//! Wraith MCP Bridge — stdio JSON-RPC proxy to Wraith's HTTP API. //! //! This binary is spawned by AI CLIs (Claude Code, Gemini CLI) as an MCP //! server. It reads JSON-RPC requests from stdin, translates them to HTTP //! calls against the running Wraith instance, and writes responses to stdout. //! //! The Wraith instance's MCP HTTP port is read from the data directory's //! `mcp-port` file. use std::io::{self, BufRead, Write}; use serde::{Deserialize, Serialize}; use serde_json::Value; #[derive(Deserialize)] #[allow(dead_code)] struct JsonRpcRequest { jsonrpc: String, id: Value, method: String, #[serde(default)] params: Value, } #[derive(Serialize)] struct JsonRpcResponse { jsonrpc: String, id: Value, #[serde(skip_serializing_if = "Option::is_none")] result: Option, #[serde(skip_serializing_if = "Option::is_none")] error: Option, } #[derive(Serialize)] struct JsonRpcError { code: i32, message: String, } fn get_mcp_port() -> Result { // Check standard locations for the port file let port_file = if let Ok(appdata) = std::env::var("APPDATA") { std::path::PathBuf::from(appdata).join("Wraith").join("mcp-port") } else if let Ok(home) = std::env::var("HOME") { if cfg!(target_os = "macos") { std::path::PathBuf::from(home).join("Library").join("Application Support").join("Wraith").join("mcp-port") } else { std::path::PathBuf::from(home).join(".local").join("share").join("wraith").join("mcp-port") } } else { return Err("Cannot determine data directory".to_string()); }; let port_str = std::fs::read_to_string(&port_file) .map_err(|e| format!("Cannot read MCP port file at {}: {} — is Wraith running?", port_file.display(), e))?; port_str.trim().parse::() .map_err(|e| format!("Invalid port in MCP port file: {}", e)) } fn handle_initialize(id: Value) -> JsonRpcResponse { JsonRpcResponse { jsonrpc: "2.0".to_string(), id, result: Some(serde_json::json!({ "protocolVersion": "2024-11-05", "capabilities": { "tools": {} }, "serverInfo": { "name": "wraith-terminal", "version": "1.0.0" } })), error: None, } } fn handle_tools_list(id: Value) -> JsonRpcResponse { JsonRpcResponse { jsonrpc: "2.0".to_string(), id, result: Some(serde_json::json!({ "tools": [ { "name": "terminal_read", "description": "Read recent terminal output from an active SSH or PTY session (ANSI codes stripped)", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string", "description": "The session ID to read from. Use list_sessions to find IDs." }, "lines": { "type": "number", "description": "Number of recent lines to return (default: 50)" } }, "required": ["session_id"] } }, { "name": "terminal_execute", "description": "Execute a command in an active SSH session and return the output", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string", "description": "The SSH session ID to execute in" }, "command": { "type": "string", "description": "The command to run" }, "timeout_ms": { "type": "number", "description": "Max wait time in ms (default: 5000)" } }, "required": ["session_id", "command"] } }, { "name": "terminal_screenshot", "description": "Capture a screenshot of an active RDP session as a base64-encoded PNG image for visual analysis", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string", "description": "The RDP session ID to screenshot" } }, "required": ["session_id"] } }, { "name": "sftp_list", "description": "List files in a directory on a remote host via SFTP", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string", "description": "The SSH session ID" }, "path": { "type": "string", "description": "Remote directory path" } }, "required": ["session_id", "path"] } }, { "name": "sftp_read", "description": "Read a file from a remote host via SFTP", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string", "description": "The SSH session ID" }, "path": { "type": "string", "description": "Remote file path" } }, "required": ["session_id", "path"] } }, { "name": "sftp_write", "description": "Write content to a file on a remote host via SFTP", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string", "description": "The SSH session ID" }, "path": { "type": "string", "description": "Remote file path" }, "content": { "type": "string", "description": "File content to write" } }, "required": ["session_id", "path", "content"] } }, { "name": "network_scan", "description": "Discover all devices on a remote network subnet via ARP + ping sweep", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "subnet": { "type": "string", "description": "First 3 octets, e.g. 192.168.1" } }, "required": ["session_id", "subnet"] } }, { "name": "port_scan", "description": "Scan TCP ports on a target host through an SSH session", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "target": { "type": "string" }, "ports": { "type": "array", "items": { "type": "number" }, "description": "Specific ports. Omit for quick scan of 24 common ports." } }, "required": ["session_id", "target"] } }, { "name": "ping", "description": "Ping a host through an SSH session", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "target": { "type": "string" } }, "required": ["session_id", "target"] } }, { "name": "traceroute", "description": "Traceroute to a host through an SSH session", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "target": { "type": "string" } }, "required": ["session_id", "target"] } }, { "name": "dns_lookup", "description": "DNS lookup for a domain through an SSH session", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "domain": { "type": "string" }, "record_type": { "type": "string", "description": "A, AAAA, MX, NS, TXT, CNAME, SOA, SRV, PTR" } }, "required": ["session_id", "domain"] } }, { "name": "whois", "description": "Whois lookup for a domain or IP through an SSH session", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "target": { "type": "string" } }, "required": ["session_id", "target"] } }, { "name": "wake_on_lan", "description": "Send Wake-on-LAN magic packet through an SSH session to wake a device", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "mac_address": { "type": "string", "description": "MAC address (AA:BB:CC:DD:EE:FF)" } }, "required": ["session_id", "mac_address"] } }, { "name": "bandwidth_test", "description": "Run an internet speed test on a remote host through SSH", "inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" } }, "required": ["session_id"] } }, { "name": "subnet_calc", "description": "Calculate subnet details from CIDR notation (no SSH needed)", "inputSchema": { "type": "object", "properties": { "cidr": { "type": "string", "description": "e.g. 192.168.1.0/24" } }, "required": ["cidr"] } }, { "name": "generate_ssh_key", "description": "Generate an SSH key pair (ed25519 or RSA)", "inputSchema": { "type": "object", "properties": { "key_type": { "type": "string", "description": "ed25519 or rsa" }, "comment": { "type": "string" } }, "required": ["key_type"] } }, { "name": "generate_password", "description": "Generate a cryptographically secure random password", "inputSchema": { "type": "object", "properties": { "length": { "type": "number" }, "uppercase": { "type": "boolean" }, "lowercase": { "type": "boolean" }, "digits": { "type": "boolean" }, "symbols": { "type": "boolean" } } } }, { "name": "list_sessions", "description": "List all active Wraith sessions (SSH, RDP, PTY) with connection details", "inputSchema": { "type": "object", "properties": {} } } ] })), error: None, } } fn call_wraith(port: u16, endpoint: &str, body: Value) -> Result { let url = format!("http://127.0.0.1:{}{}", port, endpoint); let body_str = serde_json::to_string(&body).unwrap_or_default(); let mut resp = ureq::post(url) .header("Content-Type", "application/json") .send(body_str.as_bytes()) .map_err(|e| format!("HTTP request to Wraith failed: {}", e))?; let resp_str = resp.body_mut().read_to_string() .map_err(|e| format!("Failed to read Wraith response: {}", e))?; let json: Value = serde_json::from_str(&resp_str) .map_err(|e| format!("Failed to parse Wraith response: {}", e))?; if json.get("ok").and_then(|v| v.as_bool()) == Some(true) { Ok(json.get("data").cloned().unwrap_or(Value::Null)) } else { let err_msg = json.get("error").and_then(|e| e.as_str()).unwrap_or("Unknown error"); Err(err_msg.to_string()) } } fn handle_tool_call(id: Value, port: u16, tool_name: &str, args: &Value) -> JsonRpcResponse { let result = match tool_name { "list_sessions" => call_wraith(port, "/mcp/sessions", serde_json::json!({})), "terminal_read" => call_wraith(port, "/mcp/terminal/read", args.clone()), "terminal_execute" => call_wraith(port, "/mcp/terminal/execute", args.clone()), "sftp_list" => call_wraith(port, "/mcp/sftp/list", args.clone()), "sftp_read" => call_wraith(port, "/mcp/sftp/read", args.clone()), "sftp_write" => call_wraith(port, "/mcp/sftp/write", args.clone()), "network_scan" => call_wraith(port, "/mcp/tool/scan-network", args.clone()), "port_scan" => call_wraith(port, "/mcp/tool/scan-ports", args.clone()), "ping" => call_wraith(port, "/mcp/tool/ping", args.clone()), "traceroute" => call_wraith(port, "/mcp/tool/traceroute", args.clone()), "dns_lookup" => call_wraith(port, "/mcp/tool/dns", args.clone()), "whois" => call_wraith(port, "/mcp/tool/whois", args.clone()), "wake_on_lan" => call_wraith(port, "/mcp/tool/wol", args.clone()), "bandwidth_test" => call_wraith(port, "/mcp/tool/bandwidth", args.clone()), "subnet_calc" => call_wraith(port, "/mcp/tool/subnet", args.clone()), "generate_ssh_key" => call_wraith(port, "/mcp/tool/keygen", args.clone()), "generate_password" => call_wraith(port, "/mcp/tool/passgen", args.clone()), "terminal_screenshot" => { let result = call_wraith(port, "/mcp/screenshot", args.clone()); // Screenshot returns base64 PNG — wrap as image content for multimodal AI return match result { Ok(b64) => JsonRpcResponse { jsonrpc: "2.0".to_string(), id, result: Some(serde_json::json!({ "content": [{ "type": "image", "data": b64, "mimeType": "image/png" }] })), error: None, }, Err(e) => JsonRpcResponse { jsonrpc: "2.0".to_string(), id, result: None, error: Some(JsonRpcError { code: -32000, message: e }), }, }; } _ => Err(format!("Unknown tool: {}", tool_name)), }; match result { Ok(data) => JsonRpcResponse { jsonrpc: "2.0".to_string(), id, result: Some(serde_json::json!({ "content": [{ "type": "text", "text": if data.is_string() { data.as_str().unwrap().to_string() } else { serde_json::to_string_pretty(&data).unwrap_or_default() } }] })), error: None, }, Err(e) => JsonRpcResponse { jsonrpc: "2.0".to_string(), id, result: None, error: Some(JsonRpcError { code: -32000, message: e }), }, } } fn main() { let port = match get_mcp_port() { Ok(p) => p, Err(e) => { eprintln!("wraith-mcp-bridge: {}", e); std::process::exit(1); } }; let stdin = io::stdin(); let mut stdout = io::stdout(); for line in stdin.lock().lines() { let line = match line { Ok(l) => l, Err(_) => break, }; if line.trim().is_empty() { continue; } let request: JsonRpcRequest = match serde_json::from_str(&line) { Ok(r) => r, Err(e) => { let err_resp = JsonRpcResponse { jsonrpc: "2.0".to_string(), id: Value::Null, result: None, error: Some(JsonRpcError { code: -32700, message: format!("Parse error: {}", e) }), }; let _ = writeln!(stdout, "{}", serde_json::to_string(&err_resp).unwrap()); let _ = stdout.flush(); continue; } }; let response = match request.method.as_str() { "initialize" => handle_initialize(request.id), "tools/list" => handle_tools_list(request.id), "tools/call" => { let tool_name = request.params.get("name") .and_then(|v| v.as_str()) .unwrap_or(""); let args = request.params.get("arguments") .cloned() .unwrap_or(Value::Object(serde_json::Map::new())); handle_tool_call(request.id, port, tool_name, &args) } "notifications/initialized" | "notifications/cancelled" => { // Notifications don't get responses continue; } _ => JsonRpcResponse { jsonrpc: "2.0".to_string(), id: request.id, result: None, error: Some(JsonRpcError { code: -32601, message: format!("Method not found: {}", request.method) }), }, }; let _ = writeln!(stdout, "{}", serde_json::to_string(&response).unwrap()); let _ = stdout.flush(); } }