#!/usr/bin/env python3
"""Full packet capture via tcpdump with post-rotation compression and encryption.

Manages tcpdump as a supervised subprocess. After each PCAP rotation:
  1. Compress with zstd (level configurable per platform)
  2. Encrypt with AES-256-GCM
  3. Remove plaintext PCAP
  4. Publish PCAP_ROTATED event

Disk monitoring auto-purges oldest encrypted PCAPs at 85% threshold.
"""

import glob
import logging
import os
import shutil
import signal
import struct
import subprocess
import threading
import time
from pathlib import Path
from typing import Optional

from modules.base import BaseModule
from utils.networking import detect_interface_with_retry

logger = logging.getLogger(__name__)


class PacketCapture(BaseModule):
    """Manage tcpdump for continuous full packet capture with rotation."""

    name = "packet_capture"
    module_type = "passive"
    priority = 50
    requires_root = True
    requires_capture_bus = True

    # Defaults
    DEFAULT_INTERFACE = "eth0"
    DEFAULT_SNAP_LEN = 65535
    DEFAULT_ROTATION_SECS = 3600
    DEFAULT_MAX_FILES = 168  # 7 days at 1h rotation
    DEFAULT_COMPRESSION_LEVEL = 3
    DEFAULT_DISK_THRESHOLD = 85  # percent

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._proc: Optional[subprocess.Popen] = None
        self._watcher_thread: Optional[threading.Thread] = None
        self._rotation_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()
        self._pcap_dir = ""
        self._pcap_count = 0
        self._bytes_written = 0
        self._encryption_key = b""

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

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

        iface = self.config.get("interface") or detect_interface_with_retry(
            max_retries=3,
            retry_delay=5,
            exponential=True,
            config_interface=self.DEFAULT_INTERFACE
        )
        snap_len = self.config.get("snap_length", self.DEFAULT_SNAP_LEN)
        rotation = self.config.get("rotation_minutes", 60) * 60
        if rotation <= 0:
            rotation = self.DEFAULT_ROTATION_SECS
        compress_level = self.config.get("compression_level", self.DEFAULT_COMPRESSION_LEVEL)
        disk_threshold = self.config.get("disk_threshold", self.DEFAULT_DISK_THRESHOLD)

        # Directories
        base_dir = self.config.get("data_dir", os.path.expanduser("~/.implant"))
        self._pcap_dir = os.path.join(base_dir, "pcaps")
        Path(self._pcap_dir).mkdir(parents=True, exist_ok=True)

        # Encryption key from config (derived by crypto module at startup)
        self._encryption_key = self.config.get("encryption_key", b"")
        if isinstance(self._encryption_key, str):
            self._encryption_key = self._encryption_key.encode()

        # Build tcpdump command
        pcap_template = os.path.join(self._pcap_dir, "capture_%Y%m%d_%H%M%S.pcap")
        cmd = [
            "tcpdump",
            "-i", iface,
            "-G", str(rotation),
            "-s", str(snap_len),
            "-w", pcap_template,
            "--time-stamp-precision=nano",
            "-Z", "root",  # don't drop privs (we need to read output files)
        ]

        try:
            self._proc = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                preexec_fn=os.setsid,
            )
        except FileNotFoundError:
            logger.error("tcpdump binary not found")
            return
        except Exception as e:
            logger.error("Failed to start tcpdump: %s", e)
            return

        self._running = True
        self._pid = self._proc.pid
        self._start_time = time.time()

        # Watcher thread — monitors tcpdump stderr and detects crashes
        self._watcher_thread = threading.Thread(
            target=self._watch_tcpdump, daemon=True, name="sensor-pcap-watch"
        )
        self._watcher_thread.start()

        # Rotation thread — scans for completed PCAPs to compress/encrypt
        self._rotation_thread = threading.Thread(
            target=self._rotation_loop,
            args=(compress_level, disk_threshold),
            daemon=True, name="sensor-pcap-rotate",
        )
        self._rotation_thread.start()

        self.state.set_module_status(self.name, "running", pid=self._proc.pid)
        logger.info(
            "PacketCapture started — tcpdump PID %d, iface=%s, rotation=%ds, snap=%d",
            self._proc.pid, iface, rotation, snap_len,
        )

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

        # Graceful shutdown of tcpdump
        if self._proc and self._proc.poll() is None:
            try:
                os.killpg(os.getpgid(self._proc.pid), signal.SIGTERM)
            except (OSError, ProcessLookupError):
                pass
            try:
                self._proc.wait(timeout=5.0)
            except subprocess.TimeoutExpired:
                try:
                    os.killpg(os.getpgid(self._proc.pid), signal.SIGKILL)
                    self._proc.wait(timeout=2.0)
                except Exception:
                    pass

        self._proc = None
        self._pid = None
        self.state.set_module_status(self.name, "stopped")
        logger.info("PacketCapture stopped — %d PCAPs rotated", self._pcap_count)

    def status(self) -> dict:
        alive = self._proc is not None and self._proc.poll() is None
        return {
            "running": alive,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "pcap_dir": self._pcap_dir,
            "pcap_count": self._pcap_count,
            "bytes_written": self._bytes_written,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)
        # Live reconfig requires restart
        if self._running:
            logger.info("PacketCapture config changed — restarting tcpdump")
            self.stop()
            self.start()

    # ------------------------------------------------------------------
    # tcpdump watcher
    # ------------------------------------------------------------------

    def _watch_tcpdump(self) -> None:
        """Read tcpdump stderr for stats and detect exit."""
        try:
            for line in iter(self._proc.stderr.readline, b""):
                if not self._running:
                    break
                decoded = line.decode("utf-8", errors="replace").rstrip()
                if decoded:
                    logger.debug("tcpdump: %s", decoded)
        except Exception:
            pass

        # tcpdump exited
        if self._running:
            rc = self._proc.poll() if self._proc else -1
            logger.warning("tcpdump exited unexpectedly (rc=%s)", rc)
            self.bus.emit("TOOL_CRASHED", {
                "tool": "tcpdump", "exit_code": rc,
            }, source_module=self.name)

    # ------------------------------------------------------------------
    # Rotation loop — compress, encrypt, purge
    # ------------------------------------------------------------------

    def _rotation_loop(self, compress_level: int, disk_threshold: int) -> None:
        """Periodically scan for completed PCAPs and process them."""
        while self._running:
            time.sleep(30)  # Check every 30 seconds
            try:
                self._process_completed_pcaps(compress_level)
                self._check_disk_usage(disk_threshold)
            except Exception:
                logger.exception("Rotation loop error")

    def _process_completed_pcaps(self, compress_level: int) -> None:
        """Find completed (not actively written) PCAPs and compress+encrypt."""
        pattern = os.path.join(self._pcap_dir, "capture_*.pcap")
        pcap_files = sorted(glob.glob(pattern))

        for pcap_path in pcap_files:
            # Skip the file tcpdump is currently writing to (most recent)
            if self._is_file_locked(pcap_path):
                continue

            # Skip tiny files (likely incomplete)
            try:
                size = os.path.getsize(pcap_path)
                if size < 24:  # pcap global header minimum
                    continue
            except OSError:
                continue

            compressed_path = pcap_path + ".zst"
            encrypted_path = compressed_path + ".enc"

            try:
                # Step 1: Compress with zstd
                self._compress_zstd(pcap_path, compressed_path, compress_level)

                # Step 2: Encrypt with AES-256-GCM
                if self._encryption_key:
                    self._encrypt_file(compressed_path, encrypted_path)
                    os.unlink(compressed_path)
                    final_path = encrypted_path
                else:
                    final_path = compressed_path

                # Step 3: Remove plaintext PCAP
                os.unlink(pcap_path)

                self._pcap_count += 1
                final_size = os.path.getsize(final_path)
                self._bytes_written += final_size

                # Publish rotation event
                self.bus.emit("PCAP_ROTATED", {
                    "path": final_path,
                    "original_size": size,
                    "final_size": final_size,
                    "compressed": True,
                    "encrypted": bool(self._encryption_key),
                }, source_module=self.name)

                logger.info(
                    "PCAP rotated: %s -> %s (%.1f%% ratio)",
                    os.path.basename(pcap_path),
                    os.path.basename(final_path),
                    (final_size / size * 100) if size > 0 else 0,
                )

            except Exception:
                logger.exception("Failed to process PCAP %s", pcap_path)

    @staticmethod
    def _is_file_locked(path: str) -> bool:
        """Check if a file is still being written by tcpdump (via fuser)."""
        try:
            result = subprocess.run(
                ["fuser", path], capture_output=True, timeout=2
            )
            return result.returncode == 0
        except Exception:
            # If fuser not available, check if file was modified in last 5 seconds
            try:
                mtime = os.path.getmtime(path)
                return (time.time() - mtime) < 5
            except OSError:
                return True

    @staticmethod
    def _compress_zstd(src: str, dst: str, level: int) -> None:
        """Compress a file with zstd."""
        try:
            import zstandard as zstd
            cctx = zstd.ZstdCompressor(level=level)
            with open(src, "rb") as fin, open(dst, "wb") as fout:
                cctx.copy_stream(fin, fout)
        except ImportError:
            # Fallback to CLI zstd
            subprocess.run(
                ["zstd", f"-{level}", "-f", "-o", dst, src],
                check=True, capture_output=True,
            )

    def _encrypt_file(self, src: str, dst: str) -> None:
        """Encrypt a file with AES-256-GCM."""
        try:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        except ImportError:
            logger.warning("cryptography not available — skipping encryption")
            os.rename(src, dst)
            return

        key = self._encryption_key[:32].ljust(32, b"\x00")
        nonce = os.urandom(12)
        aesgcm = AESGCM(key)

        with open(src, "rb") as f:
            plaintext = f.read()

        ciphertext = aesgcm.encrypt(nonce, plaintext, None)

        with open(dst, "wb") as f:
            # Header: magic(4) + nonce(12) + ciphertext
            f.write(b"BB01")
            f.write(nonce)
            f.write(ciphertext)

    def _check_disk_usage(self, threshold: int) -> None:
        """Auto-purge oldest encrypted PCAPs when disk usage exceeds threshold."""
        try:
            usage = shutil.disk_usage(self._pcap_dir)
            pct = (usage.used / usage.total) * 100
        except OSError:
            return

        if pct < threshold:
            return

        logger.warning("Disk usage %.1f%% exceeds threshold %d%% — purging oldest PCAPs", pct, threshold)

        # Gather all encrypted/compressed PCAPs sorted by mtime
        enc_files = sorted(
            glob.glob(os.path.join(self._pcap_dir, "capture_*.pcap.zst*")),
            key=os.path.getmtime,
        )

        purged = 0
        for path in enc_files:
            if not enc_files:
                break
            try:
                os.unlink(path)
                purged += 1
            except OSError:
                continue

            # Re-check usage
            try:
                usage = shutil.disk_usage(self._pcap_dir)
                if (usage.used / usage.total) * 100 < threshold - 5:
                    break
            except OSError:
                break

        if purged:
            logger.info("Purged %d old PCAPs to reclaim disk space", purged)
            self.bus.emit("RESOURCE_WARNING", {
                "type": "disk_purge",
                "purged_count": purged,
                "disk_pct": pct,
            }, source_module=self.name)
