diff --git a/arango_cve_processor/tools/cpe_db.py b/arango_cve_processor/tools/cpe_db.py index a814b8f..3ce05df 100644 --- a/arango_cve_processor/tools/cpe_db.py +++ b/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() diff --git a/arango_cve_processor/tools/nvd.py b/arango_cve_processor/tools/nvd.py index be416a1..c3bc9c6 100644 --- a/arango_cve_processor/tools/nvd.py +++ b/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 ) diff --git a/pyproject.toml b/pyproject.toml index 9e21e31..6914ba4 100644 --- a/pyproject.toml +++ b/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] diff --git a/requirements.txt b/requirements.txt index ea43938..167652c 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file + # 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 diff --git a/tests/unit/tools/test_nvd.py b/tests/unit/tools/test_nvd.py index e587932..78b5f77 100644 --- a/tests/unit/tools/test_nvd.py +++ b/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,7 +53,7 @@ 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 @@ -61,3 +61,36 @@ def test_fetch_nvd_with_backoff(): # 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