diff --git a/censys/common/base.py b/censys/common/base.py index d67410d..b094916 100644 --- a/censys/common/base.py +++ b/censys/common/base.py @@ -98,6 +98,7 @@ def __init__( self._session.proxies.update(proxies) if cookies: self._session.cookies.update(cookies) + self.request_id = kwargs.get("request_id") self._session.headers.update( { "accept": "application/json, */8", @@ -112,6 +113,29 @@ def __init__( } ) + @property + def request_id(self) -> Optional[str]: + """The x-request-id header value for API requests. + + The x-request-id header is not set when the value is None. + Value is None by default + + Returns: + Type[Optional[str]]: The value of the header. + """ + value = self._session.headers.get("x-request-id") + if not isinstance(value, str): + return None + return value + + @request_id.setter + def request_id(self, value: Optional[str]): + if value is None: + self._session.headers.pop("x-request-id", None) + return + + self._session.headers["x-request-id"] = value + @staticmethod def _get_exception_class(_: Response) -> Type[CensysAPIException]: """Maps HTTP status code or ASM error code to exception. diff --git a/pyproject.toml b/pyproject.toml index 6c6ed71..74073c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "censys" -version = "2.2.11" +version = "2.2.12" description = "An easy-to-use and lightweight API wrapper for Censys APIs (censys.io)." authors = ["Censys, Inc. "] license = "Apache-2.0" diff --git a/tests/test_base.py b/tests/test_base.py index 9bc2de2..8df60f4 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -96,6 +96,19 @@ def test_user_agent(self): == requests.utils.default_user_agent() + " test" ) + def test_request_id(self): + id_value = "my-request-id" + + # Test request id value is present + base = CensysAPIBase(TEST_URL, request_id=id_value) + assert base.request_id == id_value + assert base._session.headers.get("x-request-id") == id_value + + # Test request id value is not present + base.request_id = None + assert base.request_id is None + assert base._session.headers.get("x-request-id") is None + @pytest.mark.filterwarnings("ignore:HTTP proxies will not be used.") def test_proxies(self): # Mock/actual call