Use pipeline_id for created_by field

- For now pipeline_id should be module name of pipeline

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

vulnerabilities/importers/__init__.py

@@ -39,6 +39,7 @@
 from vulnerabilities.importers import ubuntu_usn
 from vulnerabilities.importers import vulnrichment
 from vulnerabilities.importers import xen
+from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipeline
 from vulnerabilities.pipelines import npm_importer
 from vulnerabilities.pipelines import pypa_importer
 
@@ -79,4 +80,7 @@
     npm_importer.NpmImporterPipeline,
 ]
 
-IMPORTERS_REGISTRY = {x.qualified_name: x for x in IMPORTERS_REGISTRY}
+IMPORTERS_REGISTRY = {
+    x.pipeline_id if issubclass(x, VulnerableCodeBaseImporterPipeline) else x.qualified_name: x
+    for x in IMPORTERS_REGISTRY
+}

vulnerabilities/improvers/__init__.py

@@ -10,6 +10,7 @@
 from vulnerabilities.improvers import valid_versions
 from vulnerabilities.improvers import vulnerability_kev
 from vulnerabilities.improvers import vulnerability_status
+from vulnerabilities.pipelines import VulnerableCodePipeline
 from vulnerabilities.pipelines import flag_ghost_packages
 
 IMPROVERS_REGISTRY = [
@@ -34,4 +35,7 @@
     flag_ghost_packages.FlagGhostPackagePipeline,
 ]
 
-IMPROVERS_REGISTRY = {x.qualified_name: x for x in IMPROVERS_REGISTRY}
+IMPROVERS_REGISTRY = {
+    x.pipeline_id if issubclass(x, VulnerableCodePipeline) else x.qualified_name: x
+    for x in IMPROVERS_REGISTRY
+}

vulnerabilities/management/commands/import.py

@@ -57,14 +57,15 @@ def import_data(self, importers):
         failed_importers = []
 
         for importer in importers:
-            self.stdout.write(f"Importing data using {importer.qualified_name}")
             if issubclass(importer, VulnerableCodeBaseImporterPipeline):
+                self.stdout.write(f"Importing data using {importer.pipeline_id}")
                 status, error = importer().execute()
                 if status != 0:
                     self.stdout.write(error)
-                    failed_importers.append(importer.qualified_name)
+                    failed_importers.append(importer.pipeline_id)
                 continue
 
+            self.stdout.write(f"Importing data using {importer.qualified_name}")
             try:
                 ImportRunner(importer).run()
                 self.stdout.write(

vulnerabilities/management/commands/improve.py

@@ -56,14 +56,15 @@ def improve_data(self, improvers):
         failed_improvers = []
 
         for improver in improvers:
-            self.stdout.write(f"Improving data using {improver.qualified_name}")
             if issubclass(improver, VulnerableCodePipeline):
+                self.stdout.write(f"Improving data using {improver.pipeline_id}")
                 status, error = improver().execute()
                 if status != 0:
                     self.stdout.write(error)
-                    failed_improvers.append(improver.qualified_name)
+                    failed_improvers.append(improver.pipeline_id)
                 continue
 
+            self.stdout.write(f"Improving data using {improver.qualified_name}")
             try:
                 ImproveRunner(improver_class=improver).run()
                 self.stdout.write(

vulnerabilities/pipelines/__init__.py

@@ -27,6 +27,8 @@
 
 
 class VulnerableCodePipeline(BasePipeline):
+    pipeline_id = None  # Unique Pipeline ID
+
     def log(self, message, level=logging.INFO):
         """Log the given `message` to the current module logger and execution_log."""
         now_local = datetime.now(timezone.utc).astimezone()
@@ -36,11 +38,12 @@ def log(self, message, level=logging.INFO):
         self.append_to_log(message)
 
     @classproperty
-    def qualified_name(cls):
-        """
-        Fully qualified name prefixed with the module name of the pipeline used in logging.
-        """
-        return f"{cls.__module__}.{cls.__qualname__}"
+    def pipeline_id(cls):
+        """Return unique pipeline_id set in cls.pipeline_id"""
+
+        if cls.pipeline_id is None or cls.pipeline_id == "":
+            raise NotImplementedError("pipeline_id is not defined or is empty")
+        return cls.pipeline_id
 
 
 class VulnerableCodeBaseImporterPipeline(VulnerableCodePipeline):
@@ -52,6 +55,7 @@ class VulnerableCodeBaseImporterPipeline(VulnerableCodePipeline):
         Also override the ``steps`` and ``advisory_confidence`` as needed.
     """
 
+    pipeline_id = None  # Unique Pipeline ID, this should be the name of pipeline module.
     license_url = None
     spdx_license_expression = None
     repo_url = None
@@ -89,7 +93,7 @@ def collect_and_store_advisories(self):
         for advisory in progress.iter(self.collect_advisories()):
             if _obj := insert_advisory(
                 advisory=advisory,
-                pipeline_name=self.qualified_name,
+                pipeline_id=self.pipeline_id,
                 logger=self.log,
             ):
                 collected_advisory_count += 1
@@ -98,7 +102,7 @@ def collect_and_store_advisories(self):
 
     def import_new_advisories(self):
         new_advisories = Advisory.objects.filter(
-            created_by=self.qualified_name,
+            created_by=self.pipeline_id,
             date_imported__isnull=True,
         )
 
@@ -119,7 +123,7 @@ def import_advisory(self, advisory: Advisory) -> int:
         try:
             import_advisory(
                 advisory=advisory,
-                pipeline_name=self.qualified_name,
+                pipeline_id=self.pipeline_id,
                 confidence=self.advisory_confidence,
                 logger=self.log,
             )

vulnerabilities/pipelines/flag_ghost_packages.py

@@ -23,6 +23,8 @@
 class FlagGhostPackagePipeline(VulnerableCodePipeline):
     """Detect and flag packages that do not exist upstream."""
 
+    pipeline_id = "flag_ghost_packages"
+
     @classmethod
     def steps(cls):
         return (cls.flag_ghost_packages,)

vulnerabilities/pipelines/pypa_importer.py

@@ -14,7 +14,6 @@
 from fetchcode.vcs import fetch_via_vcs
 
 from vulnerabilities.importer import AdvisoryData
-from vulnerabilities.importers.osv import parse_advisory_data
 from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipeline
 from vulnerabilities.utils import get_advisory_url
 
@@ -24,6 +23,8 @@
 class PyPaImporterPipeline(VulnerableCodeBaseImporterPipeline):
     """Collect advisories from PyPA GitHub repository."""
 
+    pipeline_id = "pypa_importer"
+
     spdx_license_expression = "CC-BY-4.0"
     license_url = "https://github.com/pypa/advisory-database/blob/main/LICENSE"
     repo_url = "git+https://github.com/pypa/advisory-database"
@@ -47,6 +48,8 @@ def advisories_count(self):
         return sum(1 for _ in vulns_directory.rglob("*.yaml"))
 
     def collect_advisories(self) -> Iterable[AdvisoryData]:
+        from vulnerabilities.importers.osv import parse_advisory_data
+
         base_directory = Path(self.vcs_response.dest_dir)
         vulns_directory = base_directory / "vulns"
 

vulnerabilities/pipes/advisory.py

@@ -26,7 +26,7 @@
 from vulnerabilities.models import Weakness
 
 
-def insert_advisory(advisory: AdvisoryData, pipeline_name: str, logger: Callable = None):
+def insert_advisory(advisory: AdvisoryData, pipeline_id: str, logger: Callable = None):
     obj = None
     try:
         obj, _ = Advisory.objects.get_or_create(
@@ -38,7 +38,7 @@ def insert_advisory(advisory: AdvisoryData, pipeline_name: str, logger: Callable
             weaknesses=advisory.weaknesses,
             url=advisory.url,
             defaults={
-                "created_by": pipeline_name,
+                "created_by": pipeline_id,
                 "date_collected": datetime.now(timezone.utc),
             },
         )
@@ -55,7 +55,7 @@ def insert_advisory(advisory: AdvisoryData, pipeline_name: str, logger: Callable
 @transaction.atomic
 def import_advisory(
     advisory: Advisory,
-    pipeline_name: str,
+    pipeline_id: str,
     confidence: int = MAX_CONFIDENCE,
     logger: Callable = None,
 ):
@@ -141,7 +141,7 @@ def import_advisory(
         PackageRelatedVulnerability(
             vulnerability=vulnerability,
             package=vulnerable_package,
-            created_by=pipeline_name,
+            created_by=pipeline_id,
             confidence=confidence,
             fix=False,
         ).update_or_create(advisory=advisory)
@@ -151,7 +151,7 @@ def import_advisory(
         PackageRelatedVulnerability(
             vulnerability=vulnerability,
             package=fixed_package,
-            created_by=pipeline_name,
+            created_by=pipeline_id,
             confidence=confidence,
             fix=True,
         ).update_or_create(advisory=advisory)

vulnerabilities/tests/pipelines/test_base_pipeline.py

@@ -62,6 +62,8 @@ def test_collect_and_store_advisories(self, mock_advisories_count, mock_collect_
         self.assertEqual(0, models.Advisory.objects.count())
 
         base_pipeline = VulnerableCodeBaseImporterPipeline()
+        base_pipeline.pipeline_id = "test_pipeline"
+
         base_pipeline.collect_and_store_advisories()
 
         mock_advisories_count.assert_called_once()
@@ -74,13 +76,14 @@ def test_collect_and_store_advisories(self, mock_advisories_count, mock_collect_
         expected_aliases = advisory_data1.aliases
 
         self.assertEqual(expected_aliases, result_aliases)
-        self.assertEqual(base_pipeline.qualified_name, collected_advisory.created_by)
+        self.assertEqual(base_pipeline.pipeline_id, collected_advisory.created_by)
 
     def test_import_new_advisories(self):
         self.assertEqual(0, models.Vulnerability.objects.count())
 
         base_pipeline = VulnerableCodeBaseImporterPipeline()
-        advisory1 = get_advisory1(created_by=base_pipeline.qualified_name)
+        base_pipeline.pipeline_id = "test_pipeline"
+        advisory1 = get_advisory1()
         base_pipeline.import_new_advisories()
 
         self.assertEqual(1, models.Vulnerability.objects.count())

vulnerabilities/tests/pipelines/test_pipeline_id.py

@@ -0,0 +1,66 @@
+#
+# Copyright (c) nexB Inc. and others. All rights reserved.
+# VulnerableCode is a trademark of nexB Inc.
+# SPDX-License-Identifier: Apache-2.0
+# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
+# See https://github.com/aboutcode-org/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+
+import importlib
+import inspect
+import unittest
+from pathlib import Path
+
+from vulnerabilities.pipelines import VulnerableCodePipeline
+
+PIPELINE_DIR = Path(__file__).parent.parent.parent / "pipelines"
+
+
+class PipelineTests(unittest.TestCase):
+    def setUp(self):
+        self.pipeline_dict = self.collect_pipeline_ids()
+
+    def collect_pipeline_ids(self):
+        """Return pipeline_ids from all the VulnerableCodePipeline."""
+        pipeline_dict = {}
+
+        for pipeline in PIPELINE_DIR.glob("*.py"):
+            if pipeline.name == "__init__.py":
+                continue
+
+            module_name = pipeline.stem
+            module = importlib.import_module(f"vulnerabilities.pipelines.{module_name}")
+
+            for _, obj in inspect.getmembers(module, inspect.isclass):
+                if issubclass(obj, VulnerableCodePipeline) and obj is not VulnerableCodePipeline:
+                    pipeline_id = obj.pipeline_id
+                    pipeline_dict[obj] = pipeline_id
+                    break
+
+        return pipeline_dict
+
+    def test_no_empty_pipeline_ids(self):
+        empty_pipeline_ids = [cls for cls, pid in self.pipeline_dict.items() if pid == ""]
+
+        if empty_pipeline_ids:
+            error_messages = [
+                f"{cls.__name__} has empty pipeline_id." for cls in empty_pipeline_ids
+            ]
+            error_message = "`pipeline_id` should not be empty string:\n" + "\n".join(
+                error_messages
+            )
+            assert False, error_message
+
+    def test_no_none_pipeline_ids(self):
+        none_pipeline_ids = [cls for cls, pid in self.pipeline_dict.items() if pid == None]
+
+        if none_pipeline_ids:
+            error_messages = [f"{cls.__name__} has None pipeline_id." for cls in none_pipeline_ids]
+            error_message = "`pipeline_id` should not be None:\n" + "\n".join(error_messages)
+            assert False, error_message
+
+    def test_unique_pipeline_ids(self):
+        pipeline_ids = self.pipeline_dict.values()
+        unique_ids = set(pipeline_ids)
+        assert len(pipeline_ids) == len(unique_ids), "`pipeline_id` should be unique."

vulnerabilities/tests/pipes/test_advisory.py

@@ -60,16 +60,16 @@ def get_all_vulnerability_relationships_objects():
 @pytest.mark.django_db
 def test_vulnerability_pipes_importer_import_advisory():
     advisory1 = get_advisory1(created_by="test_importer_pipeline")
-    import_advisory(advisory=advisory1, pipeline_name="test_importer_pipeline")
+    import_advisory(advisory=advisory1, pipeline_id="test_importer_pipeline")
     all_vulnerability_relation_objects = get_all_vulnerability_relationships_objects()
-    import_advisory(advisory=advisory1, pipeline_name="test_importer_pipeline")
+    import_advisory(advisory=advisory1, pipeline_id="test_importer_pipeline")
     assert all_vulnerability_relation_objects == get_all_vulnerability_relationships_objects()
 
 
 @pytest.mark.django_db
 def test_vulnerability_pipes_importer_import_advisory_different_pipelines():
     advisory1 = get_advisory1(created_by="test_importer_pipeline")
-    import_advisory(advisory=advisory1, pipeline_name="test_importer1_pipeline")
+    import_advisory(advisory=advisory1, pipeline_id="test_importer1_pipeline")
     all_vulnerability_relation_objects = get_all_vulnerability_relationships_objects()
-    import_advisory(advisory=advisory1, pipeline_name="test_importer2_pipeline")
+    import_advisory(advisory=advisory1, pipeline_id="test_importer2_pipeline")
     assert all_vulnerability_relation_objects == get_all_vulnerability_relationships_objects()