Skip to content

Commit 7f5b546

Browse files
mariotmariot
committed
fix(signatures): send canonical signature types and skip invalid expectations (#352)
SignatureManager.build_payload sent the ExecutionSignature field names (start_time, source_ipv4, target_hostname, ...) as signature types. The platform stores them as is, so every collector failed to build the whole expectation batch and no expectation was matched. - build_payload maps the fields to their SignatureTypes value and leaves out unknown keys and None, list or dict values, with one warning - expectations_models_for_source builds expectations one by one and skips the invalid ones with a warning instead of failing the whole fetch
1 parent 8f4f9d6 commit 7f5b546

4 files changed

Lines changed: 446 additions & 14 deletions

File tree

‎pyoaev/apis/inject_expectation/inject_expectation.py‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
1+
import logging
12
from typing import Any, Dict, List
23

4+
from pydantic import ValidationError
5+
36
from pyoaev import exceptions as exc
47
from pyoaev.apis.inject_expectation.model import (
58
DetectionExpectation,
@@ -10,6 +13,8 @@
1013
from pyoaev.mixins import ListMixin, UpdateMixin
1114
from pyoaev.utils import RequiredOptional
1215

16+
logger = logging.getLogger(__name__)
17+
1318

1419
class InjectExpectation(RESTObject):
1520
_id_attr = "inject_expectation_id"
@@ -98,19 +103,21 @@ def expectations_models_for_source(self, source_id: str, **kwargs: Any):
98103
expectation_dict["inject_expectation_type"]
99104
== ExpectationTypeEnum.Detection.value
100105
):
101-
expectations.append(
102-
DetectionExpectation(**expectation_dict, api_client=self)
103-
)
104-
elif (
105-
expectation_dict["inject_expectation_type"]
106-
== ExpectationTypeEnum.Prevention.value
107-
):
108-
expectations.append(
109-
PreventionExpectation(**expectation_dict, api_client=self)
110-
)
106+
expectation_class = DetectionExpectation
111107
else:
108+
expectation_class = PreventionExpectation
109+
# One invalid expectation (e.g. an unknown signature type) must not
110+
# hide the valid ones from the collector.
111+
try:
112112
expectations.append(
113-
PreventionExpectation(**expectation_dict, api_client=self)
113+
expectation_class(**expectation_dict, api_client=self)
114+
)
115+
except ValidationError as err:
116+
logger.warning(
117+
"Skipping expectation %s, it does not match the %s model: %s",
118+
expectation_dict.get("inject_expectation_id"),
119+
expectation_class.__name__,
120+
err,
114121
)
115122
return expectations
116123

‎pyoaev/signatures/signature_manager.py‎

Lines changed: 60 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,28 @@
2525
TargetSignatures,
2626
ToolOutput,
2727
)
28+
from pyoaev.signatures.types import SignatureTypes
2829

2930
if TYPE_CHECKING:
3031
from pyoaev.client import OpenAEV
3132

33+
logger = logging.getLogger(__name__)
34+
35+
# ExecutionSignature fields whose name differs from the signature type the platform
36+
# and collectors expect. The cloud fields already use the canonical names.
37+
EXECUTION_SIGNATURE_TYPES = {
38+
"start_time": SignatureTypes.SIG_TYPE_START_DATE.value,
39+
"end_time": SignatureTypes.SIG_TYPE_END_DATE.value,
40+
"source_ipv4": SignatureTypes.SIG_TYPE_SOURCE_IPV4_ADDRESS.value,
41+
"source_ipv6": SignatureTypes.SIG_TYPE_SOURCE_IPV6_ADDRESS.value,
42+
"target_ipv4": SignatureTypes.SIG_TYPE_TARGET_IPV4_ADDRESS.value,
43+
"target_ipv6": SignatureTypes.SIG_TYPE_TARGET_IPV6_ADDRESS.value,
44+
"target_hostname": SignatureTypes.SIG_TYPE_TARGET_HOSTNAME_ADDRESS.value,
45+
}
46+
SIGNATURE_TYPE_VALUES = frozenset(
47+
signature_type.value for signature_type in SignatureTypes
48+
)
49+
3250

3351
class SignatureManager:
3452
"""End-to-end signature pipeline: compile, merge, transmit. One class, three jobs."""
@@ -165,19 +183,29 @@ def build_payload(
165183
into the base post_signatures.
166184
167185
Returns:
168-
A payload dict ready for send_signatures.
186+
A payload dict ready for send_signatures. Signature types are the canonical
187+
``SignatureTypes`` values (e.g. ``start_time`` is sent as ``start_date``);
188+
keys that are not a ``SignatureTypes`` value, and non-scalar values, are
189+
left out because collectors cannot parse them.
169190
"""
170191
if not isinstance(execution_signatures, list):
171192
execution_signatures = [execution_signatures]
172193
if isinstance(targets_meta, dict):
173194
targets_meta = [targets_meta] * len(execution_signatures)
174195

196+
skipped: set[str] = set()
175197
targets = []
176198
for signature, target in zip(execution_signatures, targets_meta):
177199
signature_values = []
178200
for expectation_type in expectation_types:
179-
signature_data = signature.model_dump(exclude_none=True)
180-
signature_data.update(extra_signatures.get_extra(expectation_type))
201+
signature_data = SignatureManager._canonical_signatures(
202+
signature.model_dump(exclude_none=True), skipped
203+
)
204+
signature_data.update(
205+
SignatureManager._canonical_signatures(
206+
extra_signatures.get_extra(expectation_type), skipped
207+
)
208+
)
181209
values = [
182210
SignatureValue(signature_type=key, signature_value=value)
183211
for key, value in signature_data.items()
@@ -199,8 +227,37 @@ def build_payload(
199227
)
200228
)
201229

230+
if skipped:
231+
logger.warning(
232+
"Signatures left out of the payload, not a known signature type "
233+
"or not a single value: %s",
234+
", ".join(sorted(skipped)),
235+
)
202236
return SignaturePayload(targets=targets).model_dump()
203237

238+
@staticmethod
239+
def _canonical_signatures(
240+
signature_data: dict[str, Any], skipped: set[str]
241+
) -> dict[str, Any]:
242+
"""Rename signature keys to their ``SignatureTypes`` value and drop the rest.
243+
244+
Collectors parse each signature as a ``SignatureTypes`` and a string, and the
245+
platform discards empty and non-scalar values, so unknown keys and ``None``,
246+
list or dict values are added to ``skipped`` instead.
247+
"""
248+
canonical = {}
249+
for key, value in signature_data.items():
250+
signature_type = EXECUTION_SIGNATURE_TYPES.get(key, key)
251+
if (
252+
signature_type not in SIGNATURE_TYPE_VALUES
253+
or value is None
254+
or isinstance(value, (list, dict))
255+
):
256+
skipped.add(key)
257+
continue
258+
canonical[signature_type] = value
259+
return canonical
260+
204261
def send_signatures(
205262
self,
206263
inject_id: str,

‎test/apis/inject_expectation/test_inject_expectation.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
from unittest import TestCase, main, mock
2+
from uuid import uuid4
23

34
from pyoaev import OpenAEV
5+
from pyoaev.apis.inject_expectation.model import (
6+
DetectionExpectation,
7+
PreventionExpectation,
8+
)
49
from pyoaev.exceptions import OpenAEVHttpError, OpenAEVListError, OpenAEVParsingError
510

611

@@ -65,5 +70,97 @@ def test_http_error_is_mapped_to_list_error(self):
6570
api_client.inject_expectation.ai_expectations_for_source("collector-1")
6671

6772

73+
def make_expectation(expectation_type, signatures):
74+
return {
75+
"inject_expectation_id": str(uuid4()),
76+
"inject_expectation_type": expectation_type,
77+
"inject_expectation_signatures": signatures,
78+
}
79+
80+
81+
class TestExpectationsModelsForSource(TestCase):
82+
def setUp(self):
83+
self.api_client = OpenAEV("url", "token")
84+
85+
def models_for(self, expectations):
86+
with mock.patch.object(
87+
self.api_client.inject_expectation,
88+
"expectations_assets_for_source",
89+
return_value=expectations,
90+
):
91+
return self.api_client.inject_expectation.expectations_models_for_source(
92+
"collector-1"
93+
)
94+
95+
def test_builds_the_model_matching_each_expectation_type(self):
96+
signatures = [{"type": "process_name", "value": "obas-implant.exe"}]
97+
98+
models = self.models_for(
99+
[
100+
make_expectation("DETECTION", signatures),
101+
make_expectation("PREVENTION", signatures),
102+
make_expectation("VULNERABILITY", signatures),
103+
]
104+
)
105+
106+
self.assertEqual(
107+
[DetectionExpectation, PreventionExpectation, PreventionExpectation],
108+
[type(model) for model in models],
109+
)
110+
111+
def test_an_invalid_expectation_does_not_hide_the_valid_ones(self):
112+
# Regression for client-python#352: signature types sent by SignatureManager
113+
# before the fix (start_time, source_ipv4, ...) are not SignatureTypes.
114+
valid_detection = make_expectation(
115+
"DETECTION", [{"type": "process_name", "value": "eicar.com"}]
116+
)
117+
legacy = make_expectation(
118+
"PREVENTION",
119+
[
120+
{"type": "start_time", "value": "2026-10-16T08:00:00Z"},
121+
{"type": "source_ipv4", "value": "10.0.0.5"},
122+
],
123+
)
124+
valid_prevention = make_expectation(
125+
"PREVENTION", [{"type": "target_ipv4_address", "value": "192.0.2.11"}]
126+
)
127+
128+
with self.assertLogs(
129+
"pyoaev.apis.inject_expectation.inject_expectation", level="WARNING"
130+
) as logs:
131+
models = self.models_for([valid_detection, legacy, valid_prevention])
132+
133+
self.assertEqual(
134+
[
135+
valid_detection["inject_expectation_id"],
136+
valid_prevention["inject_expectation_id"],
137+
],
138+
[str(model.inject_expectation_id) for model in models],
139+
)
140+
self.assertEqual(1, len(logs.records))
141+
self.assertIn(legacy["inject_expectation_id"], logs.output[0])
142+
self.assertIn("PreventionExpectation", logs.output[0])
143+
144+
def test_returns_an_empty_list_when_every_expectation_is_invalid(self):
145+
with self.assertLogs(
146+
"pyoaev.apis.inject_expectation.inject_expectation", level="WARNING"
147+
) as logs:
148+
models = self.models_for(
149+
[
150+
make_expectation("DETECTION", [{"type": "end_time", "value": "x"}]),
151+
make_expectation("PREVENTION", [{"type": "process_name"}]),
152+
]
153+
)
154+
155+
self.assertEqual([], models)
156+
self.assertEqual(2, len(logs.records))
157+
158+
def test_expectations_with_null_signatures_are_still_returned(self):
159+
models = self.models_for([make_expectation("DETECTION", None)])
160+
161+
self.assertEqual(1, len(models))
162+
self.assertEqual([], models[0].inject_expectation_signatures)
163+
164+
68165
if __name__ == "__main__":
69166
main()

0 commit comments

Comments
 (0)