#!/usr/bin/env python3
"""
souleyez.plugins.impacket_mssqlclient - MSSQL client for database access and command execution
"""

import subprocess
import time
from typing import List

from .impacket_common import find_impacket_command
from .plugin_base import PluginBase

HELP = {
    "name": "Impacket mssqlclient - MSSQL Database Client",
    "description": (
        "Need to connect to Microsoft SQL Server?\n\n"
        "mssqlclient provides an interactive SQL shell for Microsoft SQL Server. "
        "It supports Windows authentication, SQL authentication, and Kerberos. "
        "With xp_cmdshell enabled, you can execute operating system commands.\n\n"
        "Use mssqlclient to:\n"
        "- Connect to MSSQL with stolen credentials\n"
        "- Execute SQL queries\n"
        "- Enable and use xp_cmdshell for RCE\n"
        "- Perform Silver Ticket attacks against MSSQL\n\n"
        "Quick tips:\n"
        "- Use -windows-auth for domain credentials\n"
        "- Use -k for Kerberos authentication\n"
        "- xp_cmdshell may need to be enabled via sp_configure\n"
    ),
    "usage": 'souleyez jobs enqueue impacket-mssqlclient <target> --args "DOMAIN/user:pass@host"',
    "examples": [
        'souleyez jobs enqueue impacket-mssqlclient 10.0.0.82 --args "sa:Password123@10.0.0.82"',
        'souleyez jobs enqueue impacket-mssqlclient 10.0.0.82 --args "CONTOSO/svc_mssql:Trustno1@10.0.0.82 -windows-auth"',
        'souleyez jobs enqueue impacket-mssqlclient 10.0.0.82 --args "-k breachdc.breach.vl"',
    ],
    "flags": [
        ["-windows-auth", "Use Windows Authentication (NTLM)"],
        ["-k", "Use Kerberos authentication (requires valid ccache)"],
        ["-hashes <LM:NT>", "Pass-the-hash authentication"],
        ["-port <port>", "MSSQL port (default 1433)"],
        ["-db <database>", "Connect to specific database"],
    ],
    "preset_categories": {
        "authentication": [
            {
                "name": "SQL Authentication",
                "args": ["sa:password@<target>"],
                "desc": "Connect with SQL Server authentication (sa account)",
            },
            {
                "name": "Windows Authentication",
                "args": ["DOMAIN/user:pass@<target>", "-windows-auth"],
                "desc": "Connect with domain credentials via NTLM",
            },
            {
                "name": "Kerberos Authentication",
                "args": ["-k", "-no-pass", "<target>"],
                "desc": "Connect using Kerberos ticket (Silver Ticket)",
            },
            {
                "name": "Pass-the-Hash",
                "args": [
                    "DOMAIN/user@<target>",
                    "-hashes",
                    ":<ntlm_hash>",
                    "-windows-auth",
                ],
                "desc": "Authenticate with NTLM hash instead of password",
            },
        ],
        "exploitation": [
            {
                "name": "Enable xp_cmdshell",
                "args": [],
                "desc": "After connecting, run: enable_xp_cmdshell",
            },
            {
                "name": "Execute OS Command",
                "args": [],
                "desc": "After connecting, run: xp_cmdshell whoami",
            },
        ],
    },
    "presets": [],
}

# Flatten presets
for category_presets in HELP["preset_categories"].values():
    HELP["presets"].extend(category_presets)

HELP["help_sections"] = [
    {
        "title": "What is mssqlclient?",
        "color": "cyan",
        "content": [
            {
                "title": "Overview",
                "desc": "mssqlclient is an interactive MSSQL client from Impacket. "
                "It connects to Microsoft SQL Server using SQL auth, Windows auth, or Kerberos.",
            },
            {
                "title": "Use Cases",
                "desc": "Database exploitation and lateral movement",
                "tips": [
                    "Connect with stolen/cracked credentials",
                    "Execute xp_cmdshell for RCE",
                    "Silver Ticket attacks (impersonate users)",
                    "Extract data from databases",
                ],
            },
        ],
    },
    {
        "title": "How to Use",
        "color": "green",
        "content": [
            {
                "title": "Basic Workflow",
                "desc": "1. Obtain MSSQL service account credentials\n"
                "     2. Connect with mssqlclient\n"
                "     3. Enable xp_cmdshell if needed\n"
                "     4. Execute commands or queries",
            },
            {
                "title": "Built-in Commands",
                "desc": "mssqlclient shell commands",
                "tips": [
                    "enable_xp_cmdshell - Enable xp_cmdshell",
                    "disable_xp_cmdshell - Disable xp_cmdshell",
                    "xp_cmdshell <cmd> - Execute OS command",
                    "sp_start_job <job> - Start SQL Agent job",
                    "lcd/!<cmd> - Local shell command",
                ],
            },
        ],
    },
    {
        "title": "Tips & Best Practices",
        "color": "yellow",
        "content": [
            (
                "Best Practices:",
                [
                    "Use -windows-auth for domain accounts",
                    "Export KRB5CCNAME for Kerberos auth",
                    "Check for sysadmin role first",
                    "xp_cmdshell runs as SQL service account",
                ],
            ),
            (
                "Silver Ticket Attack:",
                [
                    "Get svc_mssql NTLM hash via Kerberoasting",
                    "Create ticket: ticketer.py -spn MSSQLSvc/host:1433 ...",
                    "Export: export KRB5CCNAME=Administrator.ccache",
                    "Connect: mssqlclient.py -k -no-pass host",
                ],
            ),
        ],
    },
]


class ImpacketMssqlclientPlugin(PluginBase):
    name = "Impacket mssqlclient"
    tool = "impacket-mssqlclient"
    category = "lateral_movement"
    HELP = HELP

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

        # Replace <target> placeholder
        args = [arg.replace("<target>", target) for arg in args]

        # Check for --ccache arg (Silver Ticket support)
        # This sets KRB5CCNAME environment variable for Kerberos auth
        ccache_path = None
        filtered_args = []
        i = 0
        while i < len(args):
            if args[i] == "--ccache" and i + 1 < len(args):
                ccache_path = args[i + 1]
                i += 2  # Skip both --ccache and path
            else:
                filtered_args.append(args[i])
                i += 1
        args = filtered_args

        # Find the correct command (varies by install: apt vs pipx)
        mssqlclient_cmd = find_impacket_command("mssqlclient")
        if not mssqlclient_cmd:
            return None  # Tool not installed

        # Build command (args should include credentials)
        cmd = [mssqlclient_cmd] + args

        # Set environment for Kerberos if ccache provided
        env = None
        if ccache_path:
            import os

            env = {"KRB5CCNAME": ccache_path}

        return {"cmd": cmd, "timeout": 600, "env": env}

    def run(
        self, target: str, args: List[str] = None, label: str = "", log_path: str = None
    ) -> int:
        """Execute impacket-mssqlclient and write output to log_path."""

        args = args or []

        # Replace <target> placeholder
        args = [arg.replace("<target>", target) for arg in args]

        # Find the correct command (varies by install: apt vs pipx)
        mssqlclient_cmd = find_impacket_command("mssqlclient")
        if not mssqlclient_cmd:
            if log_path:
                with open(log_path, "w", encoding="utf-8") as fh:
                    fh.write("ERROR: mssqlclient not found. Install with:\n")
                    fh.write("  Kali: sudo apt install python3-impacket\n")
                    fh.write("  Ubuntu: pipx install impacket\n")
            return 1

        # Build command
        cmd = [mssqlclient_cmd]

        # Add target/credentials (should be in args like "DOMAIN/user:pass@host")
        cmd.extend(args)

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

        try:
            # Create metadata header
            with open(log_path, "w", encoding="utf-8", errors="replace") as fh:
                fh.write(f"=== Plugin: Impacket mssqlclient ===\n")
                fh.write(f"Target: {target}\n")
                fh.write(f"Args: {args}\n")
                fh.write(f"Label: {label}\n")
                fh.write(
                    f"Started: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())}\n"
                )
                fh.write(f"Command: {' '.join(cmd)}\n\n")

            # Run mssqlclient with timeout
            # Note: mssqlclient is interactive, so we run with short timeout for initial connection test
            proc = subprocess.run(
                cmd, capture_output=True, timeout=60, check=False, text=True
            )

            # Write output
            with open(log_path, "a", encoding="utf-8", errors="replace") as fh:
                if proc.stdout:
                    fh.write(proc.stdout)

                if proc.stderr:
                    fh.write(f"\n\n# Error output:\n{proc.stderr}\n")

                fh.write(
                    f"\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:
            with open(log_path, "a", encoding="utf-8", errors="replace") as fh:
                fh.write("\n\nERROR: mssqlclient timed out after 60 seconds\n")
            return 124

        except FileNotFoundError:
            with open(log_path, "w", encoding="utf-8", errors="replace") as fh:
                fh.write("ERROR: impacket-mssqlclient not found in PATH\n")
                fh.write("Install: apt install python3-impacket\n")
            return 127

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


plugin = ImpacketMssqlclientPlugin()
