Add progress bar to NVDImporter

Signed-off-by: Harsh Mishra <hmisraji07@gmail.com>

Author: harsh098

requirements.txt

@@ -120,3 +120,4 @@ drf-spectacular==0.24.2
 coreapi==2.3.3
 coreschema==0.0.4
 itypes==1.2.0
+progress==1.6

vulnerabilities/importers/nvd.py

@@ -10,6 +10,7 @@
 import gzip
 import json
 from datetime import date
+from progress.bar import ChargingBar
 
 import attr
 import requests
@@ -78,10 +79,18 @@ def fetch_cve_data_1_1(starting_year=2002):
     year since ``starting_year`` defaulting to 2002.
     """
     current_year = date.today().year
+    progress_bar_for_records_fetched = ChargingBar("\tRecords fetched", max=(current_year-starting_year)+1)
+    progress_bar_for_records_fetched.start()
     # NVD json feeds start from 2002.
-    for year in range(starting_year, current_year + 1):
-        download_url = f"https://nvd.nist.gov/feeds/json/cve/1.1/nvdcve-1.1-{year}.json.gz"
-        yield year, fetch(url=download_url)
+    try:
+        for year in range(starting_year, current_year + 1):
+            try:
+                download_url = f"https://nvd.nist.gov/feeds/json/cve/1.1/nvdcve-1.1-{year}.json.gz"
+                yield year, fetch(url=download_url)
+            finally:
+                progress_bar_for_records_fetched.next()
+    finally:
+        progress_bar_for_records_fetched.finish()
 
 
 def to_advisories(cve_data):

vulnerabilities/management/commands/import.py

@@ -8,9 +8,10 @@
 #
 import traceback
 
+import progress
 from django.core.management.base import BaseCommand
 from django.core.management.base import CommandError
-
+from progress.bar import  IncrementalBar
 from vulnerabilities.import_runner import ImportRunner
 from vulnerabilities.importers import IMPORTERS_REGISTRY
 
@@ -54,9 +55,10 @@ def import_data(self, importers):
         names for the importers.
         """
         failed_importers = []
-
+        progress_bar_for_import = IncrementalBar("Fetching Data from Databases", max=len(importers))
+        progress_bar_for_import.start()
         for importer in importers:
-            self.stdout.write(f"Importing data using {importer.qualified_name}")
+            self.stdout.write(f"\nImporting data using {importer.qualified_name}")
             try:
                 ImportRunner(importer).run()
                 self.stdout.write(
@@ -72,6 +74,9 @@ def import_data(self, importers):
                         f"Failed to run importer {importer.qualified_name}. Continuing..."
                     )
                 )
+            finally:
+                progress_bar_for_import.next()
+        progress_bar_for_import.finish()
 
         if failed_importers:
             raise CommandError(f"{len(failed_importers)} failed!: {','.join(failed_importers)}")