🏗️ Add option to whether create custom vulcodes or not

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

Author: sbs2001

vulnerabilities/import_runner.py

@@ -98,7 +98,7 @@ def __init__(self, importer: models.Importer, batch_size: int):
         self.importer = importer
         self.batch_size = batch_size
 
-    def run(self, cutoff_date: datetime.datetime = None) -> None:
+    def run(self, create_vulcodes=True, cutoff_date: datetime.datetime = None) -> None:
         """
         Create a data source for the given importer and store the data retrieved in the database.
 
@@ -112,7 +112,7 @@ def run(self, cutoff_date: datetime.datetime = None) -> None:
         logger.info(f"Starting import for {self.importer.name}.")
         data_source = self.importer.make_data_source(self.batch_size, cutoff_date=cutoff_date)
         with data_source:
-            process_advisories(data_source)
+            process_advisories(data_source, create_vulcodes)
         self.importer.last_run = datetime.datetime.now(tz=datetime.timezone.utc)
         self.importer.data_source_cfg = dataclasses.asdict(data_source.config)
         self.importer.save()
@@ -133,7 +133,7 @@ def get_vuln_pkg_refs(vulnerability, package):
     )
 
 
-def process_advisories(data_source: DataSource) -> None:
+def process_advisories(data_source: DataSource, create_vulcodes) -> None:
     bulk_create_vuln_refs = set()
     bulk_create_vuln_pkg_refs = set()
     # Treat updated_advisories and added_advisories as same. Eventually
@@ -143,6 +143,9 @@ def process_advisories(data_source: DataSource) -> None:
     for batch in advisory_batches:
         for advisory in batch:
             try:
+                if not advisory.identifier and not create_vulcodes:
+                    continue
+
                 vuln, vuln_created = _get_or_create_vulnerability(advisory)
                 for vuln_ref in advisory.vuln_references:
                     ref = VulnerabilityReferenceInserter(
@@ -151,11 +154,9 @@ def process_advisories(data_source: DataSource) -> None:
                         reference_id=vuln_ref.reference_id,
                     )
 
-                    if vuln_created or not vuln_ref_exists(
-                        vuln, vuln_ref.url, vuln_ref.reference_id
-                    ):
-                        # A vulnerability reference can't exist if the
-                        # vulnerability is just created so insert it
+                    if vuln_created or not vuln_ref_exists(vuln, vuln_ref.url, vuln_ref.reference_id):
+                        # A vulnerability reference can't exist if the vulnerability is just created so
+                        # insert it
                         bulk_create_vuln_refs.add(ref)
 
                 for purl in chain(advisory.impacted_package_urls, advisory.resolved_package_urls):
@@ -175,12 +176,20 @@ def process_advisories(data_source: DataSource) -> None:
                         existing_ref = get_vuln_pkg_refs(vuln, pkg)
                         if not existing_ref:
                             bulk_create_vuln_pkg_refs.add(pkg_vuln_ref)
+                            # A vulnerability-package relationship does not exist already if either the
+                            # vulnerability or the package is just created.
 
                         else:
-                            # This handles conflicts between existing data and obtained data
-                            if existing_ref[0].is_vulnerable != pkg_vuln_ref.is_vulnerable:
-                                handle_conflicts([existing_ref[0], pkg_vuln_ref.to_model_object()])
-                                existing_ref.delete()
+                            # insert only if it there is no existing vulnerability-package relationship.
+                            existing_ref = get_vuln_pkg_refs(vuln, pkg)
+                            if not existing_ref:
+                                bulk_create_vuln_pkg_refs.add(pkg_vuln_ref)
+
+                            else:
+                                # This handles conflicts between existing data and obtained data
+                                if existing_ref[0].is_vulnerable != pkg_vuln_ref.is_vulnerable:
+                                    handle_conflicts([existing_ref[0], pkg_vuln_ref.to_model_object()])
+                                    existing_ref.delete()
             except Exception:
                 # TODO: store error but continue
                 logger.error(
@@ -249,20 +258,13 @@ def handle_conflicts(conflicts):
 def _get_or_create_vulnerability(
     advisory: Advisory,
 ) -> Tuple[models.Vulnerability, bool]:
-    if advisory.identifier:
-        query_kwargs = {"identifier": advisory.identifier}
-    elif advisory.summary:
-        query_kwargs = {"summary": advisory.summary}
-    else:
-        return models.Vulnerability.objects.create(), True
-
-    try:
-        vuln, created = models.Vulnerability.objects.get_or_create(**query_kwargs)
-        # Eventually we only want to keep summary from NVD and ignore other descriptions.
-        if advisory.summary and vuln.summary != advisory.summary:
-            vuln.summary = advisory.summary
-            vuln.save()
-        return vuln, created
+
+    vuln, created = models.Vulnerability.objects.get_or_create(identifier=advisory.identifier)
+
+    # Eventually we only want to keep summary from NVD and ignore other descriptions.
+    if advisory.summary and vuln.summary != advisory.summary:
+        vuln.summary = advisory.summary
+        vuln.save()
 
     except Exception:
         logger.error(

vulnerabilities/management/commands/import.py

@@ -32,55 +32,64 @@
 
 
 class Command(BaseCommand):
-    help = 'Import vulnerability data'
+    help = "Import vulnerability data"
 
     def add_arguments(self, parser):
         parser.add_argument(
-            '--list',
-            action='store_true',
-            help='List available data sources',
+            "--list",
+            action="store_true",
+            help="List available data sources",
+        )
+        parser.add_argument(
+            "--all", action="store_true", help="Import data from all available sources"
         )
-        parser.add_argument('--all', action='store_true',
-                            help='Import data from all available sources')
 
         parser.add_argument(
-            '--cutoff-date',
+            "--cutoff-date",
             type=datetime.fromisoformat,
-            help='ISO8601 formatted timestamp denoting the maximum age of vulnerability '
-                 'information to import.',
+            help="ISO8601 formatted timestamp denoting the maximum age of vulnerability "
+            "information to import.",
         )
-        parser.add_argument('sources', nargs='*',
-                            help='Data sources from which to import')
+        parser.add_argument("sources", nargs="*", help="Data sources from which to import")
 
         parser.add_argument(
-            '--batch_size', help='The batch size to be used for bulk inserting data')
+            "--batch_size", help="The batch size to be used for bulk inserting data"
+        )
+
+        parser.add_argument(
+            "--cv",
+            action="store_true",
+            help="This will import and assign id's to vulnerabilities without any identifiers",
+        )
 
     def handle(self, *args, **options):
         # load_importers() seeds the DB with Importers
         load_importers()
-        if options['list']:
+        if options["list"]:
             self.list_sources()
             return
 
-        if options['batch_size']:
-            self.batch_size = options['batch_size']
+        if options["batch_size"]:
+            self.batch_size = options["batch_size"]
+
+        self.create_vulcodes = options["cv"]
 
-        if options['all']:
-            self._import_data(Importer.objects.all(), options['cutoff_date'])
+        if options["all"]:
+            self._import_data(Importer.objects.all(), options["cutoff_date"])
             return
 
-        sources = options['sources']
+        sources = options["sources"]
         if not sources:
             raise CommandError(
-                'Please provide at least one data source to import from or use "--all".')
+                'Please provide at least one data source to import from or use "--all".'
+            )
 
-        self.import_data(sources, options['cutoff_date'])
+        self.import_data(sources, options["cutoff_date"])
 
     def list_sources(self):
         importers = Importer.objects.all()
-        self.stdout.write(
-            'Vulnerability data can be imported from the following sources:')
-        self.stdout.write(', '.join([i.name for i in importers]))
+        self.stdout.write("Vulnerability data can be imported from the following sources:")
+        self.stdout.write(", ".join([i.name for i in importers]))
 
     def import_data(self, names, cutoff_date):
         importers = []
@@ -93,15 +102,18 @@ def import_data(self, names, cutoff_date):
                 unknown_importers.add(name)
 
         if unknown_importers:
-            unknown_importers = ', '.join(unknown_importers)
-            raise CommandError(f'Unknown data sources: {unknown_importers}')
+            unknown_importers = ", ".join(unknown_importers)
+            raise CommandError(f"Unknown data sources: {unknown_importers}")
 
         self._import_data(importers, cutoff_date)
 
     def _import_data(self, importers, cutoff_date):
         for importer in importers:
-            self.stdout.write(f'Importing data from {importer.name}')
-            batch_size = int(getattr(self, 'batch_size', 10))
-            ImportRunner(importer, batch_size).run(cutoff_date=cutoff_date)
+            self.stdout.write(f"Importing data from {importer.name}")
+            batch_size = int(getattr(self, "batch_size", 10))
+            ImportRunner(importer, batch_size).run(
+                cutoff_date=cutoff_date, create_vulcodes=self.create_vulcodes
+            )
             self.stdout.write(
-                self.style.SUCCESS(f'Successfully imported data from {importer.name}'))
+                self.style.SUCCESS(f"Successfully imported data from {importer.name}")
+            )