Update (exploitdb, kev, metasploit ) pipelines to do single db query

Signed-off-by: ziad hany <ziadhany2016@gmail.com>

Author: ziadhany

vulnerabilities/pipelines/enhance_with_exploitdb.py

@@ -78,20 +78,16 @@ def add_exploit(self):
 
 
 def add_vulnerability_exploit(row, logger):
-    vulnerabilities = set()
-
     aliases = row["codes"].split(";") if row["codes"] else []
 
     if not aliases:
         return 0
 
-    for raw_alias in aliases:
-        try:
-            if alias := Alias.objects.get(alias=raw_alias):
-                if alias.vulnerability:
-                    vulnerabilities.add(alias.vulnerability)
-        except Alias.DoesNotExist:
-            continue
+    vulnerabilities = (
+        Alias.objects.filter(alias__in=aliases, vulnerability__isnull=False)
+        .values_list("vulnerability_id", flat=True)
+        .distinct()
+    )
 
     if not vulnerabilities:
         logger(f"No vulnerability found for aliases {aliases}")
@@ -105,7 +101,7 @@ def add_vulnerability_exploit(row, logger):
         add_exploit_references(row["codes"], row["source_url"], row["file"], vulnerability, logger)
         try:
             Exploit.objects.update_or_create(
-                vulnerability=vulnerability,
+                vulnerability_id=vulnerability,
                 data_source="Exploit-DB",
                 defaults={
                     "date_added": date_added,
@@ -126,7 +122,7 @@ def add_vulnerability_exploit(row, logger):
     return 1
 
 
-def add_exploit_references(ref_id, direct_url, path, vul, logger):
+def add_exploit_references(ref_id, direct_url, path, vul_id, logger):
     url_map = {
         "file_url": f"https://gitlab.com/exploit-database/exploitdb/-/blob/main/{path}",
         "direct_url": direct_url,
@@ -145,7 +141,7 @@ def add_exploit_references(ref_id, direct_url, path, vul, logger):
 
                 if created:
                     VulnerabilityRelatedReference.objects.get_or_create(
-                        vulnerability=vul,
+                        vulnerability_id=vul_id,
                         reference=ref,
                     )
 

vulnerabilities/pipelines/enhance_with_kev.py

@@ -71,31 +71,29 @@ def add_vulnerability_exploit(kev_vul, logger):
     if not cve_id:
         return 0
 
-    vulnerability = None
-    try:
-        if alias := Alias.objects.get(alias=cve_id):
-            if alias.vulnerability:
-                vulnerability = alias.vulnerability
-    except Alias.DoesNotExist:
-        logger(f"No vulnerability found for aliases {cve_id}")
-        return 0
+    vulnerabilities = (
+        Alias.objects.filter(alias=cve_id, vulnerability__isnull=False)
+        .values_list("vulnerability", flat=True)
+        .distinct()
+    )
 
-    if not vulnerability:
+    if not vulnerabilities:
         logger(f"No vulnerability found for aliases {cve_id}")
         return 0
 
-    Exploit.objects.update_or_create(
-        vulnerability=vulnerability,
-        data_source="KEV",
-        defaults={
-            "description": kev_vul["shortDescription"],
-            "date_added": kev_vul["dateAdded"],
-            "required_action": kev_vul["requiredAction"],
-            "due_date": kev_vul["dueDate"],
-            "notes": kev_vul["notes"],
-            "known_ransomware_campaign_use": True
-            if kev_vul["knownRansomwareCampaignUse"] == "Known"
-            else False,
-        },
-    )
+    for vulnerability in vulnerabilities:
+        Exploit.objects.update_or_create(
+            vulnerability_id=vulnerability,
+            data_source="KEV",
+            defaults={
+                "description": kev_vul["shortDescription"],
+                "date_added": kev_vul["dateAdded"],
+                "required_action": kev_vul["requiredAction"],
+                "due_date": kev_vul["dueDate"],
+                "notes": kev_vul["notes"],
+                "known_ransomware_campaign_use": True
+                if kev_vul["knownRansomwareCampaignUse"] == "Known"
+                else False,
+            },
+        )
     return 1

vulnerabilities/pipelines/enhance_with_metasploit.py

@@ -66,7 +66,6 @@ def add_vulnerability_exploits(self):
 
 
 def add_vulnerability_exploit(record, logger):
-    vulnerabilities = set()
     references = record.get("references", [])
 
     interesting_references = [
@@ -76,13 +75,11 @@ def add_vulnerability_exploit(record, logger):
     if not interesting_references:
         return 0
 
-    for ref in interesting_references:
-        try:
-            if alias := Alias.objects.get(alias=ref):
-                if alias.vulnerability:
-                    vulnerabilities.add(alias.vulnerability)
-        except Alias.DoesNotExist:
-            continue
+    vulnerabilities = (
+        Alias.objects.filter(alias__in=interesting_references, vulnerability__isnull=False)
+        .values_list("vulnerability", flat=True)
+        .distinct()
+    )
 
     if not vulnerabilities:
         logger(f"No vulnerability found for aliases {interesting_references}")
@@ -108,7 +105,7 @@ def add_vulnerability_exploit(record, logger):
 
     for vulnerability in vulnerabilities:
         Exploit.objects.update_or_create(
-            vulnerability=vulnerability,
+            vulnerability_id=vulnerability,
             data_source="Metasploit",
             defaults={
                 "description": description,