app.py

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MODEL_ID = "CIRCL/vulnerability-attack-technique-classification-roberta-base"

# MITRE ATT&CK Enterprise technique names for the 53 parent techniques
# the model was trained on (sub-techniques collapsed to parents).
TECHNIQUE_NAMES = {
    "T1003": "OS Credential Dumping",
    "T1005": "Data from Local System",
    "T1021": "Remote Services",
    "T1036": "Masquerading",
    "T1040": "Network Sniffing",
    "T1041": "Exfiltration Over C2 Channel",
    "T1046": "Network Service Discovery",
    "T1055": "Process Injection",
    "T1059": "Command and Scripting Interpreter",
    "T1068": "Exploitation for Privilege Escalation",
    "T1070": "Indicator Removal",
    "T1071": "Application Layer Protocol",
    "T1078": "Valid Accounts",
    "T1082": "System Information Discovery",
    "T1083": "File and Directory Discovery",
    "T1087": "Account Discovery",
    "T1091": "Replication Through Removable Media",
    "T1098": "Account Manipulation",
    "T1105": "Ingress Tool Transfer",
    "T1106": "Native API",
    "T1110": "Brute Force",
    "T1133": "External Remote Services",
    "T1136": "Create Account",
    "T1185": "Browser Session Hijacking",
    "T1189": "Drive-by Compromise",
    "T1190": "Exploit Public-Facing Application",
    "T1202": "Indirect Command Execution",
    "T1203": "Exploitation for Client Execution",
    "T1204": "User Execution",
    "T1210": "Exploitation of Remote Services",
    "T1211": "Exploitation for Defense Evasion",
    "T1212": "Exploitation for Credential Access",
    "T1485": "Data Destruction",
    "T1486": "Data Encrypted for Impact",
    "T1496": "Resource Hijacking",
    "T1497": "Virtualization/Sandbox Evasion",
    "T1498": "Network Denial of Service",
    "T1499": "Endpoint Denial of Service",
    "T1505": "Server Software Component",
    "T1528": "Steal Application Access Token",
    "T1542": "Pre-OS Boot",
    "T1543": "Create or Modify System Process",
    "T1548": "Abuse Elevation Control Mechanism",
    "T1550": "Use Alternate Authentication Material",
    "T1552": "Unsecured Credentials",
    "T1555": "Credentials from Password Stores",
    "T1557": "Adversary-in-the-Middle",
    "T1563": "Remote Service Session Hijacking",
    "T1565": "Data Manipulation",
    "T1566": "Phishing",
    "T1574": "Hijack Execution Flow",
    "T1608": "Stage Capabilities",
    "T1685": "Disable or Modify Tools",
}

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID).to("cuda")
model.eval()
id2label = model.config.id2label


@spaces.GPU(duration=30)
def classify_cve(description: str, top_k: int = 10, threshold: float = 0.5) -> dict:
    """Classify a CVE description into MITRE ATT&CK techniques.

    Given a free-text vulnerability description, this function predicts the
    most relevant MITRE ATT&CK Enterprise techniques using a RoBERTa-base
    multi-label classifier trained on a curated gold set of 1,207 CVEs.

    Args:
        description: A free-text vulnerability (CVE) description.
        top_k: Maximum number of techniques to return (sorted by probability).
        threshold: Minimum sigmoid probability to include a technique in the
            "high-confidence" set. All predicted probabilities are returned
            regardless; this only affects the summary text.
    """
    if not description or not description.strip():
        return {}, "Please enter a CVE description."

    inputs = tokenizer(
        description,
        truncation=True,
        max_length=512,
        return_tensors="pt",
    ).to("cuda")

    with torch.no_grad():
        probs = torch.sigmoid(model(**inputs).logits)[0]

    # Sort by descending probability
    sorted_indices = probs.argsort(descending=True)
    top_k = min(top_k, len(sorted_indices))

    results = {}
    high_conf = []
    for idx in sorted_indices[:top_k]:
        tid = id2label[int(idx)]
        prob = float(probs[idx])
        label_full = f"{tid} — {TECHNIQUE_NAMES.get(tid, 'Unknown')}"
        results[label_full] = prob
        if prob >= threshold:
            high_conf.append(f"{tid} ({prob:.3f})")

    summary = (
        f"High-confidence techniques (≥ {threshold:.1f}): {', '.join(high_conf)}"
        if high_conf
        else f"No techniques above threshold {threshold:.1f}. Top: {id2label[int(sorted_indices[0])]} ({float(probs[sorted_indices[0]]):.3f})"
    )

    return results, summary


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    [
        "Zoho ManageEngine ServiceDesk Plus before 11306, ServiceDesk Plus MSP "
        "before 10530, and SupportCenter Plus before 11014 are vulnerable to "
        "unauthenticated remote code execution."
    ],
    [
        "A buffer overflow vulnerability in the Simple Network Management "
        "Protocol (SNMP) service of Cisco IOS allows remote attackers to "
        "execute arbitrary code on the affected device via crafted SNMP packets."
    ],
    [
        "An SQL injection vulnerability in the login page of a web application "
        "allows unauthenticated remote attackers to bypass authentication and "
        "access the administrator dashboard by injecting malicious SQL "
        "queries through the username field."
    ],
    [
        "A path traversal vulnerability in a file download component allows "
        "remote attackers to read arbitrary files on the system via specially "
        "crafted requests containing dot-dot sequences in the filename "
        "parameter."
    ],
    [
        "A cross-site scripting (XSS) vulnerability in the comment section "
        "of a popular forum allows attackers to inject malicious JavaScript "
        "that executes in the browser of any visitor viewing the affected "
        "page, enabling session token theft."
    ],
]

with gr.Blocks(title="CVE to MITRE ATT&CK Mapper") as demo:
    gr.Markdown(
        "# CVE to MITRE ATT&CK Technique Mapper\n"
        "Enter a vulnerability (CVE) description and get the most relevant "
        "MITRE ATT&CK Enterprise techniques, predicted by a RoBERTa-base "
        "multi-label classifier trained on a curated gold set of 1,207 CVEs "
        "with expert MITRE CTID labels.\n\n"
        "📖 [Paper](https://arxiv.org/abs/2607.25572) | "
        "🤖 [Model](https://huggingface.co/CIRCL/vulnerability-attack-technique-classification-roberta-base) | "
        "💻 [Code (VulnTrain)](https://github.com/vulnerability-lookup/VulnTrain)"
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            description = gr.Textbox(
                label="CVE / Vulnerability Description",
                placeholder="Paste a CVE description here...",
                lines=6,
                scale=4,
            )
            run_btn = gr.Button("Classify", variant="primary", scale=1)

        with gr.Accordion("Advanced settings", open=False):
            top_k = gr.Slider(
                minimum=1, maximum=53, value=10, step=1,
                label="Top-K techniques to show",
            )
            threshold = gr.Slider(
                minimum=0.0, maximum=1.0, value=0.5, step=0.05,
                label="Confidence threshold (for high-confidence summary)",
            )

        label_output = gr.Label(
            label="Predicted MITRE ATT&CK Techniques",
            num_top_classes=10,
        )
        summary_output = gr.Textbox(
            label="Summary",
            interactive=False,
            lines=2,
        )

    run_btn.click(
        fn=classify_cve,
        inputs=[description, top_k, threshold],
        outputs=[label_output, summary_output],
        api_name="classify",
    )
    description.submit(
        fn=classify_cve,
        inputs=[description, top_k, threshold],
        outputs=[label_output, summary_output],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[description],
        outputs=[label_output, summary_output],
        fn=classify_cve,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)