diff --git a/reptor/api/APIClient.py b/reptor/api/APIClient.py index 76c5220..490a990 100644 --- a/reptor/api/APIClient.py +++ b/reptor/api/APIClient.py @@ -1,6 +1,8 @@ import typing import requests +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry import reptor.settings as settings from reptor.lib.console import reptor_console @@ -28,6 +30,8 @@ def __init__(self, require_project_id=True, **kwargs) -> None: urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) # type: ignore + self._session = self._build_session() + try: self._project_id = self.reptor.get_active_project_id() self.debug(f"Project ID is {self._project_id}") @@ -56,10 +60,36 @@ def _get_headers(self, json_content=True) -> typing.Dict: self.debug(f"HTTP Headers: {headers_debug}") return headers + def _build_session(self) -> requests.Session: + """Builds a requests Session with automatic retries for transient + failures. Retries cover connection/TLS handshake errors (e.g. a reset + socket producing SSLEOFError) and the retryable status codes in + settings.API_RETRY_STATUS_FORCELIST. Only idempotent methods are + retried on a status code (urllib3's default allowed_methods), so + POST/PATCH are not replayed against the server; connection-establishment + failures are still retried for every method. + """ + retry = Retry( + total=settings.API_MAX_RETRIES, + connect=settings.API_MAX_RETRIES, + read=settings.API_MAX_RETRIES, + status=settings.API_MAX_RETRIES, + backoff_factor=settings.API_RETRY_BACKOFF_FACTOR, + status_forcelist=settings.API_RETRY_STATUS_FORCELIST, + raise_on_status=False, + respect_retry_after_header=True, + ) + adapter = HTTPAdapter(max_retries=retry) + session = requests.Session() + session.mount("http://", adapter) + session.mount("https://", adapter) + return session + def _prepare_kwargs(self, kwargs, json_content=True): return kwargs | { 'headers': kwargs.get('headers', {}) | self._get_headers(json_content=json_content), 'verify': kwargs.get('verify', self.verify), + 'timeout': kwargs.get('timeout', settings.API_TIMEOUT), 'allow_redirects': False, } @@ -153,11 +183,11 @@ def _do_request( self, url, method: str = "GET", json_content: bool = True, **kwargs ) -> requests.models.Response: methods = { - "GET": requests.get, - "POST": requests.post, - "PUT": requests.put, - "PATCH": requests.patch, - "DELETE": requests.delete, + "GET": self._session.get, + "POST": self._session.post, + "PUT": self._session.put, + "PATCH": self._session.patch, + "DELETE": self._session.delete, } method = method.upper() if method not in methods.keys(): diff --git a/reptor/api/tests/test_apiclient.py b/reptor/api/tests/test_apiclient.py new file mode 100644 index 0000000..f03a8cb --- /dev/null +++ b/reptor/api/tests/test_apiclient.py @@ -0,0 +1,64 @@ +from unittest.mock import MagicMock + +import pytest +from urllib3.util.retry import Retry + +import reptor.settings as settings +from reptor.lib.reptor import Reptor + +from ..APIClient import APIClient + + +class TestAPIClientSession: + @pytest.fixture(autouse=True) + def setUp(self): + self.reptor = Reptor() + self.reptor._config._raw_config["server"] = "https://demo.sysre.pt" + self.reptor._config._raw_config["token"] = "sysreptor_test" + self.client = APIClient(reptor=self.reptor, require_project_id=False) + + def test_build_session_mounts_retry_on_both_schemes(self): + for scheme in ("http://demo.sysre.pt", "https://demo.sysre.pt"): + adapter = self.client._session.get_adapter(scheme) + retry = adapter.max_retries + assert isinstance(retry, Retry) + assert retry.total == settings.API_MAX_RETRIES + assert retry.connect == settings.API_MAX_RETRIES + assert retry.read == settings.API_MAX_RETRIES + assert retry.status == settings.API_MAX_RETRIES + assert retry.backoff_factor == settings.API_RETRY_BACKOFF_FACTOR + assert list(retry.status_forcelist) == settings.API_RETRY_STATUS_FORCELIST + # Never turn an exhausted status-retry into an exception here; the + # existing response.raise_for_status() in _do_request owns that. + assert retry.raise_on_status is False + assert retry.respect_retry_after_header is True + + def test_prepare_kwargs_sets_defaults(self): + prepared = self.client._prepare_kwargs({}) + assert prepared["timeout"] == settings.API_TIMEOUT + assert prepared["verify"] is True # insecure not set -> verify on + assert prepared["allow_redirects"] is False + + def test_prepare_kwargs_respects_caller_overrides(self): + prepared = self.client._prepare_kwargs({"timeout": 5, "verify": False}) + assert prepared["timeout"] == 5 + assert prepared["verify"] is False + + def test_insecure_config_disables_verify(self): + self.reptor._config._raw_config["insecure"] = True + client = APIClient(reptor=self.reptor, require_project_id=False) + assert client._prepare_kwargs({})["verify"] is False + + def test_do_request_routes_through_session_with_defaults(self): + response = MagicMock() + response.headers = {} + response.content = b"{}" + self.client._session.get = MagicMock(return_value=response) + + self.client.get("https://demo.sysre.pt/api/v1/pentestprojects/") + + self.client._session.get.assert_called_once() + _, kwargs = self.client._session.get.call_args + assert kwargs["timeout"] == settings.API_TIMEOUT + assert kwargs["verify"] is True + assert kwargs["allow_redirects"] is False diff --git a/reptor/settings.py b/reptor/settings.py index 8c3c937..72ddbe8 100644 --- a/reptor/settings.py +++ b/reptor/settings.py @@ -37,6 +37,12 @@ NEWLINE = "\n" USER_AGENT = "reptor CLI v0.1.0" # TODO dynamic version + +# HTTP client defaults (see reptor/api/APIClient.py) +API_TIMEOUT = 30 # per-request (connect, read) timeout in seconds; prevents indefinite hangs +API_MAX_RETRIES = 3 # retries for transient connection/TLS failures and retryable status codes +API_RETRY_BACKOFF_FACTOR = 0.5 # exponential backoff: 0s, 0.5s, 1s, 2s, ... +API_RETRY_STATUS_FORCELIST = [429, 500, 502, 503, 504] # only these statuses are retried LANGUAGE_CODE = "en" FORMAT_MODULE_PATH = [] LOCALE_PATHS = (