diff --git a/requirements.txt b/requirements.txt index 042f541da..7a26d552b 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file diff --git a/vulnerabilities/importers/nvd.py b/vulnerabilities/importers/nvd.py index b22d3db68..3c72d0206 100644 --- a/vulnerabilities/importers/nvd.py +++ b/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): diff --git a/vulnerabilities/management/commands/import.py b/vulnerabilities/management/commands/import.py index 5ae885299..9bbb2394f 100644 --- a/vulnerabilities/management/commands/import.py +++ b/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)}")