#!/usr/bin/env python3
"""METATRON — lightweight Ollama terminal harness.

Talks to Ollama on nightmare (hardcoded). Launched by typing `metatron`:
  1. model selection screen (fetches live models + sizes from nightmare)
  2. chat REPL with streaming
  3. system access: `!command` + native run_command tool
  4. internet: `?query` (search) / `?fetch <url>` (read a page) + web_search tool

Stdlib-only (urllib + subprocess). Hardcoded: Ollama URL, SearXNG URL, system
prompt, tools. Only the model list is dynamic.
"""
import json
import os
import re
import subprocess
import sys
import time
import urllib.parse
import urllib.request
from html import unescape

# Windows console is cp1252 — replace unencodable chars instead of crashing on
# arbitrary Unicode the model emits (arrows, emoji, etc.).
for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(errors="replace")
    except Exception:
        pass

OLLAMA = "http://10.30.20.29:11434"   # nightmare
SEARXNG = "http://10.30.20.35:6969"   # self-hosted meta-search (CT 516)
HERE = os.path.dirname(os.path.abspath(__file__))
LAST_MODEL = os.path.join(HERE, ".last_model")

SYSTEM_PROMPT = (
    "You are METATRON, a system-access AI harness on a Windows Commando VM with a "
    "large tool catalog organized by category (recon, web, exploit, post, crack, "
    "wireless, sniff, osint, forensics, utils). DISCOVER LAZILY: call list_categories "
    "first, then list_tools(<category>) to see a category's tools, then run_command to "
    "execute one. Never enumerate everything at once — drill in only as needed. Use "
    "web_search for internet facts. Be concise and direct."
)

TOOLS = [{
    "type": "function",
    "function": {
        "name": "run_command",
        "description": "Execute a system command (cmd/PowerShell) and return its output. "
                       "Use for file ops, process management, recon, anything on the box.",
        "parameters": {
            "type": "object",
            "properties": {
                "command": {"type": "string", "description": "The command to run"},
            },
            "required": ["command"],
        },
    },
}, {
    "type": "function",
    "function": {
        "name": "web_search",
        "description": "Search the internet and return the top results (title, URL, snippet). "
                       "Use for current events, facts, anything you don't already know.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The search query"},
            },
            "required": ["query"],
        },
    },
}, {
    "type": "function",
    "function": {
        "name": "list_categories",
        "description": "List the available tool categories (recon, web, exploit, post, crack, "
                       "wireless, sniff, osint, forensics, utils) with descriptions. Use this "
                       "first to discover what tools exist.",
        "parameters": {"type": "object", "properties": {}},
    },
}, {
    "type": "function",
    "function": {
        "name": "list_tools",
        "description": "List the tools within a category (name, description, usage). Use after "
                       "list_categories to drill into a category.",
        "parameters": {
            "type": "object",
            "properties": {
                "category": {"type": "string", "description": "Category name (e.g. 'recon', 'post')"},
            },
            "required": ["category"],
        },
    },
}]

C = {"r": "\033[0m", "b": "\033[1m", "c": "\033[36m", "g": "\033[32m",
     "y": "\033[33m", "m": "\033[35m"}

# ---- Commando tool catalog, organized by category (lazy discovery) ----
# The model sees list_categories -> list_tools -> run_command, never the whole tree.
CATALOG = {
    "recon": {"desc": "Network recon & scanning", "tools": [
        {"name": "nmap", "desc": "Port/service/OS scanner",
         "usage": "nmap -sV -sC -p- <target>"},
        {"name": "masscan", "desc": "Ultra-fast port scanner",
         "usage": "masscan -p1-65535 <target> --rate=1000"},
        {"name": "rustscan", "desc": "Fast port scanner feeding nmap",
         "usage": "rustscan -a <target>"},
        {"name": "netdiscover", "desc": "ARP-based LAN discovery",
         "usage": "netdiscover -r <subnet>"},
        {"name": "arp-scan", "desc": "ARP host discovery",
         "usage": "arp-scan -l"},
        {"name": "amass", "desc": "Subdomain enumeration",
         "usage": "amass enum -d <domain>"},
        {"name": "sublist3r", "desc": "Subdomain enumeration",
         "usage": "sublist3r -d <domain>"},
        {"name": "theHarvester", "desc": "Email/domain OSINT",
         "usage": "theHarvester -d <domain> -b all"},
        {"name": "dnsrecon", "desc": "DNS reconnaissance",
         "usage": "dnsrecon -d <domain>"},
        {"name": "whatweb", "desc": "Web tech fingerprinting",
         "usage": "whatweb <url>"},
    ]},
    "web": {"desc": "Web app testing", "tools": [
        {"name": "nikto", "desc": "Web server vuln scanner",
         "usage": "nikto -h <url>"},
        {"name": "gobuster", "desc": "Directory/file brute force",
         "usage": "gobuster dir -u <url> -w <wordlist>"},
        {"name": "dirb", "desc": "Directory brute force",
         "usage": "dirb <url>"},
        {"name": "ffuf", "desc": "Fast web fuzzer",
         "usage": "ffuf -u <url>/FUZZ -w <wordlist>"},
        {"name": "wpscan", "desc": "WordPress scanner",
         "usage": "wpscan --url <url>"},
        {"name": "sqlmap", "desc": "SQL injection",
         "usage": "sqlmap -u <url> --dbs"},
        {"name": "commix", "desc": "Command injection",
         "usage": "commix --url <url>"},
    ]},
    "exploit": {"desc": "Exploitation frameworks", "tools": [
        {"name": "msfconsole", "desc": "Metasploit framework",
         "usage": "msfconsole"},
        {"name": "searchsploit", "desc": "Exploit-DB search",
         "usage": "searchsploit <term>"},
        {"name": "msfvenom", "desc": "Payload generation",
         "usage": "msfvenom -p <payload> LHOST=<ip> LPORT=<port> -f exe"},
    ]},
    "post": {"desc": "Post-exploitation & lateral movement", "tools": [
        {"name": "impacket-secretsdump", "desc": "Remote credential dump",
         "usage": "impacket-secretsdump <domain>/<user>:<pass>@<target>"},
        {"name": "impacket-psexec", "desc": "Remote shell (SMB)",
         "usage": "impacket-psexec <domain>/<user>:<pass>@<target>"},
        {"name": "impacket-wmiexec", "desc": "Remote shell (WMI)",
         "usage": "impacket-wmiexec <domain>/<user>:<pass>@<target>"},
        {"name": "impacket-GetNPUsers", "desc": "AS-REP roast (no creds)",
         "usage": "impacket-GetNPUsers <domain>/ -usersfile <file> -no-pass"},
        {"name": "impacket-GetUserSPNs", "desc": "Kerberoast",
         "usage": "impacket-GetUserSPNs <domain>/<user>:<pass>"},
        {"name": "crackmapexec", "desc": "SMB/WinRM brute & exec (also nxc)",
         "usage": "crackmapexec smb <target>"},
        {"name": "evil-winrm", "desc": "WinRM shell",
         "usage": "evil-winrm -i <target> -u <user> -p <pass>"},
        {"name": "mimikatz", "desc": "Credential dumping",
         "usage": "mimikatz"},
        {"name": "bloodhound-python", "desc": "AD enumeration",
         "usage": "bloodhound-python -d <domain> -u <user> -p <pass> -ns <dc>"},
        {"name": "responder", "desc": "LLMNR/NBT-NS poisoning",
         "usage": "responder -I <interface>"},
    ]},
    "crack": {"desc": "Password cracking & brute force", "tools": [
        {"name": "hashcat", "desc": "GPU password cracking",
         "usage": "hashcat -m <mode> <hash> <wordlist>"},
        {"name": "john", "desc": "Password cracking",
         "usage": "john --wordlist=<wordlist> <hashfile>"},
        {"name": "hydra", "desc": "Network login brute force",
         "usage": "hydra -l <user> -P <wordlist> <target> <service>"},
        {"name": "medusa", "desc": "Parallel brute force",
         "usage": "medusa -h <target> -u <user> -P <wordlist> -M <module>"},
    ]},
    "wireless": {"desc": "Wireless attacks", "tools": [
        {"name": "aircrack-ng", "desc": "WEP/WPA cracking",
         "usage": "aircrack-ng -w <wordlist> <capture.cap>"},
        {"name": "airmon-ng", "desc": "Monitor mode",
         "usage": "airmon-ng start <interface>"},
        {"name": "bettercap", "desc": "MITM/recon framework",
         "usage": "bettercap"},
    ]},
    "sniff": {"desc": "Sniffing & MITM", "tools": [
        {"name": "tshark", "desc": "CLI packet analysis",
         "usage": "tshark -i <interface>"},
        {"name": "tcpdump", "desc": "Packet capture",
         "usage": "tcpdump -i <interface> -w out.pcap"},
        {"name": "mitmproxy", "desc": "HTTP MITM proxy",
         "usage": "mitmproxy"},
        {"name": "ettercap", "desc": "MITM attacks",
         "usage": "ettercap -T -M arp"},
    ]},
    "osint": {"desc": "Open-source intelligence", "tools": [
        {"name": "theHarvester", "desc": "Email/domain harvesting",
         "usage": "theHarvester -d <domain> -b all"},
        {"name": "sherlock", "desc": "Username search across sites",
         "usage": "sherlock <username>"},
        {"name": "holehe", "desc": "Check email against services",
         "usage": "holehe <email>"},
        {"name": "spiderfoot", "desc": "OSINT automation",
         "usage": "spiderfoot -s <target>"},
    ]},
    "forensics": {"desc": "Forensics & analysis", "tools": [
        {"name": "volatility", "desc": "Memory forensics",
         "usage": "volatility -f <dump> <plugin>"},
        {"name": "binwalk", "desc": "Firmware/embedded analysis",
         "usage": "binwalk <file>"},
        {"name": "exiftool", "desc": "Metadata extraction",
         "usage": "exiftool <file>"},
        {"name": "strings", "desc": "String extraction",
         "usage": "strings <file>"},
    ]},
    "utils": {"desc": "Utilities & networking", "tools": [
        {"name": "nc", "desc": "Netcat — raw TCP/UDP",
         "usage": "nc -lvnp <port>"},
        {"name": "socat", "desc": "Networking relay",
         "usage": "socat TCP-LISTEN:<port>,reuseaddr,fork TCP:<host>:<port>"},
        {"name": "proxychains", "desc": "Proxy chaining",
         "usage": "proxychains <cmd>"},
        {"name": "curl", "desc": "HTTP client",
         "usage": "curl <url>"},
        {"name": "certutil", "desc": "Download/encode (Windows)",
         "usage": "certutil -urlcache -f <url> <out>"},
        {"name": "powershell", "desc": "Scripting / one-liners",
         "usage": "powershell -c <cmd>"},
    ]},
}


def list_categories():
    lines = [f"{cat} — {info['desc']} ({len(info['tools'])} tools)"
             for cat, info in CATALOG.items()]
    return "\n".join(lines)


def list_tools(category):
    info = CATALOG.get(category)
    if not info:
        return (f"unknown category '{category}'. "
                f"Categories: {', '.join(CATALOG.keys())}")
    return "\n".join(
        f"- {t['name']}: {t['desc']}\n    {t['usage']}"
        for t in info["tools"])


def get(path):
    req = urllib.request.Request(OLLAMA + path)
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)


def post_stream(path, payload):
    data = json.dumps(payload).encode()
    req = urllib.request.Request(OLLAMA + path, data=data,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as r:
        for line in r:
            line = line.strip()
            if line:
                yield json.loads(line)


def list_models():
    d = get("/api/tags")
    out = []
    for m in d.get("models", []):
        ps = (m.get("details") or {}).get("parameter_size", "") or ""
        out.append((m["name"], ps))
    return out


def load_last():
    try:
        return open(LAST_MODEL).read().strip()
    except Exception:
        return None


def save_last(model):
    try:
        open(LAST_MODEL, "w").write(model)
    except Exception:
        pass


def pick_model(models):
    last = load_last()
    names = [n for n, _ in models]
    print(f"\n{C['b']}{C['m']}  METATRON{C['r']} — select a model "
          f"{C['y']}(smaller = faster){C['r']}\n")
    for i, (name, ps) in enumerate(models, 1):
        mark = f" {C['m']}*{C['r']}" if name == last else ""
        print(f"  {C['c']}[{i}]{C['r']} {name}  {C['g']}{ps}{C['r']}{mark}")
    tail = f"  {C['c']}[0]{C['r']} {C['y']}quit{C['r']}"
    if last and last in names:
        tail += f"  {C['y']}(Enter = {last}){C['r']}"
    print(tail + "\n")
    while True:
        try:
            sel = input(f"{C['g']}model>{C['r']} ").strip()
            if sel == "" and last and last in names:
                return last
            if sel in ("0", "q", "quit", "exit"):
                sys.exit(0)
            idx = int(sel) - 1
            if 0 <= idx < len(models):
                return names[idx]
        except ValueError:
            pass
        print(f"{C['y']}  pick a number 1-{len(models)}{C['r']}")


def warm_model(model):
    sys.stdout.write(f"{C['y']}  warming {model} ...{C['r']}")
    sys.stdout.flush()
    payload = {"model": model, "messages": [{"role": "user", "content": "hi"}],
               "stream": False, "think": False, "keep_alive": "30m",
               "options": {"num_predict": 1}}
    try:
        data = json.dumps(payload).encode()
        req = urllib.request.Request(OLLAMA + "/api/chat", data=data,
                                     headers={"Content-Type": "application/json"})
        urllib.request.urlopen(req, timeout=120)
        sys.stdout.write(f"\r{C['g']}  ready.{C['r']}      \n")
    except Exception as e:
        sys.stdout.write(f"\r{C['y']}  warm failed ({e}){C['r']}      \n")


def supports_tools(model):
    try:
        data = json.dumps({"model": model}).encode()
        req = urllib.request.Request(OLLAMA + "/api/show", data=data,
                                     headers={"Content-Type": "application/json"})
        d = json.load(urllib.request.urlopen(req, timeout=15))
        return "tools" in (d.get("capabilities") or [])
    except Exception:
        return False


def run_command(cmd):
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True,
                           timeout=120)
        out = (r.stdout or "") + (r.stderr or "")
        return out.strip() or "(no output)"
    except Exception as e:
        return f"error: {e}"


def web_search(query):
    url = f"{SEARXNG}/search?q={urllib.parse.quote(query)}&format=json"
    try:
        d = json.load(urllib.request.urlopen(url, timeout=25))
    except Exception as e:
        return f"search error: {e}"
    results = d.get("results", [])
    if not results:
        return "no results"
    lines = []
    for r in results[:8]:
        lines.append(f"- {r.get('title', '')}\n  {r.get('url', '')}\n"
                     f"  {(r.get('content') or '')[:220]}")
    return "\n".join(lines)


def fetch_page(url):
    if not re.match(r"^https?://", url):
        url = "https://" + url
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        raw = urllib.request.urlopen(req, timeout=25).read().decode("utf-8", "ignore")
    except Exception as e:
        return f"fetch error: {e}"
    t = re.sub(r"<script.*?</script>", " ", raw, flags=re.S)
    t = re.sub(r"<style.*?</style>", " ", t, flags=re.S)
    t = re.sub(r"<[^>]+>", " ", t)
    t = unescape(t)
    t = re.sub(r"\s+", " ", t)
    return t.strip()[:2000]


def truncate(text, max_lines=50):
    lines = text.split("\n")
    if len(lines) > max_lines:
        return "\n".join(lines[:max_lines]) + f"\n  [... {len(lines) - max_lines} more lines]"
    return text


def print_help():
    print(f"""
  {C['b']}commands{C['r']}
    {C['c']}!<cmd>{C['r']}        run a system command
    {C['c']}?<query>{C['r']}      search the web (SearXNG)
    {C['c']}?fetch <url>{C['r']}  read a page's text
    {C['c']}/tools{C['r']}        list tool categories
    {C['c']}/tools <cat>{C['r']}  list tools in a category
    {C['c']}/save{C['r']}         save this session to markdown
    {C['c']}/help{C['r']}         this list
    {C['c']}exit{C['r']}          quit
""")


def save_transcript(messages, model):
    fn = os.path.join(os.getcwd(),
                      f"metatron-{time.strftime('%Y%m%d-%H%M%S')}.md")
    try:
        with open(fn, "w", encoding="utf-8") as f:
            f.write(f"# METATRON session — {model}\n\n")
            for m in messages:
                role, content = m.get("role"), m.get("content", "")
                if role == "user":
                    f.write(f"**you:** {content}\n\n")
                elif role == "assistant" and content:
                    f.write(f"**metatron:** {content}\n\n")
                elif role == "tool":
                    f.write(f"```\n{content}\n```\n\n")
        return fn
    except Exception as e:
        return f"save failed: {e}"


def chat(model):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    use_tools = supports_tools(model)
    print(f"\n{C['b']}  METATRON {C['c']}:: {model}{C['r']}  "
          f"{C['y']}(!cmd, ?search, ?fetch, /tools, /save, exit){C['r']}\n")
    while True:
        try:
            user = input(f"{C['g']}you>{C['r']} ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not user:
            continue
        if user.lower() in ("exit", "quit", "/q", "/exit"):
            break
        if user in ("/help", "help", "?"):
            print_help()
            continue
        if user == "/save":
            fn = save_transcript(messages, model)
            print(f"  {C['g']}saved: {fn}{C['r']}\n")
            continue
        if user == "/tools":
            print(f"{C['c']}  [categories]{C['r']}\n  {list_categories()}\n")
            continue
        if user.startswith("/tools "):
            cat = user[7:].strip()
            print(f"{C['c']}  [tools: {cat}]{C['r']}\n  {list_tools(cat)}\n")
            continue
        if user.startswith("!"):
            cmd = user[1:].strip()
            print(f"{C['y']}  $ {cmd}{C['r']}")
            print(f"  {truncate(run_command(cmd))}\n")
            continue
        if user.startswith("?"):
            q = user[1:].strip()
            if q.startswith("fetch "):
                url = q[6:].strip()
                print(f"{C['c']}  [fetch] {url}{C['r']}")
                print(f"  {fetch_page(url)}\n")
            elif re.match(r"^https?://|^www\.", q):
                print(f"{C['c']}  [fetch] {q}{C['r']}")
                print(f"  {fetch_page(q)}\n")
            else:
                print(f"{C['c']}  [search] {q}{C['r']}")
                print(f"  {web_search(q)}\n")
            continue

        messages.append({"role": "user", "content": user})
        for _ in range(6):
            payload = {"model": model, "messages": messages, "stream": True,
                       "think": False, "keep_alive": "30m"}
            if use_tools:
                payload["tools"] = TOOLS
            buf = ""
            tool_calls = []
            got_content = False
            sys.stdout.write("  ...")
            sys.stdout.flush()
            try:
                for obj in post_stream("/api/chat", payload):
                    msg = obj.get("message", {})
                    piece = msg.get("content") or ""
                    if piece:
                        if not got_content:
                            sys.stdout.write("\r  ")
                            got_content = True
                        buf += piece
                        sys.stdout.write(piece)
                        sys.stdout.flush()
                    if obj.get("done") and msg.get("tool_calls"):
                        tool_calls = msg["tool_calls"]
            except Exception as e:
                print(f"\n{C['y']}  [ollama error: {e}]{C['r']}")
                break
            if not got_content:
                sys.stdout.write("\r")
                sys.stdout.flush()

            if tool_calls:
                messages.append({"role": "assistant", "content": buf,
                                 "tool_calls": tool_calls})
                for tc in tool_calls:
                    fn = tc.get("function", {})
                    name = fn.get("name")
                    raw = fn.get("arguments") or {}
                    if isinstance(raw, str):
                        try:
                            args = json.loads(raw)
                        except Exception:
                            args = {}
                    else:
                        args = raw
                    if name == "run_command":
                        cmd = args.get("command", "")
                        print(f"\n{C['y']}  $ {cmd}{C['r']}")
                        out = run_command(cmd)
                        print(f"  {truncate(out)}")
                        messages.append({"role": "tool", "content": out})
                    elif name == "web_search":
                        q = args.get("query", "")
                        print(f"\n{C['c']}  [search] {q}{C['r']}")
                        out = web_search(q)
                        print(f"  {out[:600]}")
                        messages.append({"role": "tool", "content": out})
                    elif name == "list_categories":
                        out = list_categories()
                        print(f"\n{C['c']}  [categories]{C['r']}")
                        print(f"  {out}")
                        messages.append({"role": "tool", "content": out})
                    elif name == "list_tools":
                        cat = args.get("category", "")
                        out = list_tools(cat)
                        print(f"\n{C['c']}  [tools: {cat}]{C['r']}")
                        print(f"  {out[:900]}")
                        messages.append({"role": "tool", "content": out})
                continue
            else:
                messages.append({"role": "assistant", "content": buf})
                break
        print("\n")


def main():
    try:
        models = list_models()
    except Exception as e:
        print(f"{C['y']}  can't reach Ollama at {OLLAMA}: {e}{C['r']}")
        sys.exit(1)
    if not models:
        print(f"{C['y']}  no models on nightmare{C['r']}")
        sys.exit(1)
    model = pick_model(models)
    save_last(model)
    warm_model(model)
    chat(model)


if __name__ == "__main__":
    main()