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

Handles parsing and display for MSSQL client connection results.
"""

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 ImpacketMssqlclientHandler(BaseToolHandler):
    """Handler for Impacket mssqlclient MSSQL connection jobs."""

    tool_name = "impacket-mssqlclient"
    display_name = "MSSQL Client"

    has_error_handler = True
    has_warning_handler = True
    has_no_results_handler = True
    has_done_handler = True

    # Error patterns for MSSQL connections
    ERROR_PATTERNS = [
        (r"Login failed", "Authentication failed - invalid credentials"),
        (r"Connection refused", "Connection refused - MSSQL service may be down"),
        (r"timed out", "Connection timed out"),
        (r"No route to host", "No route to host - network unreachable"),
        (r"Errno 113", "Network unreachable"),
        (r"KDC_ERR", "Kerberos authentication error"),
    ]

    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 mssqlclient connection results."""
        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

            # Import managers if not provided
            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()

            target = job.get("target", "")

            # Read log file
            if not log_path or not os.path.exists(log_path):
                return {
                    "tool": self.tool_name,
                    "status": STATUS_ERROR,
                    "error": "Log file not found",
                }

            with open(log_path, "r", encoding="utf-8", errors="replace") as f:
                output = f.read()

            # Parse connection results
            connected = False
            sql_version = ""
            current_user = ""
            current_db = ""
            is_sysadmin = False
            xp_cmdshell_enabled = False
            xp_cmdshell_output = []

            # Check for successful connection
            # mssqlclient shows "SQL>" prompt when connected
            if "SQL>" in output or "SQL (" in output:
                connected = True

            # Extract SQL Server version
            # Format: [*] Microsoft SQL Server 2019 (RTM) ...
            version_match = re.search(
                r"\[\*\]\s*Microsoft SQL Server\s+(\d+[^\n]*)", output
            )
            if version_match:
                sql_version = version_match.group(1).strip()

            # Extract current user from connection info
            # Format: [*] ENVCHANGE(DATABASE): Old Value: master, New Value: master
            # Or from SQL> SELECT SYSTEM_USER output
            user_match = re.search(r"Logged in as:\s*([^\n]+)", output, re.IGNORECASE)
            if user_match:
                current_user = user_match.group(1).strip()

            # Check for sysadmin role
            if "is_srvrolemember" in output.lower() and "1" in output:
                is_sysadmin = True
            if "sysadmin" in output.lower():
                is_sysadmin = True

            # Check for xp_cmdshell usage and RCE success
            has_rce = False
            if "xp_cmdshell" in output:
                xp_cmdshell_enabled = True

                # Look for xp_cmdshell whoami output
                # Format: SQL> xp_cmdshell whoami\noutput\n------\ndomain\user\nNULL
                # Extract lines after xp_cmdshell whoami until === or SQL>
                whoami_match = re.search(
                    r"xp_cmdshell\s+whoami\s*\n(.*?)(?====|SQL[\s>]|$)",
                    output,
                    re.DOTALL | re.IGNORECASE,
                )
                if whoami_match:
                    whoami_block = whoami_match.group(1)
                    # Look for domain\user pattern in the output
                    for line in whoami_block.split("\n"):
                        line = line.strip()
                        # Skip header/separator lines
                        if (
                            not line
                            or line.startswith("-")
                            or line.lower() in ["null", "output"]
                        ):
                            continue
                        # Check for valid user output (domain\user or nt authority\...)
                        if "\\" in line:
                            xp_cmdshell_output.append(line)
                            has_rce = True
                            is_sysadmin = True  # xp_cmdshell working = sysadmin
                            break  # Found the user, stop

            # Get host for findings
            host_id = None
            ip_match = re.search(r"(\d+\.\d+\.\d+\.\d+)", target)
            if ip_match:
                host_ip = ip_match.group(1)
                host = host_manager.get_host_by_ip(engagement_id, host_ip)
                if host:
                    host_id = host["id"]
                else:
                    host_id = host_manager.add_or_update_host(
                        engagement_id, {"ip": host_ip, "status": "up"}
                    )

            # Create finding for successful MSSQL access
            if connected and host_id:
                severity = "critical" if (is_sysadmin or has_rce) else "high"
                if has_rce:
                    title = "MSSQL RCE Achieved (Pwn3d!)"
                    description = f"Successfully authenticated to MSSQL Server and achieved remote code execution via xp_cmdshell"
                    finding_type = "admin_access"  # Triggers 🔑 in findings view
                elif is_sysadmin:
                    title = "MSSQL Sysadmin Access Achieved"
                    description = f"Successfully authenticated to MSSQL Server"
                    finding_type = "mssql_access"
                else:
                    title = "MSSQL Authentication Successful"
                    description = f"Successfully authenticated to MSSQL Server"
                    finding_type = "mssql_access"

                if sql_version:
                    description += f" ({sql_version})"
                if is_sysadmin and not has_rce:
                    description += ". User has sysadmin privileges - xp_cmdshell can be used for RCE."
                if xp_cmdshell_output:
                    description += f"\nCommand output: {xp_cmdshell_output[0]}"

                evidence = f"Version: {sql_version}\nSysadmin: {is_sysadmin}"
                if xp_cmdshell_output:
                    evidence += (
                        f"\nxp_cmdshell output: {chr(10).join(xp_cmdshell_output[:5])}"
                    )

                findings_manager.add_finding(
                    engagement_id=engagement_id,
                    host_id=host_id,
                    finding_type=finding_type,
                    severity=severity,
                    title=title,
                    description=description,
                    evidence=evidence,
                    tool="impacket-mssqlclient",
                    port=1433,
                )

            # Check for errors
            tool_error = detect_tool_error(output, "impacket_mssqlclient")

            # Determine status - findings ALWAYS take priority
            if connected:
                status = STATUS_DONE
            elif any(re.search(p[0], output, re.I) for p in self.ERROR_PATTERNS):
                # Find the specific error
                error_msg = "Connection failed"
                for pattern, msg in self.ERROR_PATTERNS:
                    if re.search(pattern, output, re.I):
                        error_msg = msg
                        break
                return {
                    "tool": self.tool_name,
                    "status": STATUS_ERROR,
                    "target": target,
                    "error": error_msg,
                    "summary": error_msg,
                }
            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

            # Build summary
            summary_parts = []
            if connected:
                summary_parts.append("Connected")
            if sql_version:
                summary_parts.append(f"SQL Server {sql_version.split()[0]}")
            if has_rce:
                summary_parts.append("RCE (Pwn3d!)")
            elif is_sysadmin:
                summary_parts.append("SYSADMIN!")
            if xp_cmdshell_enabled and not has_rce:
                summary_parts.append("xp_cmdshell")
            summary = " | ".join(summary_parts) if summary_parts else "No connection"

            return {
                "tool": self.tool_name,
                "status": status,
                "target": target,
                "summary": summary,
                "connected": connected,
                "sql_version": sql_version,
                "current_user": current_user,
                "current_db": current_db,
                "is_sysadmin": is_sysadmin,
                "xp_cmdshell_enabled": xp_cmdshell_enabled,
                "xp_cmdshell_output": xp_cmdshell_output,
                "has_rce": has_rce,
                "has_admin_access": has_rce,  # Triggers 🔑 emoji in UI
            }

        except Exception as e:
            logger.error(f"Error parsing mssqlclient 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 MSSQL connection results."""
        parse_result = job.get("parse_result", {})

        target = parse_result.get("target", job.get("target", ""))
        sql_version = parse_result.get("sql_version", "")
        is_sysadmin = parse_result.get("is_sysadmin", False)
        has_rce = parse_result.get("has_rce", False)
        xp_cmdshell_enabled = parse_result.get("xp_cmdshell_enabled", False)
        xp_cmdshell_output = parse_result.get("xp_cmdshell_output", [])

        click.echo()

        # Different banner for RCE vs just connection
        if has_rce:
            click.echo(click.style("=" * 70, fg="red"))
            click.echo(click.style("MSSQL RCE ACHIEVED! (Pwn3d!)", fg="red", bold=True))
            click.echo(click.style("=" * 70, fg="red"))
        else:
            click.echo(click.style("=" * 70, fg="green"))
            click.echo(
                click.style("MSSQL CONNECTION SUCCESSFUL!", fg="green", bold=True)
            )
            click.echo(click.style("=" * 70, fg="green"))

        click.echo()
        click.echo(f"  Target: {click.style(target, fg='cyan', bold=True)}")

        if sql_version:
            click.echo(f"  Version: Microsoft SQL Server {sql_version}")

        if has_rce:
            click.secho("  Access: SYSADMIN + xp_cmdshell RCE", fg="red", bold=True)
            # Show whoami output
            if xp_cmdshell_output:
                # Filter to show actual user output (contains backslash)
                for line in xp_cmdshell_output:
                    if "\\" in line:
                        click.echo(
                            f"  Running as: {click.style(line, fg='yellow', bold=True)}"
                        )
                        break
        elif is_sysadmin:
            click.secho(
                "  Access: SYSADMIN (xp_cmdshell available)", fg="yellow", bold=True
            )

        click.echo()
        click.echo(click.style("=" * 70, fg="red" if has_rce else "green"))

        # Show spawn shell instructions
        click.echo()
        click.echo(click.style("  Use with:", bold=True))

        # Check if this is Kerberos auth
        args = job.get("args", [])
        ccache_path = None
        target_ip = None
        kerberos_hostname = None
        if isinstance(args, list):
            for i, arg in enumerate(args):
                if arg == "--ccache" and i + 1 < len(args):
                    ccache_path = args[i + 1]
                elif arg == "-target-ip" and i + 1 < len(args):
                    target_ip = args[i + 1]
                if arg and not arg.startswith("-") and "." in arg:
                    kerberos_hostname = arg

        if ccache_path:
            click.echo(f"    export KRB5CCNAME={ccache_path}")
            cmd = "impacket-mssqlclient -k -no-pass"
            if target_ip:
                cmd += f" -target-ip {target_ip}"
            cmd += f" {kerberos_hostname or target}"
            click.echo(f"    {cmd}")
        else:
            click.echo(f"    impacket-mssqlclient <creds>@{target} -windows-auth")

        click.echo()
        click.echo(click.style("=" * 70, fg="red" if has_rce else "green"))
        click.echo()

    def display_error(
        self,
        job: Dict[str, Any],
        log_path: str,
        show_all: bool = False,
    ) -> None:
        """Display MSSQL connection error."""
        click.echo()
        click.echo(click.style("=" * 70, fg="red"))
        click.echo(click.style("MSSQL CONNECTION FAILED", fg="red", bold=True))
        click.echo(click.style("=" * 70, fg="red"))
        click.echo()

        parse_result = job.get("parse_result", {})
        error = parse_result.get("error", "")

        if error:
            click.echo(f"  {error}")
        else:
            try:
                with open(log_path, "r", encoding="utf-8", errors="replace") as f:
                    output = f.read()

                for pattern, msg in self.ERROR_PATTERNS:
                    if re.search(pattern, output, re.I):
                        click.echo(f"  {msg}")
                        break
                else:
                    click.echo("  Connection failed - check raw logs for details")
            except Exception:
                click.echo("  Could not read error details")

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

    def display_warning(
        self,
        job: Dict[str, Any],
        log_path: str,
        log_content: Optional[str] = None,
    ) -> None:
        """Display warning status."""
        click.echo()
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo(click.style("MSSQL CLIENT - WARNING", fg="yellow", bold=True))
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()
        click.echo("  Connection completed with warnings.")
        click.echo("  Press [r] to view raw logs for details.")
        click.echo()
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()

    def display_no_results(
        self,
        job: Dict[str, Any],
        log_path: str,
        show_all: bool = False,
    ) -> None:
        """Display no connection established."""
        click.echo()
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo(click.style("MSSQL - NO CONNECTION", fg="yellow", bold=True))
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()
        click.echo("  Could not establish MSSQL connection.")
        click.echo()
        click.echo(click.style("  Possible reasons:", dim=True))
        click.echo("    - Invalid credentials")
        click.echo("    - MSSQL service not running")
        click.echo("    - Firewall blocking port 1433")
        click.echo("    - Wrong authentication method (try -windows-auth)")
        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 next steps based on MSSQL connection results.

        Smart suggestions:
        - xp_cmdshell for RCE if sysadmin
        - Database enumeration
        - Privilege escalation paths
        """
        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", "")
        connected = parse_result.get("connected", False)
        is_sysadmin = parse_result.get("is_sysadmin", False)

        if not connected:
            return steps

        # Priority 1: If sysadmin - suggest xp_cmdshell RCE
        if is_sysadmin:
            steps.append(
                create_next_step(
                    title="Enable xp_cmdshell for RCE",
                    reason="Sysadmin privileges - can enable xp_cmdshell",
                    commands=[
                        "# In mssqlclient shell:",
                        "enable_xp_cmdshell",
                        "xp_cmdshell whoami",
                        "# For reverse shell:",
                        "xp_cmdshell powershell -exec bypass -enc <base64_payload>",
                    ],
                    priority=1,
                    category="exploitation",
                    target=target,
                    tool="impacket-mssqlclient",
                )
            )

            steps.append(
                create_next_step(
                    title="Dump database credentials",
                    reason="Sysadmin access - can extract linked server passwords",
                    commands=[
                        "# In mssqlclient shell:",
                        "SELECT * FROM sys.sql_logins;",
                        "SELECT * FROM master.sys.linked_logins;",
                    ],
                    priority=2,
                    category="credential_access",
                    target=target,
                    tool="impacket-mssqlclient",
                )
            )

        # Priority 2: Database enumeration
        steps.append(
            create_next_step(
                title="Enumerate databases and sensitive data",
                reason="MSSQL access achieved - look for sensitive data",
                commands=[
                    "# List databases:",
                    "SELECT name FROM sys.databases;",
                    "# List tables:",
                    "SELECT table_name FROM information_schema.tables;",
                    "# Search for password columns:",
                    "SELECT * FROM information_schema.columns WHERE column_name LIKE '%pass%';",
                ],
                priority=2,
                category="discovery",
                target=target,
                tool="impacket-mssqlclient",
            )
        )

        return steps[:10]


# Register handler
handler = ImpacketMssqlclientHandler()
