fix(manager): build_payload accepts extra by expectation_type (#206)

Author: Megafredo

pyoaev/signatures/__init__.py

@@ -11,12 +11,13 @@
     build_network_configs,
 )
 from pyoaev.signatures.signature_manager import SignatureManager
-from pyoaev.signatures.types import MatchTypes, SignatureTypes
+from pyoaev.signatures.types import ExpectationType, MatchTypes, SignatureTypes
 
 __all__ = [
     "CloudInjectorConfig",
     "ExpectationSignatureGroup",
     "ExternalInjectorConfig",
+    "ExpectationType",
     "InjectorConfig",
     "MatchTypes",
     "NetworkInjectorConfig",

pyoaev/signatures/models.py

@@ -3,7 +3,7 @@
 import ipaddress
 from typing import Any
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, JsonValue
 
 
 class SignatureValue(BaseModel):
@@ -12,7 +12,7 @@ class SignatureValue(BaseModel):
     model_config = ConfigDict(extra="allow")
 
     signature_type: str
-    signature_value: str
+    signature_value: JsonValue
 
 
 class ExpectationSignatureGroup(BaseModel):

pyoaev/signatures/signature_manager.py

@@ -23,6 +23,7 @@
     TargetSignatures,
     ToolOutput,
 )
+from pyoaev.signatures.types import ExpectationType
 
 if TYPE_CHECKING:
     from pyoaev.client import OpenAEV
@@ -172,10 +173,11 @@ def _merge_post(
 
         return merged
 
+    @staticmethod
     def build_payload(
-        self,
         post_signatures: dict[str, Any] | list[dict[str, Any]],
         expectation_types: list[str],
+        extra_signature: dict | None = None,
     ) -> dict[str, Any]:
         """Build the nested wire payload from flat post-execution signatures.
 
@@ -185,26 +187,54 @@ def build_payload(
         Args:
             post_signatures: A single post-execution dict or a list (multi-targets).
             expectation_types: The 1+ expectation type labels (e.g. ['DETECTION', 'PREVENTION']).
+            extra_signature: Optional mapping of expectation types to additional signature fields that will be merged
+            into the base post_signatures.
 
         Returns:
             A payload dict ready for send_signatures.
         """
         if isinstance(post_signatures, dict):
             post_signatures = [post_signatures]
 
+        # Validate expectation types
+        expectation_types_valid = [
+            ExpectationType(expectation_type.upper())
+            for expectation_type in expectation_types
+        ]
+        # Normalize and validate extra_signatures keys
+        extra_signatures = {
+            ExpectationType(key.upper()): value
+            for key, value in (extra_signature or {}).items()
+        }
+
         targets = []
-        for sig in post_signatures:
-            values = [
-                SignatureValue(signature_type=k, signature_value=str(v))
-                for k, v in sig.items()
-            ]
-            signature_values = [
-                ExpectationSignatureGroup(
-                    expectation_type=expectation_type, values=values
+        for signature in post_signatures:
+            signature_values = []
+
+            for expectation_type in expectation_types_valid:
+                signature_data = signature.copy()
+                if expectation_type in extra_signatures:
+                    signature_data.update(extra_signatures[expectation_type])
+
+                values = [
+                    SignatureValue(
+                        signature_type=key,
+                        signature_value=value,
+                    )
+                    for key, value in signature_data.items()
+                ]
+
+                signature_values.append(
+                    ExpectationSignatureGroup(
+                        expectation_type=expectation_type.value,
+                        values=values,
+                    )
                 )
-                for expectation_type in expectation_types
-            ]
-            targets.append(TargetSignatures(signature_values=signature_values))
+            targets.append(
+                TargetSignatures(
+                    signature_values=signature_values,
+                )
+            )
 
         return SignaturePayload(targets=targets).model_dump()
 

pyoaev/signatures/types.py

@@ -1,6 +1,12 @@
 from enum import Enum
 
 
+class ExpectationType(str, Enum):
+    DETECTION = "DETECTION"
+    PREVENTION = "PREVENTION"
+    VULNERABILITY = "VULNERABILITY"
+
+
 class MatchTypes(str, Enum):
     MATCH_TYPE_FUZZY = "fuzzy"
     MATCH_TYPE_SIMPLE = "simple"