"""
Module 3: Analyzer

Signature-based detector that labels each injection as VULNERABLE or
NOT_VULNERABLE.  Per-type heuristics:

* SQLi: count "First name:" occurrences (DVWA dump signature) and SQL error msgs
* XSS:  check payload reflection in HTML response, dangerous tag/attr patterns
* LFI:  search for /etc/passwd content signatures (root:x:0:0, etc.)

When a pre-trained CSIC Random Forest model is supplied, the analyzer also
attaches an ML confidence score (ml_confidence) and ML label (ml_label) to
each result, giving a second opinion alongside the signature heuristics.
"""

from __future__ import annotations

import json
import logging
from dataclasses import asdict, dataclass, fields
from pathlib import Path
from typing import Any

from intelliscan.config import (
    LFI_SIGNATURES,
    RESULTS_DIR,
    SEVERITY,
    SQL_ERROR_PATTERNS,
    SQLI_DUMP_SIGNATURES,
    XSS_DANGEROUS_PATTERNS,
)

log = logging.getLogger(__name__)


@dataclass
class LabeledResult:
    """An injection result enriched with a vulnerability label and reason."""

    target_url: str
    method: str
    vuln_type: str
    param: str
    payload: str
    status_code: int
    response_len: int
    response_body: str
    label: str  # "VULNERABLE" | "NOT_VULNERABLE"
    reason: str
    severity: str  # "HIGH" | "MEDIUM" | "LOW" | "INFO"
    ml_confidence: float | None = None  # CSIC model confidence score [0–1]
    ml_label: str | None = None  # "ATTACK" | "NORMAL" from CSIC model

    def to_dict(self) -> dict:
        return asdict(self)


# Injection-schema fields accepted from raw entries; anything else (e.g. a
# previous run's label/reason when re-analyzing a labeled file) is dropped
# instead of crashing the LabeledResult constructor.
_INPUT_FIELDS: frozenset[str] = frozenset(f.name for f in fields(LabeledResult)) - {
    "label",
    "reason",
    "severity",
    "ml_confidence",
    "ml_label",
}


class Analyzer:
    """Signature-based vulnerability analyzer with optional CSIC ML scoring."""

    def __init__(
        self,
        results: list[dict],
        csic_model: Any | None = None,
        baselines: dict[str, str] | None = None,
    ) -> None:
        self.raw = results
        self.labeled: list[LabeledResult] = []
        self._csic_model = csic_model  # sklearn RandomForestClassifier or None
        # target_url -> benign baseline response body (lowercased on lookup).
        # Used to subtract static page content from signature counts so form
        # chrome ("Username:"/"Password:" labels) isn't counted as a data dump.
        self._baselines = baselines or {}

    # ── public API ────────────────────────────────────────────────────────
    def run(self) -> list[LabeledResult]:
        for entry in self.raw:
            label, reason = self._classify(entry)
            severity = SEVERITY.get(entry["vuln_type"], "INFO") if label == "VULNERABLE" else "INFO"
            self.labeled.append(
                LabeledResult(
                    **{k: v for k, v in entry.items() if k in _INPUT_FIELDS},
                    label=label,
                    reason=reason,
                    severity=severity,
                )
            )

        if self._csic_model is not None:
            self._attach_ml_scores()

        n_vuln = sum(1 for r in self.labeled if r.label == "VULNERABLE")
        log.info(
            "Analyzed %d entries: %d vulnerable, %d not vulnerable%s",
            len(self.labeled),
            n_vuln,
            len(self.labeled) - n_vuln,
            " (CSIC ML scores attached)" if self._csic_model else "",
        )
        return self.labeled

    # ── CSIC ML scoring ───────────────────────────────────────────────────
    def _attach_ml_scores(self) -> None:
        """Score every labeled result with the CSIC model in one batched call."""
        import numpy as np

        from intelliscan.modules.csic_trainer import features_from_request

        model = self._csic_model
        if model is None or not self.labeled:
            return
        try:
            X = np.array(
                [
                    features_from_request(
                        url=r.target_url,
                        content=r.payload,
                        method=r.method,
                    )
                    for r in self.labeled
                ],
                dtype=np.float32,
            )
            expected = getattr(model, "n_features_in_", X.shape[1])
            if X.shape[1] != expected:
                log.warning(
                    "CSIC model expects %d features but extractor produced %d; "
                    "skipping ML scoring — retrain with `intelliscan train-csic`",
                    expected,
                    X.shape[1],
                )
                return
            preds = model.predict(X)
            probas = model.predict_proba(X)
            for r, pred, proba in zip(self.labeled, preds, probas, strict=True):
                r.ml_label = "ATTACK" if int(pred) == 1 else "NORMAL"
                r.ml_confidence = float(proba[int(pred)])
        except Exception as exc:  # noqa: BLE001
            log.warning("CSIC ML scoring failed: %s", exc)

    # ── per-type detectors ────────────────────────────────────────────────
    def _classify(self, entry: dict) -> tuple[str, str]:
        body = entry.get("response_body", "").lower()
        payload = entry.get("payload", "")
        vt = entry.get("vuln_type", "").lower()
        baseline = self._baselines.get(entry.get("target_url", ""), "").lower()

        if vt == "sqli":
            return self._detect_sqli(body, payload, baseline)
        if vt.startswith("xss"):
            return self._detect_xss(body, payload, baseline)
        if vt == "lfi":
            return self._detect_lfi(body, payload)
        return "NOT_VULNERABLE", "Unknown vulnerability type"

    @staticmethod
    def _detect_sqli(body: str, payload: str, baseline: str = "") -> tuple[str, str]:
        # SQL error messages (multi-dialect: MySQL, PostgreSQL, MSSQL, Oracle, SQLite).
        # Only count errors the benign baseline did NOT already contain, so a page
        # that always shows an error string isn't flagged for every payload.
        for pattern in SQL_ERROR_PATTERNS:
            if pattern in body and pattern not in baseline:
                return "VULNERABLE", f"SQL error: {pattern!r}"

        # Data-dump indicators: repeated column-name labels suggest row exfiltration.
        # Subtract two sources of noise before counting:
        #   1. the benign baseline (static form labels like "Username:"), and
        #   2. the payload's OWN reflected text — a page that echoes the payload
        #      back (e.g. the stored-XSS guestbook reflecting a UNION payload that
        #      literally contains "information_schema") must not look like a dump.
        # Require at least 2 NET occurrences beyond both.
        payload_lower = payload.lower()
        reflections = body.count(payload_lower) if payload_lower else 0
        for sig in SQLI_DUMP_SIGNATURES:
            net = body.count(sig) - baseline.count(sig) - reflections * payload_lower.count(sig)
            if net >= 2:
                return "VULNERABLE", f"Data dump — {sig!r} appears {net} more time(s) than baseline"

        return "NOT_VULNERABLE", "No SQL error or data-dump signature beyond baseline"

    @staticmethod
    def _detect_xss(body: str, payload: str, baseline: str = "") -> tuple[str, str]:
        payload_lower = payload.lower()

        # Full payload reflection (strongest signal for reflected XSS). Ignore if the
        # exact payload already appears in the benign baseline (static page content).
        if payload_lower in body and payload_lower not in baseline:
            return "VULNERABLE", "Payload reflected verbatim in response"

        # Partial match: a dangerous XSS pattern present in both payload and body, and
        # not already in the baseline page.
        for pattern in XSS_DANGEROUS_PATTERNS:
            if pattern in payload_lower and pattern in body and pattern not in baseline:
                return "VULNERABLE", f"Dangerous XSS pattern reflected: {pattern!r}"

        return "NOT_VULNERABLE", "Payload not reflected in response"

    @staticmethod
    def _detect_lfi(body: str, payload: str) -> tuple[str, str]:
        for sig in LFI_SIGNATURES:
            if sig in body:
                return "VULNERABLE", f"File content signature: {sig!r}"
        return "NOT_VULNERABLE", "No file inclusion signature"

    # ── persistence ───────────────────────────────────────────────────────
    def save(self, path: str | Path) -> Path:
        path = Path(path)
        path.write_text(
            json.dumps([r.to_dict() for r in self.labeled], indent=2),
            encoding="utf-8",
        )
        return path

    # ── stats ─────────────────────────────────────────────────────────────
    def stats_by_type(self) -> dict[str, dict[str, int]]:
        out: dict[str, dict[str, int]] = {}
        for r in self.labeled:
            d = out.setdefault(r.vuln_type, {"total": 0, "vulnerable": 0, "not_vulnerable": 0})
            d["total"] += 1
            d["vulnerable" if r.label == "VULNERABLE" else "not_vulnerable"] += 1
        return out


# ── CLI helper ────────────────────────────────────────────────────────────
def main(input_file: str = "injection_results.json", output: str = "labeled_results.json") -> None:
    input_path = RESULTS_DIR / input_file
    raw = json.loads(input_path.read_text())
    analyzer = Analyzer(raw)
    analyzer.run()
    analyzer.save(RESULTS_DIR / output)
    print(json.dumps(analyzer.stats_by_type(), indent=2))


if __name__ == "__main__":
    import argparse

    p = argparse.ArgumentParser(description="IntelliScan response analyzer")
    p.add_argument("--input", default="injection_results.json")
    p.add_argument("--output", default="labeled_results.json")
    args = p.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")
    main(args.input, args.output)