✨ feat(pre_execution): add InjectorConfigModel and builder for NetworkOne

Kakudou · Kakudou · commit f7b5e24f0ea8 · 2026-05-27T18:41:17.000+02:00
diff --git a/pyoaev/signatures/__init__.py b/pyoaev/signatures/__init__.py
@@ -1,22 +1,32 @@
 from pyoaev.signatures.models import (
+    CloudInjectorConfig,
     ExpectationSignatureGroup,
+    ExternalInjectorConfig,
+    InjectorConfig,
+    NetworkInjectorConfig,
     SignatureCallbackPayload,
     SignaturePayload,
     SignatureTarget,
     SignatureValue,
     TargetSignatures,
+    build_network_configs,
 )
 from pyoaev.signatures.signature_manager import SignatureManager
 from pyoaev.signatures.types import MatchTypes, SignatureTypes
 
 __all__ = [
+    "CloudInjectorConfig",
     "ExpectationSignatureGroup",
+    "ExternalInjectorConfig",
+    "InjectorConfig",
     "MatchTypes",
+    "NetworkInjectorConfig",
     "SignatureCallbackPayload",
     "SignatureManager",
     "SignaturePayload",
     "SignatureTarget",
     "SignatureTypes",
     "SignatureValue",
     "TargetSignatures",
+    "build_network_configs",
 ]
diff --git a/pyoaev/signatures/models.py b/pyoaev/signatures/models.py
@@ -1,5 +1,6 @@
 """Pydantic schemas pinning every shape SignatureManager touches."""
 
+import ipaddress
 from typing import Any
 
 from pydantic import BaseModel, ConfigDict
@@ -122,6 +123,90 @@ class ToolOutput(BaseModel):
     extra_signatures: dict[str, Any] | None = None
 
 
+class NetworkInjectorConfig(BaseModel):
+    """A single network target. Exactly one of ``target_ipv4``, ``target_ipv6``, or ``target_hostname``."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    target_ipv4: str | None = None
+    target_ipv6: str | None = None
+    target_hostname: str | None = None
+
+
+class CloudInjectorConfig(BaseModel):
+    """A single cloud target row. One config per region; fan out by passing a list."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    cloud_provider: str
+    cloud_account_id: str
+    cloud_region: str
+    target_service: str | None = None
+
+
+class ExternalInjectorConfig(BaseModel):
+    """A single external scan target (e.g. Shodan): a query against an asset."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    query: str
+    target_ipv4: str | None = None
+    target_hostname: str | None = None
+
+
+InjectorConfig = NetworkInjectorConfig | CloudInjectorConfig | ExternalInjectorConfig
+
+
+# ---------------------------------------------------------------------------
+# Builders. Cheap helpers to turn raw injector input into typed configs.
+# ---------------------------------------------------------------------------
+
+
+def _classify_network_target(value: str) -> NetworkInjectorConfig:
+    """Decide whether ``value`` is an IPv4, IPv6, or hostname and wrap it."""
+    try:
+        addr = ipaddress.ip_address(value)
+    except ValueError:
+        return NetworkInjectorConfig(target_hostname=value)
+    if isinstance(addr, ipaddress.IPv4Address):
+        return NetworkInjectorConfig(target_ipv4=value)
+    return NetworkInjectorConfig(target_ipv6=value)
+
+
+def build_network_configs(
+    targets: list[str | dict[str, Any] | NetworkInjectorConfig],
+) -> list[NetworkInjectorConfig]:
+    """Forge a list of `NetworkInjectorConfig` from a heterogeneous target list.
+
+    Each item is one distinct asset. Accepted shapes:
+
+    - `NetworkInjectorConfig`: passed through unchanged.
+    - `dict`: validated against :class:`NetworkInjectorConfig`.
+    - `str`: auto-classified into IPv4 / IPv6 / hostname.
+
+    Args:
+        targets: Raw target list straight out of the injector.
+
+    Returns:
+        One `NetworkInjectorConfig` per input target, order preserved.
+
+    Raises:
+        TypeError: An item is not one of the accepted shapes.
+        ValidationError: A dict item fails the one-identity invariant.
+    """
+    configs: list[NetworkInjectorConfig] = []
+    for target in targets:
+        if isinstance(target, NetworkInjectorConfig):
+            configs.append(target)
+        elif isinstance(target, dict):
+            configs.append(NetworkInjectorConfig(**target))
+        elif isinstance(target, str):
+            configs.append(_classify_network_target(target))
+        else:
+            raise TypeError(f"unsupported network target type: {type(target).__name__}")
+    return configs
+
+
 __all__ = [
     "SignatureValue",
     "ExpectationSignatureGroup",
@@ -134,4 +219,9 @@ class ToolOutput(BaseModel):
     "ToolErrorInfo",
     "ToolTimeoutInfo",
     "ToolOutput",
+    "NetworkInjectorConfig",
+    "CloudInjectorConfig",
+    "ExternalInjectorConfig",
+    "InjectorConfig",
+    "build_network_configs",
 ]
diff --git a/pyoaev/signatures/signature_manager.py b/pyoaev/signatures/signature_manager.py
@@ -5,13 +5,17 @@
 import socket
 import subprocess
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any
 
 from pydantic import ValidationError
 
 from pyoaev.exceptions import OpenAEVError
 from pyoaev.signatures.models import (
+    CloudInjectorConfig,
     ExpectationSignatureGroup,
+    ExternalInjectorConfig,
+    InjectorConfig,
+    NetworkInjectorConfig,
     PostExecutionSignature,
     PreExecutionSignature,
     SignaturePayload,
@@ -48,120 +52,72 @@ def _utcnow(self) -> datetime:
 
     def compile_pre_execution_signatures(
         self,
-        inject_config: dict[str, Any],
-        category: Literal["network", "cloud", "external"],
+        config: InjectorConfig | list[InjectorConfig],
     ) -> dict[str, Any] | list[dict[str, Any]]:
-        """Build pre-execution signature dicts off the inject config and category.
+        """Build pre-execution signature dicts from one or more typed injector configs.
+
+        The category is carried by the config type itself
+        (:class:`NetworkInjectorConfig`, :class:`CloudInjectorConfig`,
+        :class:`ExternalInjectorConfig`), so no separate ``category`` flag is needed.
 
         Args:
-            inject_config: The inject payload dict.
-            category: One of 'network', 'cloud', 'external'.
+            config: A single injector config or a homogeneous list of them.
+                Multi-target injects must be expressed as a list.
 
         Returns:
-            One dict for single-target, list of dicts for multi-target configs.
+            One dict when a single config is given, otherwise a list of dicts in
+            input order.
 
         Raises:
-            ValueError: Unknown category or required fields missing.
+            ValueError: Empty list, or mixed config types in a single call.
+            TypeError: Unknown injector config type.
         """
-        now = self._utcnow()
-        start_time = now.strftime("%Y-%m-%dT%H:%M:%SZ")
-
-        if category == "network":
-            return self._compile_network_pre(inject_config, start_time)
-        elif category == "cloud":
-            return self._compile_cloud_pre(inject_config, start_time)
-        elif category == "external":
-            return self._compile_external_pre(inject_config, start_time)
-        else:
-            raise ValueError(f"Unknown category: {category!r}")
-
-    def _compile_network_pre(
-        self, config: dict[str, Any], start_time: str
-    ) -> dict[str, Any] | list[dict[str, Any]]:
-        ipv4 = self.resolve_container_ip()
-        ipv6 = self._cached_ipv6
-
-        assets = config.get("target_assets") or config.get("assets") or []
-        if not assets:
-            asset = config.get("asset")
-            if asset:
-                assets = [asset]
-
-        if not assets:
+        configs = list(config) if isinstance(config, list) else [config]
+        if not configs:
             raise ValueError(
-                "inject_config must contain 'target_assets', 'assets', or 'asset'"
+                "compile_pre_execution_signatures requires at least one config"
             )
 
-        results: list[dict[str, Any]] = []
-        for asset in assets:
-            sig = PreExecutionSignature(
-                start_time=start_time,
-                source_ipv4=ipv4,
-                source_ipv6=ipv6,
-                target_ipv4=asset["target_ipv4"],
-                target_ipv6=asset.get("target_ipv6"),
-                target_hostname=asset.get("target_hostname"),
+        first_type = type(configs[0])
+        if any(type(c) is not first_type for c in configs):
+            raise ValueError(
+                "compile_pre_execution_signatures does not mix injector config types; "
+                f"got {sorted({type(c).__name__ for c in configs})}"
             )
-            results.append(sig.model_dump(mode="json", exclude_none=True))
 
+        start_time = self._utcnow().strftime("%Y-%m-%dT%H:%M:%SZ")
+        results = [self._compile_one(cfg, start_time) for cfg in configs]
         return results[0] if len(results) == 1 else results
 
-    def _compile_cloud_pre(
-        self, config: dict[str, Any], start_time: str
-    ) -> dict[str, Any] | list[dict[str, Any]]:
-        cloud_provider = config["cloud_provider"]
-        cloud_account_id = config["cloud_account_id"]
-        target_service = config.get("target_service")
-
-        regions = config.get("regions") or []
-        if not regions:
-            region = config.get("cloud_region")
-            if region:
-                regions = [region]
-
-        if not regions:
-            raise ValueError("inject_config must contain 'regions' or 'cloud_region'")
-
-        results: list[dict[str, Any]] = []
-        for region in regions:
-            sig = PreExecutionSignature(
-                start_time=start_time,
-                cloud_provider=cloud_provider,
-                cloud_account_id=cloud_account_id,
-                cloud_region=region,
-                target_service=target_service,
-            )
-            results.append(sig.model_dump(mode="json", exclude_none=True))
+    def _compile_one(self, config: InjectorConfig, start_time: str) -> dict[str, Any]:
+        """Project a single injector config into a flat pre-execution signature dict.
 
-        return results[0] if len(results) == 1 else results
-
-    def _compile_external_pre(
-        self, config: dict[str, Any], start_time: str
-    ) -> dict[str, Any] | list[dict[str, Any]]:
-        targets = config.get("targets") or []
-        if not targets:
-            query = config.get("query")
-            if query is None:
-                raise ValueError("inject_config must contain 'query'")
-            targets = [
-                {
-                    "target_ipv4": config["target_ipv4"],
-                    "target_hostname": config.get("target_hostname"),
-                    "query": query,
-                }
-            ]
+        Common pipeline for every category:
+          1. Seed the base dict with ``start_time`` and category-specific context
+             (network gets resolved source IPs; cloud/external add nothing).
+          2. Layer the config's own fields on top.
+          3. Run it through :class:`PreExecutionSignature` for validation
+             and emit JSON-ready output stripped of ``None``\\ s.
+        """
+        base: dict[str, Any] = {"start_time": start_time}
+        base.update(self._source_context(config))
+        base.update(config.model_dump(exclude_none=True))
+        return PreExecutionSignature(**base).model_dump(mode="json", exclude_none=True)
 
-        results: list[dict[str, Any]] = []
-        for target in targets:
-            sig = PreExecutionSignature(
-                start_time=start_time,
-                target_ipv4=target["target_ipv4"],
-                target_hostname=target.get("target_hostname"),
-                query=target.get("query"),
-            )
-            results.append(sig.model_dump(mode="json", exclude_none=True))
+    def _source_context(self, config: InjectorConfig) -> dict[str, Any]:
+        """Return the source identity bits injected for the config's category.
 
-        return results[0] if len(results) == 1 else results
+        Only network signatures need the running container's source IPs;
+        cloud and external rows have no source identity to carry.
+        """
+        if isinstance(config, NetworkInjectorConfig):
+            return {
+                "source_ipv4": self.resolve_container_ip(),
+                "source_ipv6": self._cached_ipv6,
+            }
+        if isinstance(config, (CloudInjectorConfig, ExternalInjectorConfig)):
+            return {}
+        raise TypeError(f"unsupported injector config type: {type(config).__name__}")
 
     def compile_post_execution_signatures(
         self,
diff --git a/test/signatures/constraints/signature_manager_pre_execution_constraints.feature b/test/signatures/constraints/signature_manager_pre_execution_constraints.feature
@@ -3,13 +3,13 @@ Feature: SignatureManager pre-execution constraints
     I want pre-execution compilation to handle timing edge cases correctly
     So that signatures always reflect the actual moment of execution
 
-    Background:
-        Given a SignatureManager initialised with constructor SignatureManager(client, logger)
+  Background:
+    Given a SignatureManager initialised with constructor SignatureManager(client, logger)
 
-    Scenario: start_time is generated at method-call time not at class instantiation time
-        Given a SignatureManager that was instantiated at timestamp T0
-        And 5 seconds elapse after instantiation
-        And an inject_config with a single network target
-        When I call compile_pre_execution_signatures with category="network" at timestamp T1
-        Then the start_time in the returned dict equals T1 within 1 second tolerance
-        And start_time does not equal T0
+  Scenario: start_time is generated at method-call time not at class instantiation time
+    Given a SignatureManager that was instantiated at timestamp T0
+    And 5 seconds elapse after instantiation
+    And a NetworkInjectorConfig with target_ipv4="192.168.1.10"
+    When I call compile_pre_execution_signatures with the config at timestamp T1
+    Then the start_time in the returned dict equals T1 within 1 second tolerance
+    And start_time does not equal T0
diff --git a/test/signatures/features/signature_manager_pre_execution.feature b/test/signatures/features/signature_manager_pre_execution.feature
diff --git a/test/signatures/test_signature_manager_pre_execution.py b/test/signatures/test_signature_manager_pre_execution.py
