"""
MundiX CLI - System Integrator
Handles safe execution of commands on the host system (Kali Linux).
Provides command extraction from AI responses and interactive execution.
"""

import subprocess
import re
import os
import signal
import shlex
from typing import Tuple, List, Optional


# Commands that require extra confirmation (destructive or dangerous)
DANGEROUS_PATTERNS = [
    r"rm\s+-rf\s+/",
    r"mkfs\.",
    r"dd\s+if=",
    r":\(\)\{",  # fork bomb
    r"chmod\s+-R\s+777\s+/",
    r">\s*/dev/sd",
    r"shutdown",
    r"reboot",
    r"init\s+0",
    r"halt",
]

# Commands that are known pentest tools (highlight differently)
PENTEST_TOOLS = [
    "nmap", "masscan", "rustscan",
    "gobuster", "dirb", "dirbuster", "feroxbuster", "ffuf",
    "sqlmap", "nosqlmap",
    "hydra", "medusa", "john", "hashcat",
    "metasploit", "msfconsole", "msfvenom",
    "nikto", "wpscan", "nuclei",
    "burpsuite", "zaproxy",
    "responder", "impacket", "crackmapexec", "netexec",
    "enum4linux", "smbclient", "rpcclient",
    "searchsploit", "exploit-db",
    "aircrack-ng", "wifite",
    "setoolkit", "gophish",
    "bloodhound", "sharphound",
    "mimikatz", "rubeus",
    "chisel", "ligolo", "socat",
    "curl", "wget", "nc", "ncat", "netcat",
    "python", "python3", "perl", "ruby",
    "tcpdump", "wireshark", "tshark",
    "apt", "apt-get", "pip", "pip3",
    "cat", "grep", "awk", "sed", "find", "ls", "cd",
    "echo", "export", "source", "chmod", "chown",
    "ssh", "scp", "ftp", "telnet",
    "docker", "systemctl", "service",
    "iptables", "ufw", "firewall-cmd",
    "openssl", "certbot",
    "git", "make", "gcc", "g++",
    "tar", "unzip", "gzip",
    "ifconfig", "ip", "route", "arp",
    "traceroute", "ping", "dig", "nslookup", "host",
    "whois", "dnsrecon", "fierce", "sublist3r",
    "whatweb", "wafw00f",
    "cewl", "crunch",
    "evil-winrm", "psexec",
    "proxychains", "tor",
    "strace", "ltrace", "gdb",
    "binwalk", "foremost", "volatility",
    "steghide", "exiftool",
    "wfuzz", "arjun",
    "amass", "subfinder", "httpx",
    "sudo",
]

# Known shell built-ins and common first-words of real commands
VALID_COMMAND_STARTERS = set(PENTEST_TOOLS + [
    "bash", "sh", "zsh", "env", "test", "[", "if", "for", "while",
    "case", "function", "alias", "unalias", "type", "which", "where",
    "head", "tail", "sort", "uniq", "wc", "cut", "tr", "tee",
    "xargs", "parallel", "watch", "timeout", "nice", "nohup",
    "mkdir", "rmdir", "touch", "cp", "mv", "rm", "ln",
    "mount", "umount", "fdisk", "parted",
    "ps", "top", "htop", "kill", "killall", "pkill",
    "crontab", "at", "sleep",
    "man", "info", "help",
    "less", "more", "vi", "vim", "nano",
    "diff", "patch", "file", "stat",
    "base64", "md5sum", "sha256sum", "xxd", "hexdump",
    "jq", "yq", "xmllint",
    "node", "npm", "npx", "pnpm",
    "go", "cargo", "rustc",
    "java", "javac", "mvn",
    "php", "composer",
])


class SystemIntegrator:
    """Manages command execution and system interaction."""

    def __init__(self):
        self.command_history = []
        self.last_output = ""

    def extract_commands(self, ai_response: str) -> List[str]:
        """Extract executable commands from AI response.
        
        STRICT RULES:
        1. ONLY extract from ```bash/sh/shell code blocks
        2. Each line must start with a valid command/binary name
        3. Skip comments, empty lines, and output examples
        4. Skip lines that look like natural language (too many spaces, no slashes/dashes)
        """
        commands = []

        # ONLY match explicit bash/sh/shell code blocks — NOT generic ``` blocks
        pattern = r'```(?:bash|sh|shell|zsh|kali)\s*\n(.*?)```'
        matches = re.findall(pattern, ai_response, re.DOTALL)

        for match in matches:
            lines = match.strip().split("\n")
            for line in lines:
                line = line.strip()
                # Skip empty lines and comments
                if not line or line.startswith("#"):
                    continue
                # Remove common prompt prefixes
                line = re.sub(r'^[\$\#\>]\s*', '', line).strip()
                if not line:
                    continue
                # VALIDATE: line must look like a real command
                if self._is_valid_command(line):
                    commands.append(line)

        return commands

    def _is_valid_command(self, line: str) -> bool:
        """Validate that a line looks like a real shell command, not text."""
        # Reject lines that are clearly natural language
        # 1. Too long with too many spaces and no special chars = probably text
        words = line.split()
        if not words:
            return False

        first_word = words[0].lower()
        
        # Remove path prefix (e.g., /usr/bin/nmap -> nmap)
        if "/" in first_word:
            first_word = os.path.basename(first_word)

        # 2. First word must be a known command OR a path to a binary
        if first_word in VALID_COMMAND_STARTERS:
            return True

        # 3. Check if first word looks like a path to a binary
        if words[0].startswith("/") or words[0].startswith("./"):
            return True

        # 4. Check if it's a variable assignment (VAR=value command)
        if "=" in words[0] and not words[0].startswith("="):
            return True

        # 5. Check if it's a pipe chain starting with a valid command
        if "|" in line:
            pipe_first = line.split("|")[0].strip().split()[0].lower()
            pipe_first = os.path.basename(pipe_first) if "/" in pipe_first else pipe_first
            if pipe_first in VALID_COMMAND_STARTERS:
                return True

        # 6. Check if it's a command with redirection
        if ">" in line or "<" in line:
            redir_first = re.split(r'[><]', line)[0].strip().split()
            if redir_first:
                rf = redir_first[0].lower()
                rf = os.path.basename(rf) if "/" in rf else rf
                if rf in VALID_COMMAND_STARTERS:
                    return True

        # 7. Check if it's a subshell or command substitution
        if line.startswith("(") or line.startswith("$"):
            return True

        # 8. Reject: if the line has more than 6 words and no flags/special chars,
        #    it's probably natural language
        if len(words) > 6:
            has_flags = any(w.startswith("-") for w in words)
            has_special = any(c in line for c in ["|", ">", "<", ";", "&", "=", "/", "\\"])
            if not has_flags and not has_special:
                return False

        # 9. Reject: if first word starts with uppercase and has no special chars
        if words[0][0].isupper() and not any(c in words[0] for c in ["/", ".", "-", "_"]):
            return False

        # 10. Reject: lines that are clearly output (start with common output patterns)
        output_patterns = [
            r'^\d+\.\d+\.\d+\.\d+',  # IP addresses (scan output)
            r'^PORT\s+STATE',  # nmap output header
            r'^\d+/tcp',  # nmap port output
            r'^Starting\s',  # tool startup messages
            r'^Completed\s',  # tool completion messages
            r'^NSE:',  # nmap script output
            r'^Read\s+data',  # nmap data output
            r'^OS\s+and\s+Service',  # nmap summary
            r'^Nmap\s+done',  # nmap done
            r'^\[[\+\-\*\!]\]',  # tool status messages [+] [-] [*]
        ]
        for pat in output_patterns:
            if re.match(pat, line):
                return False

        return False  # Default: reject unknown patterns (safe)

    def is_dangerous(self, command: str) -> bool:
        """Check if a command matches dangerous patterns."""
        for pattern in DANGEROUS_PATTERNS:
            if re.search(pattern, command):
                return True
        return False

    def is_pentest_tool(self, command: str) -> Optional[str]:
        """Check if the command uses a known pentest tool. Returns tool name or None."""
        first_word = command.split()[0] if command.split() else ""
        # Check direct match or path match
        tool_name = os.path.basename(first_word)
        for tool in PENTEST_TOOLS:
            if tool_name == tool or tool in command.split()[0]:
                return tool
        return None

    def execute_command(self, command: str, timeout: int = 120, interactive: bool = False) -> Tuple[str, str, int]:
        """
        Execute a command and return (stdout, stderr, return_code).
        
        Args:
            command: The shell command to execute.
            timeout: Maximum execution time in seconds.
            interactive: If True, stream output in real-time.
        
        Returns:
            Tuple of (stdout, stderr, return_code)
        """
        try:
            if interactive:
                # Stream output in real-time
                process = subprocess.Popen(
                    command,
                    shell=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True,
                    preexec_fn=os.setsid,
                )

                stdout_lines = []
                try:
                    for line in iter(process.stdout.readline, ''):
                        print(line, end='')
                        stdout_lines.append(line)
                    process.wait(timeout=timeout)
                except subprocess.TimeoutExpired:
                    os.killpg(os.getpgid(process.pid), signal.SIGTERM)
                    return "".join(stdout_lines), "[TIMEOUT] Command exceeded time limit", -1

                stderr = process.stderr.read()
                stdout = "".join(stdout_lines)
                returncode = process.returncode

            else:
                result = subprocess.run(
                    command,
                    shell=True,
                    capture_output=True,
                    text=True,
                    timeout=timeout,
                )
                stdout = result.stdout
                stderr = result.stderr
                returncode = result.returncode

            # Save to history
            self.command_history.append({
                "command": command,
                "stdout": stdout[:5000],  # Truncate
                "stderr": stderr[:2000],
                "returncode": returncode,
            })
            self.last_output = stdout

            return stdout, stderr, returncode

        except subprocess.TimeoutExpired:
            return "", f"[TIMEOUT] Command exceeded {timeout}s limit", -1
        except FileNotFoundError:
            return "", f"[ERROR] Command not found: {command.split()[0]}", 127
        except Exception as e:
            return "", f"[ERROR] {str(e)}", -1

    def get_system_info(self) -> dict:
        """Gather comprehensive system information."""
        info = {}

        commands = {
            "hostname": "hostname",
            "kernel": "uname -r",
            "os": "cat /etc/os-release 2>/dev/null | head -3",
            "ip_addresses": "ip -4 addr show | grep inet | awk '{print $2}'",
            "interfaces": "ip link show | grep -E '^[0-9]+:' | awk '{print $2}' | tr -d ':'",
            "user": "whoami",
            "groups": "groups",
            "shell": "echo $SHELL",
            "python": "python3 --version 2>&1",
            "disk": "df -h / | tail -1 | awk '{print $4 \" free of \" $2}'",
            "memory": "free -h | grep Mem | awk '{print $3 \" used of \" $2}'",
        }

        for key, cmd in commands.items():
            try:
                result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=5)
                info[key] = result.stdout.strip()
            except Exception:
                info[key] = "N/A"

        return info

    def check_tool_installed(self, tool_name: str) -> bool:
        """Check if a specific tool is installed."""
        try:
            result = subprocess.run(
                f"which {tool_name}",
                shell=True,
                capture_output=True,
                text=True,
                timeout=5,
            )
            return result.returncode == 0
        except Exception:
            return False

    def install_tool(self, tool_name: str) -> Tuple[str, bool]:
        """Attempt to install a tool via apt."""
        try:
            result = subprocess.run(
                f"sudo apt-get install -y {tool_name}",
                shell=True,
                capture_output=True,
                text=True,
                timeout=120,
            )
            success = result.returncode == 0
            output = result.stdout if success else result.stderr
            return output, success
        except Exception as e:
            return str(e), False

    def get_last_output(self) -> str:
        """Return the output of the last executed command."""
        return self.last_output

    def get_command_history(self) -> List[dict]:
        """Return command execution history."""
        return self.command_history
