Skip to content

Commit de3ecf6

Browse files
committed
Update (exploitdb, kev, metasploit ) pipelines to do single db query
Signed-off-by: ziad hany <ziadhany2016@gmail.com>
1 parent 1e208cf commit de3ecf6

File tree

3 files changed

+35
-44
lines changed

3 files changed

+35
-44
lines changed

vulnerabilities/pipelines/enhance_with_exploitdb.py

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -78,20 +78,16 @@ def add_exploit(self):
7878

7979

8080
def add_vulnerability_exploit(row, logger):
81-
vulnerabilities = set()
82-
8381
aliases = row["codes"].split(";") if row["codes"] else []
8482

8583
if not aliases:
8684
return 0
8785

88-
for raw_alias in aliases:
89-
try:
90-
if alias := Alias.objects.get(alias=raw_alias):
91-
if alias.vulnerability:
92-
vulnerabilities.add(alias.vulnerability)
93-
except Alias.DoesNotExist:
94-
continue
86+
vulnerabilities = (
87+
Alias.objects.filter(alias__in=aliases, vulnerability__isnull=False)
88+
.values_list("vulnerability_id", flat=True)
89+
.distinct()
90+
)
9591

9692
if not vulnerabilities:
9793
logger(f"No vulnerability found for aliases {aliases}")
@@ -105,7 +101,7 @@ def add_vulnerability_exploit(row, logger):
105101
add_exploit_references(row["codes"], row["source_url"], row["file"], vulnerability, logger)
106102
try:
107103
Exploit.objects.update_or_create(
108-
vulnerability=vulnerability,
104+
vulnerability_id=vulnerability,
109105
data_source="Exploit-DB",
110106
defaults={
111107
"date_added": date_added,
@@ -126,7 +122,7 @@ def add_vulnerability_exploit(row, logger):
126122
return 1
127123

128124

129-
def add_exploit_references(ref_id, direct_url, path, vul, logger):
125+
def add_exploit_references(ref_id, direct_url, path, vul_id, logger):
130126
url_map = {
131127
"file_url": f"https://gitlab.com/exploit-database/exploitdb/-/blob/main/{path}",
132128
"direct_url": direct_url,
@@ -145,7 +141,7 @@ def add_exploit_references(ref_id, direct_url, path, vul, logger):
145141

146142
if created:
147143
VulnerabilityRelatedReference.objects.get_or_create(
148-
vulnerability=vul,
144+
vulnerability_id=vul_id,
149145
reference=ref,
150146
)
151147

vulnerabilities/pipelines/enhance_with_kev.py

Lines changed: 21 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -71,31 +71,29 @@ def add_vulnerability_exploit(kev_vul, logger):
7171
if not cve_id:
7272
return 0
7373

74-
vulnerability = None
75-
try:
76-
if alias := Alias.objects.get(alias=cve_id):
77-
if alias.vulnerability:
78-
vulnerability = alias.vulnerability
79-
except Alias.DoesNotExist:
80-
logger(f"No vulnerability found for aliases {cve_id}")
81-
return 0
74+
vulnerabilities = (
75+
Alias.objects.filter(alias=cve_id, vulnerability__isnull=False)
76+
.values_list("vulnerability", flat=True)
77+
.distinct()
78+
)
8279

83-
if not vulnerability:
80+
if not vulnerabilities:
8481
logger(f"No vulnerability found for aliases {cve_id}")
8582
return 0
8683

87-
Exploit.objects.update_or_create(
88-
vulnerability=vulnerability,
89-
data_source="KEV",
90-
defaults={
91-
"description": kev_vul["shortDescription"],
92-
"date_added": kev_vul["dateAdded"],
93-
"required_action": kev_vul["requiredAction"],
94-
"due_date": kev_vul["dueDate"],
95-
"notes": kev_vul["notes"],
96-
"known_ransomware_campaign_use": True
97-
if kev_vul["knownRansomwareCampaignUse"] == "Known"
98-
else False,
99-
},
100-
)
84+
for vulnerability in vulnerabilities:
85+
Exploit.objects.update_or_create(
86+
vulnerability_id=vulnerability,
87+
data_source="KEV",
88+
defaults={
89+
"description": kev_vul["shortDescription"],
90+
"date_added": kev_vul["dateAdded"],
91+
"required_action": kev_vul["requiredAction"],
92+
"due_date": kev_vul["dueDate"],
93+
"notes": kev_vul["notes"],
94+
"known_ransomware_campaign_use": True
95+
if kev_vul["knownRansomwareCampaignUse"] == "Known"
96+
else False,
97+
},
98+
)
10199
return 1

vulnerabilities/pipelines/enhance_with_metasploit.py

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,6 @@ def add_vulnerability_exploits(self):
6666

6767

6868
def add_vulnerability_exploit(record, logger):
69-
vulnerabilities = set()
7069
references = record.get("references", [])
7170

7271
interesting_references = [
@@ -76,13 +75,11 @@ def add_vulnerability_exploit(record, logger):
7675
if not interesting_references:
7776
return 0
7877

79-
for ref in interesting_references:
80-
try:
81-
if alias := Alias.objects.get(alias=ref):
82-
if alias.vulnerability:
83-
vulnerabilities.add(alias.vulnerability)
84-
except Alias.DoesNotExist:
85-
continue
78+
vulnerabilities = (
79+
Alias.objects.filter(alias__in=interesting_references, vulnerability__isnull=False)
80+
.values_list("vulnerability", flat=True)
81+
.distinct()
82+
)
8683

8784
if not vulnerabilities:
8885
logger(f"No vulnerability found for aliases {interesting_references}")
@@ -108,7 +105,7 @@ def add_vulnerability_exploit(record, logger):
108105

109106
for vulnerability in vulnerabilities:
110107
Exploit.objects.update_or_create(
111-
vulnerability=vulnerability,
108+
vulnerability_id=vulnerability,
112109
data_source="Metasploit",
113110
defaults={
114111
"description": description,

0 commit comments

Comments
 (0)