app.py

"""Adversarial SAST — the false positive dies on screen (snippet + whole repo).

A code model (Qwen2.5-Coder) scans code for vulnerabilities (stage 1, broad), then a
second *adversarial* pass refutes each candidate (stage 2): a finding survives only if it's
a real, reachable, unsanitized sink — with a PoC. False positives get dropped.

Two modes (tabs):
- Snippet — paste code, the Verify ON/OFF toggle IS the demo.
- Whole repo — clone a public repo (or upload a .zip), scanned file-by-file with streaming.

Runs on ZeroGPU. Static analysis only (target code is never executed). Defensive / educational.
"""

import os
import shutil
import tempfile
import time

import gradio as gr
import spaces
import torch
import outlines
from huggingface_hub import InferenceClient
from transformers import AutoModelForCausalLM, AutoTokenizer

from examples import EXAMPLES
from engine import detect, detect_hf, refute, refute_hf, render_snippet
from repo_audit import (
    clone_repo, extract_zip, scan_repo, render_repo_report,
    DEFAULT_MAX_FILES, MAX_FINDINGS_PER_FILE,
)

MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-Coder-7B-Instruct")
MAX_FINDINGS = int(os.environ.get("MAX_FINDINGS", "5"))

# --- Backends for the WHOLE-REPO audit -------------------------------------------------------
# The Snippet tab always uses the local 7B on ZeroGPU (Outlines-constrained — the interactive demo).
# The repo audit runs BOTH stages on HF Inference by default:
#  - detection on a LARGE model (480B): whole-file detection needs a wide view + a long data-flow
#    (source→sink can be ~200 lines apart, e.g. an LFI through an object), which the 7B misses.
#  - refutation on a CAPABLE model (32B) over a bounded window (±REFUTE_CTX): the refuter must CREDIT
#    real sanitization (prepared statements, allow-lists, casts) — a 7B refuter can't, so it drowns the
#    report in false positives — while not being fooled by file_exists/isset guards (see REFUTE_SYS).
# Both off-GPU ⇒ a long whole-repo scan needs NO ZeroGPU reservation, so there is no proxy-token to
# expire mid-scan (the limit that capped multi-file scans). ZeroGPU is the no-token fallback.
DETECT_BACKEND = os.environ.get("DETECT_BACKEND", "hf_inference").strip().lower()
REFUTE_BACKEND = os.environ.get("REFUTE_BACKEND", "hf_inference").strip().lower()
HF_INFERENCE_MODEL = os.environ.get("HF_INFERENCE_MODEL", "Qwen/Qwen3-Coder-480B-A35B-Instruct")  # detect (large)
HF_REFUTE_MODEL = os.environ.get("HF_REFUTE_MODEL", "Qwen/Qwen2.5-Coder-32B-Instruct")  # refute (judge — needs to credit real sanitization)
HF_TOKEN = os.environ.get("HF_TOKEN")
# Rough $/1M-token figure for cost *logging* only (the real cost shows on the HF billing dashboard).
USD_PER_MTOK = float(os.environ.get("HF_INFERENCE_USD_PER_MTOK", "0.9"))
REFUTE_CHUNK = int(os.environ.get("REFUTE_CHUNK", "6"))   # candidates per ZeroGPU call (zerogpu fallback only)
REFUTE_CTX = int(os.environ.get("REFUTE_CTX", "60"))      # ± lines of context for refutation (enough to see the local sanitization, not the whole file)

_inf_client = InferenceClient(token=HF_TOKEN) if HF_TOKEN else None
if _inf_client is None:  # no token → both repo stages fall back to the local 7B on ZeroGPU
    if DETECT_BACKEND == "hf_inference":
        print("[init] no HF_TOKEN — repo detection falls back to ZeroGPU 7B"); DETECT_BACKEND = "zerogpu"
    if REFUTE_BACKEND == "hf_inference":
        print("[init] no HF_TOKEN — repo refutation falls back to ZeroGPU 7B"); REFUTE_BACKEND = "zerogpu"
print(f"[init] repo backends — detect: {DETECT_BACKEND}"
      + (f" ({HF_INFERENCE_MODEL})" if DETECT_BACKEND == "hf_inference" else "")
      + f" · refute: {REFUTE_BACKEND}"
      + (f" ({HF_REFUTE_MODEL})" if REFUTE_BACKEND == "hf_inference" else ""))

print(f"[init] loading {MODEL_ID} …")
_t0 = time.perf_counter()
_tok = AutoTokenizer.from_pretrained(MODEL_ID)
_hf = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, device_map="cuda")
model = outlines.from_transformers(_hf, _tok)
print(f"[init] model ready in {time.perf_counter() - _t0:.1f}s")


# ---------------- Snippet path (tab 1) ----------------
@spaces.GPU(duration=120)
def _analyze(code, lang, verify):
    cands = detect(model, _tok, code, lang, MAX_FINDINGS)
    if not verify:
        return {"verified": False, "candidates": cands}
    results = [{"candidate": c, "verdict": refute(model, _tok, code, c)} for c in cands]
    return {"verified": True, "results": results}


def audit(code, lang, verify):
    """Audit ONE code snippet for security vulnerabilities with adversarial verification.

    Stage 1 detects candidate vulnerabilities; if verify is on, stage 2 refutes each and keeps
    only the genuinely exploitable ones (with a PoC). False positives (sanitized input, dead
    code, guards) are dropped.

    Args:
        code: The source code snippet to audit.
        lang: Programming language (e.g. "python", "javascript", "php").
        verify: If true, run the adversarial refutation pass (kills false positives).

    Returns:
        A markdown report of confirmed and refuted findings.
    """
    code = (code or "").strip()
    if not code:
        return "Paste some code first."
    t0 = time.perf_counter()
    out = _analyze(code, lang, verify)
    return render_snippet(out, time.perf_counter() - t0)


# ---------------- Whole-repo path (tab 2) ----------------
@spaces.GPU(duration=60)
def _detect_gpu(code, lang):
    """Fallback stage-1 on ZeroGPU (7B, capped) — used only when DETECT_BACKEND=zerogpu."""
    return detect(model, _tok, code, lang, MAX_FINDINGS_PER_FILE)


def _refute_context(full_code, line):
    """Targeted, line-numbered window (±REFUTE_CTX lines) around a candidate for refutation."""
    lines = full_code.splitlines()
    lo = max(0, line - 1 - REFUTE_CTX)
    hi = min(len(lines), line + REFUTE_CTX)
    return "\n".join(f"{lo + 1 + j}| {ln}" for j, ln in enumerate(lines[lo:hi]))


@spaces.GPU(duration=90)
def _refute_gpu(code, cands):
    """Stage 2 on ZeroGPU (7B): refute a CHUNK of candidates for ONE file, each on a TARGETED
    window (args are picklable). The focused context is what lets the 7B confirm a real finding
    and kill false positives — a whole file misleads it (it excused this LFI on a file_exists guard)."""
    return [{"candidate": c, "verdict": refute(model, _tok, _refute_context(code, c["line"]), c)}
            for c in cands]


def _audit_file(code, lang, stats):
    """One file: stage 1 detect (uncapped) → stage 2 refute each candidate on a TARGETED window,
    each stage on its configured backend. With both on HF Inference the whole file is audited
    off-GPU — a long whole-repo scan needs no ZeroGPU reservation, so no proxy token expires
    mid-scan. ZeroGPU is the fallback. Accumulates token cost in `stats`."""
    if DETECT_BACKEND == "hf_inference" and _inf_client is not None:
        cands, tok = detect_hf(_inf_client, HF_INFERENCE_MODEL, code, lang)
        stats["detect_tokens"] += tok
    else:
        cands = _detect_gpu(code, lang)
    if not cands:
        return []
    stats["files"] += 1
    if REFUTE_BACKEND == "hf_inference" and _inf_client is not None:
        out = []
        for c in cands:
            verdict, tok = refute_hf(_inf_client, HF_REFUTE_MODEL, _refute_context(code, c["line"]), c)
            stats["refute_tokens"] += tok
            out.append({"candidate": c, "verdict": verdict})
        return out
    # ZeroGPU fallback — chunked so an uncapped detector can't blow the per-call GPU budget
    results = []
    for i in range(0, len(cands), REFUTE_CHUNK):
        results += _refute_gpu(code, cands[i:i + REFUTE_CHUNK])
    return results


def audit_repo(git_url, branch, zip_file, max_files):
    """Audit a WHOLE public repository for security vulnerabilities (adversarial SAST).

    Clones a public git repo (shallow) or unpacks an uploaded .zip into an ephemeral sandbox,
    selects the code files (within hard caps), then runs the two-stage audit file-by-file:
    detect candidates, then adversarially refute each — keeping only genuinely exploitable
    findings with a proof-of-concept. Static analysis only; the target code is never executed.
    Public / authorized repositories only.

    Args:
        git_url: HTTPS URL of a public git repo (e.g. https://github.com/owner/repo).
        branch: Optional branch, tag, or full commit SHA to clone (empty = default branch).
        zip_file: Optional path to an uploaded .zip of the codebase (used if git_url is empty).
        max_files: Hard cap on files scanned (protects the GPU quota).

    Yields:
        Progress lines during the scan, then the final aggregated markdown report.
    """
    git_url = (git_url or "").strip()
    try:
        max_files = max(1, min(int(max_files), 200))
    except (TypeError, ValueError):
        max_files = DEFAULT_MAX_FILES
    if not git_url and not zip_file:
        yield "Provide a public git URL or upload a .zip."
        return

    dest = tempfile.mkdtemp(prefix="sast-repo-")
    stats = {"detect_tokens": 0, "refute_tokens": 0, "files": 0}
    try:
        yield "⏳ Fetching the codebase…"
        try:
            if git_url:
                clone_repo(git_url, (branch or "").strip() or None, dest)
            else:
                extract_zip(zip_file, dest)
        except Exception as e:
            yield f"❌ Could not fetch the repo: {str(e)[:200]}"
            return

        det = (f"{HF_INFERENCE_MODEL.split('/')[-1]} (HF Inference)"
               if DETECT_BACKEND == "hf_inference" and _inf_client else "7B (ZeroGPU)")
        ref = (f"{HF_REFUTE_MODEL.split('/')[-1]} (HF Inference)"
               if REFUTE_BACKEND == "hf_inference" and _inf_client else "7B (ZeroGPU)")
        result = None
        for ev in scan_repo(lambda code, lang: _audit_file(code, lang, stats), dest, max_files=max_files):
            if ev.get("done"):
                result = ev["result"]
            else:
                yield (f"⏳ Scanning file {ev['i'] + 1}/{ev['n']} — `{ev['file']}`  ·  "
                       f"detect: {det} · refute: {ref}")
        if not result:
            yield "No code files found to scan."
            return
        report = render_repo_report(result)
        tot = stats["detect_tokens"] + stats["refute_tokens"]
        if tot:
            cost = tot / 1_000_000 * USD_PER_MTOK
            print(f"[repo-audit] detect {stats['detect_tokens']} + refute {stats['refute_tokens']} tok "
                  f"over {stats['files']} files ~${cost:.4f}")
            report += (f"\n\n---\n_Detection **{det}** · refutation **{ref}** · {stats['files']} files · "
                       f"{stats['detect_tokens']:,}+{stats['refute_tokens']:,} tokens · ~${cost:.4f} (estimate)._")
        yield report
    finally:
        shutil.rmtree(dest, ignore_errors=True)


# ---------------- UI ----------------
HEADER = "# 🛡️ Adversarial SAST — the false positive dies on screen"

SNIPPET_INTRO = """
Paste code; stage 1 detects candidate vulnerabilities, stage 2 **adversarially refutes** each.
**The Verify ON/OFF toggle is the demo**: off, the raw detector is noisy; on, the false positives
die and only real findings (with a PoC) survive. The default example hides a fake SQL injection
(a `SELECT` built from an `int`) next to a real command injection — only the real one survives.
"""

REPO_INTRO = """
Audit a **whole public repository**. It clones shallow (or unpacks your `.zip`) into an ephemeral
sandbox, ranks the code files **by security relevance**, and scans the top ones **file-by-file with
streaming** — two stages **split by model size**: detection runs on a large model
(**Qwen3-Coder-480B**) that can follow a long source→sink data-flow across a whole file, then each
candidate is **adversarially refuted by Qwen2.5-Coder-32B on a bounded context** — it credits real
sanitization (prepared statements, allow-lists) and drops false positives — both via **HF Inference**.
**Static analysis only — the target code is never executed. Public / authorized repos only.**

⏱️ A scan uses HF Inference credits (both stages, off-GPU — no scan-length limit). Try
`https://github.com/digininja/DVWA` or `https://github.com/phpipam/phpipam` (branch/tag/commit accepted).
"""

ABOUT = """
---
Snippet: **Qwen2.5-Coder-7B** on **ZeroGPU** · whole-repo: detect **Qwen3-Coder-480B** + refute
**Qwen2.5-Coder-32B**, both via **HF Inference** · static analysis, no code execution, no secrets · built by
[Ferr0](https://huggingface.co/Ferr0) · [pixelium.win](https://pixelium.win) · [GitHub](https://github.com/ferr079)
"""

with gr.Blocks(title="Adversarial SAST") as demo:
    gr.Markdown(HEADER)
    with gr.Tabs():
        with gr.Tab("Snippet"):
            gr.Markdown(SNIPPET_INTRO)
            with gr.Row():
                with gr.Column(scale=1):
                    lang = gr.Dropdown(
                        choices=["python", "javascript", "go", "java", "php", "ruby", "c"],
                        value="python", label="Language",
                    )
                    code = gr.Code(value=EXAMPLES[0][0], language="python", label="Code")
                    verify = gr.Checkbox(value=True, label="Verify (adversarial refutation)")
                    go = gr.Button("Audit", variant="primary")
                with gr.Column(scale=1):
                    out = gr.Markdown()
            gr.Examples(examples=EXAMPLES, inputs=[code, lang, verify])
            go.click(audit, inputs=[code, lang, verify], outputs=out)

        with gr.Tab("Whole repo"):
            gr.Markdown(REPO_INTRO)
            with gr.Row():
                with gr.Column(scale=1):
                    git_url = gr.Textbox(label="Public git URL", placeholder="https://github.com/owner/repo")
                    branch = gr.Textbox(label="Branch / tag / commit (optional)",
                                        placeholder="default branch, or a tag / full commit SHA")
                    zip_file = gr.File(label="…or upload a .zip", file_types=[".zip"], type="filepath")
                    max_files = gr.Slider(5, 200, value=DEFAULT_MAX_FILES, step=5, label="Max files (GPU quota guard)")
                    go_repo = gr.Button("Audit repo", variant="primary")
                with gr.Column(scale=1):
                    out_repo = gr.Markdown()
            go_repo.click(audit_repo, inputs=[git_url, branch, zip_file, max_files], outputs=out_repo)

    gr.Markdown(ABOUT)


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