🏗️ Add option to whether create custom vulcodes or not

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

Author: sbs2001

vulnerabilities/import_runner.py

@@ -100,7 +100,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.
 
@@ -114,7 +114,7 @@ def run(self, cutoff_date: datetime.datetime = None) -> None:
         logger.debug(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()
@@ -135,7 +135,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
@@ -144,6 +144,10 @@ def process_advisories(data_source: DataSource) -> None:
     advisory_batches = chain(data_source.updated_advisories(), data_source.added_advisories())
     for batch in advisory_batches:
         for advisory in batch:
+
+            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(
@@ -242,14 +246,8 @@ 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
-
-    vuln, created = models.Vulnerability.objects.get_or_create(**query_kwargs)
+
+    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:

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}")
+            )