use pypdl for download and handle chunk error

arango_cve_processor/tools/cpe_db.py

@@ -6,8 +6,9 @@
 import os
 import tempfile
 import logging
+from pathlib import Path
 
-import requests
+from pypdl import Pypdl
 
 from arango_cve_processor.tools.nvd import fetch_nvd_api
 
@@ -38,18 +39,24 @@ def get_db(cls):
             cls._class_db = SwidTitleDB()
         return cls._class_db
 
+
     def _download_zip(self):
-        logging.info("downloading cpe dictionary")
-        retry_count = 0
-        while retry_count < 5:
-            try:
-                resp = requests.get(self.archive_url)
-                resp.raise_for_status()
-                return io.BytesIO(resp.content)
-            except requests.RequestException as e:
-                logging.warning(f"Failed to download cpe dictionary: {e}")
-                retry_count += 1
-        raise RuntimeError("Failed to download cpe dictionary after 5 attempts")
+        logging.info("Downloading CPE dictionary")
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output = Path(tmpdir) / "cpe.zip"
+
+            dl = Pypdl()
+            result = dl.start(
+                url=self.archive_url,
+                file_path=str(output),
+                retries=5,
+                block=True,
+            )
+            if not result or not output.exists():
+                raise RuntimeError("Failed to download CPE dictionary after 5 attempts")
+
+            return io.BytesIO(output.read_bytes())
 
     def _create_db(self):
         self.cur = self.conn.cursor()

arango_cve_processor/tools/nvd.py

@@ -9,13 +9,13 @@
 RATE_LIMIT_MAX_SLEEP = 10*60 #10 minutes
 
 
-def fetch_nvd_api(url, query: dict):
+def fetch_nvd_api(url, query: dict, api_key=None):
     total_results = math.inf
     start_index = 0
     query.update(startIndex=0)
+    api_key = api_key or os.environ.get("NVD_API_KEY")
 
     session = requests.Session()
-    api_key = os.environ.get("NVD_API_KEY")
     requests_per_window = 5
     if api_key:
         session.headers = {"apiKey": api_key}
@@ -38,7 +38,7 @@ def fetch_nvd_api(url, query: dict):
                 logging.warning("Got response status code %d.", response.status_code)
                 raise requests.ConnectionError
 
-        except requests.ConnectionError as ex:
+        except requests.RequestException as ex:
             logging.warning(
                 "Got ConnectionError. Backing off for %d seconds.", backoff_time
             )

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "arango_cve_processor"
-version = "1.4.11"
+version = "1.4.12"
 authors = [
   { name = "dogesec" }
 ]
@@ -28,6 +28,7 @@ dependencies = [
     "arango-cve-processor",
     "stix2arango",
     "stix2extensions>=2.0.1",
+    "pypdl",
   ]
 
 [project.urls]

requirements.txt

@@ -1,27 +1,130 @@
--i https://pypi.org/simple
+#
+# This file is autogenerated by pip-compile with Python 3.13
+# by the following command:
+#
+#    pip-compile --output-file=requirements.txt pyproject.toml
+#
+aiofiles==25.1.0
+    # via pypdl
+aiohappyeyeballs==2.6.2
+    # via aiohttp
+aiohttp==3.14.1
+    # via pypdl
+aiosignal==1.4.0
+    # via aiohttp
+annotated-types==0.7.0
+    # via pydantic
 antlr4-python3-runtime==4.9.3
-attrs==23.2.0; python_version >= '3.7'
-certifi==2024.7.4; python_version >= '3.6'
-charset-normalizer==3.3.2; python_full_version >= '3.7.0'
-idna==3.7; python_version >= '3.5'
-importlib-metadata==7.1.0; python_version >= '3.8'
-jsonschema==4.22.0; python_version >= '3.8'
-jsonschema-specifications==2023.12.1; python_version >= '3.8'
-packaging==24.0; python_version >= '3.7'
-pyjwt==2.12.0; python_version >= '3.7'
-python-arango>=7.9.1; python_version >= '3.8'
-python-dotenv==1.0.1; python_version >= '3.8'
+    # via stix2-patterns
+arango-cve-processor==1.4.11
+    # via arango_cve_processor (pyproject.toml)
+attrs==23.2.0
+    # via
+    #   aiohttp
+    #   jsonschema
+    #   referencing
+certifi==2024.7.4
+    # via requests
+charset-normalizer==3.3.2
+    # via requests
+frozenlist==1.8.0
+    # via
+    #   aiohttp
+    #   aiosignal
+idna==3.7
+    # via
+    #   requests
+    #   yarl
+ijson==3.5.0
+    # via stix2arango
+jsonschema==4.22.0 ; python_version >= "3.8"
+    # via
+    #   arango-cve-processor
+    #   arango_cve_processor (pyproject.toml)
+    #   stix2arango
+jsonschema-specifications==2023.12.1
+    # via jsonschema
+multidict==6.7.1
+    # via
+    #   aiohttp
+    #   yarl
+packaging==24.0
+    # via python-arango
+propcache==0.5.2
+    # via
+    #   aiohttp
+    #   yarl
+pydantic==2.13.4
+    # via stix2extensions
+pydantic-core==2.46.4
+    # via pydantic
+pyjwt==2.12.0
+    # via python-arango
+pypdl==1.5.7
+    # via arango_cve_processor (pyproject.toml)
+python-arango==8.3.3 ; python_version >= "3.8"
+    # via
+    #   arango-cve-processor
+    #   arango_cve_processor (pyproject.toml)
+    #   stix2arango
+python-dotenv==1.0.1
+    # via stix2arango
 pytz==2024.1
-referencing==0.35.1; python_version >= '3.8'
-requests==2.32.4; python_version >= '3.7'
-requests-toolbelt==1.0.0; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3'
-rpds-py==0.18.1; python_version >= '3.8'
-setuptools==78.1.1; python_version >= '3.8'
-simplejson==3.19.2; python_version >= '2.5' and python_version not in '3.0, 3.1, 3.2, 3.3'
-six==1.16.0; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3'
-stix2==3.0.1; python_version >= '3.6'
-stix2-patterns==2.0.0; python_version >= '3.6'
-tqdm==4.66.4; python_version >= '3.7'
-urllib3==2.6.3; python_version >= '3.8'
-zipp==3.19.1; python_version >= '3.8'
-stix2arango==1.1.2
+    # via stix2
+referencing==0.35.1
+    # via
+    #   jsonschema
+    #   jsonschema-specifications
+requests==2.32.4 ; python_version >= "3.7"
+    # via
+    #   arango-cve-processor
+    #   arango_cve_processor (pyproject.toml)
+    #   python-arango
+    #   requests-toolbelt
+    #   stix2
+    #   stix2arango
+    #   stix2extensions
+requests-toolbelt==1.0.0
+    # via python-arango
+rpds-py==0.18.1
+    # via
+    #   jsonschema
+    #   referencing
+simplejson==3.19.2
+    # via stix2
+six==1.16.0
+    # via stix2-patterns
+stix2==3.0.1
+    # via
+    #   arango-cve-processor
+    #   arango_cve_processor (pyproject.toml)
+    #   stix2arango
+    #   stix2extensions
+stix2-patterns==2.0.0
+    # via stix2
+stix2arango==1.1.2
+    # via
+    #   arango-cve-processor
+    #   arango_cve_processor (pyproject.toml)
+stix2extensions==2.3.3
+    # via
+    #   arango-cve-processor
+    #   arango_cve_processor (pyproject.toml)
+tqdm==4.66.4 ; python_version >= "3.7"
+    # via
+    #   arango-cve-processor
+    #   arango_cve_processor (pyproject.toml)
+    #   stix2arango
+typing-extensions==4.15.0
+    # via
+    #   pydantic
+    #   pydantic-core
+    #   typing-inspection
+typing-inspection==0.4.2
+    # via pydantic
+urllib3==2.6.3
+    # via
+    #   python-arango
+    #   requests
+yarl==1.24.2
+    # via aiohttp

tests/unit/tools/test_nvd.py

@@ -1,6 +1,6 @@
-
 from unittest.mock import MagicMock, patch
 
+import pytest
 import requests
 
 from arango_cve_processor.tools.nvd import RATE_LIMIT_WINDOW, fetch_nvd_api
@@ -36,7 +36,7 @@ def test_fetch_nvd_success():
             make_mock_response(start_index=1),
         ]
 
-        results = list(fetch_nvd_api('', {}))
+        results = list(fetch_nvd_api("", {}))
 
         assert len(results) == 2  # 2 pages
         assert all(isinstance(group, dict) for group in results)
@@ -53,11 +53,44 @@ def test_fetch_nvd_with_backoff():
             make_mock_response(start_index=1),
         ]
 
-        results = list(fetch_nvd_api('', {}))
+        results = list(fetch_nvd_api("", {}))
 
         assert len(results) == 2
         assert mock_sleep.call_count >= 2  # One for backoff, one for pagination
         # First sleep is backoff, second is rate limit
         # first_sleep_duration = RATE_LIMIT_WINDOW / 2
         # second_sleep_duration = RATE_LIMIT_WINDOW / cpematch_manager.requests_per_window
         assert mock_sleep.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "api_key,api_key_env,expected",
+    [
+        [None, None, None],
+        ["arg", None, "arg"],
+        ["arg", "env", "arg"],
+        [None, "env", "env"],
+    ],
+)
+def test_fetch_nvd_uses_api_key(api_key, api_key_env, expected, monkeypatch):
+    monkeypatch.delenv("NVD_API_KEY", raising=False)
+    if api_key_env:
+        monkeypatch.setenv("NVD_API_KEY", api_key_env)
+
+    patched_session = requests.Session()
+    monkeypatch.setattr(requests, "Session", MagicMock(return_value=patched_session))
+    monkeypatch.setattr(
+        patched_session,
+        "get",
+        MagicMock(
+            side_effect=[
+                make_mock_response(start_index=0),
+                make_mock_response(start_index=1),
+            ]
+        ),
+    )
+    with patch("time.sleep") as mock_sleep:
+
+        # Simulate a connection error first, then success
+        results = list(fetch_nvd_api("", {}, api_key=api_key))
+    assert patched_session.headers.get("apiKey") == expected