fix(hatching-triage-sandbox): handle search API errors gracefully (#4756)

internal-enrichment/hatching-triage-sandbox/src/connector/connector.py

@@ -18,6 +18,7 @@
 )
 from stix2 import URL, DomainName, EmailAddress, IPv4Address
 from triage import Client
+from triage.client import ServerError
 
 
 class HatchingTriageSandboxConnector:
@@ -25,7 +26,7 @@
         self.config = config
         self.helper = helper
 
         self.identity = self.helper.api.identity.create(
             type="Organization",
             name="Hatching Triage",
             description="Hatching Triage",
@@ -61,12 +62,12 @@
             report_url = f"https://triage.ge/{sample_id}/"
 
         # Create external reference
         external_reference = self.helper.api.external_reference.create(
             source_name="Hatching Triage Sandbox Analysis",
             url=report_url,
             description="Hatching Triage Sandbox Analysis",
         )
         self.helper.api.stix_cyber_observable.add_external_reference(
             id=observable["id"],
             external_reference_id=external_reference["id"],
         )
@@ -87,7 +88,7 @@
                     tag_value = tag_split[1]
 
                 label = self.helper.api.label.create(value=tag_value, color=label_color)
                 self.helper.api.stix_cyber_observable.add_label(
                     id=observable["id"], label_id=label["id"]
                 )
 
@@ -244,7 +245,7 @@
                     "mime_type": mime_type,
                     "x_opencti_description": f"Extracted file from sample ID {sample_id} and task ID {task_id}.",
                 }
                 response = self.helper.api.stix_cyber_observable.upload_artifact(
                     **kwargs
                 )
                 # Create Relationship between original Observable and the extracted
@@ -411,7 +412,7 @@
 
     def _send_bundle(self, bundle_objects):
         bundle = self.helper.stix2_create_bundle(bundle_objects)
         bundles_sent = self.helper.send_stix2_bundle(bundle)
         return f"Sent {len(bundles_sent)} stix bundle(s) for worker import"
 
     def _process_file(self, observable, entity_type):
@@ -456,21 +457,27 @@
         sample_id = None
 
         if self.use_existing_analysis:
-            search_paginator = self.triage_client.search(query=search_query, max=1)
-            for search in search_paginator:
-                existing_status = search["status"]
-                if existing_status == "reported":
-                    sample_id = search["id"]
-                    self.helper.connector_logger.info(
-                        f"Found existing analysis with id {sample_id} and status {existing_status}."
-                    )
-                elif existing_status == "pending":
-                    sample_id = search["id"]
-                    self.helper.connector_logger.info(
-                        f"Found existing analysis with id {sample_id} and status {existing_status}."
-                    )
-                    self._wait_for_analysis(sample_id)
-                break
+            try:
+                search_paginator = self.triage_client.search(query=search_query, max=1)
+                for search in search_paginator:
+                    existing_status = search["status"]
+                    if existing_status == "reported":
+                        sample_id = search["id"]
+                        self.helper.connector_logger.info(
+                            f"Found existing analysis with id {sample_id} and status {existing_status}."
+                        )
+                    elif existing_status == "pending":
+                        sample_id = search["id"]
+                        self.helper.connector_logger.info(
+                            f"Found existing analysis with id {sample_id} and status {existing_status}."
+                        )
+                        self._wait_for_analysis(sample_id)
+                    break
+            except ServerError as e:
+                self.helper.connector_logger.warning(
+                    f"Search API returned an error for query '{search_query}': {e}. "
+                    "Falling back to submitting a new sample."
+                )
 
         return sample_id
 

internal-enrichment/hatching-triage-sandbox/tests/tests_connector/test_connector.py

@@ -0,0 +1,81 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from connector.connector import HatchingTriageSandboxConnector
+from triage.client import ServerError
+
+
+@pytest.fixture
+def connector(mock_opencti_connector_helper):
+    config = MagicMock()
+    config.opencti.url = "http://localhost:8080"
+    config.hatching_triage_sandbox.base_url = "https://tria.ge/api"
+    config.hatching_triage_sandbox.token.get_secret_value.return_value = "fake-token"
+    config.hatching_triage_sandbox.use_existing_analysis = True
+    config.hatching_triage_sandbox.family_color = "#0059f7"
+    config.hatching_triage_sandbox.botnet_color = "#f79e00"
+    config.hatching_triage_sandbox.campaign_color = "#7a01e5"
+    config.hatching_triage_sandbox.tag_color = "#54483b"
+    config.hatching_triage_sandbox.max_tlp = "TLP:AMBER"
+
+    helper = MagicMock()
+    helper.api.identity.create.return_value = {"standard_id": "identity--fake"}
+
+    with patch("connector.connector.Client"):
+        instance = HatchingTriageSandboxConnector(config, helper)
+
+    return instance
+
+
+class TestSearchForAnalysis:
+    def test_returns_sample_id_when_reported(self, connector):
+        """Should return sample_id when an existing reported analysis is found."""
+        connector.triage_client.search.return_value = iter(
+            [{"id": "sample-123", "status": "reported"}]
+        )
+
+        result = connector._search_for_analysis("url:https://example.com")
+
+        assert result == "sample-123"
+        connector.helper.connector_logger.info.assert_called()
+
+    def test_returns_none_when_no_results(self, connector):
+        """Should return None when no existing analysis is found."""
+        connector.triage_client.search.return_value = iter([])
+
+        result = connector._search_for_analysis("url:https://example.com")
+
+        assert result is None
+
+    def test_returns_none_on_server_error(self, connector):
+        """Should return None and log a warning when a ServerError occurs (e.g. 504 timeout)."""
+        import json
+
+        mock_response = MagicMock()
+        mock_response.status_code = 504
+        mock_response.json.side_effect = json.JSONDecodeError("Expecting value", "", 0)
+
+        http_error = MagicMock()
+        http_error.response = mock_response
+
+        connector.triage_client.search.return_value = _raise_server_error(http_error)
+
+        result = connector._search_for_analysis("url:https://example.com")
+
+        assert result is None
+        connector.helper.connector_logger.warning.assert_called_once()
+
+    def test_returns_none_when_use_existing_analysis_disabled(self, connector):
+        """Should return None without searching when use_existing_analysis is False."""
+        connector.use_existing_analysis = False
+
+        result = connector._search_for_analysis("url:https://example.com")
+
+        assert result is None
+        connector.triage_client.search.assert_not_called()
+
+
+def _raise_server_error(http_error):
+    """Generator that raises a ServerError when iterated."""
+    raise ServerError(http_error)
+    yield  # noqa: unreachable - makes this a generator