app.py

"""
╔══════════════════════════════════════════════════════════════════════════╗
║  ZKAEDI PRIME — Smart Contract Energy Signatures Auditor                 ║
║  Recursively Coupled Hamiltonian Vulnerability Detection                 ║
║  Model: zkaedi/gemma-7b-solidity-energy-signatures (GGUF Q4_K_M)        ║
╚══════════════════════════════════════════════════════════════════════════╝
"""

import os
import time
import threading
import gradio as gr
from huggingface_hub import hf_hub_download, list_repo_files

# ─────────────────────────────────────────────────────────────────────────────
# §1  MODEL LOADING
# ─────────────────────────────────────────────────────────────────────────────
MODEL_REPO   = "zkaedi/gemma-7b-solidity-energy-signatures"
_llm         = None
_load_status = {"ready": False, "error": None, "msg": "Initializing..."}
GGUF_PREF    = ["Q4_K_M", "Q4_K_S", "Q5_K_M", "Q4_0", "Q8_0", "f16", ""]

def _detect_gguf(repo_id):
    token = os.environ.get("HF_TOKEN")
    files = list(list_repo_files(repo_id, token=token))
    gguf  = [f for f in files if f.endswith(".gguf")]
    if not gguf:
        raise FileNotFoundError(f"No GGUF in {repo_id}")
    for pref in GGUF_PREF:
        for f in gguf:
            if pref.lower() in f.lower():
                return f
    return gguf[0]

def _load_model():
    global _llm
    try:
        _load_status["msg"] = "🔍 Detecting GGUF file..."
        filename = _detect_gguf(MODEL_REPO)

        _load_status["msg"] = f"📥 Downloading {filename} (~5.3 GB)..."
        token = os.environ.get("HF_TOKEN")
        model_path = hf_hub_download(
            repo_id=MODEL_REPO,
            filename=filename,
            cache_dir="/tmp/models",
            token=token,
        )

        _load_status["msg"] = "⚡ Loading into memory..."
        from llama_cpp import Llama

        n_gpu = 0
        try:
            import torch
            if torch.cuda.is_available():
                n_gpu = -1
        except ImportError:
            pass

        _llm = Llama(
            model_path=model_path,
            n_ctx=4096,
            n_threads=os.cpu_count() or 4,
            n_gpu_layers=n_gpu,
            verbose=False,
        )
        mode = "GPU" if n_gpu else "CPU"
        _load_status["ready"] = True
        _load_status["msg"]   = f"✅ Armed [{mode}] — {filename}"

    except Exception as e:
        import traceback; traceback.print_exc()
        _load_status["error"] = str(e)
        _load_status["msg"]   = f"❌ {e}"

threading.Thread(target=_load_model, daemon=True).start()


# ─────────────────────────────────────────────────────────────────────────────
# §2  PROMPTING
# ─────────────────────────────────────────────────────────────────────────────
SYSTEM = (
    "You are ZKAEDI PRIME, an expert Solidity smart contract security auditor "
    "trained on energy signature patterns.\n"
    "Analyze the provided smart contract using Hamiltonian energy field methodology:\n"
    "- Map vulnerability topology (reentrancy wells, overflow cliffs, access control gaps)\n"
    "- Classify energy signatures: CRITICAL ⚡, HIGH 🔴, MEDIUM 🟡, LOW 🟢, INFO ℹ️\n"
    "- Identify attack vectors and exploit paths\n"
    "- Provide remediation with corrected code snippets\n"
    "- Calculate overall Hamiltonian risk score (0-100)\n\n"
    "Format your response as:\n"
    "## ⚡ ENERGY SIGNATURE SCAN\n"
    "[vulnerability list with severity]\n\n"
    "## 🔬 DEEP FIELD ANALYSIS\n"
    "[detailed explanation per vulnerability]\n\n"
    "## 🛡️ REMEDIATION\n"
    "[fixes with code]\n\n"
    "## 📊 HAMILTONIAN RISK SCORE\n"
    "[score]/100 — [risk level]\n"
)

def make_prompt(code):
    return (
        "<start_of_turn>system\n" + SYSTEM + "<end_of_turn>\n"
        "<start_of_turn>user\nAudit this Solidity contract:\n\n"
        "```solidity\n" + code + "\n```<end_of_turn>\n"
        "<start_of_turn>model\n"
    )


# ─────────────────────────────────────────────────────────────────────────────
# §3  INFERENCE
# ─────────────────────────────────────────────────────────────────────────────
def audit(contract_code, temperature, max_tokens):
    if not contract_code or not contract_code.strip():
        yield "⚠️ Please paste a Solidity contract.", ""
        return

    waited = 0
    while not _load_status["ready"] and not _load_status["error"] and waited < 600:
        yield f"⏳ {_load_status['msg']}", ""
        time.sleep(3); waited += 3

    if _load_status["error"]:
        yield f"❌ Load failed:\n\n{_load_status['error']}", ""
        return
    if not _load_status["ready"]:
        yield "⏰ Timeout — refresh and retry.", ""
        return

    prompt = make_prompt(contract_code.strip())
    start  = time.time()
    out    = ""

    try:
        stream = _llm(
            prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=0.95,
            repeat_penalty=1.1,
            stream=True,
            stop=["<end_of_turn>", "<start_of_turn>"],
        )
        for chunk in stream:
            out += chunk["choices"][0]["text"]
            t    = time.time() - start
            spd  = len(out.split()) / max(t, 0.1)
            yield out, f"🔥 {len(out)} chars | {spd:.1f} tok/s | {t:.1f}s"

        yield out, f"✅ Done in {time.time()-start:.1f}s"

    except Exception as e:
        yield f"❌ Inference error: {e}", ""


# ─────────────────────────────────────────────────────────────────────────────
# §4  EXAMPLES
# ─────────────────────────────────────────────────────────────────────────────
REENTRANCY = """// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract VulnerableBank {
    mapping(address => uint256) public balances;

    function deposit() external payable {
        balances[msg.sender] += msg.value;
    }

    function withdraw(uint256 amount) external {
        require(balances[msg.sender] >= amount, "Insufficient");
        (bool ok,) = msg.sender.call{value: amount}("");
        require(ok);
        balances[msg.sender] -= amount; // state AFTER external call
    }
}"""

OVERFLOW = """// SPDX-License-Identifier: MIT
pragma solidity ^0.6.0;

contract TokenSale {
    mapping(address => uint256) public tokens;
    uint256 public constant PRICE = 0.01 ether;

    function buyTokens(uint256 n) external payable {
        require(msg.value == n * PRICE); // overflow possible
        tokens[msg.sender] += n;
    }
}"""

ACCESS = """// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Treasury {
    address public owner;

    constructor() { owner = msg.sender; }

    function initOwner(address newOwner) external {
        // Missing: require(owner == address(0))
        owner = newOwner;
    }

    function drain(address to) external {
        require(msg.sender == owner);
        payable(to).transfer(address(this).balance);
    }
}"""


# ─────────────────────────────────────────────────────────────────────────────
# §5  UI
# ─────────────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Share+Tech+Mono&family=Orbitron:wght@400;700;900&family=Exo+2:wght@300;400;600&display=swap');
:root {
  --c: #00f5ff; --m: #ff00cc; --o: #ff6b00;
  --bg0: #050810; --bg1: #0a0f1e; --bg2: #0d1428;
  --bd: rgba(0,245,255,.2); --tx: #c8d8f0;
}
* { box-sizing: border-box; }
body,.gradio-container { background:var(--bg0)!important; font-family:'Exo 2',sans-serif!important; color:var(--tx)!important; }
.gradio-container::before {
  content:''; position:fixed; inset:0;
  background:repeating-linear-gradient(0deg,transparent,transparent 2px,rgba(0,245,255,.015) 2px,rgba(0,245,255,.015) 4px);
  pointer-events:none; z-index:9999;
}
#hdr { text-align:center; padding:2rem 0 1rem; }
#hdr h1 {
  font-family:'Orbitron',monospace!important; font-size:2.2rem!important; font-weight:900!important;
  background:linear-gradient(135deg,var(--c),var(--m),var(--o));
  -webkit-background-clip:text; -webkit-text-fill-color:transparent; background-clip:text;
  letter-spacing:4px; margin:0; animation:pg 3s ease-in-out infinite;
}
@keyframes pg { 0%,100%{filter:brightness(1)} 50%{filter:brightness(1.3) drop-shadow(0 0 15px rgba(0,245,255,.6))} }
.lbl { font-family:'Orbitron',monospace!important; font-size:.7rem!important; letter-spacing:3px!important; color:var(--c)!important; text-transform:uppercase!important; margin-bottom:8px!important; }
textarea { background:var(--bg2)!important; border:1px solid var(--bd)!important; color:var(--c)!important; font-family:'Share Tech Mono',monospace!important; font-size:.85rem!important; border-radius:4px!important; }
textarea:focus { border-color:var(--c)!important; box-shadow:0 0 20px rgba(0,245,255,.4)!important; outline:none!important; }
.out textarea { color:var(--tx)!important; background:var(--bg1)!important; border-color:rgba(255,0,204,.3)!important; }
button.primary { background:linear-gradient(135deg,rgba(0,245,255,.15),rgba(255,0,204,.15))!important; border:1px solid var(--c)!important; color:var(--c)!important; font-family:'Orbitron',monospace!important; font-size:.75rem!important; letter-spacing:3px!important; text-transform:uppercase!important; border-radius:3px!important; }
button.primary:hover { background:rgba(0,245,255,.25)!important; box-shadow:0 0 20px rgba(0,245,255,.4)!important; transform:translateY(-1px)!important; }
button.secondary { background:transparent!important; border:1px solid rgba(255,0,204,.4)!important; color:var(--m)!important; font-family:'Share Tech Mono',monospace!important; font-size:.7rem!important; letter-spacing:2px!important; }
button.secondary:hover { border-color:var(--m)!important; box-shadow:0 0 20px rgba(255,0,204,.4)!important; }
#sb { font-family:'Share Tech Mono',monospace; font-size:.75rem; color:rgba(0,245,255,.6); text-align:center; padding:6px; border:1px solid rgba(0,245,255,.1); border-radius:3px; background:rgba(0,245,255,.03); letter-spacing:2px; min-height:2rem; }
#model-status-text textarea {
  height: 2.2rem !important;
  min-height: 2.2rem !important;
  max-height: 2.2rem !important;
  resize: none !important;
  overflow: hidden !important;
  font-family: 'Share Tech Mono', monospace !important;
  font-size: .75rem !important;
  color: rgba(0,245,255,.7) !important;
  background: rgba(0,245,255,.03) !important;
  border: 1px solid rgba(0,245,255,.1) !important;
  border-radius: 3px !important;
  text-align: center !important;
  letter-spacing: 2px !important;
  padding: 6px !important;
}
#model-status-text .wrap { display: none !important; }
#model-status-text label { display: none !important; }::-webkit-scrollbar-track{background:var(--bg0)}::-webkit-scrollbar-thumb{background:linear-gradient(var(--c),var(--m));border-radius:3px}
"""

def status_text():
    if _load_status["error"]:
        return f"❌  {_load_status['msg']}"
    if _load_status["ready"]:
        return f"✅  {_load_status['msg']}"
    return f"⏳  {_load_status['msg']}"

with gr.Blocks(css=CSS, title="ZKAEDI PRIME — Solidity Auditor") as demo:

    gr.HTML("""
    <div id="hdr">
      <div style="font-family:'Share Tech Mono',monospace;font-size:.65rem;color:rgba(0,245,255,.4);letter-spacing:8px;margin-bottom:8px;">▓▓▓ SYSTEM ONLINE ▓▓▓</div>
      <h1>ZKAEDI PRIME</h1>
      <div style="font-family:'Share Tech Mono',monospace;color:rgba(0,245,255,.7);font-size:.8rem;letter-spacing:6px;text-transform:uppercase;margin-top:.5rem;">⚡ Smart Contract Energy Signatures Auditor ⚡</div>
      <div style="font-family:'Share Tech Mono',monospace;font-size:.65rem;color:rgba(255,0,204,.5);letter-spacing:4px;margin-top:8px;">Recursively Coupled Hamiltonian Vulnerability Detection</div>
    </div>
    """)

    model_status = gr.Textbox(
        value=_load_status["msg"],
        label="",
        interactive=False,
        max_lines=1,
        elem_id="model-status-text",
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            gr.HTML('<div class="lbl">◈ Contract Input</div>')
            contract_in = gr.Textbox(
                label="", lines=22, max_lines=40,
                placeholder="// Paste Solidity contract here...\npragma solidity ^0.8.0;\n\ncontract YourContract {\n    ...\n}",
            )
            with gr.Accordion("⚙ Field Parameters", open=False):
                temp      = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="Temperature")
                max_tok   = gr.Slider(256, 4096, value=2048, step=128, label="Max Tokens")
            with gr.Row():
                scan_btn  = gr.Button("⚡ INITIATE SCAN", variant="primary", scale=2)
                clear_btn = gr.Button("CLEAR", variant="secondary", scale=1)

        with gr.Column(scale=1):
            gr.HTML('<div class="lbl">◈ Vulnerability Field Map</div>')
            output = gr.Textbox(
                label="", lines=22, max_lines=40,
                interactive=False, show_copy_button=True,
                elem_classes=["out"],
                placeholder="Awaiting input...\n\n// ZKAEDI PRIME maps:\n// · Reentrancy wells\n// · Overflow cliffs\n// · Access control gaps\n// · Flash loan vectors\n// · Hamiltonian risk score",
            )
            scan_status = gr.HTML('<div id="sb">IDLE — READY FOR INPUT</div>')

    gr.HTML('<div class="lbl" style="margin-top:16px;">◈ Example Contracts</div>')
    with gr.Row():
        for lbl, code in [("🔴 Reentrancy", REENTRANCY), ("🟡 Overflow", OVERFLOW), ("🔴 Access Control", ACCESS)]:
            b = gr.Button(lbl, variant="secondary", size="sm")
            b.click(lambda c=code: c, outputs=contract_in)

    gr.HTML("""
    <div style="text-align:center;margin-top:24px;padding:16px;border-top:1px solid rgba(0,245,255,.1);">
      <div style="font-family:'Share Tech Mono',monospace;font-size:.65rem;color:rgba(0,245,255,.3);letter-spacing:4px;">
        MODEL: zkaedi/gemma-7b-solidity-energy-signatures · Q4_K_M · CPU
      </div>
      <div style="font-family:'Share Tech Mono',monospace;font-size:.6rem;color:rgba(255,0,204,.25);letter-spacing:2px;margin-top:6px;">
        "The field does not find the vulnerability. The field becomes the vulnerability." — ZKAEDI PRIME
      </div>
    </div>
    """)

    demo.load(status_text, outputs=model_status, every=3)

    scan_btn.click(
        fn=audit,
        inputs=[contract_in, temp, max_tok],
        outputs=[output, scan_status],
        show_progress="hidden",
    )
    clear_btn.click(
        fn=lambda: ("", "", '<div id="sb">IDLE — READY FOR INPUT</div>'),
        outputs=[contract_in, output, scan_status],
    )

if __name__ == "__main__":
    demo.queue(max_size=5).launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )