15055aeb01
Some checks failed
Build & Sign Wraith / Build Windows + Sign (push) Failing after 6s
Every tool in Wraith is now callable by the AI through MCP: | MCP Tool | AI Use Case | |-------------------|------------------------------------------| | network_scan | "What devices are on this subnet?" | | port_scan | "Which servers have SSH open?" | | ping | "Is this host responding?" | | traceroute | "Show me the route to this server" | | dns_lookup | "What's the MX record for this domain?" | | whois | "Who owns this IP?" | | wake_on_lan | "Wake up the backup server" | | bandwidth_test | "How fast is this server's internet?" | | subnet_calc | "How many hosts in a /22?" | | generate_ssh_key | "Generate an ed25519 key pair" | | generate_password | "Generate a 32-char password" | | terminal_read | "What's on screen right now?" | | terminal_execute | "Run df -h on this server" | | terminal_screenshot| "What's that RDP error?" | | sftp_list/read/write| "Read the nginx config" | | list_sessions | "What sessions are active?" | 11 new HTTP endpoints on the MCP server. 11 new tool definitions in the bridge binary. The AI doesn't just chat — it scans, discovers, analyzes, and connects. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
387 lines
17 KiB
Rust
387 lines
17 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_mcp_port() -> Result<u16, String> {
|
|
// 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::<u16>()
|
|
.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<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")
|
|
.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();
|
|
}
|
|
}
|