fix(signatures): send canonical signature types and skip invalid expectations (#352)

pyoaev/apis/inject_expectation/inject_expectation.py

@@ -1,5 +1,8 @@
+import logging
 from typing import Any, Dict, List
 
+from pydantic import ValidationError
+
 from pyoaev import exceptions as exc
 from pyoaev.apis.inject_expectation.model import (
     DetectionExpectation,
@@ -10,6 +13,8 @@
 from pyoaev.mixins import ListMixin, UpdateMixin
 from pyoaev.utils import RequiredOptional
 
+logger = logging.getLogger(__name__)
+
 
 class InjectExpectation(RESTObject):
     _id_attr = "inject_expectation_id"
@@ -98,19 +103,21 @@ def expectations_models_for_source(self, source_id: str, **kwargs: Any):
                 expectation_dict["inject_expectation_type"]
                 == ExpectationTypeEnum.Detection.value
             ):
-                expectations.append(
-                    DetectionExpectation(**expectation_dict, api_client=self)
-                )
-            elif (
-                expectation_dict["inject_expectation_type"]
-                == ExpectationTypeEnum.Prevention.value
-            ):
-                expectations.append(
-                    PreventionExpectation(**expectation_dict, api_client=self)
-                )
+                expectation_class = DetectionExpectation
             else:
+                expectation_class = PreventionExpectation
+            # One invalid expectation (e.g. an unknown signature type) must not
+            # hide the valid ones from the collector.
+            try:
                 expectations.append(
-                    PreventionExpectation(**expectation_dict, api_client=self)
+                    expectation_class(**expectation_dict, api_client=self)
+                )
+            except ValidationError as err:
+                logger.warning(
+                    "Skipping expectation %s, it does not match the %s model: %s",
+                    expectation_dict.get("inject_expectation_id"),
+                    expectation_class.__name__,
+                    err,
                 )
         return expectations
 

pyoaev/signatures/signature_manager.py

@@ -25,10 +25,28 @@
     TargetSignatures,
     ToolOutput,
 )
+from pyoaev.signatures.types import SignatureTypes
 
 if TYPE_CHECKING:
     from pyoaev.client import OpenAEV
 
+logger = logging.getLogger(__name__)
+
+# ExecutionSignature fields whose name differs from the signature type the platform
+# and collectors expect. The cloud fields already use the canonical names.
+EXECUTION_SIGNATURE_TYPES = {
+    "start_time": SignatureTypes.SIG_TYPE_START_DATE.value,
+    "end_time": SignatureTypes.SIG_TYPE_END_DATE.value,
+    "source_ipv4": SignatureTypes.SIG_TYPE_SOURCE_IPV4_ADDRESS.value,
+    "source_ipv6": SignatureTypes.SIG_TYPE_SOURCE_IPV6_ADDRESS.value,
+    "target_ipv4": SignatureTypes.SIG_TYPE_TARGET_IPV4_ADDRESS.value,
+    "target_ipv6": SignatureTypes.SIG_TYPE_TARGET_IPV6_ADDRESS.value,
+    "target_hostname": SignatureTypes.SIG_TYPE_TARGET_HOSTNAME_ADDRESS.value,
+}
+SIGNATURE_TYPE_VALUES = frozenset(
+    signature_type.value for signature_type in SignatureTypes
+)
+
 
 class SignatureManager:
     """End-to-end signature pipeline: compile, merge, transmit. One class, three jobs."""
@@ -165,19 +183,29 @@ def build_payload(
             into the base post_signatures.
 
         Returns:
-            A payload dict ready for send_signatures.
+            A payload dict ready for send_signatures. Signature types are the canonical
+            ``SignatureTypes`` values (e.g. ``start_time`` is sent as ``start_date``);
+            keys that are not a ``SignatureTypes`` value, and non-scalar values, are
+            left out because collectors cannot parse them.
         """
         if not isinstance(execution_signatures, list):
             execution_signatures = [execution_signatures]
         if isinstance(targets_meta, dict):
             targets_meta = [targets_meta] * len(execution_signatures)
 
+        skipped: set[str] = set()
         targets = []
         for signature, target in zip(execution_signatures, targets_meta):
             signature_values = []
             for expectation_type in expectation_types:
-                signature_data = signature.model_dump(exclude_none=True)
-                signature_data.update(extra_signatures.get_extra(expectation_type))
+                signature_data = SignatureManager._canonical_signatures(
+                    signature.model_dump(exclude_none=True), skipped
+                )
+                signature_data.update(
+                    SignatureManager._canonical_signatures(
+                        extra_signatures.get_extra(expectation_type), skipped
+                    )
+                )
                 values = [
                     SignatureValue(signature_type=key, signature_value=value)
                     for key, value in signature_data.items()
@@ -199,8 +227,37 @@ def build_payload(
                 )
             )
 
+        if skipped:
+            logger.warning(
+                "Signatures left out of the payload, not a known signature type "
+                "or not a single value: %s",
+                ", ".join(sorted(skipped)),
+            )
         return SignaturePayload(targets=targets).model_dump()
 
+    @staticmethod
+    def _canonical_signatures(
+        signature_data: dict[str, Any], skipped: set[str]
+    ) -> dict[str, Any]:
+        """Rename signature keys to their ``SignatureTypes`` value and drop the rest.
+
+        Collectors parse each signature as a ``SignatureTypes`` and a string, and the
+        platform discards empty and non-scalar values, so unknown keys and ``None``,
+        list or dict values are added to ``skipped`` instead.
+        """
+        canonical = {}
+        for key, value in signature_data.items():
+            signature_type = EXECUTION_SIGNATURE_TYPES.get(key, key)
+            if (
+                signature_type not in SIGNATURE_TYPE_VALUES
+                or value is None
+                or isinstance(value, (list, dict))
+            ):
+                skipped.add(key)
+                continue
+            canonical[signature_type] = value
+        return canonical
+
     def send_signatures(
         self,
         inject_id: str,

test/apis/inject_expectation/test_inject_expectation.py

@@ -1,6 +1,11 @@
 from unittest import TestCase, main, mock
+from uuid import uuid4
 
 from pyoaev import OpenAEV
+from pyoaev.apis.inject_expectation.model import (
+    DetectionExpectation,
+    PreventionExpectation,
+)
 from pyoaev.exceptions import OpenAEVHttpError, OpenAEVListError, OpenAEVParsingError
 
 
@@ -65,5 +70,97 @@ def test_http_error_is_mapped_to_list_error(self):
                 api_client.inject_expectation.ai_expectations_for_source("collector-1")
 
 
+def make_expectation(expectation_type, signatures):
+    return {
+        "inject_expectation_id": str(uuid4()),
+        "inject_expectation_type": expectation_type,
+        "inject_expectation_signatures": signatures,
+    }
+
+
+class TestExpectationsModelsForSource(TestCase):
+    def setUp(self):
+        self.api_client = OpenAEV("url", "token")
+
+    def models_for(self, expectations):
+        with mock.patch.object(
+            self.api_client.inject_expectation,
+            "expectations_assets_for_source",
+            return_value=expectations,
+        ):
+            return self.api_client.inject_expectation.expectations_models_for_source(
+                "collector-1"
+            )
+
+    def test_builds_the_model_matching_each_expectation_type(self):
+        signatures = [{"type": "process_name", "value": "obas-implant.exe"}]
+
+        models = self.models_for(
+            [
+                make_expectation("DETECTION", signatures),
+                make_expectation("PREVENTION", signatures),
+                make_expectation("VULNERABILITY", signatures),
+            ]
+        )
+
+        self.assertEqual(
+            [DetectionExpectation, PreventionExpectation, PreventionExpectation],
+            [type(model) for model in models],
+        )
+
+    def test_an_invalid_expectation_does_not_hide_the_valid_ones(self):
+        # Regression for client-python#352: signature types sent by SignatureManager
+        # before the fix (start_time, source_ipv4, ...) are not SignatureTypes.
+        valid_detection = make_expectation(
+            "DETECTION", [{"type": "process_name", "value": "eicar.com"}]
+        )
+        legacy = make_expectation(
+            "PREVENTION",
+            [
+                {"type": "start_time", "value": "2026-10-16T08:00:00Z"},
+                {"type": "source_ipv4", "value": "10.0.0.5"},
+            ],
+        )
+        valid_prevention = make_expectation(
+            "PREVENTION", [{"type": "target_ipv4_address", "value": "192.0.2.11"}]
+        )
+
+        with self.assertLogs(
+            "pyoaev.apis.inject_expectation.inject_expectation", level="WARNING"
+        ) as logs:
+            models = self.models_for([valid_detection, legacy, valid_prevention])
+
+        self.assertEqual(
+            [
+                valid_detection["inject_expectation_id"],
+                valid_prevention["inject_expectation_id"],
+            ],
+            [str(model.inject_expectation_id) for model in models],
+        )
+        self.assertEqual(1, len(logs.records))
+        self.assertIn(legacy["inject_expectation_id"], logs.output[0])
+        self.assertIn("PreventionExpectation", logs.output[0])
+
+    def test_returns_an_empty_list_when_every_expectation_is_invalid(self):
+        with self.assertLogs(
+            "pyoaev.apis.inject_expectation.inject_expectation", level="WARNING"
+        ) as logs:
+            models = self.models_for(
+                [
+                    make_expectation("DETECTION", [{"type": "end_time", "value": "x"}]),
+                    make_expectation("PREVENTION", [{"type": "process_name"}]),
+                ]
+            )
+
+        self.assertEqual([], models)
+        self.assertEqual(2, len(logs.records))
+
+    def test_expectations_with_null_signatures_are_still_returned(self):
+        models = self.models_for([make_expectation("DETECTION", None)])
+
+        self.assertEqual(1, len(models))
+        self.assertEqual([], models[0].inject_expectation_signatures)
+
+
 if __name__ == "__main__":
     main()