sshp_mcp.py (9716 bytes)
1 #!/usr/bin/env python3 2 """Stdio MCP server that runs commands on many hosts through sshp. 3 4 Safety comes from the server config, never from the model: 5 - Only hosts listed in the hosts file can be targeted. 6 - Out of the box only read-only commands run; set "allow_write": true for a 7 general-purpose runner. A command must start with an allowed verb 8 and must not contain shell metacharacters, unless the config sets 9 "allow_write": true. 10 11 Config: $SSHP_AI_CONFIG, else ~/.config/sshp-ai/config.json. Keys: 12 hosts_file path to a newline-separated hosts file (default ~/.config/sshp-ai/hosts) 13 allow_write bool, default false 14 read_only_verbs list of first words allowed in read-only mode (optional override) 15 sshp sshp binary (default "sshp") 16 ssh_options list of "key=value" ssh options (default BatchMode=yes, ConnectTimeout=10) 17 """ 18 import json 19 import os 20 import re 21 import shlex 22 import subprocess 23 import sys 24 import tempfile 25 26 VERSION = "0.1.0" 27 CONFIG_DIR = os.path.expanduser("~/.config/sshp-ai") 28 29 READ_ONLY_VERBS = [ 30 "uptime", "uname", "hostname", "whoami", "id", "date", "df", "du", "free", 31 "cat", "head", "tail", "ls", "stat", "wc", "grep", "ps", "pgrep", "who", 32 "w", "last", "lsblk", "ip", "ss", "netstat", "systemctl", "journalctl", 33 "docker", "nproc", "lscpu", "env", "printenv", "which", "test", "echo", 34 "sha256sum", "md5sum", "find", "dmesg", "uptime", "mount", 35 ] 36 # Subcommands of multi-purpose verbs that are safe to read. 37 READ_ONLY_SUBCOMMANDS = { 38 "systemctl": {"status", "is-active", "is-enabled", "is-failed", "list-units", 39 "list-timers", "show", "cat", "--failed"}, 40 "docker": {"ps", "images", "inspect", "logs", "stats", "version", "info"}, 41 "ip": {"a", "addr", "address", "r", "route", "link", "-br", "-4", "-6"}, 42 } 43 FORBIDDEN = re.compile(r"[;&|<>`$\n\\]|\(|\)") 44 FIND_WRITE = {"-delete", "-exec", "-execdir", "-ok", "-okdir", "-fprint", "-fprintf", "-fls"} 45 EXIT_RE = re.compile(r"^\[(?P<host>[^\]]+)\] exited: (?P<code>-?\d+)") 46 LINE_RE = re.compile(r"^\[(?P<host>[^\]]+)\] ?(?P<text>.*)$") 47 48 49 def load_config(): 50 path = os.environ.get("SSHP_AI_CONFIG", os.path.join(CONFIG_DIR, "config.json")) 51 cfg = {} 52 if os.path.exists(path): 53 with open(path, encoding="utf-8") as f: 54 cfg = json.load(f) 55 cfg.setdefault("hosts_file", os.path.join(CONFIG_DIR, "hosts")) 56 cfg["hosts_file"] = os.path.expanduser(cfg["hosts_file"]) 57 cfg.setdefault("allow_write", False) 58 cfg.setdefault("read_only_verbs", READ_ONLY_VERBS) 59 cfg.setdefault("sshp", "sshp") 60 cfg.setdefault("ssh_options", ["BatchMode=yes", "ConnectTimeout=10"]) 61 return cfg 62 63 64 def read_hosts(path): 65 if not os.path.exists(path): 66 return [] 67 with open(path, encoding="utf-8") as f: 68 return [l.strip() for l in f if l.strip() and not l.lstrip().startswith("#")] 69 70 71 def check_command(command, cfg): 72 """Return an error string, or None if the command is allowed.""" 73 if not command.strip(): 74 return "empty command" 75 if cfg["allow_write"]: 76 return None 77 if FORBIDDEN.search(command): 78 return "read-only mode: shell metacharacters are not allowed" 79 try: 80 words = shlex.split(command) 81 except ValueError as e: 82 return f"cannot parse command: {e}" 83 verb = words[0] 84 if verb not in cfg["read_only_verbs"]: 85 return f"read-only mode: '{verb}' is not an allowed verb" 86 subs = READ_ONLY_SUBCOMMANDS.get(verb) 87 if subs is not None and len(words) > 1 and words[1] not in subs: 88 return f"read-only mode: '{verb} {words[1]}' is not allowed" 89 if verb == "find" and FIND_WRITE.intersection(words): 90 return "read-only mode: find actions that write or execute are not allowed" 91 if verb == "dmesg" and any(w.startswith("-C") or w in ("-c", "--clear", "--read-clear") for w in words): 92 return "read-only mode: dmesg clear is not allowed" 93 if verb == "date" and any(w.startswith("-s") or w == "--set" for w in words): 94 return "read-only mode: date --set is not allowed" 95 if verb == "hostname" and len(words) > 1 and not words[1].startswith("-"): 96 return "read-only mode: setting the hostname is not allowed" 97 return None 98 99 100 def sshp_run(args, cfg): 101 command = args.get("command", "") 102 allowed = read_hosts(cfg["hosts_file"]) 103 hosts = args.get("hosts") or allowed 104 bad = [h for h in hosts if h not in allowed] 105 if bad: 106 return {"error": f"hosts not in allowlist {cfg['hosts_file']}: {bad}"} 107 if not hosts: 108 return {"error": f"no hosts: add them to {cfg['hosts_file']}"} 109 err = check_command(command, cfg) 110 if err: 111 return {"error": err} 112 113 max_jobs = max(1, min(int(args.get("max_jobs", 10)), 50)) 114 timeout = max(1, min(int(args.get("timeout", 60)), 600)) 115 argv = [cfg["sshp"], "-e", "-c", "off", "-m", str(max_jobs)] 116 if args.get("dry_run"): 117 argv.append("-n") 118 user = args.get("user") 119 if user: 120 if not re.fullmatch(r"[A-Za-z0-9._-]+", user): 121 return {"error": "invalid user"} 122 argv += ["-l", user] 123 for opt in cfg["ssh_options"]: 124 argv += ["-o", opt] 125 126 # newline="\n": sshp keeps a CR from a Windows CRLF line as part of the host name. 127 with tempfile.NamedTemporaryFile("w", suffix=".hosts", delete=False, newline="\n") as f: 128 f.write("\n".join(hosts) + "\n") 129 hosts_path = f.name 130 try: 131 argv += ["-f", hosts_path, "--", command] 132 try: 133 p = subprocess.run(argv, capture_output=True, text=True, timeout=timeout) 134 except subprocess.TimeoutExpired: 135 return {"error": f"sshp did not finish within {timeout}s"} 136 except FileNotFoundError: 137 return {"error": f"sshp binary not found: {cfg['sshp']}"} 138 finally: 139 os.unlink(hosts_path) 140 141 results = {h: {"output": [], "exit_code": None} for h in hosts} 142 for line in p.stdout.splitlines(): 143 m = EXIT_RE.match(line) 144 if m and m["host"] in results: 145 results[m["host"]]["exit_code"] = int(m["code"]) 146 continue 147 m = LINE_RE.match(line) 148 if m and m["host"] in results: 149 results[m["host"]]["output"].append(m["text"]) 150 return { 151 "sshp_exit_code": p.returncode, 152 "dry_run": bool(args.get("dry_run")), 153 "hosts": {h: {"output": "\n".join(r["output"]), "exit_code": r["exit_code"]} 154 for h, r in results.items()}, 155 "sshp_stderr": p.stderr.strip(), 156 } 157 158 159 TOOLS = [ 160 { 161 "name": "sshp_run", 162 "description": "Run one command on many SSH hosts in parallel with sshp. " 163 "Returns per-host output (stdout+stderr merged) and exit code. " 164 "Hosts must be in the server allowlist. Without allow_write in the " 165 "server config, only read-only commands are accepted.", 166 "inputSchema": { 167 "type": "object", 168 "properties": { 169 "command": {"type": "string", "description": "Remote command, e.g. 'uptime' or 'df -h /'."}, 170 "hosts": {"type": "array", "items": {"type": "string"}, 171 "description": "Subset of allowlisted hosts. Default: all."}, 172 "max_jobs": {"type": "integer", "minimum": 1, "maximum": 50, "default": 10}, 173 "timeout": {"type": "integer", "minimum": 1, "maximum": 600, "default": 60, 174 "description": "Seconds for the whole run."}, 175 "user": {"type": "string", "description": "SSH login name (-l)."}, 176 "dry_run": {"type": "boolean", "default": False, 177 "description": "Pass -n: show what would run, start nothing."}, 178 }, 179 "required": ["command"], 180 }, 181 }, 182 { 183 "name": "sshp_hosts", 184 "description": "List allowlisted hosts and whether writes are enabled.", 185 "inputSchema": {"type": "object", "properties": {}}, 186 }, 187 ] 188 189 190 def handle(req): 191 method, rid = req.get("method"), req.get("id") 192 if rid is None: 193 return None # notification 194 if method == "initialize": 195 result = {"protocolVersion": req.get("params", {}).get("protocolVersion", "2025-06-18"), 196 "capabilities": {"tools": {}}, 197 "serverInfo": {"name": "sshp-ai", "version": VERSION}} 198 elif method == "ping": 199 result = {} 200 elif method == "tools/list": 201 result = {"tools": TOOLS} 202 elif method == "tools/call": 203 params = req.get("params", {}) 204 cfg = load_config() 205 if params.get("name") == "sshp_run": 206 out = sshp_run(params.get("arguments", {}), cfg) 207 elif params.get("name") == "sshp_hosts": 208 out = {"hosts": read_hosts(cfg["hosts_file"]), "allow_write": cfg["allow_write"]} 209 else: 210 return {"jsonrpc": "2.0", "id": rid, "error": {"code": -32602, "message": "unknown tool"}} 211 result = {"content": [{"type": "text", "text": json.dumps(out, indent=1)}], 212 "isError": "error" in out} 213 else: 214 return {"jsonrpc": "2.0", "id": rid, "error": {"code": -32601, "message": f"unknown method {method}"}} 215 return {"jsonrpc": "2.0", "id": rid, "result": result} 216 217 218 def main(): 219 for line in sys.stdin: 220 if not line.strip(): 221 continue 222 try: 223 resp = handle(json.loads(line)) 224 except Exception as e: # keep the server alive on bad input 225 resp = {"jsonrpc": "2.0", "id": None, "error": {"code": -32700, "message": str(e)}} 226 if resp is not None: 227 sys.stdout.write(json.dumps(resp) + "\n") 228 sys.stdout.flush() 229 230 231 if __name__ == "__main__": 232 main()