feat(contracts): declare expected security platform types on expectations (#313) (#314)

Author: SamuelHassine

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

test/contracts/test_contract_expectations.py

@@ -7,17 +7,24 @@
     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,
         expectation_description="",
         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()