feat: retry if NVD API Key is invalid (#1574)

Co-authored-by: Sahil <[email protected]>

Author: terriko

.github/workflows/testing.yml

@@ -13,6 +13,7 @@ on:
 env:
   ACTIONS: 1
   LONG_TESTS: 0
+  nvd_api_key: ${{ secrets.NVD_API_KEY }}
 
 jobs:
   docs:

cve_bin_tool/cli.py

@@ -351,6 +351,10 @@ def main(argv=None):
     if not args["nvd_api_key"] and os.getenv("nvd_api_key"):
         args["nvd_api_key"] = os.getenv("nvd_api_key")
 
+    # Also try the uppercase env variable, in case people prefer those
+    if not args["nvd_api_key"] and os.getenv("NVD_API_KEY"):
+        args["nvd_api_key"] = os.getenv("NVD_API_KEY")
+
     # If you're not using an NVD key, let you know how to get one
     if not args["nvd_api_key"] and not args["offline"]:
         LOGGER.info("Not using an NVD API key. Your access may be rate limited by NVD.")

cve_bin_tool/error_handler.py

@@ -88,6 +88,12 @@ class NVDServiceError(Exception):
     """
 
 
+class NVDKeyError(Exception):
+    """
+    Raised if the NVD API key is invalid.
+    """
+
+
 class SHAMismatch(Exception):
     """
     Raised if the sha of a file in the cache was not what it should be.

cve_bin_tool/nvd_api.py

@@ -19,7 +19,7 @@
 from rich.progress import Progress, track
 
 from cve_bin_tool.async_utils import RateLimiter
-from cve_bin_tool.error_handler import ErrorMode, NVDServiceError
+from cve_bin_tool.error_handler import ErrorMode, NVDKeyError, NVDServiceError
 from cve_bin_tool.log import LOGGER
 
 FEED = "https://services.nvd.nist.gov/rest/json/cves/1.0"
@@ -102,6 +102,9 @@ async def get_nvd_params(
         self.logger.debug("Fetching metadata from NVD...")
         cve_count = await self.nvd_count_metadata(self.session)
 
+        if "apiKey" in self.params:
+            await self.validate_nvd_api()
+
         if time_of_last_update:
             # Fetch all the updated CVE entries from the modified date. Subtracting 2-minute offset for updating cve entries
             self.params["modStartDate"] = self.convert_date_to_nvd_date(
@@ -125,6 +128,28 @@ async def get_nvd_params(
             self.total_results = cve_count["Total"] - cve_count["Rejected"]
         self.logger.info(f"Adding {self.total_results} CVE entries")
 
+    async def validate_nvd_api(self):
+        """
+        Validate NVD API
+        """
+        param_dict = self.params.copy()
+        param_dict["startIndex"] = 0
+        param_dict["resultsPerPage"] = 1
+        try:
+            self.logger.debug("Validating NVD API...")
+            async with await self.session.get(
+                self.feed, params=param_dict, raise_for_status=True
+            ) as response:
+                data = await response.json()
+                if data.get("error", False):
+                    self.logger.error(f"NVD API error: {data['error']}")
+                    raise NVDKeyError(self.params["apiKey"])
+        except NVDKeyError:
+            # If the API key provided is invalid, delete from params
+            # list and try the request again.
+            self.logger.error("unset api key, retrying")
+            del self.params["apiKey"]
+
     async def load_nvd_request(self, start_index):
         """Get single NVD request and update year_wise_data list which contains list of all CVEs"""
 
@@ -141,6 +166,7 @@ async def load_nvd_request(self, start_index):
                 ) as response:
                     if response.status == 200:
                         fetched_data = await response.json()
+
                         if start_index == 0:
                             # Update total results in case there is discrepancy between NVD dashboard and API
                             self.total_results = fetched_data["totalResults"]