[client] Implement timeout on HTTP requests / RabbitMQ consumption + …

…heartbeat (#726)
OpenCTI-Platform · Sep 10, 2024 · 2e01c16 · 2e01c16
1 parent e0c3a8e
commit 2e01c16
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/pycti/api/opencti_api_client.py b/pycti/api/opencti_api_client.py
@@ -330,6 +330,7 @@ def query(self, query, variables=None):
                 verify=self.ssl_verify,
                 cert=self.cert,
                 proxies=self.proxies,
+                timeout=300,
             )
         # If no
         else:
@@ -340,6 +341,7 @@ def query(self, query, variables=None):
                 verify=self.ssl_verify,
                 cert=self.cert,
                 proxies=self.proxies,
+                timeout=300,
             )
         # Build response
         if r.status_code == 200:
@@ -383,6 +385,7 @@ def fetch_opencti_file(self, fetch_uri, binary=False, serialize=False):
             verify=self.ssl_verify,
             cert=self.cert,
             proxies=self.proxies,
+            timeout=300,
         )
         if binary:
             if serialize:

diff --git a/pycti/connector/opencti_connector_helper.py b/pycti/connector/opencti_connector_helper.py
@@ -375,6 +375,8 @@ def run(self) -> None:
                 # Connect the broker
                 self.pika_credentials = pika.PlainCredentials(self.user, self.password)
                 self.pika_parameters = pika.ConnectionParameters(
+                    heartbeat=10,
+                    blocked_connection_timeout=10,
                     host=self.host,
                     port=self.port,
                     virtual_host=self.vhost,
@@ -466,6 +468,7 @@ def ping(self) -> None:
                         "Connector state has been remotely reset",
                         {"state": self.get_state()},
                     )
+
                 if self.in_error:
                     self.in_error = False
                     self.connector_logger.info("API Ping back to normal")
@@ -1714,6 +1717,8 @@ def send_stix2_bundle(self, bundle: str, **kwargs) -> list:
                     self.connector_config["connection"]["pass"],
                 )
                 pika_parameters = pika.ConnectionParameters(
+                    heartbeat=10,
+                    blocked_connection_timeout=10,
                     host=self.connector_config["connection"]["host"],
                     port=self.connector_config["connection"]["port"],
                     virtual_host=self.connector_config["connection"]["vhost"],