"""Rule-based security log classifier.

Ported from the original demo so detection works fully offline (no API key,
no model download). Given a log row it assigns a category, severity hint,
MITRE technique, risk score and human-readable reasons.
"""
from __future__ import annotations

import re
from typing import Any

# Single source of truth for category metadata (label / MITRE / risk / runbook).
CATEGORY_META: dict[str, dict[str, Any]] = {
    "benign": {
        "label": "ปกติ",
        "mitre": "-",
        "risk": 10,
        "runbook": "บันทึกไว้เป็น baseline และใช้เทียบพฤติกรรมผิดปกติในช่วงเวลาเดียวกัน",
    },
    "brute_force": {
        "label": "Brute force",
        "mitre": "T1110",
        "risk": 78,
        "runbook": "ตรวจ IP ต้นทาง, lock account ชั่วคราว, เพิ่ม MFA, block IP ถ้ามี failed login ถี่ผิดปกติ",
    },
    "port_scan": {
        "label": "Port scan",
        "mitre": "T1046",
        "risk": 62,
        "runbook": "ตรวจ firewall drop burst, จำกัด exposure ของ port สำคัญ, เพิ่ม deny rule สำหรับ source ที่ scan ซ้ำ",
    },
    "web_attack": {
        "label": "Web attack",
        "mitre": "T1190",
        "risk": 82,
        "runbook": "ตรวจ request payload, WAF rule, access log รอบเวลาเดียวกัน และ patch endpoint ที่รับ parameter เสี่ยง",
    },
    "priv_esc": {
        "label": "Privilege escalation",
        "mitre": "T1068/T1548",
        "risk": 94,
        "runbook": "ตรวจ sudo command, revoke session, review sudoers, rotate credential และดู command history",
    },
    "data_exfil": {
        "label": "Data exfiltration",
        "mitre": "T1041",
        "risk": 96,
        "runbook": "ตัด connection ปลายทาง, ตรวจปริมาณ bytes, หา process เจ้าของ traffic และ preserve evidence",
    },
    "c2_beacon": {
        "label": "C2 beacon",
        "mitre": "T1071",
        "risk": 95,
        "runbook": "isolate host, ตรวจ interval คงที่, block C2 destination และเก็บ memory/network artifact",
    },
}

# Severity derived from category when the source row does not provide one.
CATEGORY_DEFAULT_SEVERITY = {
    "benign": "info",
    "brute_force": "high",
    "port_scan": "medium",
    "web_attack": "high",
    "priv_esc": "critical",
    "data_exfil": "critical",
    "c2_beacon": "critical",
}

TOKEN_RE = re.compile(r"[a-zA-Z0-9_.:/-]+|[\u0E00-\u0E7F]+")
SQLI_RE = re.compile(r"(union\s+select|drop\s+table|or\s+'?1'?='?1|sleep\()", re.I)
XSS_RE = re.compile(r"(<script|onerror=|document\.cookie)", re.I)

QUERY_EXPANSIONS = {
    "brute_force": ["brute_force", "failed", "password", "login_failed", "ssh"],
    "port_scan": ["port_scan", "drop", "syn", "dpt", "scan"],
    "web_attack": ["web_attack", "sqli", "xss", "union", "select", "script", "payload"],
    "priv_esc": ["priv_esc", "sudo", "root", "shadow", "passwd"],
    "data_exfil": ["data_exfil", "outbound", "bytes", "exfil", "dst"],
    "c2_beacon": ["c2_beacon", "c2", "beacon", "conn", "interval"],
}

QUERY_HINTS = {
    "brute": "brute_force", "bruteforce": "brute_force", "failed login": "brute_force",
    "เดารหัส": "brute_force", "ล็อกอิน": "brute_force",
    "scan": "port_scan", "port": "port_scan", "สแกน": "port_scan",
    "sqli": "web_attack", "sql": "web_attack", "xss": "web_attack",
    "payload": "web_attack", "เว็บ": "web_attack",
    "sudo": "priv_esc", "root": "priv_esc", "privilege": "priv_esc", "ยกระดับ": "priv_esc",
    "exfil": "data_exfil", "outbound": "data_exfil", "ข้อมูลรั่ว": "data_exfil", "ส่งข้อมูล": "data_exfil",
    "c2": "c2_beacon", "beacon": "c2_beacon", "interval": "c2_beacon", "คงที่": "c2_beacon",
}


def tokenize(text: str) -> list[str]:
    return [t.lower() for t in TOKEN_RE.findall(text or "") if len(t) > 1]


def expand_query(query: str) -> tuple[list[str], set[str]]:
    lowered = query.lower()
    categories = {cat for cat in CATEGORY_META if cat in lowered}
    for hint, category in QUERY_HINTS.items():
        if hint in lowered:
            categories.add(category)
    terms = tokenize(query.replace("_", " "))
    for category in categories:
        terms.extend(QUERY_EXPANSIONS.get(category, []))
    return terms, categories


def detect_rule(row: dict[str, Any]) -> dict[str, Any]:
    """Classify a single log row purely from its message + fields.

    Returns the detected category and explanation. Does NOT trust an incoming
    `category` label — that is what we are detecting — but compares against it
    when present so callers can measure accuracy.
    """
    message = str(row.get("message", "") or "")
    category = "benign"
    reasons: list[str] = []

    if "Failed password" in message:
        category = "brute_force"
        reasons.append("พบ failed SSH login")
    if "DROP" in message and "FLAGS=SYN" in message:
        category = "port_scan"
        reasons.append("พบ SYN หลายปลายทาง/หลาย port")
    if SQLI_RE.search(message) or XSS_RE.search(message):
        category = "web_attack"
        reasons.append("พบ SQLi/XSS payload ใน request")
    if "sudo:" in message and "USER=root" in message:
        category = "priv_esc"
        reasons.append("พบ sudo ไปยัง root")
    if "OUTBOUND" in message and "bytes=" in message:
        category = "data_exfil"
        reasons.append("พบ outbound transfer ขนาดใหญ่")
    if "interval=" in message and "CONN" in message:
        category = "c2_beacon"
        reasons.append("พบ connection interval คงที่แบบ beacon")

    meta = CATEGORY_META[category]
    severity = row.get("severity") or CATEGORY_DEFAULT_SEVERITY[category]
    expected = row.get("category")
    return {
        "category": category,
        "label": meta["label"],
        "risk_score": meta["risk"],
        "severity": severity,
        "mitre": meta["mitre"],
        "reasons": reasons or ["ไม่พบ pattern เสี่ยงตาม rule พื้นฐาน"],
        "matches_label": (expected == category) if expected else None,
    }


def row_text(row: dict[str, Any]) -> str:
    parts = [
        row.get("message", ""), row.get("category", ""), row.get("severity", ""),
        row.get("source", ""), row.get("src_ip", ""), row.get("dst_ip", ""),
        row.get("path", ""), row.get("user", ""),
    ]
    return " ".join(str(p) for p in parts if p is not None)