#!/usr/bin/python3 -I
"""MausBot guard: a Claude Code and Codex PreToolUse hook that blocks the common prompt-injection exfiltration paths.

Installed root-owned (managed settings) so an agent cannot edit or disable it. It is a speed bump, not a sandbox:
it reads commands as text, so obfuscated commands can slip past it. Pair it with an outbound firewall (LuLu).

Agents launched by OpenMausBot ("bots") get hard denials. Other Claude Code sessions (you at the keyboard, other agent tools) get an "ask" prompt instead, so nothing silently breaks for a person who is present. Codex can't show
"ask" prompts from hooks, so in Codex the guard asks with a native macOS dialog instead.
"""
import json, os, re, sys, subprocess, time

HOME = os.path.expanduser("~")
STATE = os.path.join(HOME, ".openmausbot-security")
LOG = os.path.join(STATE, "guard.log")
KNOWN = os.path.join(STATE, "known-recipients.txt")
CONFIG = os.path.join(STATE, "guard-config.json")  # optional: {"extra_hosts": [...], "extra_secret_paths": [...]}

def load_config():
    try:
        with open(CONFIG) as f: return json.load(f)
    except Exception: return {}
CFG = load_config()

# ---------- what counts as a secret file ----------
SECRET_PATHS = [
    r"\.config/op/", r"\.openmausbot/(logins-key|logins\.json\.enc|mcp-oauth\.json|hook-tokens|browser-engine-key|providers)",
    r"\.ssh/(?!.*\.pub\b)(?!config\b)(?!known_hosts)(?!agent\b)[^\s'\"]+", r"\.ssh/?(\s|$|['\"])",  # private keys, or the whole folder
    r"\.aws/(credentials|config)", r"\.config/gcloud", r"\.config/gh/hosts\.yml", r"(^|[\s'\"=/])\.npmrc", r"\.docker/config\.json",
    r"\.config/stripe", r"\.codex/auth\.json", r"\.wrangler/config", r"com\.vercel\.cli/.*auth\.json", r"\.config/gws/",
    r"(^|[\s'\"=/])\.netrc", r"Library/Keychains", r"\.git-credentials", r"\.config/claude/|\.claude/\.credentials",
    r"\.openmausbot-security/(known-recipients|guard-config)",
    r"(^|[\s/'\"=])\.env(\.[A-Za-z0-9_-]+)?(?=$|[\s'\"|;&)<>])(?<!\.example)(?<!\.sample)(?<!\.template)", r"\.dev\.vars\b",
] + CFG.get("extra_secret_paths", [])
SECRET_RE = re.compile("|".join(f"(?:{p})" for p in SECRET_PATHS))
# The documented way bots read logins: OP_SERVICE_ACCOUNT_TOKEN="$(cat ~/.config/op/openmausbot-bots.token)" op item get ...
OP_OK = re.compile(r'''^\s*(export\s+)?OP_SERVICE_ACCOUNT_TOKEN="?\$\(cat\s+(~|\$HOME|/Users/[^/]+)/\.config/op/openmausbot-bots\.token\)"?\s*;?\s*(&&\s*)?op\s+(item\s+(get|list)|read|whoami|vault\s+list)\b''')
KEYCHAIN_DUMP = re.compile(r"\bsecurity\s+(find-(generic|internet)-password\b[^|;&]*\s-w\b|dump-keychain|export\b)|\bgh\s+auth\s+token\b|\bop\s+service-account\b")

# ---------- outbound ----------
ALLOW_HOSTS = [
    "localhost", "127.0.0.1", "0.0.0.0", "::1",
    "anthropic.com", "claude.ai", "github.com", "githubusercontent.com", "googleapis.com", "google.com", "gstatic.com",
    "npmjs.org", "npmjs.com", "pypi.org", "pythonhosted.org", "vercel.com", "vercel.app", "cloudflare.com",
    "openai.com", "1password.com", "slack.com", 
    "osv.dev", "cisa.gov", "nvd.nist.gov",
] + CFG.get("extra_hosts", [])
NET_TOOLS = re.compile(r"(^|[\s;&|(`$])(curl|wget|http|https|xh|httpie|nc|ncat|netcat|socat|telnet|ftp|tftp|scp|sftp|rsync)(\s|$)")
RAW_SOCKET = re.compile(r"(^|[\s;&|(`$])(nc|ncat|netcat|socat|telnet|ftp|tftp)(\s|$)|/dev/tcp/|/dev/udp/")
UPLOAD_FLAGS = re.compile(r"\s(-d|--data[a-z-]*|-F|--form[a-z-]*|-T|--upload-file|--json|-X\s*(POST|PUT|PATCH)|--request\s+(POST|PUT|PATCH)|--post-(data|file)|--method=(POST|PUT|PATCH))(\s|=|$)")
URL_RE = re.compile(r"https?://([A-Za-z0-9.-]+)", re.I)
SCRIPT_NET = re.compile(r"(requests\.(post|put)|urllib\.request|http\.client|socket\.socket|fetch\(|axios\.|XMLHttpRequest|smtplib)")

SECRET_VALUE = re.compile(
    r"AKIA[0-9A-Z]{16}|ASIA[0-9A-Z]{16}|gh[pousr]_[A-Za-z0-9]{30,}|github_pat_[A-Za-z0-9_]{40,}|sk-(ant-)?[A-Za-z0-9_-]{24,}"
    r"|sk_(live|test)_[A-Za-z0-9]{16,}|rk_live_[A-Za-z0-9]{16,}|xox[abprs]-[A-Za-z0-9-]{10,}|-----BEGIN [A-Z ]*PRIVATE KEY-----"
    r"|ops_[A-Za-z0-9_-]{40,}|AIza[0-9A-Za-z_-]{35}|gho_[A-Za-z0-9]{30,}")
CARD = re.compile(r"(?<!\d)(?:\d[ -]?){13,19}(?!\d)")

def luhn(s):
    d = [int(c) for c in s if c.isdigit()]
    if not 13 <= len(d) <= 19: return False
    t = 0
    for i, x in enumerate(reversed(d)):
        if i % 2: x = x * 2 - 9 if x * 2 > 9 else x * 2
        t += x
    return t % 10 == 0

def has_secret_value(text, cards=True):
    if SECRET_VALUE.search(text): return True
    return cards and any(luhn(m.group()) for m in CARD.finditer(text))

def host_ok(h):
    h = h.lower().rstrip(".")
    return any(h == a or h.endswith("." + a) for a in ALLOW_HOSTS) or re.match(r"^(10|192\.168|172\.(1[6-9]|2\d|3[01]))\.", h) is not None

# ---------- sends ----------
SEND_TOOL = re.compile(r"(send_message|send_email|send_draft|reply_to|create_permission|add_permission|__reply$|__forward$|GMAIL_SEND|SEND_EMAIL|GMAIL_REPLY|FORWARD_MESSAGE|slack_send|schedule_message|share_file|SEND_MESSAGE|createScheduledPost|updateScheduledPost|sendScheduledPost)", re.I)
EMAIL = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")
BASH_SEND = re.compile(r"\bgws\b[^|;&]*\bgmail\b[^|;&]*\bsend\b|\bgws\b[^|;&]*\bdrive\b[^|;&]*\bpermissions\b[^|;&]*\bcreate\b|users\(\)\.(messages|drafts)\(\)\.send|gmail/v1/users/[^\s'\"]*/(messages|drafts)/send|\+send\b|\bimsg\s+send\b|osascript[^|;&]*(Messages|iMessage)[^|;&]*\bsend\b|tell application \"Messages\"")

# Mail settings that silently leak future mail (forwarding, filters, delegates, send-as): never for bots.
MAIL_SETTINGS_TOOL = re.compile(r"(gmail|mail|outlook)\w*?(filter|forward|delegate|send_?as|alias|imap|pop|vacation|auto_?reply|settings)|(filter|forward(ing)?_address|auto_?forward|delegate|send_?as)\w*?(gmail|mail)", re.I)
GIST_POST = re.compile(r"\bgh\s+gist\s+(create|new|edit)\b|\bgh\s+api\b[^|;&]*\bgists\b|api\.github\.com/gists")
MAIL_SETTINGS_BASH = re.compile(r"\bgws\b[^|;&]*\bgmail\b[^|;&]*\bsettings\b[^|;&]*\b(create|update|patch|delete|updateAutoForwarding|updateImap|updatePop|updateVacation)\b")

CODEX = False  # set in main(): Codex hook input carries turn_id

def is_bot(cwd=None):
    if any(k.startswith("OPENMAUSBOT") for k in os.environ): return True
    base = os.path.join(HOME, ".openmausbot")
    return os.getcwd().startswith(base) or str(cwd or "").startswith(base)

def bot_name(cwd):
    m = re.search(r"\.openmausbot/(?:task-)?workspaces/([0-9a-f-]{36})", cwd or "")
    if not m: return "A bot"
    try:
        b = json.load(open(os.path.join(HOME, ".openmausbot/bots.json")))
        for x in (b if isinstance(b, list) else b.get("bots", [])):
            if isinstance(x, dict) and x.get("id") == m.group(1): return f"Bot \"{x.get('name')}\""
    except Exception: pass
    return "A bot"

def log(kind, tool, detail):
    try:
        os.makedirs(STATE, exist_ok=True)
        detail = SECRET_VALUE.sub("<redacted>", detail)
        detail = CARD.sub(lambda m: "<redacted>" if luhn(m.group()) else m.group(), detail)
        with open(LOG, "a") as f: f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')}\t{kind}\t{tool}\t{detail[:300]!r}\n")
    except Exception: pass

def person_dialog(reason):
    """Codex has no hook "ask": ask you with a native dialog. True = allow once."""
    msg = ("A Codex session wants to do something the MausBot guard flags:\n\n" + reason[:500]).replace("\\", "\\\\").replace('"', '\\"')
    script = (f'display dialog "{msg}" with title "MausBot guard" buttons {{"Block", "Allow once"}} '
              f'default button "Block" cancel button "Block" with icon caution giving up after 600')
    try:
        r = subprocess.run(["/usr/bin/osascript", "-e", script], capture_output=True, text=True, timeout=620)
        return "Allow once" in r.stdout and "gave up:true" not in r.stdout
    except Exception: return False

def out(decision, reason):
    if CODEX and decision == "ask":
        if person_dialog(reason): sys.exit(0)
        decision, reason = "deny", reason + " You blocked it (or didn't answer)."
    print(json.dumps({"hookSpecificOutput": {"hookEventName": "PreToolUse", "permissionDecision": decision, "permissionDecisionReason": reason}}))
    sys.exit(0)

def block(tool, reason, detail, bot):
    log("DENY" if bot else "ASK", tool, detail)
    out("deny" if bot else "ask", "MausBot guard: " + reason + (" Do not try another way; tell your person what you needed and why." if bot else ""))

def known_recipients():
    try: return {l.strip().lower() for l in open(KNOWN) if l.strip() and not l.startswith("#")}
    except Exception: return set()

def recipient_known(r, known):
    r = r.lower()
    return r in known or ("@" in r and "@" + r.split("@", 1)[1] in known)

def ask_sam(who, what, recipients, preview):
    """Native macOS dialog; a prompt injection cannot answer it. Returns 'block' | 'once' | 'always'."""
    msg = (f"{who} wants to {what}\n\nTo: {', '.join(recipients) or '(not shown)'}\n\n"
           f"Message starts (written by the bot, may quote outside text):\n{preview[:400]}")
    msg = msg.replace("\\", "\\\\").replace('"', '\\"')
    script = (f'display dialog "{msg}" with title "MausBot guard: new recipient" buttons {{"Block", "Allow once", "Always allow"}} '
              f'default button "Block" cancel button "Block" with icon caution giving up after 600')
    try:
        r = subprocess.run(["/usr/bin/osascript", "-e", script], capture_output=True, text=True, timeout=620)
        o = r.stdout
        if "gave up:true" in o: return "block"
        if "Always allow" in o: return "always"
        if "Allow once" in o: return "once"
    except Exception: pass
    return "block"

def check_send(tool, inp, text_blob, bot, cwd):
    if has_secret_value(text_blob):
        block(tool, "the outgoing message contains what looks like a key, token or card number.", text_blob[:200], True)
    if not bot: return
    recips = []
    for k in ("to", "cc", "bcc", "recipient_email", "emailAddress", "email_addresses", "share_with", "grantee", "extra_recipients", "recipients", "emails", "email_address", "channel", "channel_id", "user", "email"):
        v = inp.get(k) if isinstance(inp, dict) else None
        if isinstance(v, str): recips += EMAIL.findall(v) or [v]
        elif isinstance(v, list): recips += [x if isinstance(x, str) else json.dumps(x) for x in v]
    if not recips and tool.startswith("Bash"): recips = EMAIL.findall(text_blob) or re.findall(r"imsg\s+send\s+(?:--to\s+)?(\S+)", text_blob)
    known = known_recipients()
    unknown = [r for r in recips if not recipient_known(r, known)] if recips else ["(recipients not visible, e.g. sending a saved draft)"]
    if recips and not unknown: return
    what = "post publicly" if re.search(r"ScheduledPost", tool) else "send a message"
    preview = ""
    for k in ("body", "text", "message", "content", "html_body", "subject"):
        if isinstance(inp, dict) and isinstance(inp.get(k), str): preview += inp[k] + "\n"
    ans = ask_sam(bot_name(cwd), what, unknown, preview or text_blob[:400])
    log("SEND-" + ans.upper(), tool, ", ".join(unknown))
    if ans == "block":
        out("deny", "MausBot guard: Your person did not approve sending to a new recipient (blocked or no answer in 10 minutes). Save it as a draft instead and tell your person.")
    if ans == "always":
        try:
            with open(KNOWN, "a") as f:
                for r in unknown:
                    if "@" in r or not r.startswith("("): f.write(r.lower() + "\n")
        except Exception: pass

def check_bash(cmd, bot, cwd, tool="Bash"):
    if isinstance(cmd, list): cmd = " ".join(map(str, cmd))
    flat = " ".join(str(cmd).split())
    if bot and MAIL_SETTINGS_BASH.search(flat):
        block(tool, "changing mail forwarding, filters, delegates or send-as settings is blocked for bots.", flat, bot)
    if bot and GIST_POST.search(flat):
        block(tool, "posting GitHub gists is blocked for bots (an easy place to leak data).", flat, bot)
    if SECRET_RE.search(flat) and not (OP_OK.search(flat) and not NET_TOOLS.search(flat) and len(SECRET_RE.findall(flat)) == 1):
        block(tool, "this command touches a secret file (keys, tokens, saved logins or .env).", flat, bot)
    if KEYCHAIN_DUMP.search(flat):
        block(tool, "reading raw passwords or tokens out of the keychain or gh is blocked.", flat, bot)
    if bot and RAW_SOCKET.search(flat):
        block(tool, "raw network tools (nc, socat, telnet, /dev/tcp) are blocked for bots.", flat, bot)
    if NET_TOOLS.search(flat) or SCRIPT_NET.search(flat):
        hosts = URL_RE.findall(flat)
        m = re.search(r"\b(scp|rsync|sftp)\b.*?\s(?:[\w.-]+@)?([\w.-]+):", flat)
        if m: hosts.append(m.group(2))
        bad = [h for h in hosts if not host_ok(h)]
        uploading = UPLOAD_FLAGS.search(" " + flat) or SCRIPT_NET.search(flat) or re.search(r"\b(scp|rsync|sftp)\b", flat)
        if has_secret_value(flat, cards=False):
            block(tool, "this network command carries what looks like a key or token.", flat, bot)
        if bad and (uploading or "$(" in flat or "`" in flat):
            block(tool, f"sending data to {', '.join(sorted(set(bad)))} is blocked (not a known destination). Reading web pages with a plain GET is fine.", flat, bot)
    if BASH_SEND.search(flat):
        check_send(tool, {}, flat, bot, cwd)

def check_file(tool, p, body, bot, cwd):
    p = p.replace(HOME, "~")
    if SECRET_RE.search(" " + p):
        block(tool, "this file holds secrets (keys, tokens, saved logins or env files).", p, bot)
    if body and re.search(r"\.(sh|py|js|mjs|cjs|ts|rb|command|zsh|bash|applescript|scpt)\b", p):
        if SECRET_RE.search(body.replace(HOME, "~")) and (NET_TOOLS.search(body) or SCRIPT_NET.search(body)):
            block(tool, "this script would read a secret file and use the network.", p, bot)
        if bot and MAIL_SETTINGS_BASH.search(body):
            block(tool, "this script would change mail forwarding or filter settings.", p, bot)
        if bot and GIST_POST.search(body):
            block(tool, "this script would post a GitHub gist.", p, bot)
        if BASH_SEND.search(body):  # a script that sends mail or messages gets the same recipient check, at write time
            check_send(tool, {}, body, bot, cwd)

def main():
    global CODEX
    try: data = json.load(sys.stdin)
    except Exception: sys.exit(0)
    tool = data.get("tool_name", "") or ""; inp = data.get("tool_input") or {}; cwd = data.get("cwd") or os.getcwd()
    CODEX = "turn_id" in data
    bot = is_bot(cwd)
    if bot and tool.startswith("mcp__") and MAIL_SETTINGS_TOOL.search(tool) and not re.search(r"(^|_)(get|list|read|search|fetch)(_|$)", tool.split("__")[-1], re.I):
        block(tool, "changing mail forwarding, filters, delegates or send-as settings is blocked for bots.", json.dumps(inp)[:200], bot)
    if tool == "Bash":
        check_bash(inp.get("command", ""), bot, cwd)
    elif tool in ("Read", "Edit", "Write", "MultiEdit", "NotebookEdit", "Grep", "Glob"):
        p = " ".join(str(inp.get(k, "")) for k in ("file_path", "path", "notebook_path", "pattern"))
        body = str(inp.get("content", "")) + str(inp.get("new_string", "")) + json.dumps(inp.get("edits", ""))
        check_file(tool, p, body if tool in ("Write", "Edit", "MultiEdit") else "", bot, cwd)
    elif tool == "apply_patch":  # Codex file edits
        patch = inp.get("command", "") if isinstance(inp, dict) else str(inp)
        if isinstance(patch, list): patch = "\n".join(map(str, patch))
        for m in re.finditer(r"^\*\*\* (?:Add File|Update File|Delete File|Move to): (.+)$", patch, re.M):
            check_file(tool, m.group(1), patch, bot, cwd)
    elif tool == "WebFetch" or tool.endswith("agent_browser_open") or tool.endswith("browser_navigate"):
        if has_secret_value(json.dumps(inp), cards=False):
            block(tool, "the address carries what looks like a key or token.", json.dumps(inp), True)
    elif SEND_TOOL.search(tool):
        check_send(tool, inp, json.dumps(inp), bot, cwd)
    elif tool.startswith("mcp__") and "browser" not in tool and re.search(r"(create|post|update|share|upload)", tool, re.I):
        if has_secret_value(json.dumps(inp), cards=False):
            block(tool, "this would send what looks like a key or token to an outside service.", json.dumps(inp)[:200], True)
    sys.exit(0)

if __name__ == "__main__":
    main()