"""
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 in UI)
PENTEST_TOOLS = [
    # ── Network Scanning & Discovery ──
    "nmap", "masscan", "rustscan", "unicornscan", "zmap",
    "arp-scan", "arping", "fping", "hping3", "netdiscover",
    "nbtscan", "onesixtyone", "snmpwalk", "snmpcheck",
    # ── Web Scanning & Fuzzing ──
    "gobuster", "dirb", "dirbuster", "feroxbuster", "ffuf",
    "wfuzz", "arjun", "dirsearch", "katana", "hakrawler",
    "gospider", "gau", "waybackurls", "paramspider",
    # ── Vulnerability Scanning ──
    "nikto", "wpscan", "nuclei", "openvas", "nessus",
    "arachni", "skipfish", "w3af", "vuls", "trivy",
    # ── SQL Injection ──
    "sqlmap", "nosqlmap", "ghauri", "jsql",
    # ── Brute Force & Password ──
    "hydra", "medusa", "john", "hashcat", "patator",
    "crowbar", "ncrack", "cewl", "crunch", "cupp",
    "hash-identifier", "hashid", "ophcrack",
    # ── Exploitation Frameworks ──
    "metasploit", "msfconsole", "msfvenom", "msfdb",
    "searchsploit", "exploitdb",
    # ── Web Proxies & Interception ──
    "burpsuite", "zaproxy", "mitmproxy", "bettercap",
    # ── Active Directory & Windows ──
    "responder", "impacket", "crackmapexec", "netexec",
    "enum4linux", "enum4linux-ng", "smbclient", "rpcclient",
    "smbmap", "ldapsearch", "kerbrute", "getTGT.py",
    "getST.py", "secretsdump.py", "psexec.py", "wmiexec.py",
    "smbexec.py", "atexec.py", "dcomexec.py",
    "bloodhound", "sharphound", "bloodhound-python",
    "mimikatz", "rubeus", "certipy", "adidnsdump",
    "evil-winrm", "psexec", "pth-winexe",
    # ── Wireless ──
    "aircrack-ng", "airmon-ng", "airodump-ng", "aireplay-ng",
    "wifite", "reaver", "bully", "pixiewps", "wash",
    # ── Social Engineering ──
    "setoolkit", "gophish", "king-phisher",
    # ── Tunneling & Pivoting ──
    "chisel", "ligolo", "socat", "proxychains", "proxychains4",
    "sshuttle", "revsocks", "rpivot",
    # ── OSINT & Recon ──
    "whois", "dnsrecon", "fierce", "sublist3r", "subfinder",
    "amass", "httpx", "theHarvester", "theharvester",
    "shodan", "censys", "recon-ng", "maltego",
    "whatweb", "wafw00f", "wappalyzer",
    # ── Reverse Engineering & Forensics ──
    "strace", "ltrace", "gdb", "radare2", "r2",
    "binwalk", "foremost", "volatility", "autopsy",
    "steghide", "exiftool", "strings", "objdump",
    "ghidra", "ida", "ollydbg",
    # ── Networking Utilities ──
    "curl", "wget", "nc", "ncat", "netcat",
    "ssh", "scp", "ftp", "telnet", "rdesktop", "xfreerdp",
    "tcpdump", "wireshark", "tshark", "ettercap",
    "traceroute", "tracepath", "ping", "dig", "nslookup", "host",
    "ifconfig", "ip", "route", "arp", "ss", "netstat",
    "tor", "torify", "torsocks",
    # ── Container & Cloud ──
    "docker", "kubectl", "aws", "az", "gcloud",
    "terraform", "ansible", "vagrant",
    # ── Programming & Scripting ──
    "python", "python3", "python2", "perl", "ruby", "lua",
    "php", "node", "npm", "npx", "pnpm",
    "go", "cargo", "rustc", "gcc", "g++", "make", "cmake",
    "java", "javac", "mvn", "gradle",
    "composer", "gem", "pip", "pip3",
    # ── System Administration ──
    "sudo", "su", "doas",
    "apt", "apt-get", "dpkg", "yum", "dnf", "pacman", "snap",
    "systemctl", "service", "journalctl",
    "iptables", "nftables", "ufw", "firewall-cmd",
    "openssl", "certbot", "gpg",
    "crontab", "at",
    "useradd", "usermod", "passwd", "chpasswd",
    "mount", "umount", "fdisk", "parted", "lsblk",
    # ── File & Text Utilities ──
    "cat", "grep", "egrep", "fgrep", "awk", "sed", "find",
    "ls", "cd", "pwd", "echo", "printf", "export", "source",
    "chmod", "chown", "chgrp", "chattr",
    "head", "tail", "sort", "uniq", "wc", "cut", "tr", "tee",
    "xargs", "parallel", "watch", "timeout", "nice", "nohup",
    "mkdir", "rmdir", "touch", "cp", "mv", "rm", "ln",
    "tar", "unzip", "zip", "gzip", "gunzip", "bzip2", "xz", "7z",
    "diff", "patch", "file", "stat", "tree",
    "less", "more", "vi", "vim", "nano", "ed",
    "base64", "md5sum", "sha256sum", "sha1sum", "xxd", "hexdump",
    "jq", "yq", "xmllint", "csvtool",
    # ── Process & System ──
    "ps", "top", "htop", "kill", "killall", "pkill",
    "sleep", "date", "uptime", "uname", "hostname", "id", "who", "w",
    "free", "df", "du", "lsof", "stty", "tput", "clear", "reset",
    "env", "printenv", "set", "unset", "alias", "unalias",
    "type", "which", "where", "whereis", "locate", "updatedb",
    "man", "info", "help", "apropos",
    # ── Version Control ──
    "git", "svn", "hg",
    # ── Shell ──
    "bash", "sh", "zsh", "dash", "fish",
    "test", "[", "[[", "if", "for", "while", "until",
    "case", "select", "function", "eval", "exec",
    "true", "false", "yes",
]

# Build the set for fast lookup
VALID_COMMAND_STARTERS = set(PENTEST_TOOLS)


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

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

    # Shell keywords that indicate multi-line constructs (must NOT be split)
    _SHELL_BLOCK_OPENERS = re.compile(
        r'^\s*(?:if\s|for\s|while\s|until\s|case\s|function\s|\w+\s*\(\)\s*\{|'
        r'do\b|then\b|else\b|elif\s|select\s)'
    )
    _SHELL_BLOCK_CLOSERS = {'fi', 'done', 'esac', '}'}
    _SHEBANG = re.compile(r'^#!\s*/(?:usr/)?bin/')

    def _is_multiline_script(self, block: str) -> bool:
        """Detect if a code block is a multi-line shell script that must run as one unit."""
        lines = [l.strip() for l in block.strip().split('\n') if l.strip() and not l.strip().startswith('#')]
        if len(lines) <= 1:
            return False
        # Shebang → always a script
        first_line = block.strip().split('\n')[0].strip()
        if self._SHEBANG.match(first_line):
            return True
        # Shell block keywords
        for line in lines:
            first_word = line.split()[0].rstrip(';') if line.split() else ''
            if self._SHELL_BLOCK_OPENERS.match(line) or first_word in self._SHELL_BLOCK_CLOSERS:
                return True
        return False

    def extract_commands(self, ai_response: str) -> List[str]:
        """Extract executable commands from AI response.

        PHILOSOPHY: If the AI wrapped it in a ```bash block, TRUST IT.
        Multi-line scripts (if/fi, for/done, shebang) are kept as a single
        script and executed via a temp file to avoid broken line-by-line splits.
        """
        commands = []

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

        for match in matches:
            block = match.strip()
            if not block:
                continue

            # If block is a multi-line script, keep it whole
            if self._is_multiline_script(block):
                # Clean leading prompt prefixes but keep structure
                cleaned_lines = []
                for line in block.split('\n'):
                    cleaned = re.sub(r'^[\$\>]{1,2}\s', '', line)
                    cleaned_lines.append(cleaned)
                script = '\n'.join(cleaned_lines)
                commands.append(script)
                continue

            # Simple commands: split line-by-line
            for line in block.split("\n"):
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                line = re.sub(r'^[\$\#\>]{1,2}\s*', '', line).strip()
                if not line:
                    continue
                if self._is_output_or_text(line):
                    continue
                commands.append(line)

        return commands

    def _is_output_or_text(self, line: str) -> bool:
        """Return True if the line is clearly command output or natural language, NOT a command."""
        words = line.split()
        if not words:
            return True

        # ── SKIP / meta-directives from AI ──
        if line.upper().startswith('SKIP') and (':' in line[:10]):
            return True

        # ── Definite output patterns ──
        output_patterns = [
            r'^PORT\s+STATE',           # nmap output header
            r'^\d+/tcp\s+',             # nmap port output (e.g., "80/tcp open http")
            r'^\d+/udp\s+',             # nmap UDP output
            r'^Starting\s+\w+\s+\d',    # "Starting Nmap 7.93..."
            r'^Completed\s',            # tool completion messages
            r'^NSE:',                   # nmap script engine output
            r'^Nmap\s+done',            # nmap done message
            r'^Nmap\s+scan\s+report',   # nmap scan report
            r'^Host\s+is\s+',           # "Host is up"
            r'^MAC\s+Address:',         # MAC address output
            r'^Service\s+Info:',        # nmap service info
            r'^OS\s+and\s+Service',     # nmap OS detection
            r'^Read\s+data\s+files',    # nmap data files
            r'^\[[\+\-\*\!]\]\s',       # tool status messages [+] [-] [*] [!]
            r'^Total\s+hosts:',         # arp-scan total
            r'^Interface:',             # arp-scan interface
            r'^Ending\s+',             # tool ending messages
            r'^\d+\s+hosts?\s+up',      # "256 hosts up"
            r'^Not\s+shown:',           # nmap "Not shown: 998 filtered ports"
            r'^Warning:',               # tool warnings
            r'^Error:',                 # tool errors
            r'^Usage:',                 # tool usage messages
        ]
        for pat in output_patterns:
            if re.match(pat, line):
                return True

        # ── Natural language detection ──
        # If line has 8+ words, no flags (-), no special chars, and starts with uppercase
        # it's almost certainly natural language text
        if len(words) >= 8:
            has_flags = any(w.startswith("-") for w in words[1:])
            has_special = any(c in line for c in ["|", ">", "<", ";", "&", "=", "/", "\\", "`"])
            if not has_flags and not has_special and words[0][0].isupper():
                return True

        # ── Lines that are just IP addresses or numbers (scan output) ──
        if re.match(r'^\d+\.\d+\.\d+\.\d+\s+\w', line) and len(words) >= 3:
            # Could be output like "192.168.1.1  00:11:22:33:44:55  Intel"
            # But NOT a command like "192.168.1.1" (which wouldn't be a command anyway)
            if not any(w.startswith("-") for w in words):
                return True

        # ── Blank-ish lines ──
        if line in ("---", "===", "***", "..."):
            return True

        # Everything else: TRUST the bash block
        return False

    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
