# SPDX-License-Identifier: GPL-3.0-or-later
"""Incident grouping — collapse related alerts into one incident.

A single intrusion usually trips several signatures: a web service spawns a
shell (`reverse_shell`), the shell opens a port (`new_listener`) and writes
persistence (`persistence`). Reported as three unrelated alerts that is noise;
reported as one incident it is a story.

Grouping is **process-lineage first, time second**. Each alert batch carries the
PIDs it flagged; we walk their ancestry (`/proc` PPid chain) into a *lineage
set*. Two batches join the same incident when their lineage sets intersect —
i.e. they share a process or a common ancestor (the web server, the SSH session).
Batches with no live PID (FIM drift, package tamper, hidden modules) can't be
placed by lineage, so they fall back to joining the most recently active open
incident within a time window. Distinct process trees stay distinct incidents.

The incident index is a JSON file in ``log_dir`` updated as snapshots are
captured; ``incident_id`` is added to each snapshot report (the per-alert JSON
schema is untouched). The matching/lineage cores are pure functions so they test
without ``/proc`` or a daemon.
"""
from __future__ import annotations

import json
import os
import threading
import time
from collections.abc import Callable, Iterable
from datetime import datetime

from .alert import Alert, Severity
from .config import Config

INDEX_NAME = "incidents.json"
MAX_INCIDENTS = 1000          # trim the index to the most recent N on save
_DEFAULT_EXCLUDE = frozenset({0, 1})   # never correlate via pid 0 or init

# The index read-modify-write happens from capture threads (sweep + eBPF), so
# serialize it within this process.
_LOCK = threading.RLock()


# --- pure cores (tested without /proc or a daemon) ------------------------

def lineage_of(
    pids: Iterable[int],
    parent_of: Callable[[int], int],
    depth: int = 8,
    exclude: frozenset[int] = _DEFAULT_EXCLUDE,
) -> set[int]:
    """Flagged PIDs plus their ancestors, walking PPid via ``parent_of``.

    ``parent_of(pid)`` returns the parent PID (0 when unknown/root). Stops at
    ``exclude`` (pid 0/1 by default, so everything doesn't correlate through
    init) and at ``depth`` ancestors.
    """
    seen: set[int] = set()
    for start in pids:
        cur = start
        steps = 0
        while cur and cur not in exclude and steps <= depth:
            if cur in seen:
                break
            seen.add(cur)
            cur = parent_of(cur)
            steps += 1
    return seen


def _sev_rank(name: str) -> int:
    try:
        return int(Severity[name])
    except KeyError:
        return 0


def max_severity(a: str, b: str) -> str:
    return a if _sev_rank(a) >= _sev_rank(b) else b


def assign(index: dict, lineage: set[int], when: float, host: str,
           window: int) -> str | None:
    """Return the id of an open incident this batch belongs to, or None.

    Lineage-bearing batches join an open incident whose lineage they intersect.
    PID-less batches (empty lineage) join the most recently active open incident.
    "Open" means last activity is within ``window`` seconds. Ties and PID-less
    matches resolve to the most recently active incident.
    """
    best: str | None = None
    best_ts = -1.0
    for inc in index.values():
        if inc.get("host") != host:
            continue
        if (when - inc.get("last_ts", 0.0)) > window:
            continue  # closed
        if lineage:
            if not (set(inc.get("lineage", ())) & lineage):
                continue
        # else: PID-less batch — eligible to join any open incident on this host
        if inc["last_ts"] > best_ts:
            best, best_ts = inc["id"], inc["last_ts"]
    return best


# --- index persistence ----------------------------------------------------

def index_path(cfg: Config):
    return cfg.log_dir / INDEX_NAME


def load_index(cfg: Config) -> dict:
    try:
        data = json.loads(index_path(cfg).read_text())
        return data if isinstance(data, dict) else {}
    except (OSError, ValueError):
        return {}


def save_index(cfg: Config, index: dict) -> None:
    # Keep only the most recent incidents so the file stays bounded.
    if len(index) > MAX_INCIDENTS:
        keep = sorted(index.values(), key=lambda i: i.get("last_ts", 0.0),
                      reverse=True)[:MAX_INCIDENTS]
        index = {i["id"]: i for i in keep}
    tmp = index_path(cfg).with_suffix(".json.tmp")
    tmp.write_text(json.dumps(index, indent=2))
    os.replace(tmp, index_path(cfg))
    try:
        index_path(cfg).chmod(0o640)
    except OSError:
        pass


def _new_id(when: float) -> str:
    stamp = datetime.fromtimestamp(when).astimezone().strftime("%Y%m%d-%H%M%S")
    return f"inc-{stamp}-{os.urandom(2).hex()}"


# --- the recorder (called from snapshot.capture) --------------------------

def record(cfg: Config, snapshot_name: str, alerts: list[Alert],
           lineage: set[int], when: float, host: str) -> str | None:
    """Attach this snapshot to an incident (new or existing); return its id.

    Best-effort: any failure returns None so a snapshot is never lost to an
    incident-index problem.
    """
    if not getattr(cfg, "incident_tracking", True):
        return None
    try:
        with _LOCK:
            index = load_index(cfg)
            iid = assign(index, lineage, when, host, cfg.incident_window)
            iso = datetime.fromtimestamp(when).astimezone().isoformat()
            severity = str(max((a.severity for a in alerts), default=Severity.HIGH))
            sigs = list(dict.fromkeys(a.signature for a in alerts))
            sids = list(dict.fromkeys(a.sid for a in alerts if a.sid))
            pids = sorted({p for a in alerts for p in a.pids})

            if iid is None:
                iid = _new_id(when)
                index[iid] = {
                    "id": iid, "host": host,
                    "first_ts": when, "last_ts": when,
                    "first_seen": iso, "last_seen": iso,
                    "severity": severity,
                    "signatures": [], "sids": [], "pids": [],
                    "lineage": [], "snapshots": [], "alert_count": 0,
                }
            inc = index[iid]
            inc["last_ts"] = when
            inc["last_seen"] = iso
            inc["severity"] = max_severity(inc.get("severity", severity), severity)
            inc["signatures"] = list(dict.fromkeys(inc["signatures"] + sigs))
            inc["sids"] = list(dict.fromkeys(inc["sids"] + sids))
            inc["pids"] = sorted(set(inc["pids"]) | set(pids))
            inc["lineage"] = sorted(set(inc["lineage"]) | lineage)
            inc["snapshots"].append(snapshot_name)
            inc["alert_count"] += len(alerts)
            save_index(cfg, index)
            return iid
    except Exception:
        return None


def lineage_from_state(pids: Iterable[int], state, cfg: Config) -> set[int]:
    """Compute a lineage set for ``pids`` using a live SystemState for PPids."""
    def parent_of(pid: int) -> int:
        proc = state.process(pid)
        return proc.ppid if proc else 0
    return lineage_of(pids, parent_of,
                      depth=getattr(cfg, "incident_lineage_depth", 8))