Change models

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/pipelines/fill_vulnerability_summary_pipeline.py

@@ -17,51 +17,43 @@
 from vulnerabilities.pipelines import VulnerableCodePipeline
 
 
-class FillVulnerabilitySummariesPipeline(VulnerableCodePipeline):
-    """Pipeline to fill missing vulnerability summaries from advisories."""
+class PopulateVulnerabilitySummariesPipeline(VulnerableCodePipeline):
+    """Pipeline to populate missing vulnerability summaries from advisories."""
 
-    pipeline_id = "fill_vulnerability_summaries"
+    pipeline_id = "populate_vulnerability_summaries"
 
     @classmethod
     def steps(cls):
-        return (cls.fill_missing_summaries,)
+        return (cls.populate_missing_summaries,)
 
-    def fill_missing_summaries(self):
-        """Find vulnerabilities without summaries and fill them using advisories with the same aliases."""
+    def populate_missing_summaries(self):
+        """Find vulnerabilities with missing summaries and populate them using advisories with the same aliases."""
         vulnerabilities_qs = Vulnerability.objects.filter(summary="")
         self.log(
             f"Processing {vulnerabilities_qs.count()} vulnerabilities without summaries",
             level=logging.INFO,
         )
-        nvd_importer_advisories = Advisory.objects.filter(
-            created_by="nvd_importer", summary__isnull=False
-        ).exclude(summary="")
-        self.log(
-            f"Found {nvd_importer_advisories.count()} advisories with summaries from NVD importer",
-            level=logging.INFO,
-        )
+        # nvd_importer_advisories = Advisory.objects.filter(created_by="nvd_importer").exclude(summary="")
+        # self.log(
+        #     f"Found {nvd_importer_advisories.count()} advisories with summaries from NVD importer",
+        #     level=logging.INFO,
+        # )
 
         progress = LoopProgress(total_iterations=vulnerabilities_qs.count(), logger=self.log)
 
-        for vulnerability in progress.iter(vulnerabilities_qs.paginated()):
-            aliases = vulnerability.aliases.values_list("alias", flat=True)
-            # get alias that start with CVE- with filter
-            alias = aliases.filter(alias__startswith="CVE-").first()
+        for vulnerability in progress.iter(vulnerabilities_qs.iterator()):
+            cve_alias = vulnerability.aliases.filter(alias__startswith="CVE-").first()
 
-            # check if the vulnerability has an alias
-            if not alias:
+            if not cve_alias:
                 self.log(
-                    f"Vulnerability {vulnerability.vulnerability_id} has no alias",
+                    f"Vulnerability {vulnerability.vulnerability_id} has no CVE alias",
                     level=logging.INFO,
                 )
                 continue
 
-            # check if the vulnerability has an alias that matches an advisory
-            matching_advisories = nvd_importer_advisories.filter(Q(aliases__contains=alias))
+            matching_advisories = Advisory.objects.filter(aliases=cve_alias)
 
             if matching_advisories.exists():
-                # Take the first matching advisory with a summary
-                # get the advisory that was collected the most recently
                 best_advisory = matching_advisories.order_by("-date_collected").first()
                 # Note: we filtered above to only get non-empty summaries
                 vulnerability.summary = best_advisory.summary

vulnerabilities/tests/pipelines/test_fill_vulnerability_summary_pipeline.py

@@ -17,7 +17,7 @@
 from vulnerabilities.models import Alias
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.pipelines.fill_vulnerability_summary_pipeline import (
-    FillVulnerabilitySummariesPipeline,
+    PopulateVulnerabilitySummariesPipeline,
 )
 
 
@@ -46,7 +46,7 @@ def test_fill_missing_summaries_from_nvd(self):
         )
 
         # Run the pipeline
-        pipeline = FillVulnerabilitySummariesPipeline()
+        pipeline = PopulateVulnerabilitySummariesPipeline()
         pipeline.fill_missing_summaries()
 
         # Check that the vulnerability now has a summary
@@ -65,7 +65,7 @@ def test_no_matching_advisory(self):
         Alias.objects.create(alias="CVE-2024-1234", vulnerability=vulnerability)
 
         # Run the pipeline
-        pipeline = FillVulnerabilitySummariesPipeline()
+        pipeline = PopulateVulnerabilitySummariesPipeline()
         pipeline.fill_missing_summaries()
 
         # Check that the vulnerability still has no summary
@@ -84,7 +84,7 @@ def test_vulnerability_without_alias(self):
         )
 
         # Run the pipeline
-        pipeline = FillVulnerabilitySummariesPipeline()
+        pipeline = PopulateVulnerabilitySummariesPipeline()
         pipeline.fill_missing_summaries()
 
         # Check that the vulnerability still has no summary
@@ -112,7 +112,7 @@ def test_non_nvd_advisory_ignored(self):
         )
 
         # Run the pipeline
-        pipeline = FillVulnerabilitySummariesPipeline()
+        pipeline = PopulateVulnerabilitySummariesPipeline()
         pipeline.fill_missing_summaries()
 
         # Check that the vulnerability still has no summary
@@ -145,7 +145,7 @@ def test_multiple_matching_advisories(self):
         )
 
         # Run the pipeline
-        pipeline = FillVulnerabilitySummariesPipeline()
+        pipeline = PopulateVulnerabilitySummariesPipeline()
         pipeline.fill_missing_summaries()
 
         # Check that the vulnerability now has the most recent summary