#!/usr/bin/env python3
"""shellmcp: an MCP server with one tool, `run`, that runs a bash command. Stdlib only, stdio.

--mode readonly (default) runs every command in a bubblewrap sandbox: the filesystem is
read-only except a private /tmp, there is no network, no other process is visible, and the
--hide paths (always ~/.ssh and ~/.gnupg) look empty. --mode all runs commands with the
user's full permissions; let the MCP client ask before each call.
"""

import argparse
import getpass
import json
import os
import shutil
import signal
import socket
import subprocess
import sys

ENV = {"PAGER": "cat", "GIT_PAGER": "cat", "NO_COLOR": "1", "TERM": "dumb"}


def cut(text, limit):
    """text, or its head and tail around a note when it is longer than limit."""
    if len(text) <= limit:
        return text
    half = limit // 2
    return f"{text[:half]}\n[... {len(text) - 2 * half} characters cut ...]\n{text[-half:]}"


def run(command, cwd, timeout=60, limit=8000, prefix=()):
    """The exit code and combined stdout and stderr of `bash -c command`, as one text."""
    proc = subprocess.Popen([*prefix, "bash", "-c", command], cwd=cwd, stdin=subprocess.DEVNULL,
                            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                            env={**os.environ, **ENV}, start_new_session=True)
    try:
        out, _ = proc.communicate(timeout=timeout)
        head = f"exit code {proc.returncode}"
    except subprocess.TimeoutExpired:
        os.killpg(proc.pid, signal.SIGKILL)
        try:
            out, _ = proc.communicate(timeout=5)
        except subprocess.TimeoutExpired:  # a process left the group and holds the pipe
            out = b""
        head = f"timed out after {timeout} s, killed"
    return f"{head}\n{cut(out.decode(errors='replace'), limit)}"

def sandbox(cwd, hide=()):
    """The bwrap argv prefix of --mode readonly."""
    argv = ["bwrap", "--ro-bind", "/", "/", "--dev", "/dev", "--proc", "/proc",
            "--tmpfs", "/tmp", "--tmpfs", "/run", "--unshare-net", "--unshare-ipc",
            "--unshare-pid", "--die-with-parent", "--new-session"]
    for path in hide:
        if os.path.isdir(path):
            argv += ["--tmpfs", path]
        elif os.path.exists(path):
            argv += ["--ro-bind", "/dev/null", path]
    return argv + ["--chdir", cwd]

PREFERRED = {  # named in the tool description when installed, in this order
    "rg": "search file contents (instead of grep)", "fd": "find files (instead of find)",
    "eza": "list files (instead of ls)", "sd": "find and replace (instead of sed)",
    "dust": "disk usage (instead of du)", "procs": "processes (instead of ps)",
    "jq": "query JSON", "ast-grep": "search code by syntax", "tokei": "count lines of code"}


def describe(mode, cwd, timeout, limit):
    """The tool description: where commands run, the limits, and the modern tools installed."""
    text = (f"Run a bash command on {socket.gethostname()} as {getpass.getuser()} in {cwd}. "
            f"Returns its exit code and combined stdout and stderr (at most {limit} characters; "
            f"the middle of longer output is cut). Commands are killed after {timeout} s. "
            "There is no terminal: interactive programs, pagers and prompts fail.")
    if mode == "readonly":
        text += (" Read-only sandbox: nothing can be written except /tmp, which is emptied "
                 "after every command; no network; other processes are not visible. To change "
                 "something, give the user the command to run instead.")
    else:
        text += (" Commands run with the user's full permissions: they can change files, "
                 "processes and the network.")
    tools = [f"{name}: {what}" for name, what in PREFERRED.items() if shutil.which(name)]
    if tools:
        text += " Installed, prefer these: " + "; ".join(tools) + "."
    return text

PROTOCOLS = ("2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05")  # newest first
HIDDEN = ("~/.ssh", "~/.gnupg")


def error(mid, code, message):
    return {"jsonrpc": "2.0", "id": mid, "error": {"code": code, "message": message}}


def handle(message, tool, call):
    """The JSON-RPC response to one message; None for a notification."""
    mid, method, params = message.get("id"), message.get("method"), message.get("params") or {}
    if mid is None:
        return None
    if method == "initialize":
        asked = params.get("protocolVersion")
        result = {"protocolVersion": asked if asked in PROTOCOLS else PROTOCOLS[0],
                  "capabilities": {"tools": {}},
                  "serverInfo": {"name": "shellmcp", "version": "0.1.0"}}
    elif method == "ping":
        result = {}
    elif method == "tools/list":
        result = {"tools": [tool]}
    elif method == "tools/call":
        command = (params.get("arguments") or {}).get("command")
        if params.get("name") != "run" or not isinstance(command, str):
            return error(mid, -32602, "call the tool `run` with a string `command`")
        result = {"content": [{"type": "text", "text": call(command)}], "isError": False}
    else:
        return error(mid, -32601, f"unknown method {method}")
    return {"jsonrpc": "2.0", "id": mid, "result": result}


def main(argv=None):
    p = argparse.ArgumentParser(description="MCP server with one tool, run: a bash command.")
    p.add_argument("--mode", choices=("readonly", "all"), default="readonly",
                   help="readonly: bwrap sandbox (default); all: the user's full permissions")
    p.add_argument("--cwd", default=os.getcwd(), help="directory commands run in")
    p.add_argument("--timeout", type=int, default=60, help="seconds before a command is killed")
    p.add_argument("--max-output", type=int, default=8000, help="characters of output kept")
    p.add_argument("--hide", action="append", default=[], metavar="PATH",
                   help="readonly: also make PATH look empty (repeatable; ~/.ssh and ~/.gnupg "
                        "always are)")
    a = p.parse_args(argv)
    cwd = os.path.abspath(a.cwd)
    if a.mode == "readonly" and not shutil.which("bwrap"):
        sys.exit("shellmcp: --mode readonly needs bwrap (bubblewrap)")
    hide = [os.path.expanduser(h) for h in (*HIDDEN, *a.hide)]
    prefix = sandbox(cwd, hide) if a.mode == "readonly" else ()
    tool = {"name": "run", "description": describe(a.mode, cwd, a.timeout, a.max_output),
            "inputSchema": {"type": "object", "required": ["command"], "properties": {
                "command": {"type": "string", "description": "a bash command line"}}}}
    # ponytail: one request at a time; a long command delays the next call (no cancellation)
    for line in sys.stdin:
        if not line.strip():
            continue
        try:
            message = json.loads(line)
        except ValueError:
            message = None
        if not isinstance(message, dict):
            reply = error(None, -32700, "parse error")
        else:
            reply = handle(message, tool, lambda c: run(c, cwd, a.timeout, a.max_output, prefix))
        if reply is not None:
            print(json.dumps(reply), flush=True)


if __name__ == "__main__":
    main()