prime_swarm_engine.py

"""
🔱 ZKAEDI PRIME — 12-Agent Vulnerability Hunter Engine
Self-contained module for HF Space deployment.
Vectorized numpy, zero external dependencies beyond numpy.
"""
from __future__ import annotations
import time
import numpy as np
from dataclasses import dataclass
from typing import Optional

N_AGENTS = 12
N_PARAMS = 10
P_ETA, P_GAMMA, P_BETA, P_SIGMA, P_KAPPA, P_DELTA = 0, 1, 2, 3, 4, 5
P_A, P_B, P_EPS_FHN, P_DIFFUSION = 6, 7, 8, 9
COUPLING_RADIUS_SQ = 2 * 20.0**2

ROLE_NAMES = [
    "RECON", "SENTINEL", "ARCHITECT", "ORACLE", "HEALER", "FORGE",
    "PHANTOM", "NEXUS", "CIPHER", "VANGUARD", "ECHO", "APEX"
]

ROLE_PRESET_MATRIX = np.array([
    [0.30,0.20,0.20,0.12,0.15,0.05,0.5,0.5,0.04,0.4],
    [0.50,0.50,0.05,0.02,0.40,0.05,0.5,0.5,0.04,0.4],
    [0.60,0.40,0.08,0.04,0.30,0.05,0.5,0.5,0.04,0.4],
    [0.55,0.45,0.10,0.03,0.35,0.05,0.5,0.5,0.04,0.4],
    [0.45,0.25,0.06,0.05,0.50,0.05,0.5,0.5,0.04,0.4],
    [0.40,0.35,0.15,0.08,0.20,0.05,0.5,0.5,0.04,0.4],
    [0.20,0.15,0.25,0.20,0.10,0.05,0.5,0.5,0.04,0.4],
    [0.48,0.38,0.07,0.04,0.60,0.05,0.5,0.5,0.04,0.4],
    [0.52,0.42,0.09,0.03,0.25,0.05,0.5,0.5,0.04,0.4],
    [0.58,0.48,0.10,0.05,0.35,0.05,0.5,0.5,0.04,0.4],
    [0.65,0.30,0.08,0.04,0.30,0.05,0.5,0.5,0.04,0.4],
    [0.50,0.40,0.12,0.06,0.45,0.05,0.5,0.5,0.04,0.4],
], dtype=np.float64)


@dataclass
class VulnSignature:
    id: str
    vuln_type: str       # e.g. "reentrancy", "integer_overflow"
    swc: str             # e.g. "SWC-107"
    severity: str        # "CRITICAL" | "HIGH" | "MEDIUM" | "LOW"
    position: list       # [x, y] in state space
    energy: float        # Well depth
    radius: float        # Influence radius
    description: str = ""
    compound_partner: Optional[str] = None
    compound_type: str = ""


class WellEntryEvent:
    __slots__ = ('agent_id', 'sig_idx', 'step', 'energy_drop', 'position')
    def __init__(self, agent_id, sig_idx, step, energy_drop, position):
        self.agent_id = agent_id
        self.sig_idx = sig_idx
        self.step = step
        self.energy_drop = energy_drop
        self.position = position.copy()


class TemporalCorrelationDetector:
    def __init__(self, signatures: list[VulnSignature], temporal_window: int = 100):
        self.signatures = signatures
        self.temporal_window = temporal_window

        self.compound_pairs: list[tuple[int, int, str]] = []
        seen = set()
        for i, sig_a in enumerate(signatures):
            if sig_a.compound_partner is None:
                continue
            for j, sig_b in enumerate(signatures):
                if sig_b.id == sig_a.compound_partner:
                    pair_key = tuple(sorted([i, j]))
                    if pair_key not in seen:
                        seen.add(pair_key)
                        self.compound_pairs.append((i, j, sig_a.compound_type))

        self.well_entries: list[WellEntryEvent] = []
        self.solo_detections: dict[int, WellEntryEvent] = {}
        self.compound_detections: list[dict] = []
        self.detected_compound_names: set[str] = set()

    def record_well_entry(self, agent_id, sig_idx, step, energy_drop, position):
        event = WellEntryEvent(agent_id, sig_idx, step, energy_drop, position)
        self.well_entries.append(event)
        if sig_idx not in self.solo_detections:
            self.solo_detections[sig_idx] = event

    def check_compounds(self, current_step):
        new_detections = []
        for idx_a, idx_b, compound_type in self.compound_pairs:
            if compound_type in self.detected_compound_names:
                continue
            events_a = [e for e in self.well_entries if e.sig_idx == idx_a]
            events_b = [e for e in self.well_entries if e.sig_idx == idx_b]
            if not events_a or not events_b:
                continue
            for ea in events_a:
                for eb in events_b:
                    dt = abs(ea.step - eb.step)
                    if dt <= self.temporal_window:
                        detection = {
                            "compound_type": compound_type,
                            "sig_a_id": self.signatures[idx_a].id,
                            "sig_b_id": self.signatures[idx_b].id,
                            "step_detected": current_step,
                            "event_a_step": ea.step,
                            "event_b_step": eb.step,
                            "temporal_gap": dt,
                            "agent_a": ea.agent_id,
                            "agent_b": eb.agent_id,
                            "agent_a_role": ROLE_NAMES[ea.agent_id],
                            "agent_b_role": ROLE_NAMES[eb.agent_id],
                            "same_agent": ea.agent_id == eb.agent_id,
                        }
                        new_detections.append(detection)
                        self.compound_detections.append(detection)
                        self.detected_compound_names.add(compound_type)
                        break
                else:
                    continue
                break
        return new_detections


def build_vulnerability_field(signatures, field_size=(100, 100)):
    sx, sy = field_size
    X, Y = np.meshgrid(np.arange(sx, dtype=np.float64),
                        np.arange(sy, dtype=np.float64), indexing='ij')
    H = np.zeros((sx, sy))
    rng = np.random.default_rng(12345)
    for _ in range(5):
        cx, cy = rng.uniform(0, sx), rng.uniform(0, sy)
        amp, width = rng.uniform(-0.5, 0.5), rng.uniform(15, 30)
        H += amp * np.exp(-((X - cx)**2 + (Y - cy)**2) / (2 * width**2))
    for sig in signatures:
        pos = np.array(sig.position)
        d_sq = (X - pos[0])**2 + (Y - pos[1])**2
        H -= sig.energy * np.exp(-d_sq / (2 * sig.radius**2))
    return H


class PrimeSwarmHunter:
    """
    12-agent PRIME vulnerability hunter with temporal compound detection.

    INPUT: List of VulnSignature dicts
    OUTPUT: Structured detection results (solo + compound)

    Designed for model tool-use: JSON in → JSON out.
    """

    def __init__(self, signatures: list[VulnSignature],
                 steps: int = 300, temporal_window: int = 150, seed: int = 42):
        self.rng = np.random.default_rng(seed)
        self.field_size = (100, 100)
        self.signatures = signatures
        self.total_steps = steps
        self.timestep = 0

        # Agent state
        self.positions = self.rng.uniform(5, 95, size=(N_AGENTS, 2))
        self.velocities = np.zeros((N_AGENTS, 2))
        self.H = np.zeros(N_AGENTS)
        self.V = np.zeros(N_AGENTS)
        self.H_prev = np.zeros(N_AGENTS)
        self.params = ROLE_PRESET_MATRIX.copy()

        # History
        self.hist_len = 200
        self.energy_hist = np.zeros((N_AGENTS, self.hist_len))
        self.hist_ptr = 0
        self.hist_count = 0

        # Field
        self.H_field = build_vulnerability_field(signatures, self.field_size)
        self._grad_x = np.zeros_like(self.H_field)
        self._grad_y = np.zeros_like(self.H_field)
        self._grad_x[1:-1, :] = (self.H_field[2:, :] - self.H_field[:-2, :]) / 2.0
        self._grad_y[:, 1:-1] = (self.H_field[:, 2:] - self.H_field[:, :-2]) / 2.0

        self.sig_positions = np.array([s.position for s in signatures])
        self.sig_radii = np.array([s.radius for s in signatures])

        # Detector
        self.detector = TemporalCorrelationDetector(signatures, temporal_window)
        self.entered_wells: set[tuple[int, int]] = set()

        # Trajectory recording for visualization
        self.trajectory_log: list[np.ndarray] = []

    def _step(self):
        self.timestep += 1
        dt = 0.1
        N = N_AGENTS
        p = self.params

        diff = self.positions[None, :, :] - self.positions[:, None, :]
        dist_sq = np.sum(diff**2, axis=-1)
        dist = np.sqrt(dist_sq) + 1e-6
        C_mat = np.exp(-dist_sq / COUPLING_RADIUS_SQ)
        np.fill_diagonal(C_mat, 0.0)

        ix = np.clip(self.positions[:, 0].astype(int), 0, 99)
        iy = np.clip(self.positions[:, 1].astype(int), 0, 99)
        H_base = self.H_field[ix, iy]
        H_prev = self.H.copy()

        eta = p[:, P_ETA]; gamma = p[:, P_GAMMA]; beta = p[:, P_BETA]
        sigma = p[:, P_SIGMA]; kappa = p[:, P_KAPPA]; delta = p[:, P_DELTA]
        eps_fhn = p[:, P_EPS_FHN]; a_fhn = p[:, P_A]; b_fhn = p[:, P_B]

        sig_val = 1.0 / (1.0 + np.exp(np.clip(-gamma * H_prev, -500, 500)))
        noise = self.rng.normal(0, 1, size=N) * (1.0 + beta * np.abs(H_prev))
        coupling_sum = C_mat @ self.H
        H_grad = H_prev - self.H_prev

        H_new = (H_base + eta * H_prev * sig_val + sigma * noise
                 + delta * H_grad + kappa * coupling_sum / max(1, N-1))
        V_new = self.V + dt * eps_fhn * (H_prev + a_fhn - b_fhn * self.V)
        H_new += dt * (-np.clip(H_prev, -10, 10)**3 / 3.0 - self.V)

        self.H_prev = H_prev
        self.H = np.clip(H_new, -50, 50)
        self.V = np.clip(V_new, -50, 50)

        # Movement
        grad = np.column_stack([self._grad_x[ix, iy], self._grad_y[ix, iy]])
        unit_diff = diff / dist[:, :, None]
        coupling_forces = np.einsum('ij,ijd->id', C_mat, unit_diff) * kappa[:, None]
        exploration = self.rng.normal(0, 1, size=(N, 2)) * sigma[:, None]

        role_bias = np.zeros((N, 2))
        role_bias[0] = self.rng.normal(0, sigma[0] * 5, size=2)
        role_bias[9] = -grad[9] * 1.0
        if self.rng.random() < 0.2:
            role_bias[6] = self.rng.normal(0, 12.0, size=2)

        self.velocities = (0.7 * self.velocities - 0.8 * grad
                           + 0.3 * coupling_forces + 0.2 * role_bias + exploration)
        speeds = np.linalg.norm(self.velocities, axis=1)
        too_fast = speeds > 3.0
        self.velocities[too_fast] *= (3.0 / speeds[too_fast])[:, None]
        self.positions += self.velocities
        np.clip(self.positions, 1, 98, out=self.positions)

        ptr = self.hist_ptr % self.hist_len
        self.energy_hist[:, ptr] = self.H
        self.hist_ptr += 1
        self.hist_count = min(self.hist_count + 1, self.hist_len)

        # Record trajectory every 5 steps
        if self.timestep % 5 == 0:
            self.trajectory_log.append(self.positions.copy())

        # Well-entry detection
        agent_sig_diff = self.positions[:, None, :] - self.sig_positions[None, :, :]
        agent_sig_dist = np.sqrt(np.sum(agent_sig_diff**2, axis=-1))
        proximity_mask = agent_sig_dist < (self.sig_radii[None, :] * 1.5)

        for agent_id in range(N):
            for sig_idx in range(len(self.signatures)):
                if proximity_mask[agent_id, sig_idx]:
                    key = (agent_id, sig_idx)
                    if key not in self.entered_wells:
                        self.entered_wells.add(key)
                        self.detector.record_well_entry(
                            agent_id, sig_idx, self.timestep,
                            abs(float(self.H[agent_id])),
                            self.positions[agent_id]
                        )

        self.detector.check_compounds(self.timestep)

    def run(self) -> dict:
        """Execute full analysis. Returns structured results dict."""
        t0 = time.time()
        for _ in range(self.total_steps):
            self._step()
        elapsed = time.time() - t0

        # Build results
        solo_findings = []
        for sig_idx, event in self.detector.solo_detections.items():
            sig = self.signatures[sig_idx]
            solo_findings.append({
                "id": sig.id,
                "vuln_type": sig.vuln_type,
                "swc": sig.swc,
                "severity": sig.severity,
                "description": sig.description,
                "detected_at_step": event.step,
                "detected_by_agent": event.agent_id,
                "detected_by_role": ROLE_NAMES[event.agent_id],
                "energy_drop": round(event.energy_drop, 4),
                "position": sig.position,
            })

        compound_findings = []
        for d in self.detector.compound_detections:
            compound_findings.append({
                "compound_type": d["compound_type"],
                "components": [d["sig_a_id"], d["sig_b_id"]],
                "detected_at_step": d["step_detected"],
                "temporal_gap": d["temporal_gap"],
                "agents_involved": {
                    "agent_a": {"id": d["agent_a"], "role": d["agent_a_role"]},
                    "agent_b": {"id": d["agent_b"], "role": d["agent_b_role"]},
                },
                "same_agent_discovery": d["same_agent"],
                "severity": "CRITICAL",
            })

        n_sigs = len(self.signatures)
        n_compound_patterns = len(self.detector.compound_pairs)

        return {
            "engine": "ZKAEDI PRIME 12-Agent Hamiltonian Swarm v3",
            "config": {
                "agents": N_AGENTS,
                "steps": self.total_steps,
                "field_size": list(self.field_size),
                "temporal_window": self.detector.temporal_window,
            },
            "execution": {
                "elapsed_seconds": round(elapsed, 3),
                "steps_per_second": round(self.total_steps / elapsed, 0),
                "total_well_entries": len(self.detector.well_entries),
            },
            "summary": {
                "total_signatures": n_sigs,
                "solo_detected": len(solo_findings),
                "solo_detection_rate": round(len(solo_findings) / max(n_sigs, 1) * 100, 1),
                "compound_patterns": n_compound_patterns,
                "compounds_detected": len(compound_findings),
                "compound_detection_rate": round(
                    len(compound_findings) / max(n_compound_patterns, 1) * 100, 1
                ),
                "risk_score": sum(
                    {"CRITICAL": 10, "HIGH": 5, "MEDIUM": 2, "LOW": 1}.get(f["severity"], 0)
                    for f in solo_findings
                ) + len(compound_findings) * 15,
            },
            "solo_findings": solo_findings,
            "compound_findings": compound_findings,
            "agent_summary": [
                {
                    "id": i,
                    "role": ROLE_NAMES[i],
                    "final_position": self.positions[i].round(2).tolist(),
                    "final_energy": round(float(self.H[i]), 4),
                    "wells_entered": sum(1 for k in self.entered_wells if k[0] == i),
                }
                for i in range(N_AGENTS)
            ],
        }


def parse_signatures(sig_dicts: list[dict]) -> list[VulnSignature]:
    """Parse JSON signature dicts into VulnSignature objects."""
    sigs = []
    for d in sig_dicts:
        sigs.append(VulnSignature(
            id=d["id"],
            vuln_type=d.get("vuln_type", "unknown"),
            swc=d.get("swc", ""),
            severity=d.get("severity", "MEDIUM"),
            position=d.get("position", [50, 50]),
            energy=d.get("energy", 5.0),
            radius=d.get("radius", 10.0),
            description=d.get("description", ""),
            compound_partner=d.get("compound_partner"),
            compound_type=d.get("compound_type", ""),
        ))
    return sigs


# ── Preset scenarios for quick testing ────────────────────────
PRESET_SCENARIOS = {
    "defi_lending_pool": {
        "name": "DeFi Lending Pool (Compound-like)",
        "signatures": [
            {"id": "v1", "vuln_type": "reentrancy", "swc": "SWC-107", "severity": "CRITICAL",
             "position": [20, 80], "energy": 9.5, "radius": 10,
             "description": "withdraw() external call before balance update",
             "compound_partner": "v2", "compound_type": "reentrancy_chain"},
            {"id": "v2", "vuln_type": "unchecked_call", "swc": "SWC-104", "severity": "HIGH",
             "position": [55, 40], "energy": 6.0, "radius": 9,
             "description": "_doWithdraw() unchecked low-level call",
             "compound_partner": "v1", "compound_type": "reentrancy_chain"},
            {"id": "v3", "vuln_type": "flash_loan", "swc": "DEFI-01", "severity": "HIGH",
             "position": [15, 45], "energy": 7.5, "radius": 10,
             "description": "flashLoan() no single-block borrow limit",
             "compound_partner": "v4", "compound_type": "price_manipulation"},
            {"id": "v4", "vuln_type": "oracle_manipulation", "swc": "DEFI-02", "severity": "CRITICAL",
             "position": [50, 75], "energy": 8.5, "radius": 10,
             "description": "getPrice() spot oracle without TWAP",
             "compound_partner": "v3", "compound_type": "price_manipulation"},
            {"id": "v5", "vuln_type": "access_control", "swc": "SWC-115", "severity": "MEDIUM",
             "position": [80, 20], "energy": 4.0, "radius": 8,
             "description": "setInterestRate() missing onlyOwner"},
            {"id": "v6", "vuln_type": "integer_overflow", "swc": "SWC-101", "severity": "LOW",
             "position": [90, 90], "energy": 2.5, "radius": 6,
             "description": "Interest accrual overflow on extreme durations"},
        ]
    },
    "nft_marketplace": {
        "name": "NFT Marketplace (OpenSea-like)",
        "signatures": [
            {"id": "v1", "vuln_type": "reentrancy", "swc": "SWC-107", "severity": "HIGH",
             "position": [40, 60], "energy": 7.0, "radius": 9,
             "description": "fulfillOrder() ETH send before listing clear"},
            {"id": "v2", "vuln_type": "front_running", "swc": "SWC-114", "severity": "MEDIUM",
             "position": [75, 50], "energy": 5.0, "radius": 10,
             "description": "cancelListing/fulfillOrder race condition"},
            {"id": "v3", "vuln_type": "delegatecall", "swc": "SWC-112", "severity": "CRITICAL",
             "position": [15, 85], "energy": 9.0, "radius": 10,
             "description": "upgradeProxy() unprotected delegatecall",
             "compound_partner": "v4", "compound_type": "proxy_takeover"},
            {"id": "v4", "vuln_type": "storage_collision", "swc": "SWC-124", "severity": "HIGH",
             "position": [55, 30], "energy": 7.5, "radius": 9,
             "description": "_implementation slot overlaps owner storage",
             "compound_partner": "v3", "compound_type": "proxy_takeover"},
            {"id": "v5", "vuln_type": "access_control", "swc": "SWC-115", "severity": "MEDIUM",
             "position": [85, 15], "energy": 4.5, "radius": 7,
             "description": "setRoyaltyReceiver() callable by any seller"},
        ]
    },
    "token_bridge": {
        "name": "Cross-Chain Token Bridge (Wormhole-like)",
        "signatures": [
            {"id": "v1", "vuln_type": "access_control", "swc": "SWC-115", "severity": "CRITICAL",
             "position": [25, 85], "energy": 9.0, "radius": 10,
             "description": "validateSignatures() accepts 2/5 multisig",
             "compound_partner": "v2", "compound_type": "governance_takeover"},
            {"id": "v2", "vuln_type": "selfdestruct", "swc": "SWC-106", "severity": "CRITICAL",
             "position": [70, 35], "energy": 8.5, "radius": 9,
             "description": "emergencyShutdown() callable after gov takeover",
             "compound_partner": "v1", "compound_type": "governance_takeover"},
            {"id": "v3", "vuln_type": "unchecked_call", "swc": "SWC-104", "severity": "HIGH",
             "position": [80, 75], "energy": 7.0, "radius": 10,
             "description": "processMessage() doesn't verify execution",
             "compound_partner": "v4", "compound_type": "double_spend"},
            {"id": "v4", "vuln_type": "reentrancy", "swc": "SWC-107", "severity": "HIGH",
             "position": [30, 40], "energy": 7.5, "radius": 9,
             "description": "claimTokens() external call before burn",
             "compound_partner": "v3", "compound_type": "double_spend"},
            {"id": "v5", "vuln_type": "oracle_manipulation", "swc": "DEFI-02", "severity": "MEDIUM",
             "position": [50, 15], "energy": 5.0, "radius": 8,
             "description": "Cross-chain price feed 30min staleness"},
            {"id": "v6", "vuln_type": "integer_overflow", "swc": "SWC-101", "severity": "LOW",
             "position": [10, 10], "energy": 2.0, "radius": 5,
             "description": "Fee truncation on sub-wei amounts"},
        ]
    },
}


def run_preset(scenario_key: str, steps: int = 300,
               temporal_window: int = 150, seed: int = 42) -> dict:
    """Run a preset scenario. Returns full results dict."""
    if scenario_key not in PRESET_SCENARIOS:
        return {"error": f"Unknown scenario. Available: {list(PRESET_SCENARIOS.keys())}"}
    scenario = PRESET_SCENARIOS[scenario_key]
    sigs = parse_signatures(scenario["signatures"])
    hunter = PrimeSwarmHunter(sigs, steps=steps, temporal_window=temporal_window, seed=seed)
    results = hunter.run()
    results["scenario"] = scenario["name"]
    return results


def run_custom(signatures_json: list[dict], steps: int = 300,
               temporal_window: int = 150, seed: int = 42) -> dict:
    """Run with custom vulnerability signatures. Returns full results dict."""
    sigs = parse_signatures(signatures_json)
    hunter = PrimeSwarmHunter(sigs, steps=steps, temporal_window=temporal_window, seed=seed)
    results = hunter.run()
    results["scenario"] = "custom"
    return results