Use advisory alias relation in pipelines and importer

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/import_runner.py

@@ -15,6 +15,7 @@
 
 from django.core.exceptions import ValidationError
 from django.db import transaction
+from django.db.models.query import QuerySet
 
 from vulnerabilities.importer import AdvisoryData
 from vulnerabilities.importer import Importer
@@ -96,12 +97,14 @@ def process_advisories(
         Insert advisories into the database
         Return the number of inserted advisories.
         """
+        from vulnerabilities.pipes.advisory import get_or_create_aliases
+
         count = 0
         advisories = []
         for data in advisory_datas:
             try:
+                aliases = get_or_create_aliases(aliases=data.aliases)
                 obj, created = Advisory.objects.get_or_create(
-                    aliases=data.aliases,
                     summary=data.summary,
                     affected_packages=[pkg.to_dict() for pkg in data.affected_packages],
                     references=[ref.to_dict() for ref in data.references],
@@ -113,6 +116,7 @@ def process_advisories(
                     },
                     url=data.url,
                 )
+                obj.aliases.add(*aliases)
                 if not obj.date_imported:
                     advisories.append(obj)
             except Exception as e:
@@ -148,6 +152,8 @@ def process_inferences(inferences: List[Inference], advisory: Advisory, improver
     erroneous. Also, the atomic transaction for every advisory and its
     inferences makes sure that date_imported of advisory is consistent.
     """
+    from vulnerabilities.pipes.advisory import get_or_create_aliases
+
     inferences_processed_count = 0
 
     if not inferences:
@@ -157,9 +163,10 @@ def process_inferences(inferences: List[Inference], advisory: Advisory, improver
     logger.info(f"Improving advisory id: {advisory.id}")
 
     for inference in inferences:
+        aliases = get_or_create_aliases(inference.aliases)
         vulnerability = get_or_create_vulnerability_and_aliases(
             vulnerability_id=inference.vulnerability_id,
-            aliases=inference.aliases,
+            aliases=aliases,
             summary=inference.summary,
             advisory=advisory,
         )
@@ -265,14 +272,13 @@ def create_valid_vulnerability_reference(url, reference_id=None):
 
 
 def get_or_create_vulnerability_and_aliases(
-    aliases: List[str], vulnerability_id=None, summary=None, advisory=None
+    aliases: QuerySet, vulnerability_id=None, summary=None, advisory=None
 ):
     """
     Get or create vulnerabilitiy and aliases such that all existing and new
     aliases point to the same vulnerability
     """
-    aliases = set(alias.strip() for alias in aliases if alias and alias.strip())
-    new_alias_names, existing_vulns = get_vulns_for_aliases_and_get_new_aliases(aliases)
+    new_aliases, existing_vulns = get_vulns_for_aliases_and_get_new_aliases(aliases)
 
     # All aliases must point to the same vulnerability
     vulnerability = None
@@ -310,11 +316,11 @@ def get_or_create_vulnerability_and_aliases(
         #         f"Inconsistent summary for {vulnerability.vulnerability_id}. "
         #         f"Existing: {vulnerability.summary!r}, provided: {summary!r}"
         #     )
-        associate_vulnerability_with_aliases(vulnerability=vulnerability, aliases=new_alias_names)
+        associate_vulnerability_with_aliases(vulnerability=vulnerability, aliases=new_aliases)
     else:
         try:
             vulnerability = create_vulnerability_and_add_aliases(
-                aliases=new_alias_names, summary=summary
+                aliases=new_aliases, summary=summary
             )
             importer_name = get_importer_name(advisory)
             VulnerabilityChangeLog.log_import(
@@ -324,24 +330,22 @@ def get_or_create_vulnerability_and_aliases(
             )
         except Exception as e:
             logger.error(
-                f"Cannot create vulnerability with summary {summary!r} and {new_alias_names!r} {e!r}.\n{traceback_format_exc()}."
+                f"Cannot create vulnerability with summary {summary!r} and {new_aliases!r} {e!r}.\n{traceback_format_exc()}."
             )
             return
 
     return vulnerability
 
 
-def get_vulns_for_aliases_and_get_new_aliases(aliases):
+def get_vulns_for_aliases_and_get_new_aliases(aliases: QuerySet):
     """
     Return ``new_aliases`` that are not in the database and
     ``existing_vulns`` that point to the given ``aliases``.
     """
-    new_aliases = set(aliases)
-    existing_vulns = set()
-    for alias in Alias.objects.filter(alias__in=aliases):
-        existing_vulns.add(alias.vulnerability)
-        new_aliases.remove(alias.alias)
-    return new_aliases, existing_vulns
+    new_aliases = aliases.filter(vulnerability__isnull=True)
+    existing_vulns = [alias.vulnerability for alias in aliases.filter(vulnerability__isnull=False)]
+
+    return new_aliases, list(set(existing_vulns))
 
 
 @transaction.atomic
@@ -360,7 +364,5 @@ def create_vulnerability_and_add_aliases(aliases, summary):
 
 
 def associate_vulnerability_with_aliases(aliases, vulnerability):
-    for alias_name in aliases:
-        alias = Alias(alias=alias_name, vulnerability=vulnerability)
-        alias.save()
-        logger.info(f"New alias for {vulnerability!r}: {alias_name}")
+    aliases.update(vulnerability=vulnerability)
+    logger.info(f"New alias for {vulnerability!r}: {aliases}")

vulnerabilities/improvers/vulnerability_status.py

@@ -37,11 +37,7 @@ class VulnerabilityStatusImprover(Improver):
 
     @property
     def interesting_advisories(self) -> QuerySet:
-        return (
-            Advisory.objects.filter(Q(created_by=NVDImporterPipeline.pipeline_id))
-            .distinct("aliases")
-            .paginated()
-        )
+        return Advisory.objects.filter(Q(created_by=NVDImporterPipeline.pipeline_id)).paginated()
 
     def get_inferences(self, advisory_data: AdvisoryData) -> Iterable[Inference]:
         """

vulnerabilities/pipelines/add_cvss31_to_CVEs.py

@@ -12,6 +12,7 @@
 
 from vulnerabilities import severity_systems
 from vulnerabilities.models import Advisory
+from vulnerabilities.models import Alias
 from vulnerabilities.models import VulnerabilitySeverity
 from vulnerabilities.pipelines import VulnerableCodePipeline
 
@@ -48,7 +49,6 @@ def process_cve_advisory_mapping(self):
         results = []
 
         for severity in progress.iter(nvd_severities.paginated(per_page=batch_size)):
-            print(severity.url)
             cve_pattern = re.compile(r"(CVE-\d{4}-\d{4,7})").search
             cve_match = cve_pattern(severity.url)
             if cve_match:
@@ -57,12 +57,10 @@ def process_cve_advisory_mapping(self):
                 self.log(f"Could not find CVE ID in URL: {severity.url}")
                 continue
 
-            matching_advisories = Advisory.objects.filter(
-                aliases=[cve_id],
-                created_by="nvd_importer",
-            )
+            if matching_alias := Alias.objects.get(alias=cve_id):
+                matching_advisories = matching_alias.advisories.filter(created_by="nvd_importer")
 
-            for advisory in matching_advisories:
+            for advisory in matching_advisories or []:
                 for reference in advisory.references:
                     for sev in reference.get("severities", []):
                         if sev.get("system") == "cvssv3.1":
@@ -76,7 +74,6 @@ def process_cve_advisory_mapping(self):
                             )
 
         if results:
-            print(results)
             self._process_batch(results)
 
         self.log(f"Completed processing CVE to Advisory mappings")

vulnerabilities/pipes/advisory.py

@@ -12,13 +12,15 @@
 from datetime import timezone
 from traceback import format_exc as traceback_format_exc
 from typing import Callable
+from typing import List
 
 from django.db import transaction
 
 from vulnerabilities.importer import AdvisoryData
 from vulnerabilities.improver import MAX_CONFIDENCE
 from vulnerabilities.models import Advisory
 from vulnerabilities.models import AffectedByPackageRelatedVulnerability
+from vulnerabilities.models import Alias
 from vulnerabilities.models import FixingPackageRelatedVulnerability
 from vulnerabilities.models import Package
 from vulnerabilities.models import VulnerabilityReference
@@ -27,11 +29,17 @@
 from vulnerabilities.models import Weakness
 
 
+def get_or_create_aliases(aliases: List) -> List:
+    for alias in aliases:
+        Alias.objects.get_or_create(alias=alias)
+    return Alias.objects.filter(alias__in=aliases)
+
+
 def insert_advisory(advisory: AdvisoryData, pipeline_id: str, logger: Callable = None):
-    obj = None
+    advisory_obj = None
+    aliases = get_or_create_aliases(aliases=advisory.aliases)
     try:
-        obj, _ = Advisory.objects.get_or_create(
-            aliases=advisory.aliases,
+        advisory_obj, _ = Advisory.objects.get_or_create(
             summary=advisory.summary,
             affected_packages=[pkg.to_dict() for pkg in advisory.affected_packages],
             references=[ref.to_dict() for ref in advisory.references],
@@ -43,14 +51,15 @@ def insert_advisory(advisory: AdvisoryData, pipeline_id: str, logger: Callable =
                 "date_collected": datetime.now(timezone.utc),
             },
         )
+        advisory_obj.aliases.add(*aliases)
     except Exception as e:
         if logger:
             logger(
                 f"Error while processing {advisory!r} with aliases {advisory.aliases!r}: {e!r} \n {traceback_format_exc()}",
                 level=logging.ERROR,
             )
 
-    return obj
+    return advisory_obj
 
 
 @transaction.atomic
@@ -82,9 +91,10 @@ def import_advisory(
         affected_purls.extend(package_affected_purls)
         fixed_purls.extend(package_fixed_purls)
 
+    aliases = get_or_create_aliases(advisory_data.aliases)
     vulnerability = import_runner.get_or_create_vulnerability_and_aliases(
         vulnerability_id=None,
-        aliases=advisory_data.aliases,
+        aliases=aliases,
         summary=advisory_data.summary,
         advisory=advisory,
     )