From 31c6fd0bd6a13f459c187f2c96ce397d0808951d Mon Sep 17 00:00:00 2001 From: Matthijs Vos Date: Thu, 27 Jun 2024 11:37:29 +0200 Subject: [PATCH 1/2] Add option to set API key header to X-IRIS-AUTH --- dfir_iris_client/session.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/dfir_iris_client/session.py b/dfir_iris_client/session.py index e95013f..03bbed0 100644 --- a/dfir_iris_client/session.py +++ b/dfir_iris_client/session.py @@ -52,7 +52,7 @@ class ClientSession(object): Returns: """ - def __init__(self, apikey=None, host=None, agent="iris-client", ssl_verify=True, proxy=None, timeout=120): + def __init__(self, apikey=None, host=None, agent="iris-client", ssl_verify=True, proxy=None, timeout=120, api_key_header=None): """ Initialize the ClientSession. APIKey validity is verified as well as API compatibility between the client and the server. @@ -70,6 +70,7 @@ def __init__(self, apikey=None, host=None, agent="iris-client", ssl_verify=True, ssl_verify: Set or unset SSL verification proxy: Proxy parameters - For future use only timeout: Default timeout for requests + api_key_header: Header to use for the API key, can either be 'Authorization' or 'X-IRIS-AUTH'. Defaults to 'Authorization' when not set """ self._apikey = apikey self._host = host @@ -84,6 +85,11 @@ def __init__(self, apikey=None, host=None, agent="iris-client", ssl_verify=True, if not self._ssl_verify: requests.packages.urllib3.disable_warnings(InsecureRequestWarning) + + if api_key_header and api_key_header not in ['Authorization', 'X-IRIS-AUTH']: + raise Exception("api_key_header can only be 'Authorization' or 'X-IRIS-AUTH'") + self._api_key_header = api_key_header or 'Authorization' + self._check_apikey_validity() self._check_api_compatibility() @@ -216,7 +222,7 @@ def _pi_request(self, uri: str, type: str = None, data: dict = None, headers = { 'Content-Type': "application/json", - 'Authorization': "Bearer " + self._apikey, + self._api_key_header: "Bearer " + self._apikey, 'User-Agent': self._agent } if type == "POST": From 167a7e2a6fc90dd7f8475c7760cf6ea7c202aa36 Mon Sep 17 00:00:00 2001 From: Matthijs Vos Date: Thu, 27 Jun 2024 11:45:10 +0200 Subject: [PATCH 2/2] Add option to specify additional headers --- dfir_iris_client/session.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/dfir_iris_client/session.py b/dfir_iris_client/session.py index 03bbed0..bffccb2 100644 --- a/dfir_iris_client/session.py +++ b/dfir_iris_client/session.py @@ -52,7 +52,7 @@ class ClientSession(object): Returns: """ - def __init__(self, apikey=None, host=None, agent="iris-client", ssl_verify=True, proxy=None, timeout=120, api_key_header=None): + def __init__(self, apikey=None, host=None, agent="iris-client", ssl_verify=True, proxy=None, timeout=120, api_key_header=None, additional_headers=None): """ Initialize the ClientSession. APIKey validity is verified as well as API compatibility between the client and the server. @@ -71,6 +71,7 @@ def __init__(self, apikey=None, host=None, agent="iris-client", ssl_verify=True, proxy: Proxy parameters - For future use only timeout: Default timeout for requests api_key_header: Header to use for the API key, can either be 'Authorization' or 'X-IRIS-AUTH'. Defaults to 'Authorization' when not set + additional_headers: Additional headers to pass in the request """ self._apikey = apikey self._host = host @@ -89,6 +90,7 @@ def __init__(self, apikey=None, host=None, agent="iris-client", ssl_verify=True, if api_key_header and api_key_header not in ['Authorization', 'X-IRIS-AUTH']: raise Exception("api_key_header can only be 'Authorization' or 'X-IRIS-AUTH'") self._api_key_header = api_key_header or 'Authorization' + self._addition_headers = additional_headers self._check_apikey_validity() @@ -225,6 +227,9 @@ def _pi_request(self, uri: str, type: str = None, data: dict = None, self._api_key_header: "Bearer " + self._apikey, 'User-Agent': self._agent } + if self._addition_headers: + headers = {**headers, **self._addition_headers} + if type == "POST": log.debug(f'POST : {self._pi_uri(uri)}') @@ -277,9 +282,11 @@ def pi_post_files(self, uri: str, files: dict = None, data: dict = None, cid: in ApiResponse object """ headers = { - 'Authorization': "Bearer " + self._apikey, + self._api_key_header: "Bearer " + self._apikey, 'User-Agent': self._agent } + if self._addition_headers: + headers = {**headers, **self._addition_headers} if cid is None: raise ValueError('cid is mandatory when uploading files')