#!/usr/bin/env python3
"""
security_overhead_monitor.py

Measure direct CPU, RAM, and disk-I/O usage of common hosting/security agents
over a defined observation window using Linux /proc only (no third-party deps).

Run as root for the most complete PSS and I/O accounting.
"""

from __future__ import annotations

import argparse
import csv
import math
import os
import re
import signal
import statistics
import sys
import time
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

CLK_TCK = os.sysconf(os.sysconf_names["SC_CLK_TCK"])
PAGE_SIZE = os.sysconf("SC_PAGE_SIZE")
NCPU = os.cpu_count() or 1


@dataclass(frozen=True)
class ProductRule:
    name: str
    patterns: Tuple[str, ...]
    security_total: bool = True


# Ordered: first match wins. Put specific products before generic audit/AV rules.
RULES: Tuple[ProductRule, ...] = (
    ProductRule("Imunify360", (
        r"imunify",
        r"\bim360:",
    )),
    ProductRule("Rustbolit", (
        r"rustbolit",
    )),
    ProductRule("Acronis", (
        r"/acronis/",
        r"\bacronis\b",
    )),
    ProductRule("OSSEC", (
        r"/var/ossec/",
        r"\bossec(?:-|_)"
    )),
    ProductRule("Linux audit", (
        r"(?:^|/)auditd(?:\s|$)",
        r"\bkauditd\b",
        r"\baudit_prune_tre\b",
    )),
    ProductRule("ClamAV", (
        r"(?:^|/)clamd(?:\s|$)",
        r"\bclamav\b",
    )),

    # Security-adjacent/provider overhead found in the supplied process list.
    # These are measured but deliberately excluded from SECURITY TOTAL.
    ProductRule("WebPros 360 monitoring", (
        r"/agent360(?:\s|$)",
    ), security_total=False),
    ProductRule("cPanel p0f", (
        r"(?:^|/)p0f(?:\s|$)",
        r"cpanelconnecttrack",
    ), security_total=False),
)

COMPILED_RULES = tuple(
    (rule, tuple(re.compile(p, re.IGNORECASE) for p in rule.patterns))
    for rule in RULES
)


@dataclass
class Proc:
    pid: int
    ppid: int
    start_ticks: int
    cpu_ticks: int
    rss_bytes: int
    pss_bytes: Optional[int]
    read_bytes: Optional[int]
    write_bytes: Optional[int]
    comm: str
    cmdline: str
    exe: str
    product: Optional[str] = None
    direct_match: bool = False

    @property
    def identity(self) -> Tuple[int, int]:
        # PID + process start time protects against PID reuse.
        return (self.pid, self.start_ticks)

    @property
    def display_command(self) -> str:
        if self.cmdline:
            return self.cmdline
        return f"[{self.comm}]"

    @property
    def signature(self) -> str:
        if self.exe:
            return self.exe
        if self.cmdline:
            return self.cmdline.split()[0]
        return f"[{self.comm}]"


def parse_duration(text: str) -> Optional[float]:
    text = text.strip().lower()
    if text in {"0", "forever", "inf", "infinite"}:
        return None
    m = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)\s*([smhd]?)", text)
    if not m:
        raise argparse.ArgumentTypeError(
            "duration must look like 30s, 15m, 24h, 2d, or 0 for until Ctrl-C"
        )
    value = float(m.group(1))
    unit = m.group(2) or "s"
    factor = {"s": 1, "m": 60, "h": 3600, "d": 86400}[unit]
    return value * factor


def safe_read_text(path: Path) -> Optional[str]:
    try:
        return path.read_text(errors="replace")
    except (FileNotFoundError, ProcessLookupError, PermissionError, OSError):
        return None


def read_cmdline(procdir: Path) -> str:
    try:
        raw = (procdir / "cmdline").read_bytes()
        if not raw:
            return ""
        return raw.replace(b"\0", b" ").decode(errors="replace").strip()
    except (FileNotFoundError, ProcessLookupError, PermissionError, OSError):
        return ""


def read_exe(procdir: Path) -> str:
    try:
        return os.readlink(procdir / "exe")
    except (FileNotFoundError, ProcessLookupError, PermissionError, OSError):
        return ""


def read_pss_bytes(procdir: Path) -> Optional[int]:
    text = safe_read_text(procdir / "smaps_rollup")
    if text is None:
        return None
    for line in text.splitlines():
        if line.startswith("Pss:"):
            try:
                return int(line.split()[1]) * 1024
            except (IndexError, ValueError):
                return None
    return None


def read_io(procdir: Path) -> Tuple[Optional[int], Optional[int]]:
    text = safe_read_text(procdir / "io")
    if text is None:
        return None, None
    values = {}
    for line in text.splitlines():
        if ":" not in line:
            continue
        k, v = line.split(":", 1)
        try:
            values[k.strip()] = int(v.strip())
        except ValueError:
            pass
    return values.get("read_bytes"), values.get("write_bytes")


def read_proc(pid: int, with_pss: bool = True) -> Optional[Proc]:
    procdir = Path("/proc") / str(pid)
    stat = safe_read_text(procdir / "stat")
    if stat is None:
        return None

    # comm can contain spaces and parentheses, so split after the final ')'.
    rparen = stat.rfind(")")
    lparen = stat.find("(")
    if lparen < 0 or rparen < 0 or rparen <= lparen:
        return None

    comm = stat[lparen + 1:rparen]
    rest = stat[rparen + 2:].split()
    try:
        ppid = int(rest[1])          # field 4
        utime = int(rest[11])        # field 14
        stime = int(rest[12])        # field 15
        start_ticks = int(rest[19])  # field 22
        rss_pages = int(rest[21])    # field 24
    except (IndexError, ValueError):
        return None

    cmdline = read_cmdline(procdir)
    exe = read_exe(procdir)
    pss = read_pss_bytes(procdir) if with_pss else None
    rb, wb = read_io(procdir)

    return Proc(
        pid=pid,
        ppid=ppid,
        start_ticks=start_ticks,
        cpu_ticks=utime + stime,
        rss_bytes=max(0, rss_pages) * PAGE_SIZE,
        pss_bytes=pss,
        read_bytes=rb,
        write_bytes=wb,
        comm=comm,
        cmdline=cmdline,
        exe=exe,
    )


def classify_direct(proc: Proc) -> Optional[str]:
    haystack = " ".join(x for x in (proc.exe, proc.cmdline, proc.comm) if x)
    for rule, regexes in COMPILED_RULES:
        if any(rx.search(haystack) for rx in regexes):
            return rule.name
    return None


def scan_processes(with_pss: bool = True) -> Dict[int, Proc]:
    procs: Dict[int, Proc] = {}
    try:
        entries = os.scandir("/proc")
    except OSError:
        return procs

    with entries:
        for entry in entries:
            if not entry.name.isdigit():
                continue
            p = read_proc(int(entry.name), with_pss=with_pss)
            if p is None:
                continue
            product = classify_direct(p)
            if product:
                p.product = product
                p.direct_match = True
            procs[p.pid] = p

    # Attribute descendants to a matched security product even when a child has
    # a generic executable name (shell, python, sleep, etc.).
    changed = True
    while changed:
        changed = False
        for p in procs.values():
            if p.product is not None:
                continue
            parent = procs.get(p.ppid)
            if parent and parent.product:
                p.product = parent.product
                p.direct_match = False
                changed = True

    return procs


def read_host_cpu() -> Optional[Tuple[int, int]]:
    text = safe_read_text(Path("/proc/stat"))
    if text is None:
        return None
    first = text.splitlines()[0].split()
    if not first or first[0] != "cpu":
        return None
    try:
        vals = [int(x) for x in first[1:]]
    except ValueError:
        return None
    total = sum(vals)
    idle = (vals[3] if len(vals) > 3 else 0) + (vals[4] if len(vals) > 4 else 0)
    return total, idle


def percentile(values: List[float], p: float) -> float:
    if not values:
        return 0.0
    s = sorted(values)
    if len(s) == 1:
        return s[0]
    pos = (len(s) - 1) * p
    lo = math.floor(pos)
    hi = math.ceil(pos)
    if lo == hi:
        return s[lo]
    frac = pos - lo
    return s[lo] * (1 - frac) + s[hi] * frac


def mib(n: float) -> float:
    return n / (1024 * 1024)


def mb(n: float) -> float:
    return n / 1_000_000


def product_is_security(name: str) -> bool:
    for rule in RULES:
        if rule.name == name:
            return rule.security_total
    return False


def human_seconds(seconds: float) -> str:
    if seconds >= 86400:
        return f"{seconds / 86400:.2f} d"
    if seconds >= 3600:
        return f"{seconds / 3600:.2f} h"
    if seconds >= 60:
        return f"{seconds / 60:.2f} min"
    return f"{seconds:.1f} s"


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Measure CPU/RAM/I/O used by hosting security agents over time."
    )
    parser.add_argument(
        "--duration", type=parse_duration, default=parse_duration("24h"),
        help="observation window: 30s, 15m, 24h, 2d; 0 = until Ctrl-C (default: 24h)",
    )
    parser.add_argument(
        "--interval", type=float, default=2.0,
        help="sample interval in seconds (default: 2)",
    )
    parser.add_argument(
        "--output", default=None,
        help="output directory (default: security-overhead-YYYYmmdd-HHMMSS)",
    )
    parser.add_argument(
        "--no-pss", action="store_true",
        help="skip smaps_rollup/PSS reads; faster but RAM totals use RSS only",
    )
    args = parser.parse_args()

    if args.interval <= 0:
        parser.error("--interval must be > 0")

    started_wall = datetime.now().astimezone()
    outdir = Path(args.output or started_wall.strftime("security-overhead-%Y%m%d-%H%M%S"))
    outdir.mkdir(parents=True, exist_ok=True)
    samples_path = outdir / "samples.csv"
    summary_path = outdir / "summary.txt"
    process_path = outdir / "process-summary.csv"

    stop = False

    def on_signal(signum, frame):
        nonlocal stop
        stop = True

    signal.signal(signal.SIGINT, on_signal)
    signal.signal(signal.SIGTERM, on_signal)

    product_names = [r.name for r in RULES]
    cpu_seconds = defaultdict(float)
    io_read = defaultdict(int)
    io_write = defaultdict(int)
    pss_samples: Dict[str, List[float]] = defaultdict(list)
    rss_samples: Dict[str, List[float]] = defaultdict(list)
    cpu_pct_one_core_samples: Dict[str, List[float]] = defaultdict(list)
    process_count_samples: Dict[str, List[int]] = defaultdict(list)

    # Per executable/signature evidence.
    proc_cpu_seconds = defaultdict(float)  # (product, signature)
    proc_pss_sum_mib = defaultdict(float)
    proc_pss_count = defaultdict(int)
    proc_pss_peak_mib = defaultdict(float)
    proc_rss_sum_mib = defaultdict(float)
    proc_rss_count = defaultdict(int)
    proc_rss_peak_mib = defaultdict(float)
    proc_last_cmd = {}

    previous: Dict[Tuple[int, int], Proc] = {}
    prev_host_cpu = read_host_cpu()
    first_mono = time.monotonic()
    prev_mono = first_mono
    sample_no = 0
    pss_permission_failures = 0

    print(f"Output:   {outdir}")
    print(f"CPUs:     {NCPU}")
    print(f"Interval: {args.interval:g}s")
    print(f"Duration: {'until Ctrl-C' if args.duration is None else human_seconds(args.duration)}")
    print("Run as root for complete PSS/I/O accounting. Ctrl-C is safe; a summary is still written.")
    print()

    with samples_path.open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow([
            "timestamp", "elapsed_s", "group", "security_total_member",
            "processes", "cpu_pct_one_core", "cpu_pct_machine",
            "cpu_seconds_this_sample", "pss_mib", "rss_mib",
            "read_mb_this_sample", "write_mb_this_sample", "host_busy_pct",
        ])

        # Establish baseline. CPU deltas become meaningful on the next sample.
        baseline = scan_processes(with_pss=not args.no_pss)
        previous = {p.identity: p for p in baseline.values() if p.product}
        prev_mono = time.monotonic()
        prev_host_cpu = read_host_cpu()

        while not stop:
            target = prev_mono + args.interval
            now = time.monotonic()
            if target > now:
                time.sleep(target - now)

            now_mono = time.monotonic()
            elapsed = now_mono - prev_mono
            total_elapsed = now_mono - first_mono
            if elapsed <= 0:
                continue

            current_all = scan_processes(with_pss=not args.no_pss)
            current = {p.identity: p for p in current_all.values() if p.product}
            current_host_cpu = read_host_cpu()

            host_busy_pct = None
            if prev_host_cpu and current_host_cpu:
                dtot = current_host_cpu[0] - prev_host_cpu[0]
                didle = current_host_cpu[1] - prev_host_cpu[1]
                if dtot > 0:
                    host_busy_pct = 100.0 * max(0, dtot - didle) / dtot

            per_group = {
                name: {
                    "cpu_sec": 0.0,
                    "rss": 0,
                    "pss": 0,
                    "pss_known": 0,
                    "count": 0,
                    "read": 0,
                    "write": 0,
                }
                for name in product_names
            }

            for ident, p in current.items():
                assert p.product is not None
                g = per_group[p.product]
                g["count"] += 1
                g["rss"] += p.rss_bytes
                if p.pss_bytes is not None:
                    g["pss"] += p.pss_bytes
                    g["pss_known"] += 1
                elif not args.no_pss:
                    pss_permission_failures += 1

                sigkey = (p.product, p.signature)
                proc_last_cmd[sigkey] = p.display_command
                proc_rss_sum_mib[sigkey] += mib(p.rss_bytes)
                proc_rss_count[sigkey] += 1
                proc_rss_peak_mib[sigkey] = max(proc_rss_peak_mib[sigkey], mib(p.rss_bytes))
                if p.pss_bytes is not None:
                    proc_pss_sum_mib[sigkey] += mib(p.pss_bytes)
                    proc_pss_count[sigkey] += 1
                    proc_pss_peak_mib[sigkey] = max(proc_pss_peak_mib[sigkey], mib(p.pss_bytes))

                old = previous.get(ident)
                if old is None:
                    # New process: no prior counter, so don't invent CPU/I/O usage.
                    continue

                dticks = p.cpu_ticks - old.cpu_ticks
                if dticks >= 0:
                    dcpu = dticks / CLK_TCK
                    g["cpu_sec"] += dcpu
                    proc_cpu_seconds[sigkey] += dcpu

                if p.read_bytes is not None and old.read_bytes is not None:
                    d = p.read_bytes - old.read_bytes
                    if d >= 0:
                        g["read"] += d
                if p.write_bytes is not None and old.write_bytes is not None:
                    d = p.write_bytes - old.write_bytes
                    if d >= 0:
                        g["write"] += d

            timestamp = datetime.now().astimezone().isoformat(timespec="seconds")

            security_tot = {k: 0 for k in ("cpu_sec", "rss", "pss", "pss_known", "count", "read", "write")}
            adjacent_tot = {k: 0 for k in ("cpu_sec", "rss", "pss", "pss_known", "count", "read", "write")}

            for name in product_names:
                g = per_group[name]
                cpu_one_core_pct = 100.0 * g["cpu_sec"] / elapsed
                cpu_machine_pct = cpu_one_core_pct / NCPU
                pss_mib = mib(g["pss"]) if g["pss_known"] else None
                rss_mib = mib(g["rss"])

                cpu_seconds[name] += g["cpu_sec"]
                io_read[name] += g["read"]
                io_write[name] += g["write"]
                cpu_pct_one_core_samples[name].append(cpu_one_core_pct)
                rss_samples[name].append(rss_mib)
                if pss_mib is not None:
                    pss_samples[name].append(pss_mib)
                process_count_samples[name].append(g["count"])

                dest = security_tot if product_is_security(name) else adjacent_tot
                for k in dest:
                    dest[k] += g[k]

                writer.writerow([
                    timestamp, f"{total_elapsed:.3f}", name,
                    "yes" if product_is_security(name) else "no",
                    g["count"], f"{cpu_one_core_pct:.4f}", f"{cpu_machine_pct:.6f}",
                    f"{g['cpu_sec']:.6f}",
                    "" if pss_mib is None else f"{pss_mib:.3f}",
                    f"{rss_mib:.3f}", f"{mb(g['read']):.6f}", f"{mb(g['write']):.6f}",
                    "" if host_busy_pct is None else f"{host_busy_pct:.4f}",
                ])

            for total_name, total, is_security in (
                ("SECURITY TOTAL", security_tot, True),
                ("ADJACENT TOTAL", adjacent_tot, False),
            ):
                cpu_one_core_pct = 100.0 * total["cpu_sec"] / elapsed
                cpu_machine_pct = cpu_one_core_pct / NCPU
                pss_mib = mib(total["pss"]) if total["pss_known"] else None
                rss_mib = mib(total["rss"])

                cpu_seconds[total_name] += total["cpu_sec"]
                io_read[total_name] += total["read"]
                io_write[total_name] += total["write"]
                cpu_pct_one_core_samples[total_name].append(cpu_one_core_pct)
                rss_samples[total_name].append(rss_mib)
                if pss_mib is not None:
                    pss_samples[total_name].append(pss_mib)
                process_count_samples[total_name].append(total["count"])

                writer.writerow([
                    timestamp, f"{total_elapsed:.3f}", total_name,
                    "yes" if is_security else "no",
                    total["count"], f"{cpu_one_core_pct:.4f}", f"{cpu_machine_pct:.6f}",
                    f"{total['cpu_sec']:.6f}",
                    "" if pss_mib is None else f"{pss_mib:.3f}",
                    f"{rss_mib:.3f}", f"{mb(total['read']):.6f}", f"{mb(total['write']):.6f}",
                    "" if host_busy_pct is None else f"{host_busy_pct:.4f}",
                ])

            f.flush()
            sample_no += 1

            if sample_no % max(1, int(round(30 / args.interval))) == 0:
                sec_cpu = cpu_pct_one_core_samples["SECURITY TOTAL"][-1]
                sec_pss = pss_samples["SECURITY TOTAL"][-1] if pss_samples["SECURITY TOTAL"] else None
                mem_text = f"{sec_pss:.0f} MiB PSS" if sec_pss is not None else f"{rss_samples['SECURITY TOTAL'][-1]:.0f} MiB RSS"
                print(
                    f"{timestamp}  security={sec_cpu:.1f}% of one core "
                    f"({sec_cpu / NCPU:.3f}% machine), {mem_text}, "
                    f"procs={process_count_samples['SECURITY TOTAL'][-1]}"
                )

            previous = current
            prev_mono = now_mono
            prev_host_cpu = current_host_cpu

            if args.duration is not None and total_elapsed >= args.duration:
                break

    ended_wall = datetime.now().astimezone()
    actual_elapsed = max(0.000001, time.monotonic() - first_mono)

    report_names = product_names + ["SECURITY TOTAL", "ADJACENT TOTAL"]

    with summary_path.open("w") as out:
        out.write("Security / provider overhead measurement\n")
        out.write("========================================\n\n")
        out.write(f"Started:       {started_wall.isoformat(timespec='seconds')}\n")
        out.write(f"Ended:         {ended_wall.isoformat(timespec='seconds')}\n")
        out.write(f"Elapsed:       {human_seconds(actual_elapsed)} ({actual_elapsed:.1f} s)\n")
        out.write(f"Logical CPUs:  {NCPU}\n")
        out.write(f"Sample period: {args.interval:g} s\n")
        out.write(f"Samples:       {sample_no}\n")
        out.write("RAM metric:    PSS when readable; RSS is also recorded\n\n")
        out.write("CPU interpretation: 100% of one core = one logical CPU continuously busy.\n")
        out.write(f"On this {NCPU}-CPU host, 100% of one core = {100.0/NCPU:.4f}% of total machine CPU capacity.\n\n")

        header = (
            f"{'Group':27} {'CPU sec':>12} {'Avg core %':>11} {'Avg mach %':>11} "
            f"{'P95 core %':>11} {'Peak core %':>12} {'Avg PSS MiB':>12} {'P95 PSS':>10} {'Peak PSS':>10} {'Read MB':>11} {'Write MB':>11}\n"
        )
        out.write(header)
        out.write("-" * (len(header.rstrip())) + "\n")

        for name in report_names:
            cs = cpu_seconds[name]
            avg_core_pct = 100.0 * cs / actual_elapsed
            avg_machine_pct = avg_core_pct / NCPU
            cpus = cpu_pct_one_core_samples[name]
            psss = pss_samples[name]
            out.write(
                f"{name:27} {cs:12.1f} {avg_core_pct:11.2f} {avg_machine_pct:11.4f} "
                f"{percentile(cpus, .95):11.2f} {(max(cpus) if cpus else 0):12.2f} "
                f"{(statistics.fmean(psss) if psss else 0):12.1f} {percentile(psss, .95):10.1f} {(max(psss) if psss else 0):10.1f} "
                f"{mb(io_read[name]):11.1f} {mb(io_write[name]):11.1f}\n"
            )

        out.write("\nNotes\n-----\n")
        out.write("* SECURITY TOTAL includes Imunify360, Rustbolit, Acronis, OSSEC, Linux audit, and ClamAV.\n")
        out.write("* WebPros 360 monitoring and cPanel p0f are reported separately as adjacent provider overhead.\n")
        out.write("* PSS is preferred for adding memory across processes because shared pages are proportionally counted.\n")
        out.write("* RSS is retained in samples.csv, but summing RSS can overstate memory because shared pages may be counted repeatedly.\n")
        out.write("* Very short-lived processes that start and finish entirely between samples can be missed. Use --interval 1 for a more aggressive capture.\n")
        out.write("* This measures direct process/kernel-audit resource use. Security hooks that charge CPU to nginx/httpd/MySQL or kernel paths may not appear here.\n")
        if pss_permission_failures and not args.no_pss:
            out.write(f"* PSS could not be read in {pss_permission_failures} process observations; run as root for complete PSS accounting.\n")

    with process_path.open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow([
            "product", "signature", "cpu_seconds", "avg_cpu_pct_one_core_over_run",
            "avg_pss_mib_when_present", "peak_pss_mib", "avg_rss_mib", "peak_rss_mib",
            "last_seen_command",
        ])
        keys = sorted(proc_last_cmd.keys(), key=lambda k: proc_cpu_seconds[k], reverse=True)
        for key in keys:
            product, signature = key
            pcount = proc_pss_count[key]
            rcount = proc_rss_count[key]
            w.writerow([
                product,
                signature,
                f"{proc_cpu_seconds[key]:.6f}",
                f"{100.0 * proc_cpu_seconds[key] / actual_elapsed:.4f}",
                "" if not pcount else f"{proc_pss_sum_mib[key] / pcount:.3f}",
                f"{proc_pss_peak_mib[key]:.3f}",
                "" if not rcount else f"{proc_rss_sum_mib[key] / rcount:.3f}",
                f"{proc_rss_peak_mib[key]:.3f}",
                proc_last_cmd[key],
            ])

    print()
    print(f"Wrote {samples_path}")
    print(f"Wrote {summary_path}")
    print(f"Wrote {process_path}")
    print()
    print(summary_path.read_text())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())