[client] Add capability to add custom headers in opencti http client (#844)

Author: richard-julien

pycti/api/opencti_api_client.py

@@ -3,7 +3,7 @@
 import datetime
 import io
 import json
-from typing import Union
+from typing import Dict, Optional, Tuple, Union
 
 import magic
 import requests
@@ -71,6 +71,25 @@
 from pycti.utils.opencti_stix2_utils import OpenCTIStix2Utils
 
 
+def build_request_headers(token: str, custom_headers: str, app_logger):
+    headers_dict = {
+        "User-Agent": "pycti/" + __version__,
+        "Authorization": "Bearer " + token,
+    }
+    # Build and add custom headers
+    if custom_headers is not None:
+        for header_pair in custom_headers.strip().split(";"):
+            if header_pair:  # Skip empty header pairs
+                try:
+                    key, value = header_pair.split(":", 1)
+                    headers_dict[key.strip()] = value.strip()
+                except ValueError:
+                    app_logger.warning(
+                        "Ignored invalid header pair", {"header_pair": header_pair}
+                    )
+    return headers_dict
+
+
 class File:
     def __init__(self, name, data, mime="text/plain"):
         self.name = name
@@ -107,16 +126,16 @@ class OpenCTIApiClient:
 
     def __init__(
         self,
-        url,
-        token,
-        log_level="info",
-        ssl_verify=False,
-        proxies=None,
-        json_logging=False,
-        bundle_send_to_queue=True,
-        cert=None,
-        auth=None,
-        perform_health_check=True,
+        url: str,
+        token: str,
+        log_level: str = "info",
+        ssl_verify: bool = False,
+        proxies: Optional[Dict[str, str]] = None,
+        json_logging: Optional[bool] = False,
+        bundle_send_to_queue: Optional[bool] = True,
+        cert: Optional[Union[str, Tuple[str, str]]] = None,
+        custom_headers: Optional[str] = None,
+        perform_health_check: Optional[bool] = True,
     ):
         """Constructor method"""
 
@@ -138,17 +157,10 @@ def __init__(
         # Define API
         self.api_token = token
         self.api_url = url + "/graphql"
-        self.request_headers = {
-            "User-Agent": "pycti/" + __version__,
-            "Authorization": "Bearer " + token,
-        }
-
-        if auth is not None:
-            self.session = requests.session()
-            self.session.auth = auth
-        else:
-            self.session = requests.session()
-
+        self.request_headers = build_request_headers(
+            token, custom_headers, self.app_logger
+        )
+        self.session = requests.session()
         # Define the dependencies
         self.work = OpenCTIApiWork(self)
         self.playbook = OpenCTIApiPlaybook(self)

pycti/connector/opencti_connector_helper.py

@@ -777,6 +777,12 @@ def __init__(self, config: Dict, playbook_compatible=False) -> None:
         self.opencti_token = get_config_variable(
             "OPENCTI_TOKEN", ["opencti", "token"], config
         )
+        self.opencti_custom_headers = get_config_variable(
+            "OPENCTI_CUSTOM_HEADERS",
+            ["opencti", "custom_headers"],
+            config,
+            default=None,
+        )
         self.opencti_ssl_verify = get_config_variable(
             "OPENCTI_SSL_VERIFY", ["opencti", "ssl_verify"], config, False, False
         )
@@ -922,6 +928,7 @@ def __init__(self, config: Dict, playbook_compatible=False) -> None:
             self.log_level,
             self.opencti_ssl_verify,
             json_logging=self.opencti_json_logging,
+            custom_headers=self.opencti_custom_headers,
             bundle_send_to_queue=self.bundle_send_to_queue,
         )
         # - Impersonate API that will use applicant id
@@ -932,6 +939,7 @@ def __init__(self, config: Dict, playbook_compatible=False) -> None:
             self.log_level,
             self.opencti_ssl_verify,
             json_logging=self.opencti_json_logging,
+            custom_headers=self.opencti_custom_headers,
             bundle_send_to_queue=self.bundle_send_to_queue,
         )
         self.connector_logger = self.api.logger_class(self.connect_name)