#!/usr/bin/env python3
"""Kill switch — full evidence wipe and shutdown.

Sequence:
  1. Stop all modules + subprocesses (bettercap, tcpdump, Responder)
  2. Send corrective ARP if ARP spoof was active
  3. Destroy LUKS header (milliseconds)
  4. Shred encryption keys + credential database
  5. dd zero-fill data partition
  6. Clear RAM
  7. Reboot (boot flag auto-completes if interrupted)

Trigger methods: CLI, SSH, BLE, dead man's switch.
Boot flag at /var/run/.bb_wipe ensures wipe completes if interrupted.
"""

import os
import sys
import time
import signal
import shutil
import struct
import socket
import logging
import subprocess
from pathlib import Path
from typing import Optional

from core.bus import EventBus

logger = logging.getLogger(__name__)

# Boot flag — if this file exists on boot, resume wipe
WIPE_FLAG = "/var/run/.bb_wipe"

# Paths to wipe (relative to install root or absolute)
DEFAULT_SENSITIVE_PATHS = [
    "storage/creds/",
    "storage/config/",
    "storage/intel/",
    "storage/audit/",
    "storage/dns_logs/",
    "storage/sni_logs/",
    "storage/logs/",
    "storage/tool_logs/",
    "storage/pcaps/",
    "storage/baseline/",
]

SENSITIVE_FILES = [
    "storage/config/encryption.key",
    "storage/creds/credentials.db",
    "storage/config/luks_header.bak",
]


class KillSwitch:
    """Full wipe orchestrator.

    Call execute() to perform the complete wipe sequence.
    Call check_boot_flag() on startup to resume an interrupted wipe.
    """

    def __init__(self, bus: EventBus, config: dict,
                 engine=None, tool_manager=None):
        self.bus = bus
        self.config = config
        self.engine = engine
        self.tool_manager = tool_manager
        self._install_path = config.get("device", {}).get(
            "install_path", "/opt/.cache/bb"
        )
        self._luks_device = config.get("security", {}).get(
            "luks_device", "/dev/mapper/sensor-data"
        )
        self._luks_partition = config.get("security", {}).get(
            "luks_partition", ""
        )
        self._arp_spoof_active = False
        self._spoof_gateway = None
        self._spoof_targets = []
        self._interface = config.get("network", {}).get(
            "primary_interface", "eth0"
        )

    # ------------------------------------------------------------------
    # Main entry
    # ------------------------------------------------------------------

    def execute(self, reason: str = "manual") -> None:
        """Execute the full kill switch sequence. Does not return."""
        logger.critical("KILL SWITCH ACTIVATED — reason: %s", reason)

        # Set boot flag first so reboot completes the wipe
        self._set_boot_flag()

        self.bus.emit("KILL_SWITCH",
                      {"reason": reason, "stage": "initiated"},
                      source_module="kill_switch")

        # Phase 1: Stop everything
        self._stop_all_processes()

        # Phase 2: Corrective ARP
        self._send_corrective_arp()

        # Phase 3: Destroy LUKS header
        self._destroy_luks()

        # Phase 4: Shred sensitive files
        self._shred_sensitive_files()

        # Phase 5: Zero-fill data partition
        self._zero_fill_partition()

        # Phase 6: Clear RAM
        self._clear_ram()

        # Phase 7: Remove boot flag and reboot
        self._clear_boot_flag()
        self._reboot()

    def check_boot_flag(self) -> bool:
        """Check if a wipe was interrupted and should resume.

        Call this on boot. Returns True if wipe was resumed.
        """
        if os.path.exists(WIPE_FLAG):
            logger.critical("Boot flag detected — resuming interrupted wipe")
            self.execute(reason="interrupted_resume")
            return True
        return False

    # ------------------------------------------------------------------
    # ARP state (called by active modules to register spoof state)
    # ------------------------------------------------------------------

    def set_arp_spoof_state(self, active: bool, gateway: str = None,
                             targets: list = None, interface: str = None) -> None:
        """Register current ARP spoof state for corrective ARP on kill."""
        self._arp_spoof_active = active
        self._spoof_gateway = gateway
        self._spoof_targets = targets or []
        if interface:
            self._interface = interface

    # ------------------------------------------------------------------
    # Phase 1: Stop processes
    # ------------------------------------------------------------------

    def _stop_all_processes(self) -> None:
        """Stop all modules and managed subprocesses."""
        logger.info("Kill switch phase 1: stopping all processes")

        if self.engine:
            try:
                self.engine.stop_all(timeout=3.0)
            except Exception:
                logger.exception("Error stopping modules")

        if self.tool_manager:
            try:
                self.tool_manager.stop_all()
            except Exception:
                logger.exception("Error stopping tools")

            # Kill any remaining tool PIDs
            for pid in self.tool_manager.get_all_pids():
                try:
                    os.kill(pid, signal.SIGKILL)
                except (OSError, ProcessLookupError):
                    pass

    # ------------------------------------------------------------------
    # Phase 2: Corrective ARP
    # ------------------------------------------------------------------

    def _send_corrective_arp(self) -> None:
        """Send corrective gratuitous ARPs if ARP spoof was active."""
        if not self._arp_spoof_active or not self._spoof_gateway:
            return

        logger.info("Kill switch phase 2: sending corrective ARP")
        try:
            gw_ip = self._spoof_gateway
            gw_mac = self._resolve_mac(gw_ip)
            if not gw_mac:
                logger.warning("Could not resolve gateway MAC for corrective ARP")
                return

            for target_ip in self._spoof_targets:
                target_mac = self._resolve_mac(target_ip)
                if target_mac:
                    self._send_arp_reply(
                        src_mac=gw_mac, src_ip=gw_ip,
                        dst_mac=target_mac, dst_ip=target_ip,
                        interface=self._interface
                    )
            # Also restore gateway's ARP cache
            for target_ip in self._spoof_targets:
                target_mac = self._resolve_mac(target_ip)
                if target_mac:
                    self._send_arp_reply(
                        src_mac=target_mac, src_ip=target_ip,
                        dst_mac=gw_mac, dst_ip=gw_ip,
                        interface=self._interface
                    )
        except Exception:
            logger.exception("Corrective ARP failed")

    @staticmethod
    def _resolve_mac(ip: str) -> Optional[str]:
        """Resolve IP to MAC from /proc/net/arp."""
        try:
            with open("/proc/net/arp", "r") as f:
                for line in f:
                    parts = line.split()
                    if len(parts) >= 4 and parts[0] == ip:
                        mac = parts[3]
                        if mac != "00:00:00:00:00:00":
                            return mac
        except (FileNotFoundError, PermissionError):
            pass
        return None

    @staticmethod
    def _send_arp_reply(src_mac: str, src_ip: str, dst_mac: str, dst_ip: str,
                         interface: str) -> None:
        """Send a single ARP reply using raw socket."""
        try:
            sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(0x0806))
            sock.bind((interface, 0))

            src_mac_bytes = bytes.fromhex(src_mac.replace(":", ""))
            dst_mac_bytes = bytes.fromhex(dst_mac.replace(":", ""))
            src_ip_bytes = socket.inet_aton(src_ip)
            dst_ip_bytes = socket.inet_aton(dst_ip)

            # Ethernet header + ARP reply
            frame = (
                dst_mac_bytes +                    # Dst MAC
                src_mac_bytes +                    # Src MAC
                b"\x08\x06" +                      # EtherType: ARP
                b"\x00\x01" +                      # Hardware type: Ethernet
                b"\x08\x00" +                      # Protocol type: IPv4
                b"\x06" +                          # Hardware size: 6
                b"\x04" +                          # Protocol size: 4
                b"\x00\x02" +                      # Opcode: reply
                src_mac_bytes +                    # Sender MAC
                src_ip_bytes +                     # Sender IP
                dst_mac_bytes +                    # Target MAC
                dst_ip_bytes                       # Target IP
            )

            # Send 3 times for reliability
            for _ in range(3):
                sock.send(frame)
                time.sleep(0.1)

            sock.close()
        except Exception as e:
            logger.warning("ARP reply send failed: %s", e)

    # ------------------------------------------------------------------
    # Phase 3: LUKS header destruction
    # ------------------------------------------------------------------

    def _destroy_luks(self) -> None:
        """Destroy the LUKS header, making encrypted data unrecoverable."""
        logger.info("Kill switch phase 3: destroying LUKS header")

        if not self._luks_partition:
            logger.warning("No LUKS partition configured, skipping")
            return

        # Close the LUKS device first
        try:
            subprocess.run(
                ["cryptsetup", "close", self._luks_device.split("/")[-1]],
                timeout=5, capture_output=True
            )
        except Exception:
            pass

        # Overwrite LUKS header (first 16MB covers all key slots)
        try:
            subprocess.run(
                ["dd", "if=/dev/urandom", f"of={self._luks_partition}",
                 "bs=1M", "count=16", "conv=notrunc"],
                timeout=30, capture_output=True, check=True
            )
            logger.info("LUKS header destroyed")
        except Exception as e:
            logger.error("LUKS header destruction failed: %s", e)

    # ------------------------------------------------------------------
    # Phase 4: Shred sensitive files
    # ------------------------------------------------------------------

    def _shred_sensitive_files(self) -> None:
        """Securely delete sensitive files and directories."""
        logger.info("Kill switch phase 4: shredding sensitive files")

        base = Path(self._install_path)

        # Shred specific files first
        for rel_path in SENSITIVE_FILES:
            fpath = base / rel_path
            if fpath.exists():
                self._shred_file(str(fpath))

        # Then directories
        for rel_path in DEFAULT_SENSITIVE_PATHS:
            dpath = base / rel_path
            if dpath.exists():
                for fpath in dpath.rglob("*"):
                    if fpath.is_file():
                        self._shred_file(str(fpath))
                try:
                    shutil.rmtree(str(dpath), ignore_errors=True)
                except Exception:
                    pass

        # State database
        state_db = os.path.expanduser("~/.implant/state.db")
        for f in [state_db, state_db + "-wal", state_db + "-shm"]:
            if os.path.exists(f):
                self._shred_file(f)

    @staticmethod
    def _shred_file(path: str) -> None:
        """Overwrite file with random data, then unlink."""
        try:
            size = os.path.getsize(path)
            with open(path, "r+b") as f:
                # 3 passes: random, zeros, random
                for pattern in [None, b"\x00", None]:
                    f.seek(0)
                    remaining = size
                    while remaining > 0:
                        chunk = min(remaining, 65536)
                        if pattern is None:
                            data = os.urandom(chunk)
                        else:
                            data = pattern * chunk
                        f.write(data)
                        remaining -= chunk
                    f.flush()
                    os.fsync(f.fileno())
            os.unlink(path)
        except Exception as e:
            # Fallback: just delete
            try:
                os.unlink(path)
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Phase 5: Zero-fill
    # ------------------------------------------------------------------

    def _zero_fill_partition(self) -> None:
        """Zero-fill the data partition."""
        logger.info("Kill switch phase 5: zero-filling data partition")

        storage_dir = os.path.join(self._install_path, "storage")
        if not os.path.isdir(storage_dir):
            storage_dir = self._install_path

        try:
            # Write zeros until disk full
            zero_file = os.path.join(storage_dir, ".wipe")
            subprocess.run(
                ["dd", "if=/dev/zero", f"of={zero_file}",
                 "bs=1M", "status=none"],
                timeout=600, capture_output=True
            )
        except Exception:
            pass
        finally:
            try:
                os.unlink(os.path.join(storage_dir, ".wipe"))
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Phase 6: Clear RAM
    # ------------------------------------------------------------------

    @staticmethod
    def _clear_ram() -> None:
        """Drop caches and clear as much RAM as possible."""
        logger.info("Kill switch phase 6: clearing RAM")
        try:
            # Drop page cache, dentries, inodes
            with open("/proc/sys/vm/drop_caches", "w") as f:
                f.write("3")
        except (PermissionError, FileNotFoundError):
            pass

        # smem or sdmem would be better but may not be installed
        try:
            subprocess.run(["sdmem", "-f", "-ll"], timeout=60,
                           capture_output=True)
        except (FileNotFoundError, subprocess.TimeoutExpired):
            pass

    # ------------------------------------------------------------------
    # Phase 7: Reboot
    # ------------------------------------------------------------------

    @staticmethod
    def _reboot() -> None:
        """Reboot the system."""
        logger.info("Kill switch phase 7: rebooting")
        try:
            subprocess.run(["reboot"], timeout=5)
        except Exception:
            try:
                os.system("reboot")
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Boot flag
    # ------------------------------------------------------------------

    @staticmethod
    def _set_boot_flag() -> None:
        """Create boot flag to resume wipe if interrupted."""
        try:
            Path(WIPE_FLAG).parent.mkdir(parents=True, exist_ok=True)
            Path(WIPE_FLAG).touch()
        except Exception:
            pass

    @staticmethod
    def _clear_boot_flag() -> None:
        """Remove boot flag after successful wipe."""
        try:
            os.unlink(WIPE_FLAG)
        except FileNotFoundError:
            pass

    # ------------------------------------------------------------------
    # Dead man's switch
    # ------------------------------------------------------------------

    def start_dead_man_switch(self, interval_hours: int = 24) -> None:
        """Start dead man's switch timer.

        If not reset within interval_hours, trigger kill switch.
        Intended to be reset by operator check-in via C2.
        """
        self._dms_interval = interval_hours * 3600
        self._dms_last_checkin = time.time()

        import threading
        def _dms_loop():
            while True:
                time.sleep(60)
                elapsed = time.time() - self._dms_last_checkin
                if elapsed > self._dms_interval:
                    logger.critical("Dead man's switch triggered (no check-in for %dh)",
                                    interval_hours)
                    self.execute(reason="dead_man_switch")

        t = threading.Thread(target=_dms_loop, daemon=True, name="sensor-dms")
        t.start()

    def reset_dead_man_switch(self) -> None:
        """Reset the dead man's switch timer (operator check-in)."""
        self._dms_last_checkin = time.time()