""" 🔱 ZKAEDI PRIME — 12-Agent Vulnerability Hunter Hugging Face Space: Interactive UI + API for model tool-use. API Usage (from your models): from gradio_client import Client client = Client("zkaedi/prime-swarm-hunter") result = client.predict( scenario="defi_lending_pool", steps=300, api_name="/run_preset" ) """ import gradio as gr import json import numpy as np import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from prime_swarm_engine import ( run_preset, run_custom, parse_signatures, PrimeSwarmHunter, PRESET_SCENARIOS, ROLE_NAMES, N_AGENTS, build_vulnerability_field, ) # ══════════════════════════════════════════════════════════════ # VISUALIZATION # ══════════════════════════════════════════════════════════════ ROLE_COLORS = [ '#00FFFF', '#FFD700', '#4169E1', '#FF00FF', '#00FF80', '#FF8C00', '#808080', '#00FFFF', '#FF69B4', '#FF4444', '#00CED1', '#FFD700', ] def render_field_plot(results: dict, signatures: list[dict]) -> plt.Figure: """Render energy landscape with agent trajectories and vuln wells.""" sigs = parse_signatures(signatures) H_field = build_vulnerability_field(sigs) fig, ax = plt.subplots(1, 1, figsize=(8, 8), facecolor='#0D0D0D') ax.set_facecolor('#0D0D0D') # Energy field heatmap ax.imshow(H_field.T, extent=[0, 100, 0, 100], origin='lower', cmap='inferno', alpha=0.6, aspect='auto') # Vulnerability wells for sig in sigs: color = '#FF0000' if sig.severity == 'CRITICAL' else ( '#FF8C00' if sig.severity == 'HIGH' else '#FFD700' if sig.severity == 'MEDIUM' else '#808080') circle = plt.Circle(sig.position, sig.radius * 1.5, fill=False, color=color, linewidth=1.5, linestyle='--', alpha=0.7) ax.add_patch(circle) ax.plot(sig.position[0], sig.position[1], 'x', color=color, markersize=10, markeredgewidth=2) ax.annotate(f'{sig.id}\n{sig.vuln_type}', xy=sig.position, fontsize=6, color='white', ha='center', va='bottom', fontweight='bold') # Agent final positions for agent in results.get("agent_summary", []): pos = agent["final_position"] color = ROLE_COLORS[agent["id"] % len(ROLE_COLORS)] ax.plot(pos[0], pos[1], 'o', color=color, markersize=8, alpha=0.9, markeredgecolor='white', markeredgewidth=0.5) ax.annotate(f'{agent["id"]}', xy=(pos[0], pos[1]), fontsize=5, color='white', ha='center', va='center', fontweight='bold') # Compound detection lines for comp in results.get("compound_findings", []): sig_a = next((s for s in sigs if s.id == comp["components"][0]), None) sig_b = next((s for s in sigs if s.id == comp["components"][1]), None) if sig_a and sig_b: ax.plot([sig_a.position[0], sig_b.position[0]], [sig_a.position[1], sig_b.position[1]], color='#FF00FF', linewidth=2, alpha=0.8, linestyle='-', marker='D', markersize=6) mid = [(sig_a.position[0] + sig_b.position[0]) / 2, (sig_a.position[1] + sig_b.position[1]) / 2] ax.annotate(f'⚠ {comp["compound_type"]}', xy=mid, fontsize=7, color='#FF00FF', ha='center', fontweight='bold', bbox=dict(boxstyle='round,pad=0.3', facecolor='#1a0a1a', alpha=0.8)) ax.set_xlim(0, 100); ax.set_ylim(0, 100) ax.set_xlabel('Call Depth / Control Flow Index', color='#888', fontsize=8) ax.set_ylabel('State Mutation Intensity', color='#888', fontsize=8) ax.set_title('🔱 PRIME Swarm Vulnerability Landscape', color='#00FFFF', fontsize=11, fontweight='bold') ax.tick_params(colors='#555', labelsize=7) # Legend legend_elements = [ plt.Line2D([0], [0], marker='x', color='#FF0000', label='CRITICAL', markersize=8, linestyle='None'), plt.Line2D([0], [0], marker='x', color='#FF8C00', label='HIGH', markersize=8, linestyle='None'), plt.Line2D([0], [0], marker='x', color='#FFD700', label='MEDIUM', markersize=8, linestyle='None'), plt.Line2D([0], [0], color='#FF00FF', label='Compound Link', linewidth=2), plt.Line2D([0], [0], marker='o', color='#00FFFF', label='Agent', markersize=6, linestyle='None'), ] ax.legend(handles=legend_elements, loc='upper right', fontsize=6, facecolor='#1a1a1a', edgecolor='#333', labelcolor='white') plt.tight_layout() return fig def format_results_markdown(results: dict) -> str: """Format results as readable markdown for the UI.""" if "error" in results: return f"❌ **Error:** {results['error']}" s = results["summary"] lines = [ f"## 🔱 ZKAEDI PRIME Analysis Complete", f"**Scenario:** {results.get('scenario', 'custom')}", f"**Engine:** {results['engine']}", f"", f"### Execution", f"- ⏱️ {results['execution']['elapsed_seconds']}s " f"({results['execution']['steps_per_second']:.0f} steps/s)", f"- 🔍 {results['execution']['total_well_entries']} well-entry events", f"", f"### Detection Summary", f"- **Solo vulnerabilities:** {s['solo_detected']}/{s['total_signatures']} " f"({s['solo_detection_rate']}%)", f"- **Compound patterns:** {s['compounds_detected']}/{s['compound_patterns']} " f"({s['compound_detection_rate']}%)", f"- **Risk score:** {s['risk_score']}", f"", f"### Solo Findings", ] severity_emoji = {"CRITICAL": "🔴", "HIGH": "🟠", "MEDIUM": "🟡", "LOW": "⚪"} for f in results["solo_findings"]: emoji = severity_emoji.get(f["severity"], "⚪") lines.append( f"- {emoji} **{f['severity']}** `{f['swc']}` {f['vuln_type']} — " f"{f['description']} *(step {f['detected_at_step']}, " f"agent {f['detected_by_role']})*" ) if results["compound_findings"]: lines.append(f"\n### ⚠️ Compound Vulnerabilities Detected") for c in results["compound_findings"]: lines.append( f"- **{c['compound_type']}** ({' + '.join(c['components'])}) — " f"temporal gap={c['temporal_gap']} steps, " f"agents: {c['agents_involved']['agent_a']['role']} → " f"{c['agents_involved']['agent_b']['role']}" ) else: lines.append(f"\n### Compound Vulnerabilities") lines.append(f"*No compound patterns detected in {results['config']['steps']} steps.*") lines.append(f"\n### Agent Activity") for a in results["agent_summary"]: if a["wells_entered"] > 0: lines.append(f"- Agent {a['id']} ({a['role']}): {a['wells_entered']} wells entered") return "\n".join(lines) # ══════════════════════════════════════════════════════════════ # GRADIO INTERFACE # ══════════════════════════════════════════════════════════════ def analyze_preset(scenario: str, steps: int, temporal_window: int, seed: int): """Run a preset scenario and return results + visualization.""" results = run_preset(scenario, steps=steps, temporal_window=temporal_window, seed=seed) sigs = PRESET_SCENARIOS[scenario]["signatures"] fig = render_field_plot(results, sigs) md = format_results_markdown(results) return md, fig, json.dumps(results, indent=2, default=str) def analyze_custom(signatures_json: str, steps: int, temporal_window: int, seed: int): """Run with custom vulnerability signatures.""" try: sigs = json.loads(signatures_json) except json.JSONDecodeError as e: return f"❌ Invalid JSON: {e}", None, json.dumps({"error": str(e)}) results = run_custom(sigs, steps=steps, temporal_window=temporal_window, seed=seed) fig = render_field_plot(results, sigs) md = format_results_markdown(results) return md, fig, json.dumps(results, indent=2, default=str) # API-only endpoints (for model tool-use via gradio_client) def api_run_preset(scenario: str, steps: int = 300, temporal_window: int = 150, seed: int = 42) -> str: """API endpoint: Run preset scenario, return JSON results.""" results = run_preset(scenario, steps=steps, temporal_window=temporal_window, seed=seed) return json.dumps(results, default=str) def api_run_custom(signatures_json: str, steps: int = 300, temporal_window: int = 150, seed: int = 42) -> str: """API endpoint: Run custom signatures, return JSON results.""" try: sigs = json.loads(signatures_json) except json.JSONDecodeError as e: return json.dumps({"error": str(e)}) results = run_custom(sigs, steps=steps, temporal_window=temporal_window, seed=seed) return json.dumps(results, default=str) EXAMPLE_CUSTOM = json.dumps(PRESET_SCENARIOS["defi_lending_pool"]["signatures"], indent=2) CSS = """ .gradio-container { max-width: 1200px !important; } .dark { background-color: #0D0D0D; } h1 { color: #00FFFF !important; } """ with gr.Blocks(theme=gr.themes.Base(primary_hue="cyan", neutral_hue="gray"), css=CSS, title="🔱 PRIME Swarm Vulnerability Hunter") as demo: gr.Markdown(""" # 🔱 ZKAEDI PRIME — 12-Agent Vulnerability Hunter **Recursively Coupled Hamiltonian Swarm** with temporal correlation compound detection. 12 specialized agents navigate an energy landscape built from vulnerability signatures. **For model tool-use:** Use the API endpoints via `gradio_client`: ```python from gradio_client import Client client = Client("zkaedi/prime-swarm-hunter") result = client.predict("defi_lending_pool", 300, 150, 42, api_name="/api_preset") ``` """) with gr.Tabs(): # ── Tab 1: Preset Scenarios ─────────────────────── with gr.TabItem("🎯 Preset Scenarios"): with gr.Row(): with gr.Column(scale=1): scenario_dropdown = gr.Dropdown( choices=list(PRESET_SCENARIOS.keys()), value="defi_lending_pool", label="Contract Scenario" ) steps_slider = gr.Slider(50, 1000, value=300, step=50, label="Steps") window_slider = gr.Slider(20, 300, value=150, step=10, label="Temporal Window") seed_input = gr.Number(value=42, label="Random Seed", precision=0) run_btn = gr.Button("🔱 Run Swarm Analysis", variant="primary", size="lg") with gr.Column(scale=2): results_md = gr.Markdown(label="Analysis Results") results_plot = gr.Plot(label="Energy Landscape") results_json = gr.Code(label="Raw JSON (for model consumption)", language="json") run_btn.click( analyze_preset, inputs=[scenario_dropdown, steps_slider, window_slider, seed_input], outputs=[results_md, results_plot, results_json] ) # ── Tab 2: Custom Signatures ────────────────────── with gr.TabItem("🔧 Custom Signatures"): gr.Markdown("Provide vulnerability signatures as JSON. Your models can generate these.") with gr.Row(): with gr.Column(scale=1): custom_json = gr.Code( value=EXAMPLE_CUSTOM, label="Vulnerability Signatures (JSON array)", language="json", lines=20, ) custom_steps = gr.Slider(50, 1000, value=300, step=50, label="Steps") custom_window = gr.Slider(20, 300, value=150, step=10, label="Temporal Window") custom_seed = gr.Number(value=42, label="Seed", precision=0) custom_btn = gr.Button("🔱 Run Custom Analysis", variant="primary") with gr.Column(scale=2): custom_md = gr.Markdown() custom_plot = gr.Plot() custom_json_out = gr.Code(label="Raw JSON", language="json") custom_btn.click( analyze_custom, inputs=[custom_json, custom_steps, custom_window, custom_seed], outputs=[custom_md, custom_plot, custom_json_out] ) # ── Tab 3: API Docs ─────────────────────────────── with gr.TabItem("📡 API / Tool-Use"): gr.Markdown(""" ## Model Tool-Use Integration This Space exposes two API endpoints your HF models can call: ### `/api_preset` — Run a preset scenario ```python from gradio_client import Client client = Client("zkaedi/prime-swarm-hunter") result_json = client.predict( scenario="defi_lending_pool", # or "nft_marketplace", "token_bridge" steps=300, temporal_window=150, seed=42, api_name="/api_preset" ) findings = json.loads(result_json) ``` ### `/api_custom` — Run with custom vulnerability signatures ```python import json from gradio_client import Client signatures = [ { "id": "v1", "vuln_type": "reentrancy", "swc": "SWC-107", "severity": "CRITICAL", "position": [20, 80], "energy": 9.5, "radius": 10, "description": "withdraw() external call before state 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" } ] client = Client("zkaedi/prime-swarm-hunter") result_json = client.predict( json.dumps(signatures), 300, # steps 150, # temporal_window 42, # seed api_name="/api_custom" ) findings = json.loads(result_json) print(f"Solo: {findings['summary']['solo_detected']}") print(f"Compound: {findings['summary']['compounds_detected']}") ``` ### Integration with gemma-7b-solidity-energy-signatures Your model generates vulnerability signatures → this tool runs the 12-agent swarm analysis → returns compound detection results your model can reason about. ```python # In your model inference pipeline: # 1. Model analyzes Solidity code → outputs vulnerability signatures # 2. Signatures fed to PRIME swarm → compound detection # 3. Results fed back to model → final audit report from gradio_client import Client swarm = Client("zkaedi/prime-swarm-hunter") # Model output: vulnerability signatures model_signatures = model.generate_signatures(solidity_code) # Swarm analysis swarm_results = json.loads( swarm.predict(json.dumps(model_signatures), 300, 150, 42, api_name="/api_custom") ) # Feed back to model for final report final_report = model.generate_report(solidity_code, swarm_results) ``` """) # ── Hidden API endpoints (for gradio_client) ────────── api_preset_iface = gr.Interface( fn=api_run_preset, inputs=[ gr.Textbox(label="scenario"), gr.Number(label="steps", value=300), gr.Number(label="temporal_window", value=150), gr.Number(label="seed", value=42), ], outputs=gr.Textbox(label="results_json"), api_name="api_preset", ) api_custom_iface = gr.Interface( fn=api_run_custom, inputs=[ gr.Textbox(label="signatures_json"), gr.Number(label="steps", value=300), gr.Number(label="temporal_window", value=150), gr.Number(label="seed", value=42), ], outputs=gr.Textbox(label="results_json"), api_name="api_custom", ) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860)