# SPDX-License-Identifier: GPL-3.0-or-later
"""A snapshot of live system state, gathered once per detector sweep.

Detectors read from a single ``SystemState`` instance instead of each shelling
out, which is what keeps a sweep cheap. Everything here is also injectable:
tests construct a ``SystemState`` with fake processes/sockets and never touch
the real ``/proc``.
"""
from __future__ import annotations

import os
import re
import stat
import subprocess
from dataclasses import dataclass, field
from functools import cached_property
from pathlib import Path

_INO_RE = re.compile(r"\bino:(\d+)")
_USER_RE = re.compile(r'users:\(\("([^"]+)",pid=(\d+),fd=\d+')


@dataclass
class Process:
    """Lazily-read view of a process via ``/proc/<pid>``."""

    pid: int

    def _read(self, name: str) -> str:
        try:
            with open(f"/proc/{self.pid}/{name}", "rb") as fh:
                return fh.read().decode("utf-8", "replace")
        except OSError:
            return ""

    @cached_property
    def comm(self) -> str:
        return self._read("comm").strip()

    @cached_property
    def cmdline(self) -> str:
        return self._read("cmdline").replace("\0", " ").strip()

    @cached_property
    def exe(self) -> str:
        try:
            return os.readlink(f"/proc/{self.pid}/exe")
        except OSError:
            return ""

    @cached_property
    def cwd(self) -> str:
        try:
            return os.readlink(f"/proc/{self.pid}/cwd")
        except OSError:
            return ""

    @cached_property
    def environ(self) -> dict[str, str]:
        raw = self._read("environ")
        out: dict[str, str] = {}
        for item in raw.split("\0"):
            if "=" in item:
                k, _, v = item.partition("=")
                out[k] = v
        return out

    @cached_property
    def status(self) -> dict[str, str]:
        out: dict[str, str] = {}
        for line in self._read("status").splitlines():
            k, _, v = line.partition(":")
            out[k] = v.strip()
        return out

    @property
    def ppid(self) -> int:
        try:
            return int(self.status.get("PPid", "0"))
        except ValueError:
            return 0

    @property
    def uid(self) -> int:
        field0 = self.status.get("Uid", "0").split()
        return int(field0[0]) if field0 else 0

    def stdio_socket_inode(self) -> int | None:
        """Socket inode if fd 0, 1, or 2 is a socket — the reverse-shell tell."""
        for fd in (0, 1, 2):
            try:
                target = os.readlink(f"/proc/{self.pid}/fd/{fd}")
            except OSError:
                continue
            if target.startswith("socket:["):
                return int(target[8:-1])
        return None


@dataclass(frozen=True)
class Socket:
    state: str
    local: str
    peer: str
    inode: int | None
    comm: str
    pid: int | None


class SystemState:
    """One sweep's worth of cached system state.

    Construct empty for the real system; pass ``processes``/``sockets`` to
    inject fakes in tests.
    """

    def __init__(
        self,
        processes: list[Process] | None = None,
        sockets: list[Socket] | None = None,
        *,
        listener_baseline: set[str] | None = None,
        suid_baseline: set[str] | None = None,
        suid_binaries: list[str] | None = None,
        persist_since: float | None = None,
    ) -> None:
        self._injected_procs = processes
        self._injected_socks = sockets
        # Context the daemon supplies for baseline-diffing detectors. Left as
        # None when not applicable (e.g. before the grace window, or in tests),
        # in which case those detectors yield nothing.
        self.listener_baseline = listener_baseline
        self.suid_baseline = suid_baseline
        self.suid_binaries = suid_binaries
        self.persist_since = persist_since

    def listener_keys(self) -> set[str]:
        """Current listening sockets as ``port/comm`` identity keys."""
        keys = set()
        for s in self.listening_sockets():
            _host, port = s.local.rpartition(":")[0], s.local.rpartition(":")[2]
            keys.add(f"{port}/{s.comm or '?'}")
        return keys

    # -- processes ---------------------------------------------------------
    @cached_property
    def processes(self) -> list[Process]:
        if self._injected_procs is not None:
            return self._injected_procs
        procs = []
        for name in os.listdir("/proc"):
            if name.isdigit():
                procs.append(Process(int(name)))
        return procs

    @cached_property
    def by_pid(self) -> dict[int, Process]:
        return {p.pid: p for p in self.processes}

    def process(self, pid: int) -> Process | None:
        return self.by_pid.get(pid)

    # -- sockets -----------------------------------------------------------
    @cached_property
    def sockets(self) -> list[Socket]:
        if self._injected_socks is not None:
            return self._injected_socks
        out: list[Socket] = []
        for args in (["-tanep"], ["-uanep"]):
            out.extend(self._parse_ss(self._run_ss(args)))
        return out

    @staticmethod
    def _run_ss(extra: list[str]) -> str:
        try:
            res = subprocess.run(
                ["ss", "-H", *extra],
                capture_output=True, text=True, timeout=10,
            )
            return res.stdout
        except (OSError, subprocess.SubprocessError):
            return ""

    @staticmethod
    def _parse_ss(text: str) -> list[Socket]:
        socks: list[Socket] = []
        for line in text.splitlines():
            parts = line.split()
            if len(parts) < 5:
                continue
            state, _rq, _sq, local, peer = parts[:5]
            rest = line
            ino_m = _INO_RE.search(rest)
            usr_m = _USER_RE.search(rest)
            socks.append(Socket(
                state=state,
                local=local,
                peer=peer,
                inode=int(ino_m.group(1)) if ino_m else None,
                comm=usr_m.group(1) if usr_m else "",
                pid=int(usr_m.group(2)) if usr_m else None,
            ))
        return socks

    @cached_property
    def net_peer_by_inode(self) -> dict[int, str]:
        """Map socket inode -> "local -> peer" for every network socket."""
        out: dict[int, str] = {}
        for s in self.sockets:
            if s.inode is not None:
                out[s.inode] = f"{s.local} -> {s.peer}"
        return out

    def listening_sockets(self) -> list[Socket]:
        return [s for s in self.sockets if s.state == "LISTEN"]

    def established_sockets(self) -> list[Socket]:
        return [s for s in self.sockets if s.state == "ESTAB"]


# -- filesystem-wide SUID scan (expensive; the daemon gates how often) ------

_SETID = stat.S_ISUID | stat.S_ISGID


def _walk_suid(root: str, *, stay_on_device: bool) -> list[str]:
    # os.scandir is much faster than os.walk + per-file lstat because DirEntry
    # caches the stat from the directory read and is_dir() often avoids a stat.
    try:
        root_dev = os.stat(root).st_dev
    except OSError:
        return []
    found: list[str] = []
    stack = [root]
    while stack:
        try:
            it = os.scandir(stack.pop())
        except OSError:
            continue
        with it:
            for entry in it:
                try:
                    if entry.is_dir(follow_symlinks=False):
                        if stay_on_device:
                            if entry.stat(follow_symlinks=False).st_dev != root_dev:
                                continue
                        stack.append(entry.path)
                    elif entry.is_file(follow_symlinks=False):
                        if entry.stat(follow_symlinks=False).st_mode & _SETID:
                            found.append(entry.path)
                except OSError:
                    continue
    return found


def scan_suid_binaries(
    root: str = "/", extra_dirs: tuple[str, ...] = ()
) -> list[str]:
    """Find SUID/SGID regular files.

    Walks ``root`` staying on its device (the ``find -xdev`` equivalent, so it
    doesn't wander into mounted media), then *also* walks each of ``extra_dirs``
    regardless of device — those are the small writable mounts (``/tmp``,
    ``/dev/shm`` …) which are usually a separate tmpfs yet are exactly where a
    malicious SUID binary gets dropped.
    """
    found = set(_walk_suid(root, stay_on_device=True))
    for d in extra_dirs:
        if os.path.isdir(d):
            found.update(_walk_suid(d, stay_on_device=False))
    return sorted(found)