#!/usr/bin/env python3
"""
Impacket psexec handler.

Consolidates parsing and display logic for Impacket psexec remote execution 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
from souleyez.handlers.base import BaseToolHandler

logger = logging.getLogger(__name__)


class ImpacketPsexecHandler(BaseToolHandler):
    """Handler for Impacket psexec remote execution jobs."""

    tool_name = "impacket-psexec"
    display_name = "PSExec"

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

    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 psexec job results.

        Checks for successful remote shell establishment.
        """
        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)",
                }

            from souleyez.engine.result_handler import detect_tool_error
            from souleyez.parsers.impacket_parser import parse_psexec

            # Import managers if not provided
            if host_manager is None:
                from souleyez.storage.hosts import HostManager

                host_manager = HostManager()

            target = job.get("target", "")
            parsed = parse_psexec(log_path, target)

            if "error" in parsed:
                return parsed

            # Get or create host
            host_id = None
            ip_match = re.search(r"@?(\d+\.\d+\.\d+\.\d+)", target)
            if ip_match:
                host_ip = ip_match.group(1)
                host_id = host_manager.add_or_update_host(
                    engagement_id, {"ip": host_ip, "status": "up"}
                )

            success = parsed.get("success", False)

            # Check for tool errors (supplement to handler's own error checks)
            with open(log_path, "r", encoding="utf-8", errors="replace") as f:
                log_content = f.read()
            tool_error = detect_tool_error(log_content, "impacket_psexec")

            # Determine status - findings ALWAYS take priority over error detection
            if success:
                status = STATUS_DONE
            elif tool_error:
                return {
                    "tool": self.tool_name,
                    "status": STATUS_ERROR,
                    "target": target,
                    "error": tool_error,
                    "summary": f"Tool error: {tool_error}",
                }
            else:
                status = STATUS_NO_RESULTS

            return {
                "tool": "impacket-psexec",
                "status": status,
                "target": target,
                "success": success,
                "output_lines": parsed.get("output_lines", 0),
            }

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

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

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

            target = job.get("target", "unknown")
            parsed = parse_psexec(log_path, target)

            success = parsed.get("success", False)
            output_lines = parsed.get("output_lines", 0)

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

            if success:
                click.echo(
                    click.style("  [SUCCESS] Remote shell established!", fg="green")
                )
                click.echo(f"  Output lines captured: {output_lines}")
                click.echo()
                click.echo("  Press 'r' to view full command output.")
            else:
                click.echo(
                    click.style(
                        "  Connection made but no shell prompt detected.", fg="yellow"
                    )
                )
                click.echo("  Press 'r' to view raw output for details.")

            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 status for psexec."""
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo(click.style("[WARNING] PSEXEC", bold=True, fg="yellow"))
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()
        click.echo("  Remote execution completed with warnings.")
        click.echo("  Check raw logs for details (press 'r').")
        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 psexec."""
        # Read log if not provided
        if log_content is None and log_path and os.path.exists(log_path):
            try:
                with open(log_path, "r", encoding="utf-8", errors="replace") as f:
                    log_content = f.read()
            except Exception:
                log_content = ""

        click.echo(click.style("=" * 70, fg="red"))
        click.echo(click.style("[ERROR] PSEXEC FAILED", bold=True, fg="red"))
        click.echo(click.style("=" * 70, fg="red"))
        click.echo()

        # Check for common psexec errors
        error_msg = None
        if log_content:
            if "Connection refused" in log_content:
                error_msg = "Connection refused - target SMB service may be down"
            elif (
                "Access denied" in log_content.lower()
                or "STATUS_ACCESS_DENIED" in log_content
            ):
                error_msg = "Access denied - need admin privileges on target"
            elif "STATUS_LOGON_FAILURE" in log_content:
                error_msg = "Logon failure - invalid credentials"
            elif "timed out" in log_content.lower() or "timeout" in log_content.lower():
                error_msg = "Connection timed out - target may be unreachable"
            elif (
                "Service installation" in log_content
                and "failed" in log_content.lower()
            ):
                error_msg = "Service installation failed - AV may be blocking"
            elif "ERROR_SERVICE" in log_content or (
                "service" in log_content.lower() and "error" in log_content.lower()
            ):
                error_msg = "Service error - may be blocked by endpoint protection"

        if error_msg:
            click.echo(f"  {error_msg}")
        else:
            click.echo(
                "  Remote execution failed - check raw logs for details (press 'r')"
            )

        click.echo()
        click.echo(click.style("=" * 70, fg="red"))
        click.echo()

    def display_no_results(
        self,
        job: Dict[str, Any],
        log_path: str,
    ) -> None:
        """Display no_results status for psexec."""
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo(click.style("PSEXEC RESULTS", bold=True, fg="yellow"))
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()
        click.echo("  No command output captured.")
        click.echo()
        click.echo(click.style("Possible reasons:", dim=True))
        click.echo("  - Connection failed before shell was established")
        click.echo("  - Insufficient privileges (need local admin)")
        click.echo("  - AV/EDR blocked the service installation")
        click.echo("  - Try smbexec or wmiexec as alternatives")
        click.echo()
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()

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

        Suggests post-exploitation actions if shell was established.
        """
        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", "")
        success = parse_result.get("success", False)

        # Extract credentials from job args if available
        args = job.get("args", "")
        cred_match = re.search(r"([^/]+)/([^:]+):([^@]+)@", args)
        domain = cred_match.group(1) if cred_match else ""
        username = cred_match.group(2) if cred_match else ""
        password = cred_match.group(3) if cred_match else ""

        if success:
            # Shell established - suggest post-exploitation
            steps.append(
                create_next_step(
                    title="Dump credentials with secretsdump",
                    reason="SYSTEM shell access - extract all secrets",
                    commands=[
                        f"secretsdump.py '{domain}/{username}:{password}'@{target}",
                    ],
                    priority=1,
                    category="post_exploit",
                    target=target,
                    tool="impacket_secretsdump",
                    requires_creds=True,
                    cred_type="smb",
                )
            )

            steps.append(
                create_next_step(
                    title="Get interactive shell",
                    reason="Verified access - spawn interactive session",
                    commands=[
                        f"psexec.py '{domain}/{username}:{password}'@{target}",
                        "# Or use evil-winrm for PowerShell:",
                        f"evil-winrm -i {target} -u {username} -p '{password}'",
                    ],
                    priority=1,
                    category="access",
                    target=target,
                    tool="impacket_psexec",
                    requires_creds=True,
                    cred_type="smb",
                )
            )

            steps.append(
                create_next_step(
                    title="Enumerate local users and groups",
                    reason="Admin access - enumerate for persistence options",
                    commands=[
                        "# In shell: net user",
                        "# In shell: net localgroup administrators",
                        "# In shell: whoami /all",
                    ],
                    priority=2,
                    category="post_exploit",
                    target=target,
                    tool="manual",
                )
            )

            steps.append(
                create_next_step(
                    title="Check for domain controller",
                    reason="Determine if this is a DC for further attacks",
                    commands=[
                        "# In shell: nltest /dclist:",
                        '# In shell: systeminfo | findstr /B /C:"Domain"',
                    ],
                    priority=2,
                    category="post_exploit",
                    target=target,
                    tool="manual",
                )
            )

        else:
            # Shell failed - suggest alternatives
            steps.append(
                create_next_step(
                    title="Try alternative execution methods",
                    reason="PSExec failed - try other Impacket tools",
                    commands=[
                        f"wmiexec.py '{domain}/{username}:{password}'@{target}",
                        f"smbexec.py '{domain}/{username}:{password}'@{target}",
                        f"atexec.py '{domain}/{username}:{password}'@{target} 'whoami'",
                    ],
                    priority=2,
                    category="access",
                    target=target,
                    tool="impacket_wmiexec",
                )
            )

        return steps[:10]
