""" 🔱 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