diff --git a/pyoaev/contracts/contract_config.py b/pyoaev/contracts/contract_config.py index b2ea50f..8de69ad 100644 --- a/pyoaev/contracts/contract_config.py +++ b/pyoaev/contracts/contract_config.py @@ -74,6 +74,25 @@ class ExpectationType(str, Enum): vulnerability: str = "VULNERABILITY" +class SecurityPlatformType(str, Enum): + """Categories of security platform expected to fulfil a technical expectation. + + When an expectation declares one or more of these, the platform focuses the + technical (DETECTION / PREVENTION / VULNERABILITY) result on collectors of + those types only (instead of every connected security platform). An empty + list means "any security platform". + """ + + EDR: str = "EDR" + XDR: str = "XDR" + SIEM: str = "SIEM" + SOAR: str = "SOAR" + NDR: str = "NDR" + ISPM: str = "ISPM" + LLM_FIREWALL: str = "LLM_FIREWALL" + AI_GATEWAY: str = "AI_GATEWAY" + + @dataclass class Expectation: expectation_type: ExpectationType @@ -83,6 +102,12 @@ class Expectation: expectation_expectation_group: bool expectation_is_predefined: bool = False expectation_is_multi_selectable: bool = False + # Security platform types expected to fulfil this expectation. Empty = any + # platform (unchanged behaviour). Typically set for technical expectations + # (DETECTION / PREVENTION / VULNERABILITY), left empty for MANUAL ones. + expectation_expected_security_platform_types: List[SecurityPlatformType] = field( + default_factory=list + ) @dataclass diff --git a/test/contracts/test_contract_expectations.py b/test/contracts/test_contract_expectations.py index f054517..7618575 100644 --- a/test/contracts/test_contract_expectations.py +++ b/test/contracts/test_contract_expectations.py @@ -7,10 +7,16 @@ ContractExpectations, Expectation, ExpectationType, + SecurityPlatformType, ) -def _expectation(expectation_type, name, is_predefined): +def _expectation(expectation_type, name, is_predefined, expected_platforms=None): + # Only pass the field when declared so the default (None) path exercises + # the dataclass default_factory instead of an explicit empty list. + kwargs = {} + if expected_platforms is not None: + kwargs["expectation_expected_security_platform_types"] = expected_platforms return Expectation( expectation_type=expectation_type, expectation_name=name, @@ -18,6 +24,7 @@ def _expectation(expectation_type, name, is_predefined): expectation_score=100, expectation_expectation_group=False, expectation_is_predefined=is_predefined, + **kwargs, ) @@ -104,6 +111,59 @@ def test_no_flag_yields_no_predefined(self): self.assertEqual([], data["predefinedExpectations"]) + def test_expected_security_platform_types_default_empty(self): + """An expectation without declared platforms serializes an empty list.""" + field = ContractExpectations( + key="expectations", + label="Expectations", + availableExpectations=[ + _expectation(ExpectationType.detection, "Detection", True), + ], + ) + + data = _serialize(field) + + self.assertEqual( + [], + data["availableExpectations"][0][ + "expectation_expected_security_platform_types" + ], + ) + + def test_expected_security_platform_types_serialize(self): + """Declared platform types serialize as their string values on both + the available and predefined arrays.""" + field = ContractExpectations( + key="expectations", + label="Expectations", + availableExpectations=[ + _expectation( + ExpectationType.detection, + "Detection", + True, + expected_platforms=[ + SecurityPlatformType.EDR, + SecurityPlatformType.XDR, + ], + ), + ], + ) + + data = _serialize(field) + + self.assertEqual( + ["EDR", "XDR"], + data["availableExpectations"][0][ + "expectation_expected_security_platform_types" + ], + ) + self.assertEqual( + ["EDR", "XDR"], + data["predefinedExpectations"][0][ + "expectation_expected_security_platform_types" + ], + ) + if __name__ == "__main__": unittest.main()