#!/usr/bin/env python3
"""
Handler for smbclient share browsing and file listing.
"""

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 SmbclientHandler(BaseToolHandler):
    """Handler for smbclient share browsing."""

    tool_name = "smbclient"
    display_name = "SMB Client"

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

    # Patterns for parsing smbclient ls output
    # Format: "  filename                          D        0  Tue Aug 22 10:17:06 2023"
    #         "  filename                          A     1234  Tue Aug 22 10:17:06 2023"
    # D = directory, A = archive (file), H = hidden, S = system, R = readonly
    DIR_ENTRY_PATTERN = (
        r"^\s{2}(\S.*?)\s{2,}([DAHSR]+)\s+(\d+)\s+\w+\s+\w+\s+\d+\s+[\d:]+\s+\d+"
    )

    # Upload success pattern: "putting file /path/to/file as \filename"
    UPLOAD_PATTERN = r"putting file\s+(\S+)\s+as\s+(\S+)"

    # Error patterns
    ERROR_PATTERNS = [
        (r"NT_STATUS_ACCESS_DENIED", "Access denied"),
        (r"NT_STATUS_LOGON_FAILURE", "Authentication failed"),
        (r"NT_STATUS_BAD_NETWORK_NAME", "Share not found"),
        (r"NT_STATUS_HOST_UNREACHABLE", "Host unreachable"),
        (r"NT_STATUS_CONNECTION_REFUSED", "Connection refused"),
        (r"Connection to .* failed", "Connection failed"),
        (r"session setup failed", "Session setup failed"),
    ]

    # Directories to skip when extracting usernames
    # Note: 'library' removed - can be a valid username in some AD environments
    SKIP_DIRS = {
        ".",
        "..",
        "public",
        "shared",
        "common",
        "scripts",
        "profiles",
        "default",
        "all users",
        "default user",
        "public documents",
    }

    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 smbclient 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

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

            if not log_path or not os.path.exists(log_path):
                return {
                    "tool": "smbclient",
                    "status": STATUS_ERROR,
                    "error": "Log file not found",
                }

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

            # Check for successful upload first
            upload_match = re.search(self.UPLOAD_PATTERN, log_content)
            if upload_match:
                return {
                    "tool": "smbclient",
                    "status": STATUS_DONE,
                    "target": target,
                    "operation": "upload",
                    "local_file": upload_match.group(1),
                    "remote_file": upload_match.group(2),
                    "summary": f"Uploaded {upload_match.group(2)}",
                }

            # Check for errors first
            for pattern, error_msg in self.ERROR_PATTERNS:
                if re.search(pattern, log_content, re.IGNORECASE):
                    return {
                        "tool": "smbclient",
                        "status": STATUS_ERROR,
                        "target": target,
                        "error": error_msg,
                    }

            # Extract share name from command or label
            share_name = ""
            share_match = re.search(r"//[^/]+/(\S+)", log_content)
            if share_match:
                share_name = share_match.group(1)
            elif "SMB_SPIDER_" in label:
                share_name = label.replace("SMB_SPIDER_", "")

            # Parse directory listing
            directories: List[Dict[str, Any]] = []
            files: List[Dict[str, Any]] = []
            extracted_usernames: List[str] = []

            for line in log_content.split("\n"):
                match = re.match(self.DIR_ENTRY_PATTERN, line)
                if match:
                    name = match.group(1).strip()
                    attrs = match.group(2)
                    size = int(match.group(3))

                    entry = {
                        "name": name,
                        "attrs": attrs,
                        "size": size,
                        "is_directory": "D" in attrs,
                    }

                    if "D" in attrs:
                        directories.append(entry)

                        # Check if this looks like a username (for homes shares)
                        if name.lower() not in self.SKIP_DIRS:
                            # Username patterns: FirstName.LastName, first.last, username
                            if re.match(r"^[A-Za-z][A-Za-z0-9._-]+$", name):
                                # Likely a username if it contains a dot (First.Last format)
                                # or is in a homes-type share
                                if "." in name or share_name.lower() in [
                                    "homes",
                                    "users",
                                    "home",
                                    "profiles",
                                ]:
                                    extracted_usernames.append(name)
                    else:
                        files.append(entry)

            # Check for interesting files
            interesting_files = []
            interesting_patterns = [
                (r"\.vbs$", "VBS script"),
                (r"\.ps1$", "PowerShell script"),
                (r"\.bat$", "Batch script"),
                (r"\.xml$", "XML config"),
                (r"Groups\.xml", "GPP file"),
                (r"web\.config", "Web config"),
                (r"\.ini$", "INI config"),
                (r"\.conf$", "Config file"),
                (r"password", "Password file"),
                (r"credential", "Credential file"),
            ]

            for f in files:
                for pattern, desc in interesting_patterns:
                    if re.search(pattern, f["name"], re.IGNORECASE):
                        interesting_files.append({"name": f["name"], "type": desc})
                        break

            # Generic error detection (supplement to handler-specific checks above)
            tool_error = detect_tool_error(log_content, "smbclient")

            # Determine status - findings take priority over error detection
            if directories or files:
                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

            # Build summary
            summary_parts = []
            if extracted_usernames:
                summary_parts.append(f"{len(extracted_usernames)} usernames extracted")
            if interesting_files:
                summary_parts.append(f"{len(interesting_files)} interesting files")
            if directories:
                summary_parts.append(f"{len(directories)} dirs")
            if files:
                summary_parts.append(f"{len(files)} files")
            summary = " | ".join(summary_parts) if summary_parts else "Empty share"

            result = {
                "tool": "smbclient",
                "status": status,
                "target": target,
                "share_name": share_name,
                "directories": directories,
                "files": files,
                "directory_count": len(directories),
                "file_count": len(files),
                "extracted_usernames": extracted_usernames,
                "interesting_files": interesting_files,
                "summary": summary,
            }

            if extracted_usernames:
                logger.info(
                    f"smbclient: Extracted {len(extracted_usernames)} potential usernames from {share_name}"
                )

                # Store extracted usernames as credentials
                if credentials_manager and host_manager:
                    try:
                        host = host_manager.get_host_by_ip(engagement_id, target)
                        if host:
                            for username in extracted_usernames:
                                try:
                                    credentials_manager.add_credential(
                                        engagement_id=engagement_id,
                                        host_id=host["id"],
                                        username=username,
                                        password="",
                                        service="smb",
                                        credential_type="username",
                                        tool="smbclient",
                                        status="untested",
                                    )
                                except Exception:
                                    pass  # Skip duplicates
                            logger.info(
                                f"smbclient: Stored {len(extracted_usernames)} usernames in credentials database"
                            )
                    except Exception as e:
                        logger.debug(f"Could not store usernames: {e}")

            if interesting_files:
                logger.info(
                    f"smbclient: Found {len(interesting_files)} interesting files in {share_name}"
                )

            return result

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

        # Handle upload operation display
        if parse_result.get("operation") == "upload":
            click.echo()
            click.echo(click.style("=" * 70, fg="green"))
            click.echo(click.style("SMB FILE UPLOAD", fg="green", bold=True))
            click.echo(click.style("=" * 70, fg="green"))
            click.echo()
            click.echo(f"  Local:  {parse_result.get('local_file', 'unknown')}")
            click.echo(
                f"  Remote: {click.style(parse_result.get('remote_file', 'unknown'), fg='cyan', bold=True)}"
            )
            click.echo()
            click.echo(click.style("  Upload successful", fg="green"))
            click.echo()
            return

        click.echo()
        click.echo(click.style("=" * 70, fg="green"))
        click.echo(click.style("SMB SHARE LISTING", fg="green", bold=True))
        click.echo(click.style("=" * 70, fg="green"))
        click.echo()
        share_name = parse_result.get("share_name", "Unknown")
        directories = parse_result.get("directories", [])
        files = parse_result.get("files", [])
        extracted_usernames = parse_result.get("extracted_usernames", [])
        interesting_files = parse_result.get("interesting_files", [])

        click.echo(f"  Share: {click.style(share_name, fg='cyan', bold=True)}")
        click.echo(f"  Directories: {len(directories)}  |  Files: {len(files)}")
        click.echo()

        # Show extracted usernames
        if extracted_usernames:
            click.echo(click.style("  EXTRACTED USERNAMES", bold=True, fg="yellow"))
            for username in extracted_usernames[:20]:  # Limit display
                click.echo(f"    {username}")
            if len(extracted_usernames) > 20:
                click.echo(f"    ... and {len(extracted_usernames) - 20} more")
            click.echo()

        # Show interesting files
        if interesting_files:
            click.echo(click.style("  INTERESTING FILES", bold=True, fg="red"))
            for f in interesting_files:
                click.echo(f"    {f['name']} ({f['type']})")
            click.echo()

        # Show directory listing if show_all
        if show_all and directories:
            click.echo(click.style("  DIRECTORIES", bold=True, fg="cyan"))
            for d in directories[:50]:
                click.echo(f"    {d['name']}/")
            click.echo()

        if show_all and files:
            click.echo(click.style("  FILES", bold=True, fg="cyan"))
            for f in files[:50]:
                size_str = f"{f['size']:,}" if f["size"] > 0 else "0"
                click.echo(f"    {f['name']:40} {size_str:>12} bytes")
            click.echo()

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

        error = job.get("parse_result", {}).get("error") or job.get("error")
        if error:
            click.echo(f"  Error: {error}")
        else:
            click.echo("  Check log for details")
        click.echo()

    def display_warning(
        self,
        job: Dict[str, Any],
        log_path: str,
        show_all: bool = False,
    ) -> None:
        """Display smbclient warning."""
        # For smbclient, warning might indicate partial results
        self.display_done(job, log_path, show_all, False)

    def display_no_results(
        self,
        job: Dict[str, Any],
        log_path: str,
        show_all: bool = False,
    ) -> None:
        """Display smbclient no results."""
        click.echo()
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo(click.style("SMB SHARE EMPTY", fg="yellow", bold=True))
        click.echo(click.style("=" * 70, fg="yellow"))
        click.echo()
        click.echo("  Share exists but contains no files or directories.")
        click.echo()

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

        Smart suggestions:
        - Password spray with extracted usernames
        - Download/examine interesting files
        - GPP credential extraction
        """
        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", "")
        share_name = parse_result.get("share_name", "")
        extracted_usernames = parse_result.get("extracted_usernames", [])
        interesting_files = parse_result.get("interesting_files", [])
        files = parse_result.get("files", [])

        # Priority 1: Password spray if usernames extracted
        if extracted_usernames:
            user_sample = ", ".join(extracted_usernames[:3])
            steps.append(
                create_next_step(
                    title=f"Password spray {len(extracted_usernames)} extracted usernames",
                    reason=f"Usernames from share: {user_sample}...",
                    commands=[
                        "# Save usernames to file, then:",
                        f"kerbrute passwordspray --dc {target} -d DOMAIN users.txt 'Password123'",
                        f"crackmapexec smb {target} -u users.txt -p 'Password123' --continue-on-success",
                    ],
                    priority=1,
                    category="auth",
                    target=target,
                    tool="kerbrute",
                )
            )

        # Priority 1: GPP extraction if Groups.xml found
        gpp_files = [f for f in interesting_files if "GPP" in f.get("type", "")]
        if gpp_files:
            steps.append(
                create_next_step(
                    title="Extract GPP credentials",
                    reason="Groups.xml found - may contain encrypted passwords",
                    commands=[
                        f"# Download the file:",
                        f"smbclient //{target}/{share_name} -N -c 'get path/to/Groups.xml'",
                        "",
                        f"# Decrypt with gpp-decrypt:",
                        "gpp-decrypt 'CPASSWORD_FROM_XML'",
                    ],
                    priority=1,
                    category="credential",
                    target=target,
                    tool="gpp_extract",
                )
            )

        # Priority 2: Download script files
        script_files = [
            f
            for f in interesting_files
            if any(
                ext in f.get("type", "").lower()
                for ext in ["script", "powershell", "batch", "vbs"]
            )
        ]
        if script_files:
            script_names = ", ".join([f["name"] for f in script_files[:3]])
            steps.append(
                create_next_step(
                    title="Examine scripts for credentials",
                    reason=f"Scripts found: {script_names}",
                    commands=[
                        f"smbclient //{target}/{share_name} -N -c 'get SCRIPT_NAME'",
                        "",
                        "# Search for credentials:",
                        "grep -i 'password\\|credential\\|secret' SCRIPT_NAME",
                    ],
                    priority=2,
                    category="credential",
                    target=target,
                    tool="manual",
                )
            )

        # Priority 2: Download config files
        config_files = [
            f
            for f in interesting_files
            if any(ext in f.get("type", "").lower() for ext in ["config", "xml", "ini"])
        ]
        if config_files and not gpp_files:
            config_names = ", ".join([f["name"] for f in config_files[:3]])
            steps.append(
                create_next_step(
                    title="Examine config files",
                    reason=f"Config files found: {config_names}",
                    commands=[
                        f"smbclient //{target}/{share_name} -N -c 'get CONFIG_FILE'",
                        "",
                        "# Search for credentials:",
                        "grep -i 'password\\|connection\\|key' CONFIG_FILE",
                    ],
                    priority=2,
                    category="credential",
                    target=target,
                    tool="manual",
                )
            )

        # Priority 3: Recursive search if many files
        if len(files) > 10:
            steps.append(
                create_next_step(
                    title="Search share recursively for sensitive data",
                    reason=f"{len(files)} files in share - search for passwords",
                    commands=[
                        f"# Spider share with smbmap:",
                        f"smbmap -H {target} -R {share_name} --depth 5",
                        "",
                        f"# Or search file contents:",
                        f"smbclient //{target}/{share_name} -N -c 'recurse on; prompt off; mget *pass*.txt *cred*'",
                    ],
                    priority=3,
                    category="enumeration",
                    target=target,
                    tool="smbmap",
                )
            )

        return steps[:10]


# Register handler
handler = SmbclientHandler()
