Skip to content

Commit

Permalink
Fixed on tags, TA and labels
Browse files Browse the repository at this point in the history
  • Loading branch information
@TheoBhang
TheoBhang committed Jul 22, 2024
1 parent ca7038e commit 1b4e0e7
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 28 deletions.
3 changes: 2 additions & 1 deletion .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,5 @@
/V2
.dockerignore
.env
.vttools.sqlite
vttools.sqlite
__pycache__/
53 changes: 35 additions & 18 deletions app/DBHandler/db_handler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
malicious_score TEXT,
total_scans TEXT,
tags TEXT,
threat_category TEXT,
threat_labels TEXT,
link TEXT,
extension TEXT,
size TEXT,
Expand Down Expand Up @@ -231,8 +233,8 @@ def insert_url_data(self, conn, url_data):
def insert_hash_data(self, conn, hash_data):
"""Insert hash data into the hashes table"""
sql = """INSERT INTO hashes(hash, malicious_score,
total_scans, tags, link, extension, size, md5, sha1, sha256, ssdeep, tlsh, names, type, type_probability)
VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"""
total_scans, tags,threat_category,threat_labels,link, extension, size, md5, sha1, sha256, ssdeep, tlsh, names, type, type_probability)
VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"""
try:
cur = conn.cursor()

Expand All @@ -250,6 +252,8 @@ def insert_hash_data(self, conn, hash_data):
hash_data.get("malicious_score"),
hash_data.get("total_scans"),
hash_data.get("tags"),
hash_data.get("threat_category"),
hash_data.get("threat_labels"),
hash_data.get("link"),
hash_data.get("extension"),
hash_data.get("size"),
Expand Down Expand Up @@ -403,7 +407,17 @@ def csv_report(self, value_type, value, report):
return csv_report

def create_object(self, value_type, value, report):
value_object = {
if value_type == "SHA-256" or value_type == "SHA-1" or value_type == "MD5":
value_object = {
"malicious_score": NOT_FOUND_ERROR,
"total_scans": NOT_FOUND_ERROR,
"tags": NOT_FOUND_ERROR,
"threat_category": NOT_FOUND_ERROR,
"threat_labels": NOT_FOUND_ERROR,
"link": NOT_FOUND_ERROR,
}
else:
value_object = {
"malicious_score": NOT_FOUND_ERROR,
"total_scans": NOT_FOUND_ERROR,
"tags": NOT_FOUND_ERROR,
Expand All @@ -413,14 +427,11 @@ def create_object(self, value_type, value, report):
total_scans = report[3]
malicious = report[2]
tags = report[4]

self.populate_tags(value_object, tags)

self.populate_scores(
value_object, total_scans, malicious
)
self.populate_link(value_object, value, value_type)

self.populate_tags(value_object, tags)
if value_type == IPV4_PUBLIC_TYPE:
self.populate_ip_data(value_object, value, report)
elif value_type == "DOMAIN":
Expand All @@ -431,7 +442,13 @@ def create_object(self, value_type, value, report):
value_type == "SHA-256" or value_type == "SHA-1" or value_type == "MD5"
):
self.populate_hash_data(value_object, value, report)


if value_type == "SHA-256" or value_type == "SHA-1" or value_type == "MD5":
value_object["threat_category"] = report[5]
value_object["threat_labels"] = report[6]



return value_object

def populate_tags( self, value_object, tags):
Expand Down Expand Up @@ -545,16 +562,16 @@ def populate_hash_data(self, value_object, value, report):
value_object.update(
{
"hash": value,
"extension": report[6],
"size": report[7],
"md5": report[8],
"sha1": report[9],
"sha256": report[10],
"ssdeep": report[11],
"tlsh": report[12],
"names": report[13],
"type": report[14],
"type_probability": report[15],
"extension": report[8],
"size": report[9],
"md5": report[10],
"sha1": report[11],
"sha256": report[12],
"ssdeep": report[13],
"tlsh": report[14],
"names": report[15],
"type": report[16],
"type_probability": report[17],
}

)
Expand Down
46 changes: 37 additions & 9 deletions app/VirusTotal/vt_reporter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,21 +60,31 @@ def create_object(self, value_type, value, report):
"""
database = "vttools.sqlite"
conn = DBHandler().create_connection(database)
value_object = {
"malicious_score": NOT_FOUND_ERROR,
"total_scans": NOT_FOUND_ERROR,
"tags": NOT_FOUND_ERROR,
"link": NO_LINK,
}

if value_type == "SHA-256" or value_type == "SHA-1" or value_type == "MD5":
value_object = {
"malicious_score": NOT_FOUND_ERROR,
"total_scans": NOT_FOUND_ERROR,
"tags": NOT_FOUND_ERROR,
"threat_category": NOT_FOUND_ERROR,
"threat_labels": NOT_FOUND_ERROR,
"link": NO_LINK,
}
else:
value_object = {
"malicious_score": NOT_FOUND_ERROR,
"total_scans": NOT_FOUND_ERROR,
"tags": NOT_FOUND_ERROR,
"link": NO_LINK,
}
if report != NOT_FOUND_ERROR and report:
total_scans = sum(report.last_analysis_stats.values())
malicious = report.last_analysis_stats.get("malicious", 0)
self.populate_scores(
value_object, total_scans, malicious
)
self.populate_tags(value_object, report)
self.populate_link(value_object, value, value_type)
self.populate_tags(value_object, report)


if value_type == IPV4_PUBLIC_TYPE:
self.populate_ip_data(value_object, value, report)
Expand All @@ -88,9 +98,27 @@ def create_object(self, value_type, value, report):
elif (
value_type == "SHA-256" or value_type == "SHA-1" or value_type == "MD5"
):
try:
if report.popular_threat_classification:
popular_threat_classification = report.get("popular_threat_classification", {})

suggested_threat_label = popular_threat_classification.get("suggested_threat_label", NOT_FOUND_ERROR)
popular_threat_category = popular_threat_classification.get('popular_threat_category', [])

# Concaténer les valeurs en une seule chaîne de caractères
categories_str = ", ".join(category['value'] for category in popular_threat_category)
if report.popular_threat_classification:
value_object["threat_category"] = categories_str
value_object["threat_labels"] = suggested_threat_label
except Exception as e:
pass

self.populate_hash_data(value_object, value, report)
DBHandler().insert_hash_data(conn, value_object)





return value_object

def populate_tags(self, value_object, report):
Expand Down
1 change: 1 addition & 0 deletions vt3_tools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,7 @@ def analyze_values(args: argparse.Namespace, value_types: list[str]) -> None:

database = "vttools.sqlite"
quota_saved = 0
error_values = 0

with init.db_handler.create_connection(database) as conn:
if conn is not None:
Expand Down

0 comments on commit 1b4e0e7

Please sign in to comment.