17973fc3dc
- New shell_escape() utility for safe command interpolation - Applied across all MCP tools, docker, scanner, network commands - MCP server generates random bearer token at startup - Token written to mcp-token file with 0600 permissions - All MCP HTTP requests require Authorization header - Bridge binary reads token and sends on every request Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
500 lines
25 KiB
Rust
500 lines
25 KiB
Rust
//! 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<Value>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
error: Option<JsonRpcError>,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct JsonRpcError {
|
|
code: i32,
|
|
message: String,
|
|
}
|
|
|
|
fn get_data_dir() -> Result<std::path::PathBuf, String> {
|
|
if let Ok(appdata) = std::env::var("APPDATA") {
|
|
Ok(std::path::PathBuf::from(appdata).join("Wraith"))
|
|
} else if let Ok(home) = std::env::var("HOME") {
|
|
if cfg!(target_os = "macos") {
|
|
Ok(std::path::PathBuf::from(home).join("Library").join("Application Support").join("Wraith"))
|
|
} else {
|
|
Ok(std::path::PathBuf::from(home).join(".local").join("share").join("wraith"))
|
|
}
|
|
} else {
|
|
Err("Cannot determine data directory".to_string())
|
|
}
|
|
}
|
|
|
|
fn get_mcp_port() -> Result<u16, String> {
|
|
let port_file = get_data_dir()?.join("mcp-port");
|
|
|
|
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::<u16>()
|
|
.map_err(|e| format!("Invalid port in MCP port file: {}", e))
|
|
}
|
|
|
|
fn get_mcp_token() -> Result<String, String> {
|
|
let token_file = get_data_dir()?.join("mcp-token");
|
|
|
|
let token = std::fs::read_to_string(&token_file)
|
|
.map_err(|e| format!("Cannot read MCP token file at {}: {} — is Wraith running?", token_file.display(), e))?;
|
|
|
|
Ok(token.trim().to_string())
|
|
}
|
|
|
|
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_type",
|
|
"description": "Type text into a terminal session (like a human typing). Optionally presses Enter after. Use this to send messages or commands without output capture.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"session_id": { "type": "string", "description": "The session ID to type into" },
|
|
"text": { "type": "string", "description": "The text to type" },
|
|
"press_enter": { "type": "boolean", "description": "Whether to press Enter after typing (default: true)" }
|
|
},
|
|
"required": ["session_id", "text"]
|
|
}
|
|
},
|
|
{
|
|
"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": "docker_ps",
|
|
"description": "List all Docker containers with status, image, and ports",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" } }, "required": ["session_id"] }
|
|
},
|
|
{
|
|
"name": "docker_action",
|
|
"description": "Perform a Docker action: start, stop, restart, remove, logs, builder-prune, system-prune",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "action": { "type": "string", "description": "start|stop|restart|remove|logs|builder-prune|system-prune" }, "target": { "type": "string", "description": "Container name (not needed for prune actions)" } }, "required": ["session_id", "action", "target"] }
|
|
},
|
|
{
|
|
"name": "docker_exec",
|
|
"description": "Execute a command inside a running Docker container",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "container": { "type": "string" }, "command": { "type": "string" } }, "required": ["session_id", "container", "command"] }
|
|
},
|
|
{
|
|
"name": "service_status",
|
|
"description": "Check systemd service status on a remote host",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "target": { "type": "string", "description": "Service name" } }, "required": ["session_id", "target"] }
|
|
},
|
|
{
|
|
"name": "process_list",
|
|
"description": "List processes on a remote host (top CPU by default, or filter by name)",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "target": { "type": "string", "description": "Process name filter (empty for top 30 by CPU)" } }, "required": ["session_id", "target"] }
|
|
},
|
|
{
|
|
"name": "git_status",
|
|
"description": "Get git status of a remote repository",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "path": { "type": "string", "description": "Path to the git repo on the remote host" } }, "required": ["session_id", "path"] }
|
|
},
|
|
{
|
|
"name": "git_pull",
|
|
"description": "Pull latest changes on a remote repository",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "path": { "type": "string" } }, "required": ["session_id", "path"] }
|
|
},
|
|
{
|
|
"name": "git_log",
|
|
"description": "Show recent commits on a remote repository",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "path": { "type": "string" } }, "required": ["session_id", "path"] }
|
|
},
|
|
{
|
|
"name": "rdp_click",
|
|
"description": "Click at a position in an RDP session (use terminal_screenshot first to see coordinates)",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "x": { "type": "number" }, "y": { "type": "number" }, "button": { "type": "string", "description": "left (default), right, or middle" } }, "required": ["session_id", "x", "y"] }
|
|
},
|
|
{
|
|
"name": "rdp_type",
|
|
"description": "Type text into an RDP session via clipboard paste",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "text": { "type": "string" } }, "required": ["session_id", "text"] }
|
|
},
|
|
{
|
|
"name": "rdp_clipboard",
|
|
"description": "Set the clipboard content on a remote RDP session",
|
|
"inputSchema": { "type": "object", "properties": { "session_id": { "type": "string" }, "text": { "type": "string" } }, "required": ["session_id", "text"] }
|
|
},
|
|
{
|
|
"name": "ssh_connect",
|
|
"description": "Open a new SSH connection through Wraith. Returns the session ID for use with other tools.",
|
|
"inputSchema": { "type": "object", "properties": {
|
|
"hostname": { "type": "string" },
|
|
"port": { "type": "number", "description": "Default: 22" },
|
|
"username": { "type": "string" },
|
|
"password": { "type": "string", "description": "Password (for password auth)" },
|
|
"private_key_path": { "type": "string", "description": "Path to SSH private key file on the local machine" }
|
|
}, "required": ["hostname", "username"] }
|
|
},
|
|
{
|
|
"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, token: &str, endpoint: &str, body: Value) -> Result<Value, String> {
|
|
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")
|
|
.header("Authorization", &format!("Bearer {}", token))
|
|
.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, token: &str, tool_name: &str, args: &Value) -> JsonRpcResponse {
|
|
let result = match tool_name {
|
|
"list_sessions" => call_wraith(port, token, "/mcp/sessions", serde_json::json!({})),
|
|
"terminal_type" => call_wraith(port, token, "/mcp/terminal/type", args.clone()),
|
|
"terminal_read" => call_wraith(port, token, "/mcp/terminal/read", args.clone()),
|
|
"terminal_execute" => call_wraith(port, token, "/mcp/terminal/execute", args.clone()),
|
|
"sftp_list" => call_wraith(port, token, "/mcp/sftp/list", args.clone()),
|
|
"sftp_read" => call_wraith(port, token, "/mcp/sftp/read", args.clone()),
|
|
"sftp_write" => call_wraith(port, token, "/mcp/sftp/write", args.clone()),
|
|
"network_scan" => call_wraith(port, token, "/mcp/tool/scan-network", args.clone()),
|
|
"port_scan" => call_wraith(port, token, "/mcp/tool/scan-ports", args.clone()),
|
|
"ping" => call_wraith(port, token, "/mcp/tool/ping", args.clone()),
|
|
"traceroute" => call_wraith(port, token, "/mcp/tool/traceroute", args.clone()),
|
|
"dns_lookup" => call_wraith(port, token, "/mcp/tool/dns", args.clone()),
|
|
"whois" => call_wraith(port, token, "/mcp/tool/whois", args.clone()),
|
|
"wake_on_lan" => call_wraith(port, token, "/mcp/tool/wol", args.clone()),
|
|
"bandwidth_test" => call_wraith(port, token, "/mcp/tool/bandwidth", args.clone()),
|
|
"subnet_calc" => call_wraith(port, token, "/mcp/tool/subnet", args.clone()),
|
|
"generate_ssh_key" => call_wraith(port, token, "/mcp/tool/keygen", args.clone()),
|
|
"generate_password" => call_wraith(port, token, "/mcp/tool/passgen", args.clone()),
|
|
"docker_ps" => call_wraith(port, token, "/mcp/docker/ps", args.clone()),
|
|
"docker_action" => call_wraith(port, token, "/mcp/docker/action", args.clone()),
|
|
"docker_exec" => call_wraith(port, token, "/mcp/docker/exec", args.clone()),
|
|
"service_status" => call_wraith(port, token, "/mcp/service/status", args.clone()),
|
|
"process_list" => call_wraith(port, token, "/mcp/process/list", args.clone()),
|
|
"git_status" => call_wraith(port, token, "/mcp/git/status", args.clone()),
|
|
"git_pull" => call_wraith(port, token, "/mcp/git/pull", args.clone()),
|
|
"git_log" => call_wraith(port, token, "/mcp/git/log", args.clone()),
|
|
"rdp_click" => call_wraith(port, token, "/mcp/rdp/click", args.clone()),
|
|
"rdp_type" => call_wraith(port, token, "/mcp/rdp/type", args.clone()),
|
|
"rdp_clipboard" => call_wraith(port, token, "/mcp/rdp/clipboard", args.clone()),
|
|
"ssh_connect" => call_wraith(port, token, "/mcp/ssh/connect", args.clone()),
|
|
"terminal_screenshot" => {
|
|
let result = call_wraith(port, token, "/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 token = match get_mcp_token() {
|
|
Ok(t) => t,
|
|
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, &token, 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();
|
|
}
|
|
}
|