app.py

"""
Hugging Face Spaces entry point.
Deploy this file as app.py in your HF Space (SDK: gradio).
"""
import os
import logging
import gradio as gr
from app.models.analyzer import analyze_contract

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ── Startup diagnostics ───────────────────────────────────────────────────────
HF_MODEL_ID = os.getenv("HF_MODEL_ID", "")
HF_TOKEN    = os.getenv("HF_TOKEN", "")

logger.info("=" * 50)
if HF_MODEL_ID:
    logger.info(f"✅ HF_MODEL_ID set: {HF_MODEL_ID}")
else:
    logger.warning("⚠️  HF_MODEL_ID not set – LLM analysis disabled")

if HF_TOKEN:
    logger.info("✅ HF_TOKEN set")
else:
    logger.warning("⚠️  HF_TOKEN not set")
logger.info("=" * 50)


# ── Helpers ───────────────────────────────────────────────────────────────────

def _fmt_vulns(vulns: list) -> str:
    if not vulns:
        return "✅ No vulnerabilities detected."
    icons = {"HIGH": "🚨", "MEDIUM": "⚠️", "LOW": "ℹ️"}
    lines = []
    for v in vulns:
        icon = icons.get(v["severity"], "•")
        name = v["name"].replace("_", " ").title()
        lines.append(f"{icon} **{name}** [{v['severity']}]")
        lines.append(f"   {v['description']}")
        if v.get("line_numbers"):
            lines.append(f"   Lines: {', '.join(map(str, v['line_numbers'][:8]))}")
        if v.get("recommendation"):
            lines.append(f"   💡 {v['recommendation']}")
        lines.append("")
    return "\n".join(lines)


def get_status() -> str:
    """Show current config – useful for debugging."""
    model_id = os.getenv("HF_MODEL_ID", "")
    token    = os.getenv("HF_TOKEN", "")
    if model_id and token:
        return f"🟢 **LLM ready** — `{model_id}`"
    elif model_id and not token:
        return f"🟡 **HF_MODEL_ID set** but no HF_TOKEN — may fail on private models"
    else:
        return "🔴 **HF_MODEL_ID not set** — pattern analysis only"


def audit(solidity_code: str):
    if not solidity_code.strip():
        return "⚠️ Please paste some Solidity code.", "", ""

    logger.info("Running analysis...")
    result = analyze_contract(solidity_code)
    logger.info(f"Analysis type: {result['analysis_type']}, risk: {result['risk_level']}")

    risk_icons = {"CRITICAL": "🚨", "MEDIUM": "⚠️", "LOW": "ℹ️", "SAFE": "✅"}
    icon = risk_icons.get(result["risk_level"], "•")
    summary = (
        f"## {icon} Risk Level: {result['risk_level']}\n\n"
        f"- **Solidity Version:** {result['solidity_version']}\n"
        f"- **Lines analyzed:** {result['total_lines']}\n"
        f"- **Severity Score:** {result['severity_score']}\n"
        f"- **Analysis type:** `{result['analysis_type']}`\n\n"
        f"| 🚨 HIGH | ⚠️ MEDIUM | ℹ️ LOW |\n|---------|----------|-------|\n"
        f"| {result['high_count']} | {result['medium_count']} | {result['low_count']} |"
    )

    vuln_text = _fmt_vulns(result["vulnerabilities"])

    if result.get("llm_analysis"):
        llm_text = result["llm_analysis"]
    else:
        llm_text = (
            "_LLM analysis not available._\n\n"
            f"**Current status:** {get_status()}\n\n"
            "Make sure `HF_MODEL_ID` and `HF_TOKEN` are set in Space **Settings → Repository secrets**."
        )

    return summary, vuln_text, llm_text


# ── UI ────────────────────────────────────────────────────────────────────────

EXAMPLES = [
    ["""pragma solidity ^0.4.24;
contract VulnerableBank {
    mapping(address => uint256) public balances;
    function withdraw(uint256 _amount) public {
        require(balances[msg.sender] >= _amount);
        msg.sender.call.value(_amount)();
        balances[msg.sender] -= _amount;
    }
}"""],
    ["""// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;
import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
import "@openzeppelin/contracts/access/Ownable.sol";
contract SafeVault is ReentrancyGuard, Ownable {
    mapping(address => uint256) private balances;
    function withdraw(uint256 amount) external nonReentrant {
        require(balances[msg.sender] >= amount);
        balances[msg.sender] -= amount;
        (bool ok,) = payable(msg.sender).call{value: amount}("");
        require(ok);
    }
}"""],
]

with gr.Blocks(title="🛡️ Smart Contract Auditor") as demo:
    gr.Markdown("# 🛡️ Smart Contract Security Auditor")
    gr.Markdown(get_status)          # live status banner – refreshes on load

    with gr.Row():
        with gr.Column():
            code_input  = gr.Code(label="Solidity Contract", language="javascript", lines=22)
            analyze_btn = gr.Button("🔍 Analyze Contract", variant="primary")

        with gr.Column():
            summary_out = gr.Markdown(label="📊 Summary")
            vuln_out    = gr.Markdown(label="🔍 Vulnerabilities")
            llm_out     = gr.Markdown(label="🤖 LLM Analysis")

    gr.Examples(examples=EXAMPLES, inputs=[code_input])

    analyze_btn.click(
        fn=audit,
        inputs=[code_input],
        outputs=[summary_out, vuln_out, llm_out],
    )

if __name__ == "__main__":
    demo.launch()