Add base importer pipeline

keshav-space · keshav-space · commit ec59d7acb072 · 2024-08-27T00:35:06.000+05:30
Signed-off-by: Keshav Priyadarshi &lt;git@keshav.space&gt;
diff --git a/vulnerabilities/importers/__init__.py b/vulnerabilities/importers/__init__.py
@@ -30,7 +30,6 @@
 from vulnerabilities.importers import oss_fuzz
 from vulnerabilities.importers import postgresql
 from vulnerabilities.importers import project_kb_msr2019
-from vulnerabilities.importers import pypa
 from vulnerabilities.importers import pysec
 from vulnerabilities.importers import redhat
 from vulnerabilities.importers import retiredotnet
@@ -40,13 +39,13 @@
 from vulnerabilities.importers import ubuntu_usn
 from vulnerabilities.importers import vulnrichment
 from vulnerabilities.importers import xen
+from vulnerabilities.pipelines import pypa_importer
 
 IMPORTERS_REGISTRY = [
     nvd.NVDImporter,
     github.GitHubAPIImporter,
     gitlab.GitLabAPIImporter,
     npm.NpmImporter,
-    pypa.PyPaImporter,
     nginx.NginxImporter,
     pysec.PyPIImporter,
     alpine_linux.AlpineImporter,
@@ -75,6 +74,7 @@
     github_osv.GithubOSVImporter,
     epss.EPSSImporter,
     vulnrichment.VulnrichImporter,
+    pypa_importer.PyPaImporterPipeline,
 ]
 
 IMPORTERS_REGISTRY = {x.qualified_name: x for x in IMPORTERS_REGISTRY}
diff --git a/vulnerabilities/management/commands/import.py b/vulnerabilities/management/commands/import.py
@@ -13,6 +13,7 @@
 
 from vulnerabilities.import_runner import ImportRunner
 from vulnerabilities.importers import IMPORTERS_REGISTRY
+from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipeline
 
 
 class Command(BaseCommand):
@@ -57,6 +58,13 @@ def import_data(self, importers):
 
         for importer in importers:
             self.stdout.write(f"Importing data using {importer.qualified_name}")
+            if issubclass(importer, VulnerableCodeBaseImporterPipeline):
+                status, error = importer().execute()
+                if status != 0:
+                    self.stdout.write(error)
+                    failed_importers.append(importer.qualified_name)
+                continue
+
             try:
                 ImportRunner(importer).run()
                 self.stdout.write(
diff --git a/vulnerabilities/pipelines/__init__.py b/vulnerabilities/pipelines/__init__.py
@@ -9,9 +9,16 @@
 import logging
 from datetime import datetime
 from datetime import timezone
+from traceback import format_exc as traceback_format_exc
+from typing import Iterable
 
 from aboutcode.pipeline import BasePipeline
+from aboutcode.pipeline import LoopProgress
 
+from vulnerabilities import import_runner
+from vulnerabilities.importer import AdvisoryData
+from vulnerabilities.improvers.default import DefaultImporter
+from vulnerabilities.models import Advisory
 from vulnerabilities.utils import classproperty
 
 module_logger = logging.getLogger(__name__)
@@ -32,3 +39,104 @@ def qualified_name(cls):
         Fully qualified name prefixed with the module name of the pipeline used in logging.
         """
         return f"{cls.__module__}.{cls.__qualname__}"
+
+
+class VulnerableCodeBaseImporterPipeline(VulnerableCodePipeline):
+    """
+    Base importer pipeline for importing advisories.
+
+    Uses:
+        Subclass this Pipeline and implement ``advisories_count`` and ``collect_advisories`` method.
+        Also override the ``steps`` if needed.
+    """
+
+    license_url = None
+    spdx_license_expression = None
+    repo_url = None
+    importer_name = None
+
+    @classmethod
+    def steps(cls):
+        return (
+            # Add step for downloading/cloning resource as required.
+            cls.collect_and_store_advisories,
+            cls.import_new_advisories,
+            # Add step for removing downloaded/cloned resource as required.
+        )
+
+    def collect_advisories(self) -> Iterable[AdvisoryData]:
+        """
+        Yield AdvisoryData for importer pipeline.
+
+        Populate the `self.collected_advisories_count` field and yield AdvisoryData
+        """
+        raise NotImplementedError
+
+    def advisories_count(self) -> int:
+        """
+        Return the estimated AdvisoryData to be yielded by ``collect_advisories``.
+
+        Used by ``collect_and_store_advisories`` to log the progress of advisory collection.
+        """
+        raise NotImplementedError
+
+    def collect_and_store_advisories(self):
+        self.new_advisories = []
+
+        collected_advisory_count = 0
+        progress = LoopProgress(total_iterations=self.advisories_count(), logger=self.log)
+        for advisory in progress.iter(self.collect_advisories()):
+            self.insert_advisory(advisory=advisory)
+            collected_advisory_count += 1
+
+        self.log(f"Successfully collected {collected_advisory_count:,d} advisories")
+
+    def insert_advisory(self, advisory: AdvisoryData):
+        try:
+            obj, created = Advisory.objects.get_or_create(
+                aliases=advisory.aliases,
+                summary=advisory.summary,
+                affected_packages=[pkg.to_dict() for pkg in advisory.affected_packages],
+                references=[ref.to_dict() for ref in advisory.references],
+                date_published=advisory.date_published,
+                weaknesses=advisory.weaknesses,
+                defaults={
+                    "created_by": self.qualified_name,
+                    "date_collected": datetime.now(timezone.utc),
+                },
+                url=advisory.url,
+            )
+            if created:
+                self.new_advisories.append(obj)
+        except Exception as e:
+            self.log(
+                f"Error while processing {advisory!r} with aliases {advisory.aliases!r}: {e!r} \n {traceback_format_exc()}",
+                level=logging.ERROR,
+            )
+
+    def import_new_advisories(self):
+        new_advisories_count = len(self.new_advisories)
+
+        imported_advisory_count = 0
+        progress = LoopProgress(total_iterations=new_advisories_count, logger=self.log)
+        for advisory in progress.iter(self.new_advisories):
+            self.import_advisory(advisory=advisory)
+            imported_advisory_count += 1
+
+        self.log(f"Successfully imported {imported_advisory_count:,d} new advisories")
+
+    def import_advisory(self, advisory) -> None:
+        if advisory.date_imported:
+            return
+        try:
+            advisory_importer = DefaultImporter(advisories=[advisory])
+            inferences = advisory_importer.get_inferences(advisory_data=advisory.to_advisory_data())
+            import_runner.process_inferences(
+                inferences=inferences,
+                advisory=advisory,
+                improver_name=self.qualified_name,
+            )
+        except Exception as e:
+            self.log(
+                f"Failed to process advisory: {advisory!r} with error {e!r}", level=logging.ERROR
+            )
diff --git a/vulnerabilities/pipelines/pypa_importer.py b/vulnerabilities/pipelines/pypa_importer.py
diff --git a/vulnerabilities/tests/pipelines/test_pypa_importer_pipeline.py b/vulnerabilities/tests/pipelines/test_pypa_importer_pipeline.py
diff --git a/vulnerabilities/tests/test_data_source.py b/vulnerabilities/tests/test_data_source.py
@@ -26,7 +26,6 @@
 from vulnerabilities.importers.istio import IstioImporter
 from vulnerabilities.importers.mozilla import MozillaImporter
 from vulnerabilities.importers.npm import NpmImporter
-from vulnerabilities.importers.pypa import PyPaImporter
 from vulnerabilities.importers.retiredotnet import RetireDotnetImporter
 from vulnerabilities.importers.ruby import RubyImporter
 from vulnerabilities.oval_parser import OvalParser
@@ -124,7 +123,6 @@ def test_git_importer(mock_clone):
         MozillaImporter,
         NpmImporter,
         RetireDotnetImporter,
-        PyPaImporter,
         RubyImporter,
         GithubOSVImporter,
     ],
