diff --git a/pyoaev/apis/inject_expectation/inject_expectation.py b/pyoaev/apis/inject_expectation/inject_expectation.py index be1c453a..082d0f69 100644 --- a/pyoaev/apis/inject_expectation/inject_expectation.py +++ b/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 diff --git a/pyoaev/signatures/signature_manager.py b/pyoaev/signatures/signature_manager.py index 28452342..37f82271 100644 --- a/pyoaev/signatures/signature_manager.py +++ b/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, diff --git a/test/apis/inject_expectation/test_inject_expectation.py b/test/apis/inject_expectation/test_inject_expectation.py index eb42ea6e..0f0435b0 100644 --- a/test/apis/inject_expectation/test_inject_expectation.py +++ b/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() diff --git a/test/signatures/test_signature_manager_build_payload.py b/test/signatures/test_signature_manager_build_payload.py new file mode 100644 index 00000000..2a5587f1 --- /dev/null +++ b/test/signatures/test_signature_manager_build_payload.py @@ -0,0 +1,271 @@ +import unittest +from types import SimpleNamespace +from uuid import uuid4 + +from pyoaev.apis.inject_expectation.model import ( + DetectionExpectation, + PreventionExpectation, +) +from pyoaev.signatures.models import ( + ExecutionSignature, + ExtraSignatureData, + SignaturePayload, +) +from pyoaev.signatures.signature_manager import SignatureManager +from pyoaev.signatures.types import SignatureTypes + +TARGET = SimpleNamespace(agent_id="agent-1", asset_id="asset-1", asset_group_id=None) +CANONICAL_TYPES = {signature_type.value for signature_type in SignatureTypes} + + +def _network_signature(**fields): + return ExecutionSignature( + start_time="2026-10-16T08:00:00Z", + end_time="2026-10-16T08:05:00Z", + source_ipv4="10.0.0.5", + source_ipv6="fd00::5", + **fields, + ) + + +def _values(payload, target_index=0, group_index=0): + group = payload["targets"][target_index]["signature_values"][group_index] + return { + value["signature_type"]: value["signature_value"] for value in group["values"] + } + + +def _as_stored_by_platform(payload): + """Signatures as the platform stores them (SignatureOutputProcessor): the type + as is, and only non-empty scalar values, as text.""" + stored = [] + for target in payload["targets"]: + for group in target["signature_values"]: + for value in group["values"]: + signature_value = value["signature_value"] + if isinstance(signature_value, (list, dict)) or signature_value in ( + None, + "", + ): + continue + stored.append( + { + "type": value["signature_type"], + "value": str(signature_value), + } + ) + return stored + + +class TestBuildPayloadSignatureTypes(unittest.TestCase): + def test_network_ipv4_signature_uses_canonical_signature_types(self): + payload = SignatureManager.build_payload( + _network_signature(target_ipv4="192.0.2.11"), [TARGET], ["DETECTION"] + ) + + self.assertEqual( + { + "start_date": "2026-10-16T08:00:00Z", + "end_date": "2026-10-16T08:05:00Z", + "source_ipv4_address": "10.0.0.5", + "source_ipv6_address": "fd00::5", + "target_ipv4_address": "192.0.2.11", + }, + _values(payload), + ) + + def test_network_ipv6_and_hostname_targets_use_canonical_signature_types(self): + payload = SignatureManager.build_payload( + [ + _network_signature(target_ipv6="2001:db8::11"), + _network_signature(target_hostname="web.example.com"), + ], + [TARGET, TARGET], + ["DETECTION"], + ) + + self.assertEqual("2001:db8::11", _values(payload, 0)["target_ipv6_address"]) + self.assertEqual( + "web.example.com", _values(payload, 1)["target_hostname_address"] + ) + + def test_cloud_signature_keeps_cloud_types_and_renames_dates(self): + signature = ExecutionSignature( + start_time="2026-10-16T08:00:00Z", + cloud_provider="aws", + cloud_account_id="123456789012", + cloud_region="eu-west-3", + target_service="s3", + ) + + payload = SignatureManager.build_payload(signature, [TARGET], ["PREVENTION"]) + + self.assertEqual( + { + "start_date": "2026-10-16T08:00:00Z", + "cloud_provider": "aws", + "cloud_account_id": "123456789012", + "cloud_region": "eu-west-3", + "target_service": "s3", + }, + _values(payload), + ) + + def test_every_signature_type_sent_is_a_signature_types_value(self): + signature = _network_signature(target_ipv4="192.0.2.11") + signature.partial_results = ["22/tcp open"] + extra = ExtraSignatureData( + detection={"ports_discovered": [22, 80], "process_name": "nmap"}, + prevention={"services_discovered": [{"port": 22}]}, + vulnerability={"cves_tested": [], "cves_found_vulnerable": []}, + ) + + payload = SignatureManager.build_payload( + signature, [TARGET], ["DETECTION", "PREVENTION", "VULNERABILITY"], extra + ) + + for group in payload["targets"][0]["signature_values"]: + for value in group["values"]: + self.assertIn(value["signature_type"], CANONICAL_TYPES) + + +class TestBuildPayloadExtraSignatures(unittest.TestCase): + def test_unknown_extra_keys_are_left_out_and_logged(self): + extra = ExtraSignatureData( + detection={"ports_discovered": [22, 80], "services_discovered": []} + ) + + with self.assertLogs( + "pyoaev.signatures.signature_manager", level="WARNING" + ) as logs: + payload = SignatureManager.build_payload( + _network_signature(target_ipv4="192.0.2.11"), + [TARGET], + ["DETECTION"], + extra, + ) + + self.assertNotIn("ports_discovered", _values(payload)) + self.assertNotIn("services_discovered", _values(payload)) + self.assertEqual(1, len(logs.records)) + self.assertIn("ports_discovered, services_discovered", logs.output[0]) + + def test_extra_keys_that_are_signature_types_are_sent(self): + extra = ExtraSignatureData( + detection={"process_name": "nmap", "command_line": "nmap -sV 192.0.2.11"} + ) + + payload = SignatureManager.build_payload( + _network_signature(target_ipv4="192.0.2.11"), [TARGET], ["DETECTION"], extra + ) + + self.assertEqual("nmap", _values(payload)["process_name"]) + self.assertEqual("nmap -sV 192.0.2.11", _values(payload)["command_line"]) + + def test_extra_keys_named_like_execution_signature_fields_are_renamed(self): + extra = ExtraSignatureData(detection={"target_hostname": "web.example.com"}) + + payload = SignatureManager.build_payload( + _network_signature(target_ipv4="192.0.2.11"), [TARGET], ["DETECTION"], extra + ) + + self.assertEqual("web.example.com", _values(payload)["target_hostname_address"]) + self.assertNotIn("target_hostname", _values(payload)) + + def test_extra_signature_overrides_the_execution_signature_value(self): + extra = ExtraSignatureData(detection={"end_date": "2026-10-16T09:00:00Z"}) + + payload = SignatureManager.build_payload( + _network_signature(target_ipv4="192.0.2.11"), [TARGET], ["DETECTION"], extra + ) + + group = payload["targets"][0]["signature_values"][0] + end_dates = [v for v in group["values"] if v["signature_type"] == "end_date"] + self.assertEqual( + [{"signature_type": "end_date", "signature_value": "2026-10-16T09:00:00Z"}], + end_dates, + ) + + def test_extras_only_apply_to_their_expectation_type(self): + extra = ExtraSignatureData(detection={"process_name": "nmap"}) + + payload = SignatureManager.build_payload( + _network_signature(target_ipv4="192.0.2.11"), + [TARGET], + ["DETECTION", "PREVENTION"], + extra, + ) + + self.assertIn("process_name", _values(payload, group_index=0)) + self.assertNotIn("process_name", _values(payload, group_index=1)) + + def test_list_dict_and_none_values_are_left_out(self): + signature = _network_signature(target_ipv4="192.0.2.11") + signature.partial_results = ["22/tcp open"] + extra = ExtraSignatureData( + detection={ + "url_hash": ["aaa", "bbb"], + "file_hash": {"sha256": "ccc"}, + "file_name": None, + } + ) + + with self.assertLogs( + "pyoaev.signatures.signature_manager", level="WARNING" + ) as logs: + payload = SignatureManager.build_payload( + signature, [TARGET], ["DETECTION"], extra + ) + + for key in ("url_hash", "file_hash", "file_name", "partial_results"): + self.assertNotIn(key, _values(payload)) + self.assertIn("file_hash, file_name, partial_results, url_hash", logs.output[0]) + + def test_nothing_is_logged_when_every_signature_is_sent(self): + with self.assertNoLogs("pyoaev.signatures.signature_manager", level="WARNING"): + SignatureManager.build_payload( + _network_signature(target_ipv4="192.0.2.11"), + [TARGET], + ["DETECTION", "PREVENTION"], + ) + + +class TestBuildPayloadIsParsedByCollectors(unittest.TestCase): + """Regression for client-python#352: a collector must be able to parse the + expectation signatures the platform stores from this payload.""" + + def test_signatures_stored_from_the_payload_build_collector_expectations(self): + signatures = [ + _network_signature(target_ipv4="192.0.2.11"), + _network_signature(target_hostname="web.example.com"), + ] + # Same extra signatures as the Nmap injector + extra_data = { + "ports_discovered": [22, 80], + "services_discovered": [{"port": 22, "service": "ssh"}], + } + extra = ExtraSignatureData(detection=extra_data, prevention=extra_data) + + payload = SignatureManager.build_payload( + signatures, [TARGET, TARGET], ["DETECTION", "PREVENTION"], extra + ) + SignaturePayload(**payload) + stored = _as_stored_by_platform(payload) + + for model_cls in (DetectionExpectation, PreventionExpectation): + expectation = model_cls( + inject_expectation_id=uuid4(), + inject_expectation_signatures=stored, + api_client=None, + ) + self.assertEqual( + len(stored), len(expectation.inject_expectation_signatures) + ) + self.assertIn( + SignatureTypes.SIG_TYPE_TARGET_IPV4_ADDRESS, + {signature.type for signature in expectation.inject_expectation_signatures}, + ) + + +if __name__ == "__main__": + unittest.main()