#!/usr/bin/env python3
"""Tested Templates Server & Secret Guardrails MCP Server (STDIO JSON-RPC 2.0).

Zero-dependency Model Context Protocol (MCP) server for Claude Code, Claude Desktop,
Cursor, Windsurf, OpenAI Codex CLI, and Google Gemini CLI.

Exposes 3 deterministic tools over JSON-RPC 2.0 (stdio):
  1. audit_server_snippet   — checks systemd units, sudoers, sshd_config, ufw, and docker-compose
                              for lockout patterns and ufw-bypassing public container ports.
  2. scan_secret_leak       — checks shell commands, configs, or git diffs for .env shell bypasses,
                              curl/wget argv token leaks (visible in ps aux), and 10+ API key formats.
  3. rewrite_safe_curl      — converts an unsafe curl command with inline tokens/headers into a
                              safe `curl -K -` stdin invocation that never leaks to /proc/<pid>/cmdline.

Usage:
  python3 mcp-ops-guardrails.py          # starts STDIO JSON-RPC 2.0 MCP loop
  python3 mcp-ops-guardrails.py --self-test  # runs built-in verification checks

Built and tested by AI agents at https://testedtemplates.com/code/mcp-server/
Full server & multi-agent ops kits:
  - Agent Ops Kit ($19): https://testedtemplatesshop.gumroad.com/l/agent-ops-kit?wanted=true
  - Multi-Vendor AI Council Ops Kit ($39): https://testedtemplatesshop.gumroad.com/l/council-ops-kit?wanted=true
"""
import hashlib
import json
import re
import sys

SERVER_INFO = {
    "name": "testedtemplates-ops-guardrails",
    "version": "1.0.0",
}

TOOLS = [
    {
        "name": "audit_server_snippet",
        "description": "Audit a systemd unit, sudoers rule, sshd_config, ufw snippet, or docker-compose.yml for server lockout and firewall-bypass bugs before applying.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "snippet": {"type": "string", "description": "Config or shell snippet to audit"}
            },
            "required": ["snippet"],
        },
    },
    {
        "name": "scan_secret_leak",
        "description": "Scan a shell command, config snippet, or git diff for .env shell bypasses (cat/python/base64), curl argv ps-table leaks, and literal API keys.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "text": {"type": "string", "description": "Command, config, or diff text to scan"}
            },
            "required": ["text"],
        },
    },
    {
        "name": "rewrite_safe_curl",
        "description": "Rewrite a curl command that exposes bearer/bot tokens in argv into a safe `curl -K -` stdin pattern so secrets never appear in `ps aux`.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "url": {"type": "string", "description": "Target URL (e.g. https://api.example.com/v1/...)"},
                "auth_header": {"type": "string", "description": "Optional Authorization header value (e.g. Authorization: Bearer $TOKEN)"},
                "data": {"type": "string", "description": "Optional request body or form data"},
            },
            "required": ["url"],
        },
    },
]


def audit_server_snippet(snippet: str):
    src = String_or_empty(snippet)
    if not src.strip():
        return [{"id": "empty", "severity": "info", "detail": "Empty snippet provided."}]
    findings = []

    def add(fid, sev, detail):
        findings.append({"id": fid, "severity": sev, "detail": detail})

    if re.search(r"NOPASSWD\s*:\s*ALL\b", src, re.I):
        add("passwordless-root", "high", "NOPASSWD: ALL grants unrestricted root execution. Scope sudoers to a single root-owned wrapper script.")
    if re.search(r"NOPASSWD\s*:\s*\S*(?:/bin/(?:ba)?sh|/usr/bin/(?:ba)?sh)\b", src, re.I):
        add("passwordless-shell", "high", "Passwordless bash/sh is equivalent to NOPASSWD: ALL.")
    if re.search(r"\bALL\s*=\s*\(\s*ALL(?::ALL)?\s*\)\s*ALL\b", src, re.I):
        add("full-sudo", "medium", "Unrestricted ALL=(ALL) ALL sudo rule; risky for unattended agents.")
    enables_ufw = bool(re.search(r"\bufw\b[^\n]{0,80}\benable\b", src, re.I))
    allows_ssh = bool(re.search(r"openssh|22/tcp|allow\s+22\b|allow\s+OpenSSH", src, re.I))
    if enables_ufw and not allows_ssh:
        add("ufw-no-ssh", "high", "Enables ufw without an explicit OpenSSH / 22/tcp allow rule; will drop active SSH sessions.")
    if re.search(r"permitrootlogin\s+yes\b", src, re.I):
        add("root-login", "high", "PermitRootLogin yes enables root password brute-forcing; use prohibit-password or no.")
    if re.search(r"passwordauthentication\s+yes\b", src, re.I):
        add("password-auth", "medium", "PasswordAuthentication yes allows password logins; prefer SSH keys only.")
    if re.search(r"systemctl\s+(?:restart|reload)\s+(?:ssh|sshd)\b|service\s+ssh(?:d)?\s+restart", src, re.I):
        add("ssh-restart", "medium", "Restarting sshd without running `sshd -t` first can lock out the only active session.")
    if re.search(r"/etc/sudoers", src, re.I) and re.search(r"\btee\b|>>|\bsed\s+-i\b", src, re.I) and not re.search(r"visudo", src, re.I):
        add("sudoers-write", "high", "Writing /etc/sudoers without `visudo -cf` can break sudo syntax for all users.")
    if re.search(r"chmod\s+777\b", src, re.I):
        add("chmod-777", "medium", "chmod 777 makes files world-writable; OpenSSH refuses authorized_keys if permissions are lax.")
    if re.search(r"\[Service\][\s\S]*?StartLimitIntervalSec\s*=", src, re.I):
        add("systemd-startlimit-section", "medium", "StartLimitIntervalSec under [Service] is silently ignored by systemd 255+; place it under [Unit].")
    if re.search(r"^\s*-\s*['\"]?\d{1,5}:\d{1,5}['\"]?\s*$", src, re.M):
        add("docker-ufw-bypass", "high", "Publishing `HOST:CONTAINER` in Docker Compose binds 0.0.0.0 and bypasses ufw via iptables nat chains. Bind `127.0.0.1:HOST:CONTAINER` instead.")
    return findings


def scan_secret_leak(text: str):
    src = String_or_empty(text)
    if not src.strip():
        return [{"id": "empty", "severity": "info", "detail": "Empty text provided."}]
    findings = []

    def add(fid, sev, detail):
        findings.append({"id": fid, "severity": sev, "detail": detail})

    if re.search(r"api\.telegram\.org/bot[0-9]{6,12}:[A-Za-z0-9_-]{20,}", src, re.I):
        add("telegram-argv-leak", "high", "Telegram bot token in URL/argv is visible to every local user via `ps aux` / `/proc/<pid>/cmdline`. Use `curl -K -`.")
    if re.search(r"\b(curl|wget)\b[^\n]*(-H|--header)\s+['\"]?Authorization:\s*(Bearer|Token)\s+[A-Za-z0-9_.\-\$]{8,}", src, re.I):
        add("curl-header-argv-leak", "high", "Authorization header on curl/wget command line leaks to `ps aux`. Pipe headers via `curl -K -` on stdin.")
    if re.search(r"\b(curl|wget)\b[^\n]*[?&](access_token|api_key|secret|token)=([A-Za-z0-9_.\-]{10,}|\$[A-Za-z_][A-Za-z0-9_]*)", src, re.I):
        add("curl-query-argv-leak", "high", "Secret token in curl/wget query string leaks to `ps aux`. Pipe URL and data via `curl -K -` on stdin.")
    if re.search(r"\b(cat|grep|awk|sed|base64|head|tail|less|more|python3?|node)\b[^\n]*(?<![\w./-])\.env(\.(?!example|sample|template)[a-zA-Z0-9_-]+)?(?![\w./-])", src, re.I):
        add("dotenv-shell-bypass", "high", "Shell command reads `.env` directly, bypassing Claude Code/Cursor `Read(.env)` deny rules and dumping secrets into LLM context.")
    if re.search(r"\btee\s+/dev/stderr\b", src):
        add("tee-dev-stderr-truncation", "medium", "`tee /dev/stderr` reopens redirected log files with O_TRUNC and wipes prior log lines. Use `>&2` instead.")

    key_patterns = [
        ("anthropic-key", "Anthropic API key (sk-ant-...)", r"\bsk-ant-[A-Za-z0-9_-]{16,}\b"),
        ("openai-key", "OpenAI API key (sk-proj-...)", r"\bsk-(?:proj|live|svcacct)-[A-Za-z0-9_-]{16,}\b"),
        ("github-token", "GitHub personal/fine-grained token", r"\b(?:ghp_[A-Za-z0-9]{20,}|github_pat_[A-Za-z0-9_]{20,})\b"),
        ("google-api-key", "Google API key (AIza...)", r"\bAIza[0-9A-Za-z_-]{30,}\b"),
        ("aws-access-key", "AWS Access Key ID (AKIA/ASIA...)", r"\b(?:AKIA|ASIA)[0-9A-Z]{16}\b"),
        ("stripe-live-key", "Stripe live secret key", r"\b(?:sk_live|rk_live)_[0-9a-zA-Z]{16,}\b"),
        ("slack-token", "Slack token (xoxb/xoxp...)", r"\bxox[baprs]-[0-9A-Za-z-]{16,}\b"),
        ("private-key", "PEM/OpenSSH private key block", r"-----BEGIN (?:RSA |EC |OPENSSH |PRIVATE )?PRIVATE KEY-----"),
    ]
    for fid, label, pat in key_patterns:
        m = re.search(pat, src)
        if m:
            val = m.group(0)
            sha = hashlib.sha256(val.encode("utf-8")).hexdigest()[:8]
            add(fid, "high", f"{label} detected [REDACTED:len={len(val)},sha256={sha}]. Never commit or paste live credentials.")
    return findings


def rewrite_safe_curl(url: str, auth_header: str = "", data: str = ""):
    u = String_or_empty(url).strip() or "https://api.example.com/v1/endpoint"
    lines = [
        "# Safe curl -K - pattern (nothing sensitive in argv / ps aux):",
        "curl -sS --max-time 30 -K - <<EOF",
        f'url = "{u}"',
    ]
    if auth_header:
        lines.append(f'header = "{auth_header}"')
    if data:
        lines.append(f'data = "{data}"')
    lines.append("EOF")
    return "\n".join(lines)


def String_or_empty(x):
    return "" if x is None else str(x)


def handle_request(msg):
    method = msg.get("method")
    req_id = msg.get("id")

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": SERVER_INFO,
            },
        }
    if method == "notifications/initialized":
        return None
    if method == "tools/list":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {"tools": TOOLS},
        }
    if method == "tools/call":
        params = msg.get("params") or {}
        name = params.get("name")
        args = params.get("arguments") or {}
        if name == "audit_server_snippet":
            res = audit_server_snippet(args.get("snippet", ""))
            text = json.dumps({"findings": res, "clean": len(res) == 0}, indent=2)
        elif name == "scan_secret_leak":
            res = scan_secret_leak(args.get("text", ""))
            text = json.dumps({"findings": res, "clean": len(res) == 0}, indent=2)
        elif name == "rewrite_safe_curl":
            text = rewrite_safe_curl(args.get("url", ""), args.get("auth_header", ""), args.get("data", ""))
        else:
            return {
                "jsonrpc": "2.0",
                "id": req_id,
                "error": {"code": -32601, "message": f"Unknown tool: {name}"},
            }
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {"content": [{"type": "text", "text": text}], "isError": False},
        }
    if req_id is not None:
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "error": {"code": -32601, "message": f"Method not found: {method}"},
        }
    return None


def main():
    if "--self-test" in sys.argv:
        r1 = audit_server_snippet("agent ALL=(ALL) NOPASSWD:ALL\nports:\n  - 8080:8080")
        assert any(x["id"] == "passwordless-root" for x in r1)
        assert any(x["id"] == "docker-ufw-bypass" for x in r1)
        fake_tg = "12345678:" + "AAHfiqksKZ8WmR2zP9xL1vK3jM5nQ7wE9rT"
        r2 = scan_secret_leak(f"cat .env\ncurl https://api.telegram.org/bot{fake_tg}/sendMessage")
        assert any(x["id"] == "dotenv-shell-bypass" for x in r2)
        assert any(x["id"] == "telegram-argv-leak" for x in r2)
        print("mcp-ops-guardrails self-test: OK")
        return

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except Exception:
            continue
        resp = handle_request(msg)
        if resp is not None:
            sys.stdout.write(json.dumps(resp) + "\n")
            sys.stdout.flush()


if __name__ == "__main__":
    main()