#!/usr/bin/env python3
"""
Metasploit exploit handler.

Consolidates parsing and display logic for msf_exploit jobs.
"""

import logging
import os
import re
from typing import Any, Dict, List, Optional

import click

from souleyez.engine.job_status import (
    STATUS_DONE,
    STATUS_ERROR,
    STATUS_NO_RESULTS,
    STATUS_WARNING,
)
from souleyez.handlers.base import BaseToolHandler

logger = logging.getLogger(__name__)


class MsfExploitHandler(BaseToolHandler):
    """Handler for Metasploit exploit module jobs."""

    tool_name = "msf_exploit"
    display_name = "Metasploit Exploit"

    # All handlers enabled
    has_error_handler = True
    has_warning_handler = True
    has_no_results_handler = True
    has_done_handler = True

    # Success patterns for exploit output
    SUCCESS_PATTERNS = [
        r"\[\*\]\s+Command shell session \d+ opened",
        r"\[\*\]\s+Meterpreter session \d+ opened",
        r"\[\+\]\s+\d+\.\d+\.\d+\.\d+:\d+\s+-\s+Session \d+ created",
        r"\[\+\].*shell.*opened",
        r"\[\+\].*session.*created",
    ]

    # Failure patterns
    FAILURE_PATTERNS = [
        r"Exploit completed, but no session was created",
        r"\[-\].*failed",
        r"\[-\].*Exploit aborted",
        r"\[-\].*not valid",
        r"\[-\].*unreachable",
        r"\[-\].*timed?\s*out",
        r"\[-\].*ConnectionTimeout",
        r"\[-\].*Connection refused",
        r"No session was created",
    ]

    def parse_job(
        self,
        engagement_id: int,
        log_path: str,
        job: Dict[str, Any],
        host_manager: Optional[Any] = None,
        findings_manager: Optional[Any] = None,
        credentials_manager: Optional[Any] = None,
    ) -> Dict[str, Any]:
        """
        Parse MSF exploit job results.

        Extracts session info, success/failure status, and creates findings.
        """
        try:
            # Guard against empty output
            if not os.path.exists(log_path) or os.path.getsize(log_path) == 0:
                return {
                    "tool": self.tool_name,
                    "status": STATUS_NO_RESULTS,
                    "summary": "Empty output (tool produced no output)",
                }

            # Import managers if not provided (for backward compatibility)
            if host_manager is None:
                from souleyez.storage.hosts import HostManager

                host_manager = HostManager()
            if findings_manager is None:
                from souleyez.storage.findings import FindingsManager

                findings_manager = FindingsManager()

            from souleyez.engine.result_handler import detect_tool_error

            # Read the log file
            with open(log_path, "r", encoding="utf-8", errors="replace") as f:
                content = f.read()

            target = job.get("target", "")
            findings_added = 0
            exploit_success = False

            # Get or create host
            host = host_manager.get_host_by_ip(engagement_id, target)
            if not host:
                host_id = host_manager.add_host(engagement_id, target)
            else:
                host_id = host["id"]

            # Extract exploit name from header or args
            exploit_match = re.search(r"Exploit:\s*(.+)$", content, re.MULTILINE)
            exploit_name = (
                exploit_match.group(1).strip() if exploit_match else "unknown"
            )

            # Check for successful exploitation
            for pattern in self.SUCCESS_PATTERNS:
                if re.search(pattern, content, re.IGNORECASE):
                    exploit_success = True
                    break

            # Check for explicit failure
            explicit_failure = False
            failure_reason = None
            for pattern in self.FAILURE_PATTERNS:
                match = re.search(pattern, content, re.IGNORECASE)
                if match:
                    explicit_failure = True
                    failure_reason = match.group(0).strip()
                    break

            # Create finding based on result
            if exploit_success:
                session_match = re.search(
                    r"session (\d+) (opened|created)", content, re.IGNORECASE
                )
                session_info = (
                    f" (Session {session_match.group(1)})" if session_match else ""
                )

                findings_manager.add_finding(
                    engagement_id=engagement_id,
                    host_id=host_id,
                    title=f'Exploit Successful: {exploit_name.split("/")[-1]}',
                    finding_type="vulnerability",
                    severity="critical",
                    description=f"Successfully exploited {target} using {exploit_name}.{session_info}\n\nThis confirms the vulnerability is exploitable.",
                    tool="msf_exploit",
                )
                findings_added += 1
            elif not explicit_failure:
                # Check for [+] lines that might indicate partial success
                plus_lines = re.findall(r"\[\+\]\s+(.+)", content)
                if plus_lines:
                    for line in plus_lines[:3]:
                        findings_manager.add_finding(
                            engagement_id=engagement_id,
                            host_id=host_id,
                            title=f'{exploit_name.split("/")[-1]}: {line[:60]}',
                            finding_type="security_issue",
                            severity="medium",
                            description=f"Exploit output: {line}",
                            tool="msf_exploit",
                        )
                        findings_added += 1

            # Check for tool-level errors (supplement to handler's own checks)
            tool_error = detect_tool_error(content, "msf_exploit")

            # Determine final status and summary
            # Findings ALWAYS take priority over error detection
            if exploit_success:
                final_status = STATUS_DONE
                summary = f"Exploit successful: {exploit_name.split('/')[-1]}"
            elif explicit_failure:
                final_status = STATUS_WARNING
                summary = self._format_failure_summary(failure_reason, exploit_name)
            elif findings_added > 0:
                final_status = STATUS_DONE
                summary = f"Exploit output captured ({findings_added} findings)"
            elif tool_error:
                final_status = STATUS_ERROR
                summary = f"Tool error: {tool_error}"
            else:
                final_status = STATUS_NO_RESULTS
                summary = f"Exploit completed: no session or findings"

            result = {
                "tool": "msf_exploit",
                "status": final_status,
                "summary": summary,
                "host": target,
                "exploit": exploit_name,
                "success": exploit_success,
                "explicit_failure": explicit_failure,
                "findings_added": findings_added,
            }
            if final_status == STATUS_ERROR and tool_error:
                result["error"] = tool_error

            return result
        except Exception as e:
            logger.error(f"Error parsing msf_exploit job: {e}")
            return {
                "tool": self.tool_name,
                "status": STATUS_ERROR,
                "error": str(e),
                "summary": f"Parse error: {e}",
            }

    def _format_failure_summary(
        self, failure_reason: Optional[str], exploit_name: str
    ) -> str:
        """Format a human-readable failure summary."""
        if not failure_reason:
            return f"Exploit failed: {exploit_name.split('/')[-1]}"

        reason_lower = failure_reason.lower()
        if "no session was created" in reason_lower:
            return "Exploit failed: no session created"
        elif "unreachable" in reason_lower or "timed out" in reason_lower:
            return "Exploit failed: target unreachable"
        elif "not valid" in reason_lower:
            return "Exploit failed: invalid configuration"
        elif "connection refused" in reason_lower:
            return "Exploit failed: connection refused"
        else:
            return f"Exploit failed: {failure_reason[:50]}"

    def display_done(
        self,
        job: Dict[str, Any],
        log_path: str,
        show_all: bool = False,
        show_passwords: bool = False,
    ) -> None:
        """Display successful exploit results."""
        try:
            from souleyez.parsers.msf_parser import parse_msf_log

            if not log_path or not os.path.exists(log_path):
                return

            parsed = parse_msf_log(log_path)
            sessions = parsed.get("sessions", [])
            findings = parsed.get("findings", [])

            if not sessions and not findings:
                # Fall back to parse_result summary
                parse_result = job.get("parse_result", {})
                if isinstance(parse_result, dict) and parse_result.get("success"):
                    click.echo(click.style("=" * 70, fg="green"))
                    click.echo(
                        click.style("METASPLOIT EXPLOIT RESULTS", bold=True, fg="green")
                    )
                    click.echo(click.style("=" * 70, fg="green"))
                    click.echo()
                    click.echo(
                        click.style(
                            f"  [SUCCESS] {parse_result.get('summary', 'Exploit successful')}",
                            fg="green",
                        )
                    )
                    click.echo()
                    click.echo(click.style("=" * 70, fg="green"))
                    click.echo()
                return

            click.echo(click.style("=" * 70, fg="green"))
            click.echo(click.style("METASPLOIT EXPLOIT RESULTS", bold=True, fg="green"))
            click.echo(click.style("=" * 70, fg="green"))
            click.echo()

            # Show sessions (most important)
            if sessions:
                click.echo(
                    click.style(
                        f"Sessions Created ({len(sessions)}):", bold=True, fg="green"
                    )
                )
                for s in sessions:
                    session_type = s.get("type", "unknown").title()
                    session_id = s.get("id", "?")
                    tunnel = s.get("tunnel", "unknown")
                    click.echo(
                        click.style(
                            f"  [{session_id}] {session_type} session", fg="green"
                        )
                        + f" - {tunnel}"
                    )
                click.echo()

            # Show findings
            if findings:
                severity_colors = {
                    "critical": "green",
                    "high": "red",
                    "medium": "yellow",
                    "low": "blue",
                    "info": "cyan",
                }
                for f in findings:
                    sev = f.get("severity", "info")
                    color = severity_colors.get(sev, "white")
                    title = f.get("title", "Unknown")
                    if "Successful" in title or "session" in title.lower():
                        click.echo(click.style(f"  [SUCCESS] {title}", fg="green"))
                    elif "Failed" in title:
                        click.echo(click.style(f"  [FAILED] {title}", fg="yellow"))
                    else:
                        click.echo(click.style(f"  [{sev.upper()}] ", fg=color) + title)
                click.echo()

            click.echo(click.style("=" * 70, fg="green"))
            click.echo()
        except Exception as e:
            logger.debug(f"Error in display_done: {e}")

    def display_warning(
        self,
        job: Dict[str, Any],
        log_path: str,
        log_content: Optional[str] = None,
    ) -> None:
        """Display warning/failure status for exploit."""
        parse_result = job.get("parse_result", {})
        summary = "Exploit failed"
        if isinstance(parse_result, dict):
            summary = parse_result.get("summary", "Exploit failed")

        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo(click.style("[FAILED] METASPLOIT EXPLOIT", bold=True, fg="yellow"))
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()
        click.echo(f"  {summary}")
        click.echo()
        click.echo(click.style("  Common causes:", fg="bright_black"))
        click.echo(
            click.style(
                "    - Target unreachable (firewall, network issue)", fg="bright_black"
            )
        )
        click.echo(
            click.style("    - Service not running on expected port", fg="bright_black")
        )
        click.echo(
            click.style("    - Connection timed out or refused", fg="bright_black")
        )
        click.echo()
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()

    def display_no_results(
        self,
        job: Dict[str, Any],
        log_path: str,
    ) -> None:
        """Display no_results status for exploit."""
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo(click.style("METASPLOIT EXPLOIT RESULTS", bold=True, fg="yellow"))
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()
        click.echo("  No session was created.")
        click.echo()
        click.echo(click.style("  Possible reasons:", fg="bright_black"))
        click.echo(
            click.style(
                "    - Target is not vulnerable to this exploit", fg="bright_black"
            )
        )
        click.echo(
            click.style(
                "    - Exploit completed but payload failed to execute",
                fg="bright_black",
            )
        )
        click.echo(
            click.style(
                "    - Target OS/version mismatch with exploit requirements",
                fg="bright_black",
            )
        )
        click.echo(
            click.style(
                "    - Security controls blocked the exploit (AV, DEP, ASLR)",
                fg="bright_black",
            )
        )
        click.echo()
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()

    def display_error(
        self,
        job: Dict[str, Any],
        log_path: str,
        log_content: Optional[str] = None,
    ) -> None:
        """Display error status for exploit."""
        parse_result = job.get("parse_result", {})
        summary = "An error occurred during exploitation"
        if isinstance(parse_result, dict):
            summary = parse_result.get("summary", summary)

        click.echo(click.style("=" * 70, fg="red"))
        click.echo(click.style("[ERROR] METASPLOIT EXPLOIT", bold=True, fg="red"))
        click.echo(click.style("=" * 70, fg="red"))
        click.echo()
        click.echo(f"  {summary}")
        click.echo()
        click.echo(click.style("=" * 70, fg="red"))
        click.echo()

    def generate_next_steps(
        self,
        job: Dict[str, Any],
        context: Any,
    ) -> List[Any]:
        """
        Generate suggested next steps based on MSF exploit results.

        Suggests post-exploitation actions if a session was obtained.
        """
        from souleyez.core.next_steps import create_next_step

        steps = []
        parse_result = job.get("parse_result", {})
        if not isinstance(parse_result, dict):
            return steps

        target = job.get("target", "")
        exploit_success = parse_result.get("success", False)
        exploit_name = parse_result.get("exploit", "unknown")

        # If exploit succeeded - suggest post-exploitation
        if exploit_success:
            steps.append(
                create_next_step(
                    title="Post-exploitation: Dump credentials",
                    reason="Session obtained - extract credentials for lateral movement",
                    commands=[
                        "# In Meterpreter session:",
                        "hashdump",
                        "load kiwi && creds_all",
                    ],
                    priority=1,
                    category="post_exploit",
                    target=target,
                    tool="meterpreter",
                )
            )

            steps.append(
                create_next_step(
                    title="Post-exploitation: Enumerate system",
                    reason="Session obtained - gather system information",
                    commands=[
                        "# In Meterpreter session:",
                        "sysinfo",
                        "getuid",
                        "ps",
                        "route",
                    ],
                    priority=2,
                    category="post_exploit",
                    target=target,
                    tool="meterpreter",
                )
            )

            steps.append(
                create_next_step(
                    title="Post-exploitation: Pivot to internal network",
                    reason="Session obtained - explore internal network",
                    commands=[
                        "# In Meterpreter session:",
                        "run autoroute -s 10.0.0.0/8",
                        "run arp_scanner -r 10.0.0.0/24",
                    ],
                    priority=3,
                    category="post_exploit",
                    target=target,
                    tool="meterpreter",
                )
            )

        # If exploit failed - suggest alternatives
        else:
            explicit_failure = parse_result.get("explicit_failure", False)
            if explicit_failure:
                steps.append(
                    create_next_step(
                        title="Try alternative exploits",
                        reason="Exploit failed - search for other vulnerabilities",
                        commands=[
                            f"searchsploit {target}",
                            "# Run vulnerability scan to find other attack vectors",
                        ],
                        priority=2,
                        category="recon",
                        target=target,
                        tool="searchsploit",
                    )
                )

        return steps[:10]
