From 56e17f99d8e1f77584c2bc3659be2ddf8bdecdf1 Mon Sep 17 00:00:00 2001 From: Nicolas Thumann Date: Thu, 5 Dec 2024 16:15:20 +0100 Subject: [PATCH 1/7] Add: Retrying of requests to NVDApi --- pontos/nvd/api.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/pontos/nvd/api.py b/pontos/nvd/api.py index 709023fa8..4debaeded 100644 --- a/pontos/nvd/api.py +++ b/pontos/nvd/api.py @@ -33,6 +33,7 @@ SLEEP_TIMEOUT = 30.0 # in seconds DEFAULT_TIMEOUT = 180.0 # three minutes DEFAULT_TIMEOUT_CONFIG = Timeout(DEFAULT_TIMEOUT) # three minutes +RETRY_DELAY = 2.0 # in seconds Headers = Dict[str, str] Params = Dict[str, Union[str, int]] @@ -342,6 +343,7 @@ def __init__( token: Optional[str] = None, timeout: Optional[Timeout] = DEFAULT_TIMEOUT_CONFIG, rate_limit: bool = True, + attempts: int = 1, ) -> None: """ Create a new instance of the CVE API. @@ -357,6 +359,7 @@ def __init__( rolling 30 second window. See https://nvd.nist.gov/developers/start-here#divRateLimits Default: True. + attempts: The number of attempts per HTTP request. Defaults to 1. """ self._url = url self._token = token @@ -370,6 +373,8 @@ def __init__( self._request_count = 0 self._last_sleep = time.monotonic() + self._attempts = attempts + def _request_headers(self) -> Headers: """ Get the default request headers @@ -409,9 +414,19 @@ async def _get( """ headers = self._request_headers() - await self._consider_rate_limit() + for retry in range(self._attempts): + if retry > 0: + delay = RETRY_DELAY**retry + await asyncio.sleep(delay) + + await self._consider_rate_limit() + response = await self._client.get( + self._url, headers=headers, params=params + ) + if not response.is_server_error: + break - return await self._client.get(self._url, headers=headers, params=params) + return response async def __aenter__(self) -> "NVDApi": # reset rate limit counter From 9ed209b7e205f0e572cef2af542dddd8ed7b894a Mon Sep 17 00:00:00 2001 From: Nicolas Thumann Date: Thu, 5 Dec 2024 16:18:02 +0100 Subject: [PATCH 2/7] Add: Usage of retries to CPEApi --- pontos/nvd/cpe/api.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pontos/nvd/cpe/api.py b/pontos/nvd/cpe/api.py index ba4df766c..aea881d46 100644 --- a/pontos/nvd/cpe/api.py +++ b/pontos/nvd/cpe/api.py @@ -61,6 +61,7 @@ def __init__( token: Optional[str] = None, timeout: Optional[Timeout] = DEFAULT_TIMEOUT_CONFIG, rate_limit: bool = True, + attempts: int = 1, ) -> None: """ Create a new instance of the CPE API. @@ -75,12 +76,14 @@ def __init__( rolling 30 second window. See https://nvd.nist.gov/developers/start-here#divRateLimits Default: True. + attempts: The number of attempts per HTTP request. Defaults to 1. """ super().__init__( DEFAULT_NIST_NVD_CPES_URL, token=token, timeout=timeout, rate_limit=rate_limit, + attempts=attempts, ) async def cpe(self, cpe_name_id: Union[str, UUID]) -> CPE: From be35f72d6c93bdb0bcd477d95849373f527d3ec8 Mon Sep 17 00:00:00 2001 From: Nicolas Thumann Date: Thu, 5 Dec 2024 16:18:13 +0100 Subject: [PATCH 3/7] Add: Usage of retries to CVEApi --- pontos/nvd/cve/api.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pontos/nvd/cve/api.py b/pontos/nvd/cve/api.py index 98f0af10f..c52991c48 100644 --- a/pontos/nvd/cve/api.py +++ b/pontos/nvd/cve/api.py @@ -65,6 +65,7 @@ def __init__( token: Optional[str] = None, timeout: Optional[Timeout] = DEFAULT_TIMEOUT_CONFIG, rate_limit: bool = True, + attempts: int = 1, ) -> None: """ Create a new instance of the CVE API. @@ -79,12 +80,14 @@ def __init__( rolling 30 second window. See https://nvd.nist.gov/developers/start-here#divRateLimits Default: True. + attempts: The number of attempts per HTTP request. Defaults to 1. """ super().__init__( DEFAULT_NIST_NVD_CVES_URL, token=token, timeout=timeout, rate_limit=rate_limit, + attempts=attempts, ) def cves( From 29c229e44e9f75742e978635637e8bdf9093a76b Mon Sep 17 00:00:00 2001 From: Nicolas Thumann Date: Thu, 5 Dec 2024 16:18:22 +0100 Subject: [PATCH 4/7] Add: Usage of retries to CVEChangesApi --- pontos/nvd/cve_changes/api.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pontos/nvd/cve_changes/api.py b/pontos/nvd/cve_changes/api.py index faa8456b9..4d6671865 100644 --- a/pontos/nvd/cve_changes/api.py +++ b/pontos/nvd/cve_changes/api.py @@ -55,6 +55,7 @@ def __init__( token: Optional[str] = None, timeout: Optional[Timeout] = DEFAULT_TIMEOUT_CONFIG, rate_limit: bool = True, + attempts: int = 1, ) -> None: """ Create a new instance of the CVE Change History API. @@ -69,12 +70,14 @@ def __init__( rolling 30 second window. See https://nvd.nist.gov/developers/start-here#divRateLimits Default: True. + attempts: The number of attempts per HTTP request. Defaults to 1. """ super().__init__( DEFAULT_NIST_NVD_CVE_HISTORY_URL, token=token, timeout=timeout, rate_limit=rate_limit, + attempts=attempts, ) def changes( From 248fd732c51c85a0285266c9a7cdf05a704d6567 Mon Sep 17 00:00:00 2001 From: Nicolas Thumann Date: Thu, 5 Dec 2024 16:47:07 +0100 Subject: [PATCH 5/7] Change: Rename attempts to request_attempts --- pontos/nvd/api.py | 12 ++++++------ pontos/nvd/cpe/api.py | 6 +++--- pontos/nvd/cve/api.py | 6 +++--- pontos/nvd/cve_changes/api.py | 6 +++--- 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/pontos/nvd/api.py b/pontos/nvd/api.py index 4debaeded..9ce9952e0 100644 --- a/pontos/nvd/api.py +++ b/pontos/nvd/api.py @@ -343,7 +343,7 @@ def __init__( token: Optional[str] = None, timeout: Optional[Timeout] = DEFAULT_TIMEOUT_CONFIG, rate_limit: bool = True, - attempts: int = 1, + request_attempts: int = 1, ) -> None: """ Create a new instance of the CVE API. @@ -359,7 +359,7 @@ def __init__( rolling 30 second window. See https://nvd.nist.gov/developers/start-here#divRateLimits Default: True. - attempts: The number of attempts per HTTP request. Defaults to 1. + request_attempts: The number of attempts per HTTP request. Defaults to 1. """ self._url = url self._token = token @@ -373,7 +373,7 @@ def __init__( self._request_count = 0 self._last_sleep = time.monotonic() - self._attempts = attempts + self._request_attempts = request_attempts def _request_headers(self) -> Headers: """ @@ -414,9 +414,9 @@ async def _get( """ headers = self._request_headers() - for retry in range(self._attempts): - if retry > 0: - delay = RETRY_DELAY**retry + for attempt in range(self._request_attempts): + if attempt > 0: + delay = RETRY_DELAY**attempt await asyncio.sleep(delay) await self._consider_rate_limit() diff --git a/pontos/nvd/cpe/api.py b/pontos/nvd/cpe/api.py index aea881d46..d397e50bb 100644 --- a/pontos/nvd/cpe/api.py +++ b/pontos/nvd/cpe/api.py @@ -61,7 +61,7 @@ def __init__( token: Optional[str] = None, timeout: Optional[Timeout] = DEFAULT_TIMEOUT_CONFIG, rate_limit: bool = True, - attempts: int = 1, + request_attempts: int = 1, ) -> None: """ Create a new instance of the CPE API. @@ -76,14 +76,14 @@ def __init__( rolling 30 second window. See https://nvd.nist.gov/developers/start-here#divRateLimits Default: True. - attempts: The number of attempts per HTTP request. Defaults to 1. + request_attempts: The number of attempts per HTTP request. Defaults to 1. """ super().__init__( DEFAULT_NIST_NVD_CPES_URL, token=token, timeout=timeout, rate_limit=rate_limit, - attempts=attempts, + request_attempts=request_attempts, ) async def cpe(self, cpe_name_id: Union[str, UUID]) -> CPE: diff --git a/pontos/nvd/cve/api.py b/pontos/nvd/cve/api.py index c52991c48..6dbce3dda 100644 --- a/pontos/nvd/cve/api.py +++ b/pontos/nvd/cve/api.py @@ -65,7 +65,7 @@ def __init__( token: Optional[str] = None, timeout: Optional[Timeout] = DEFAULT_TIMEOUT_CONFIG, rate_limit: bool = True, - attempts: int = 1, + request_attempts: int = 1, ) -> None: """ Create a new instance of the CVE API. @@ -80,14 +80,14 @@ def __init__( rolling 30 second window. See https://nvd.nist.gov/developers/start-here#divRateLimits Default: True. - attempts: The number of attempts per HTTP request. Defaults to 1. + request_attempts: The number of attempts per HTTP request. Defaults to 1. """ super().__init__( DEFAULT_NIST_NVD_CVES_URL, token=token, timeout=timeout, rate_limit=rate_limit, - attempts=attempts, + request_attempts=request_attempts, ) def cves( diff --git a/pontos/nvd/cve_changes/api.py b/pontos/nvd/cve_changes/api.py index 4d6671865..f6e37f6c6 100644 --- a/pontos/nvd/cve_changes/api.py +++ b/pontos/nvd/cve_changes/api.py @@ -55,7 +55,7 @@ def __init__( token: Optional[str] = None, timeout: Optional[Timeout] = DEFAULT_TIMEOUT_CONFIG, rate_limit: bool = True, - attempts: int = 1, + request_attempts: int = 1, ) -> None: """ Create a new instance of the CVE Change History API. @@ -70,14 +70,14 @@ def __init__( rolling 30 second window. See https://nvd.nist.gov/developers/start-here#divRateLimits Default: True. - attempts: The number of attempts per HTTP request. Defaults to 1. + request_attempts: The number of attempts per HTTP request. Defaults to 1. """ super().__init__( DEFAULT_NIST_NVD_CVE_HISTORY_URL, token=token, timeout=timeout, rate_limit=rate_limit, - attempts=attempts, + request_attempts=request_attempts, ) def changes( From d3846ad144b1ecd51c55bf2bdee3b019260af77d Mon Sep 17 00:00:00 2001 From: Nicolas Thumann Date: Fri, 6 Dec 2024 15:53:52 +0100 Subject: [PATCH 6/7] Add: Unit tests for retrying of requests in NVDApi --- tests/nvd/test_api.py | 48 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/tests/nvd/test_api.py b/tests/nvd/test_api.py index 56f2bd1b8..128227995 100644 --- a/tests/nvd/test_api.py +++ b/tests/nvd/test_api.py @@ -8,7 +8,7 @@ import unittest from datetime import datetime from typing import Any, Iterator -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch from httpx import AsyncClient, Response @@ -129,6 +129,52 @@ async def test_no_rate_limit( sleep_mock.assert_not_called() + @patch("pontos.nvd.api.asyncio.sleep", autospec=True) + @patch("pontos.nvd.api.AsyncClient", spec=AsyncClient) + async def test_retry( + self, + async_client: MagicMock, + sleep_mock: MagicMock, + ): + response_mocks = [ + MagicMock(spec=Response, status_code=500), + MagicMock(spec=Response, status_code=500), + MagicMock(spec=Response, status_code=500), + MagicMock(spec=Response, status_code=200), + ] + http_client = AsyncMock() + http_client.get.side_effect = response_mocks + async_client.return_value = http_client + + api = NVDApi("https://foo.bar/baz", request_attempts=4) + + result = await api._get() + + calls = [call(2.0), call(4.0), call(8.0)] + sleep_mock.assert_has_calls(calls) + self.assertEqual(result.status_code, 200) + + @patch("pontos.nvd.api.asyncio.sleep", autospec=True) + @patch("pontos.nvd.api.AsyncClient", spec=AsyncClient) + async def test_no_retry( + self, + async_client: MagicMock, + sleep_mock: MagicMock, + ): + response_mock = MagicMock(spec=Response) + response_mock.status_code = 200 + + http_client = AsyncMock() + http_client.get.return_value = response_mock + async_client.return_value = http_client + + api = NVDApi("https://foo.bar/baz") + + result = await api._get() + + sleep_mock.assert_not_called() + self.assertEqual(result.status_code, 200) + class Result: def __init__(self, value: int) -> None: From fd818807559af755b1c06e56ab2c555069020f95 Mon Sep 17 00:00:00 2001 From: Nicolas Thumann Date: Fri, 6 Dec 2024 16:08:29 +0100 Subject: [PATCH 7/7] Change: Fix unit test for full coverage --- tests/nvd/test_api.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/nvd/test_api.py b/tests/nvd/test_api.py index 128227995..ea27f7c9e 100644 --- a/tests/nvd/test_api.py +++ b/tests/nvd/test_api.py @@ -137,10 +137,10 @@ async def test_retry( sleep_mock: MagicMock, ): response_mocks = [ - MagicMock(spec=Response, status_code=500), - MagicMock(spec=Response, status_code=500), - MagicMock(spec=Response, status_code=500), - MagicMock(spec=Response, status_code=200), + MagicMock(spec=Response, is_server_error=True), + MagicMock(spec=Response, is_server_error=True), + MagicMock(spec=Response, is_server_error=True), + MagicMock(spec=Response, is_server_error=False), ] http_client = AsyncMock() http_client.get.side_effect = response_mocks @@ -152,7 +152,7 @@ async def test_retry( calls = [call(2.0), call(4.0), call(8.0)] sleep_mock.assert_has_calls(calls) - self.assertEqual(result.status_code, 200) + self.assertFalse(result.is_server_error) @patch("pontos.nvd.api.asyncio.sleep", autospec=True) @patch("pontos.nvd.api.AsyncClient", spec=AsyncClient) @@ -162,7 +162,7 @@ async def test_no_retry( sleep_mock: MagicMock, ): response_mock = MagicMock(spec=Response) - response_mock.status_code = 200 + response_mock.is_server_error = False http_client = AsyncMock() http_client.get.return_value = response_mock @@ -173,7 +173,7 @@ async def test_no_retry( result = await api._get() sleep_mock.assert_not_called() - self.assertEqual(result.status_code, 200) + self.assertFalse(result.is_server_error) class Result: