File size: 5,547 Bytes
f8708fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
093340d
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
"""
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()