WIP

Signed-off-by: Shivam Sandbhor <shivam.sandbhor@gmail.com>

Author: sbs2001

vulnerabilities/import_runner.py

@@ -30,11 +30,10 @@
 from typing import Set
 from typing import Tuple
 from typing import Optional
+from typing import Sequence
 
 import packageurl
 from django.db import DataError
-#remove this
-from django.db import connection
 
 from vulnerabilities import models
 from vulnerabilities.data_source import Advisory, DataSource
@@ -52,12 +51,18 @@ def __post_init__(self):
         if not any([self.reference_id,self.url]) : 
             raise TypeError("VulnerabilityReference_inserter expects either reference_id or url")
 
+    def to_model_object(self):
+        return models.VulnerabilityReference(**dataclasses.asdict(self))
+
 @dataclasses.dataclass(frozen=True)
 class Vulnerability_Package_inserter :
     vulnerability : models.Vulnerability
     is_vulnerable : bool
     package : models.Package
 
+    def to_model_object(self):
+        return models.Vulnerability_Package_Relation(**dataclasses.asdict(self))
+
 class ImportRunner:
     """
     The ImportRunner is responsible for inserting and updating data about vulnerabilities and
@@ -99,7 +104,6 @@ def run(self, cutoff_date: datetime.datetime = None) -> None:
         self.importer.save()
 
         logger.debug(f'Successfully finished import for {self.importer.name}.')
-    
 
 def _process_updated_advisories(data_source: DataSource) -> None:
     """
@@ -110,77 +114,19 @@ def _process_updated_advisories(data_source: DataSource) -> None:
     bulk_update_vuln_pkg_refs = set()
     for batch in data_source.updated_advisories():
         for advisory in batch:
-            vuln, vuln_created = _get_or_create_vulnerability(advisory)
-
-            if vuln_created:
-                # This means vulnerability didn't previously exist in the db, so bulk create 
-                # is used without any hesitation
-                for id_ in set(advisory.reference_ids):
-                    # bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, reference_id=id_))
-                    bulk_create_vuln_refs.add(VulnerabilityReference_inserter(vulnerability=vuln, reference_id=id_))
-                
-                for url in set(advisory.reference_urls):
-                    bulk_create_vuln_refs.add(VulnerabilityReference_inserter(vulnerability=vuln, url=url))
-                
-            else:
-                vuln_refs_qs = models.VulnerabilityReference.objects.filter(vulnerability=vuln)
-                # This is to avoid making additional SELECT queries, and do further filtering in python
-                # Check https://stackoverflow.com/a/5989530
-                vuln_ids = {ref.reference_id for ref in vuln_refs_qs}
-                vuln_urls = {ref.url for ref in vuln_refs_qs}
-
-                for id_ in advisory.reference_ids:
-                    # Add the item preventing duplicates pass to through.
-                    if id_ not in vuln_ids : 
-                        vuln_ids.add(id_)
-                        bulk_create_vuln_refs.add(VulnerabilityReference_inserter(vulnerability=vuln, reference_id=id_))
-                
-                for url in advisory.reference_urls:
-                    # Add the item preventing duplicates pass to through.
-                    if url not in vuln_urls:
-                        vuln_urls.add(url)
-                        bulk_create_vuln_refs.add(VulnerabilityReference_inserter(vulnerability=vuln, url=url))
-                        
-
-            for ipurl in advisory.impacted_package_urls : 
-                pkg, pkg_created = _get_or_create_package(ipurl)
-                vuln_pkgs_ref = Vulnerability_Package_inserter(package=pkg,vulnerability=vuln,is_vulnerable=True)
-
-                if pkg_created or vuln_created : 
-                    bulk_create_vuln_pkg_refs.add(vuln_pkgs_ref)
-
-                else:  
-                    qs = models.Vulnerability_Package_Relation.objects.filter(package=pkg,vulnerability=vuln)
-                    if qs:
-                        if not qs[0].is_vulnerable: 
-                            qs[0].is_vulnerable = True
-                            bulk_update_vuln_pkg_refs.add(qs[0])
-
-                    else:
-                        bulk_create_vuln_pkg_refs.add(vuln_pkgs_ref)
-
-            for rpurl in advisory.resolved_package_urls : 
-                pkg, pkg_created = _get_or_create_package(rpurl)
-                vuln_pkgs_ref = Vulnerability_Package_inserter(package=pkg,vulnerability=vuln,is_vulnerable=True)
-
-                if pkg_created or vuln_created : 
-                    # For a  `Vulnerability_Package_Relation` tp exist it needs both, the package and
-                    # vulnerability to already exist in the db.
-                    bulk_create_vuln_pkg_refs.add(vuln_pkgs_ref)
-
-                else:  
-                    qs = models.Vulnerability_Package_Relation.objects.filter(package=pkg,vulnerability=vuln)
-                    if qs:
-                        if qs[0].is_vulnerable : 
-                            qs[0].is_vulnerable = False 
-                            bulk_update_vuln_pkg_refs.add(qs[0])
-
-                    else:
-                        bulk_create_vuln_pkg_refs.add(vuln_pkgs_ref)
-
+            
+            vuln, vuln_created, references = _create_vulnerability_and_references(advisory)
+            bulk_create_vuln_refs.update(references)
+            inew_refs, iupdated_refs = _create_pkg_vuln_refs(vuln, vuln_created, advisory.impacted_package_urls, is_vulnerable=True)
+            rnew_refs, rupdated_refs = _create_pkg_vuln_refs(vuln, vuln_created, advisory.resolved_package_urls, is_vulnerable=False)
+            bulk_update_vuln_pkg_refs.update(iupdated_refs.union(rupdated_refs))
+            bulk_create_vuln_pkg_refs.update(inew_refs.union(rnew_refs))
+
+    models.VulnerabilityReference.objects.bulk_create([i.to_model_object() for i in bulk_create_vuln_refs])
     models.Vulnerability_Package_Relation.objects.bulk_update(bulk_update_vuln_pkg_refs, ['is_vulnerable'])
-    models.Vulnerability_Package_Relation.objects.bulk_create([models.Vulnerability_Package_Relation(**dataclasses.asdict(i)) for i in bulk_create_vuln_pkg_refs])
-    models.VulnerabilityReference.objects.bulk_create([models.VulnerabilityReference(**dataclasses.asdict(i)) for i in bulk_create_vuln_refs])
+    print(bulk_update_vuln_pkg_refs)
+    # print(bulk_create_vuln_pkg_refs)
+    models.Vulnerability_Package_Relation.objects.bulk_create([i.to_model_object() for i in bulk_create_vuln_pkg_refs])
 
 def _process_added_advisories(data_source: DataSource) -> None:
     for batch in data_source.added_advisories():
@@ -198,6 +144,66 @@ def _process_added_advisories(data_source: DataSource) -> None:
             logger.exception(e)
 
 
+def _create_vulnerability_and_references(advisory: Advisory) : 
+    vuln, vuln_created = _get_or_create_vulnerability(advisory)
+    vuln_references = set()
+
+    if vuln_created:
+        # print('old vuln')
+        # This means vulnerability didn't previously exist in the db, so bulk create 
+        # is used without any hesitation
+        for id_ in set(advisory.reference_ids):
+            # bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, reference_id=id_))
+            vuln_references.add(VulnerabilityReference_inserter(vulnerability=vuln, reference_id=id_))
+        
+        for url in set(advisory.reference_urls):
+            vuln_references.add(VulnerabilityReference_inserter(vulnerability=vuln, url=url))
+        
+    else:
+        # print('new vuln')
+        vuln_refs_qs = models.VulnerabilityReference.objects.filter(vulnerability=vuln)
+        # This is to avoid making additional SELECT queries, and do further filtering in python
+        # Check https://stackoverflow.com/a/5989530
+        vuln_ids = {ref.reference_id for ref in vuln_refs_qs}
+        vuln_urls = {ref.url for ref in vuln_refs_qs}
+
+        for id_ in advisory.reference_ids:
+            # Add the item preventing duplicates pass to through.
+            if id_ not in vuln_ids : 
+                vuln_ids.add(id_)
+                vuln_references.add(VulnerabilityReference_inserter(vulnerability=vuln, reference_id=id_))
+        
+        for url in advisory.reference_urls:
+            # Add the item preventing duplicates pass to through.
+            if url not in vuln_urls:
+                vuln_urls.add(url)
+                vuln_references.add(VulnerabilityReference_inserter(vulnerability=vuln, url=url))
+    
+    return vuln, vuln_created, vuln_references
+
+def _create_pkg_vuln_refs(vuln:models.Vulnerability, vuln_created:bool, purls:Sequence[PackageURL], is_vulnerable:bool) : 
+
+    new_refs, updated_refs = set(), set()            
+    for purl in purls:
+        pkg, pkg_created = _get_or_create_package(purl)
+        vuln_pkgs_ref = Vulnerability_Package_inserter(package=pkg,vulnerability=vuln,is_vulnerable=is_vulnerable)
+
+        if pkg_created or vuln_created : 
+            new_refs.add(vuln_pkgs_ref)
+
+        else:  
+            qs = models.Vulnerability_Package_Relation.objects.filter(package=pkg,vulnerability=vuln)
+            if qs:
+                if qs[0].is_vulnerable != is_vulnerable: 
+                    qs[0].is_vulnerable = is_vulnerable
+                    updated_refs.add(qs[0])
+
+            else:
+                new_refs.add(vuln_pkgs_ref)
+    
+    return new_refs, updated_refs
+
+
 def _get_or_create_vulnerability(advisory: Advisory) -> Tuple[models.Vulnerability, bool]:
     if advisory.cve_id:
         query_kwargs = {'cve_id': advisory.cve_id}
@@ -207,6 +213,7 @@ def _get_or_create_vulnerability(advisory: Advisory) -> Tuple[models.Vulnerabili
         return models.Vulnerability.objects.create(), True
 
     vuln, created = models.Vulnerability.objects.get_or_create(**query_kwargs)
+    print(vuln,"   ", created)
 
     if advisory.summary and vuln.summary != advisory.summary:
         vuln.summary = advisory.summary

vulnerabilities/models.py

@@ -116,7 +116,9 @@ class Vulnerability_Package_Relation(models.Model):
     vulnerability = models.ForeignKey(Vulnerability, on_delete=models.CASCADE)
     is_vulnerable = models.BooleanField()
 
-
+    class Meta:
+        unique_together = ('package', 'vulnerability')
+        
 
 class Importer(models.Model):
     """