add custom http headers (#110)

* add custom http headers * changelog * Update CHANGELOG.md Co-authored-by: dorschw <81086590+dorschw@users.noreply.github.com> * Update tests/mocks_test.py Co-authored-by: dorschw <81086590+dorschw@users.noreply.github.com> * validate format and add unit-test * update README + add numbers unit-test case * small typo fix * update var name * add more unit-tests cases for invalid headers * one more test-case * update regex and add unit-tests for valid forms * update test * check cases of empty strings * add comment * update changelog --------- Co-authored-by: dorschw <81086590+dorschw@users.noreply.github.com>
demisto · Apr 13, 2023 · f39a1bd · f39a1bd
1 parent f491bb9
commit f39a1bd
Show file tree

Hide file tree

Showing 4 changed files with 315 additions and 67 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,7 +3,10 @@
 [PyPI History][1]
 [1]: https://pypi.org/project/demisto-py/#history
 
-# 3.2.8
+## 3.2.9
+* Added the ability to add custom request headers.
+
+## 3.2.8
 * Fixed an issue where demisto-py used username/password authentication even when api-key was provided.
 
 ## 3.2.7

diff --git a/README.md b/README.md
@@ -18,13 +18,14 @@ Follow these instructions to generate your Demisto API Key.
 1. In Demisto, navigate to **Settings > API Keys**.
 2. Click the **Generate Your Key** button.
 
-To avoid hard coding configurations in your code, it is possible to specify configruation params
+To avoid hard coding configurations in your code, it is possible to specify configuration params
 as the following environment variables (env variables will be used if parameters are not specified):
 
 * DEMISTO_BASE_URL
 * DEMISTO_API_KEY
 * DEMISTO_USERNAME
 * DEMISTO_PASSWORD
+* DEMISTO_HTTP_HEADERS (must be in the form of: header1=value1,header2=value2,header3=value3,...headerN=valueN)
 * DEMISTO_VERIFY_SSL (true/false. Default: true)
 * XSIAM_AUTH_ID (for Cortex XSIAM only. If set, Cortex XSIAM API will be used)
 * SSL_CERT_FILE (specify an alternate certificate bundle)

diff --git a/demisto_client/__init__.py b/demisto_client/__init__.py
@@ -1,3 +1,5 @@
+import re
+
 import demisto_client.demisto_api as demisto_api
 import six
 import os
@@ -17,8 +19,11 @@
     __version__ = 'dev'
 
 
+DEMISTO_HTTP_HEADERS_REGEX_PATTERN = r'^([\w-]+=[^=,\n]+)(,[\w-]+=[^=,\n]+)*$'
+
+
 def configure(base_url=None, api_key=None, verify_ssl=None, proxy=None, username=None, password=None,
-              ssl_ca_cert=None, debug=False, connection_pool_maxsize=None, auth_id=None):
+              ssl_ca_cert=None, debug=False, connection_pool_maxsize=None, auth_id=None, additional_headers=None):
     """
     This wrapper provides an easier to use method of configuring the API client. The base
     Configuration method is still exposed if you wish to further configure the API Client.
@@ -31,6 +36,7 @@ def configure(base_url=None, api_key=None, verify_ssl=None, proxy=None, username
     * DEMISTO_USERNAME
     * DEMISTO_PASSWORD
     * DEMISTO_VERIFY_SSL (true/false. Default: true)
+    * DEMISTO_HTTP_HEADERS (must be in the form of: header1=value1,header2=value2,header3=value3,...headerN=valueN)
     * XSIAM_AUTH_ID
     * SSL_CERT_FILE (specify an alternate certificate bundle)
     * DEMISTO_CONNECTION_POOL_MAXSIZE (specify a connection pool max size)
@@ -47,6 +53,7 @@ def configure(base_url=None, api_key=None, verify_ssl=None, proxy=None, username
     :param debug: bool - Include verbose logging.
     :param connection_pool_maxsize: int - specify a connection max pool size
     :param auth_id: str - api_key_id only for the xsiam
+    :param additional_headers: dict - any additional headers to send to every http request to demisto.
     :return: Returns an API client configuration identical to the Configuration() method.
     """
     if base_url is None:
@@ -61,6 +68,18 @@ def configure(base_url=None, api_key=None, verify_ssl=None, proxy=None, username
         username = os.getenv('DEMISTO_USERNAME')
     if password is None:
         password = os.getenv('DEMISTO_PASSWORD')
+    if additional_headers is None:
+        if headers := os.getenv('DEMISTO_HTTP_HEADERS'):
+            if re.match(r'^ *$', headers):  # catch any case of empty string
+                additional_headers = {}
+            elif re.match(DEMISTO_HTTP_HEADERS_REGEX_PATTERN, headers):
+                additional_headers = dict(header.strip().split('=') for header in headers.split(','))
+            else:
+                raise ValueError(
+                    f'{headers} has invalid format, must be in the format of header1=value1,header2=value2,...headerN=valueN'
+                )
+        else:
+            additional_headers = {}
     if ssl_ca_cert is None:
         ssl_ca_cert = os.getenv('SSL_CERT_FILE')
     if verify_ssl is None:
@@ -112,15 +131,19 @@ def configure(base_url=None, api_key=None, verify_ssl=None, proxy=None, username
     if api_key or auth_id:
         api_client = ApiClient(configuration)
         api_client.user_agent = 'demisto-py/' + __version__
+        api_client.default_headers.update(additional_headers)
         api_instance = demisto_api.DefaultApi(api_client)
         return api_instance
     else:
         api_instance = login(base_url=base_url, username=username, password=password,
-                             verify_ssl=verify_ssl, proxy=proxy, debug=debug)
+                             verify_ssl=verify_ssl, proxy=proxy, debug=debug, additional_headers=additional_headers)
         return api_instance
 
 
-def login(base_url=None, username=None, password=None, verify_ssl=True, proxy=None, debug=False):
+def login(
+    base_url=None, username=None, password=None, verify_ssl=True, proxy=None, debug=False, additional_headers=None
+):
+    additional_headers = additional_headers or {}
     configuration_orig = Configuration()
     configuration_orig.host = base_url or os.getenv('DEMISTO_BASE_URL', None)
     if isinstance(configuration_orig.host, str):
@@ -139,6 +162,7 @@ def login(base_url=None, username=None, password=None, verify_ssl=True, proxy=No
             raise ValueError(err_msg)
     api_client = ApiClient(configuration_orig)
     api_client.user_agent = 'demisto-py/' + __version__
+    api_client.default_headers.update(additional_headers)
     api_instance = demisto_api.DefaultApi(api_client)
     body = {
         "user": username,
@@ -162,13 +186,15 @@ def login(base_url=None, username=None, password=None, verify_ssl=True, proxy=No
     api_client = ApiClient(configuration, header_name="X-XSRF-TOKEN", header_value=xsrf_token,
                            cookie=cookies)
     api_client.user_agent = 'demisto-py/' + __version__
+    api_client.default_headers.update(additional_headers)
     mid_client = demisto_api.DefaultApi(api_client)
     second_call = generic_request_func(self=mid_client, path='/login', method='POST', body=body,
                                        accept='application/json', content_type='application/json')
     updated_cookies = cookies + '; ' + second_call[2]['Set-Cookie']
     mid_api_client = ApiClient(configuration, header_name="X-XSRF-TOKEN", header_value=xsrf_token,
                                cookie=updated_cookies)
     mid_api_client.user_agent = 'demisto-py/' + __version__
+    mid_api_client.default_headers.update(additional_headers)
     final_api_client = demisto_api.DefaultApi(mid_api_client)
 
     return final_api_client
@@ -314,3 +340,4 @@ def get_layouts_url_for_demisto_version(api_client, params):
         if LooseVersion(server_details.get('demistoVersion')) >= LooseVersion('6.0.0'):
             url = '/layouts/import'
     return url
+