"""CVE Explainer / Summarizer.

Enter a CVE ID (e.g. CVE-2024-3094) or paste a raw vulnerability description.
Fetches public CVE metadata (CIRCL / NVD-style JSON) and produces a plain-English
summary, risk level, remediation steps and suggested MITRE ATT&CK techniques.
Uses Hugging Face Inference Providers when a hub token is available; otherwise a
structured fallback built from CVSS/CWE fields.
"""

from __future__ import annotations

import spaces  # ZeroGPU Space: required before other heavy imports

import json
import os
import re
from typing import Any

import gradio as gr
import requests

CVE_RE = re.compile(r"CVE-\d{4}-\d{4,7}", re.I)
CIRCL = "https://cve.circl.lu/api/cve/{cve_id}"
NVD = "https://services.nvd.nist.gov/rest/json/cvds/2.0?cveId={cve_id}"

LLM = os.environ.get("CVE_LLM", "Qwen/Qwen2.5-7B-Instruct")
PROVIDER = os.environ.get("CVE_INFERENCE_PROVIDER", "featherless-ai")


def _token() -> str | None:
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def parse_cve_id(text: str) -> str | None:
    m = CVE_RE.search(text.strip())
    return m.group(0).upper() if m else None


def fetch_circl(cve_id: str) -> dict[str, Any]:
    r = requests.get(CIRCL.format(cve_id=cve_id), timeout=30)
    r.raise_for_status()
    return r.json()


def fetch_nvd(cve_id: str) -> dict[str, Any] | None:
    try:
        r = requests.get(NVD.format(cve_id=cve_id), timeout=45, headers={"Accept": "application/json"})
        if r.status_code != 200:
            return None
        data = r.json()
        if not data.get("vulnerabilities"):
            return None
        return data["vulnerabilities"][0]["cve"]
    except Exception:
        return None


def cvss_from_record(rec: dict[str, Any]) -> tuple[float | None, str, dict[str, Any]]:
    """Return (score, severity_label, metric_dict)."""
    best_score, best_sev, best = None, "UNKNOWN", {}
    containers = rec.get("containers") or {}
    for block in (containers.get("cna"), *(containers.get("adp") or [])):
        if not block:
            continue
        for m in block.get("metrics") or []:
            for key in ("cvssV3_1", "cvssV3_0", "cvssV4_0"):
                if key in m:
                    d = m[key]
                    sc = float(d.get("baseScore", 0))
                    sev = str(d.get("baseSeverity", "UNKNOWN"))
                    if best_score is None or sc > best_score:
                        best_score, best_sev, best = sc, sev, d
    return best_score, best_sev, best


def descriptions_from_record(rec: dict[str, Any]) -> list[str]:
    out: list[str] = []
    cna = (rec.get("containers") or {}).get("cna") or {}
    for d in cna.get("descriptions") or []:
        if d.get("lang", "en") == "en" and d.get("value"):
            out.append(d["value"])
    nvd = rec.get("_nvd") or {}
    for d in nvd.get("descriptions") or []:
        if d.get("lang") == "en" and d.get("value"):
            out.append(d["value"])
    return out


def cwes_from_record(rec: dict[str, Any]) -> list[str]:
    ids: list[str] = []
    cna = (rec.get("containers") or {}).get("cna") or {}
    for p in cna.get("problemTypes") or []:
        for d in p.get("descriptions") or []:
            if d.get("cweId"):
                ids.append(d["cweId"])
            elif d.get("description"):
                ids.append(d["description"])
    nvd = rec.get("_nvd") or {}
    for w in nvd.get("weaknesses") or []:
        for d in w.get("description") or []:
            if d.get("value", "").startswith("CWE-"):
                ids.append(d["value"])
    return list(dict.fromkeys(ids))


def references_from_record(rec: dict[str, Any]) -> list[str]:
    refs: list[str] = []
    cna = (rec.get("containers") or {}).get("cna") or {}
    for ref in cna.get("references") or []:
        u = ref.get("url")
        if u:
            refs.append(u)
    nvd = rec.get("_nvd") or {}
    for ref in nvd.get("references") or []:
        u = ref.get("url")
        if u:
            refs.append(u)
    return list(dict.fromkeys(refs))[:12]


def risk_band(score: float | None, severity: str) -> str:
    if score is not None:
        if score >= 9.0:
            return "Critical"
        if score >= 7.0:
            return "High"
        if score >= 4.0:
            return "Medium"
        return "Low"
    sev = severity.upper()
    if "CRIT" in sev:
        return "Critical"
    if "HIGH" in sev:
        return "High"
    if "MED" in sev:
        return "Medium"
    if "LOW" in sev:
        return "Low"
    return "Unknown"


def cwe_mitre_hints(cwes: list[str], description: str) -> list[str]:
    text = " ".join(cwes) + " " + description.lower()
    hints: list[tuple[str, str]] = [
        (r"remote code execution|\brce\b|arbitrary code", "T1190 – Exploit Public-Facing Application"),
        (r"privilege escalation|elevation of privilege", "T1068 – Exploitation for Privilege Escalation"),
        (r"sql injection", "T1190 – Exploit Public-Facing Application (SQLi)"),
        (r"cross[- ]site scripting|\bxss\b", "T1189 – Drive-by Compromise (XSS delivery)"),
        (r"phishing|social engineering", "T1566 – Phishing"),
        (r"credential|password|authentication bypass", "T1078 – Valid Accounts"),
        (r"lateral movement", "T1021 – Remote Services"),
        (r"command injection|shell", "T1059 – Command and Scripting Interpreter"),
        (r"backdoor|supply chain|malicious code", "T1195 – Supply Chain Compromise"),
        (r"denial of service|\bdos\b", "T1499 – Endpoint Denial of Service"),
        (r"memory corruption|buffer overflow|use after free", "T1203 – Exploitation for Client Execution"),
    ]
    found = []
    for pat, tid in hints:
        if re.search(pat, text, re.I):
            found.append(tid)
    return list(dict.fromkeys(found))[:6]


def build_context(cve_id: str | None, raw: str, rec: dict[str, Any] | None) -> dict[str, Any]:
    if rec:
        score, sev, metrics = cvss_from_record(rec)
        descs = descriptions_from_record(rec)
        desc = descs[0] if descs else raw or "(no description in feed)"
        cwes = cwes_from_record(rec)
        refs = references_from_record(rec)
        title = ((rec.get("containers") or {}).get("cna") or {}).get("title") or cve_id
        published = (rec.get("cveMetadata") or {}).get("datePublished", "")
    else:
        score, sev, metrics, cwes, refs, title, published = None, "UNKNOWN", {}, [], [], cve_id or "Custom text", ""
        desc = raw.strip()
    band = risk_band(score, sev)
    mitre = cwe_mitre_hints(cwes, desc)
    return {
        "cve_id": cve_id or "N/A",
        "title": title,
        "published": published,
        "description": desc,
        "cvss_score": score,
        "cvss_severity": sev,
        "cvss_vector": metrics.get("vectorString"),
        "risk_band": band,
        "cwes": cwes,
        "references": refs,
        "mitre_hints": mitre,
        "attack_vector": metrics.get("attackVector"),
        "attack_complexity": metrics.get("attackComplexity"),
        "privileges_required": metrics.get("privilegesRequired"),
        "user_interaction": metrics.get("userInteraction"),
    }


def llm_explain(ctx: dict[str, Any]) -> str | None:
    token = _token()
    if not token:
        return None
    try:
        from huggingface_hub import InferenceClient
    except ImportError:
        return None
    prompt = f"""You are a senior vulnerability analyst writing for IT and security teams.

CVE: {ctx['cve_id']}
Title: {ctx['title']}
CVSS: {ctx['cvss_score']} ({ctx['cvss_severity']}) — overall risk band: {ctx['risk_band']}
CWEs: {', '.join(ctx['cwes']) or 'none listed'}
Official description:
{ctx['description'][:3500]}

Write a concise report with these sections (use markdown headings):
## Plain-English summary
(2-4 sentences, no jargon)
## Risk level
(why it is {ctx['risk_band']}; mention exploitability if known)
## Remediation steps
(numbered, practical: patch, config, detection, workarounds)
## MITRE ATT&CK mapping
(list 2-5 relevant technique IDs like T1190 with one-line justification each; include these hints if relevant: {', '.join(ctx['mitre_hints']) or 'infer from description'})

Be accurate. If something is unknown, say so. Do not invent CVE details not present above."""
    client = InferenceClient(token=token, provider=PROVIDER)
    try:
        out = client.chat_completion(
            model=LLM,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=900,
            temperature=0.2,
        )
        return out.choices[0].message.content
    except Exception:
        try:
            out = client.chat_completion(
                model=LLM,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=900,
                temperature=0.2,
            )
            return out.choices[0].message.content
        except Exception:
            return None


def fallback_report(ctx: dict[str, Any]) -> str:
    lines = [
        f"## Plain-English summary\n{ctx['description'][:1200]}",
        f"\n## Risk level\n**{ctx['risk_band']}**",
    ]
    if ctx["cvss_score"] is not None:
        lines.append(f"- CVSS base score: **{ctx['cvss_score']}** ({ctx['cvss_severity']})")
    if ctx.get("cvss_vector"):
        lines.append(f"- Vector: `{ctx['cvss_vector']}`")
    if ctx.get("attack_vector"):
        lines.append(f"- Attack vector: {ctx['attack_vector']} · Complexity: {ctx.get('attack_complexity')} · Privileges: {ctx.get('privileges_required')} · User interaction: {ctx.get('user_interaction')}")
    lines.append("\n## Remediation steps")
    steps = [
        "1. Confirm affected products/versions in your environment against vendor advisories.",
        "2. Apply the vendor patch or upgrade to a fixed release as soon as possible.",
        "3. If no patch exists, implement compensating controls (WAF rules, disable vulnerable feature, network segmentation).",
        "4. Hunt for exploitation indicators in logs (unexpected child processes, suspicious outbound connections, auth anomalies).",
        "5. Document exceptions and track residual risk until fully remediated.",
    ]
    if "supply chain" in ctx["description"].lower() or "backdoor" in ctx["description"].lower():
        steps.insert(2, "2b. For supply-chain issues: rebuild from trusted sources, rotate secrets, verify package signatures and SBOMs.")
    lines.extend(steps)
    lines.append("\n## MITRE ATT&CK mapping (heuristic)")
    if ctx["mitre_hints"]:
        lines.extend(f"- {h}" for h in ctx["mitre_hints"])
    else:
        lines.append("- No strong keyword match — review the description with your CTI team.")
    if ctx["cwes"]:
        lines.append(f"\n**CWE references:** {', '.join(ctx['cwes'])}")
    if ctx["references"]:
        lines.append("\n**References:**")
        lines.extend(f"- {u}" for u in ctx["references"][:6])
    lines.append("\n<sub>Generated from public CVE metadata (no LLM). Add `HF_TOKEN` Space secret for AI-expanded remediation.</sub>")
    return "\n".join(lines)


def _explain(cve_input: str, raw_description: str, use_llm: bool):
    cve_input = (cve_input or "").strip()
    raw_description = (raw_description or "").strip()
    cve_id = parse_cve_id(cve_input) or parse_cve_id(raw_description)
    if not cve_id and not raw_description and not cve_input:
        raise gr.Error("Enter a CVE ID (e.g. CVE-2024-3094) and/or paste a description.")

    rec: dict[str, Any] | None = None
    meta_line = ""
    if cve_id:
        try:
            rec = fetch_circl(cve_id)
            nvd = fetch_nvd(cve_id)
            if nvd:
                rec["_nvd"] = nvd
            meta_line = f"Loaded **{cve_id}** from public CVE feeds."
        except requests.HTTPError as e:
            if e.response is not None and e.response.status_code == 404:
                raise gr.Error(f"{cve_id} was not found in the CIRCL feed.")
            raise gr.Error(f"Could not fetch {cve_id}: {e}") from e
        except requests.RequestException as e:
            raise gr.Error(f"Network error fetching {cve_id}: {e}") from e
    elif cve_input and not raw_description:
        raw_description = cve_input

    ctx = build_context(cve_id, raw_description or cve_input, rec)
    report = None
    if use_llm:
        report = llm_explain(ctx)
    if not report:
        report = fallback_report(ctx)
        if use_llm and not _token():
            meta_line += " (LLM disabled: set Space secret `HF_TOKEN` with inference scope.)"
        elif use_llm:
            meta_line += " (LLM call failed — showing structured fallback.)"

    header = f"# {ctx['cve_id']}: {ctx['title']}\n"
    if ctx["published"]:
        header += f"*Published: {ctx['published'][:10]}*\n\n"
    header += meta_line + "\n\n"
    risk_md = f"**{ctx['risk_band']}**" + (f" (CVSS {ctx['cvss_score']})" if ctx["cvss_score"] is not None else "")
    json_out = json.dumps({k: v for k, v in ctx.items() if k != "description" or len(str(v)) < 500}, indent=2)
    return risk_md, header + report, json_out


@spaces.GPU(duration=90)
def explain(cve_input: str, raw_description: str, use_llm: bool):
    """Gradio entrypoint (ZeroGPU worker slot; work is mostly HTTP + optional Inference API)."""
    return _explain(cve_input, raw_description, use_llm)


EXAMPLES = [
    ["CVE-2024-3094", "", True],
    ["CVE-2021-44228", "", True],
    ["", "A buffer overflow in the XYZ VPN client allows remote attackers to execute arbitrary code via a crafted UDP packet. CVSS 9.8.", False],
]

with gr.Blocks(title="CVE Explainer") as demo:
    with gr.Column(elem_id="col"):
        gr.Markdown(
            "# 🛡️ CVE Explainer / Summarizer\n"
            "Enter a **CVE ID** or paste a **raw vulnerability description** → plain-English summary, "
            "**risk level**, **remediation steps**, and **MITRE ATT&CK** mapping.\n\n"
            "Data: [CIRCL CVE API](https://cve.circl.lu/) (+ NVD when available). "
            f"Optional AI summary via Inference Providers (`{LLM}`)."
        )
        with gr.Row():
            with gr.Column():
                cve = gr.Textbox(label="CVE ID", placeholder="CVE-2024-3094")
                raw = gr.Textbox(label="Or paste description / advisory text", lines=6, placeholder="Optional if CVE ID is set…")
                use_llm = gr.Checkbox(value=True, label="Expand with AI (requires HF_TOKEN secret on the Space)")
                btn = gr.Button("Explain", variant="primary")
            with gr.Column():
                risk = gr.Textbox(label="Risk level", interactive=False)
                out = gr.Markdown(label="Report")
        meta = gr.Code(label="Structured metadata (JSON)", language="json")
        gr.Examples(EXAMPLES, inputs=[cve, raw, use_llm], outputs=[risk, out, meta], fn=explain, cache_examples=False)
        gr.Markdown("⚠️ Advisory only — validate patches and scope with vendor bulletins and your asset inventory.")
    btn.click(explain, [cve, raw, use_llm], [risk, out, meta], api_name="explain")

if __name__ == "__main__":
    demo.queue(max_size=16).launch(theme=gr.themes.Soft(), css="#col {max-width: 960px; margin: 0 auto;}")