diff --git a/src-tauri/src/bin/wraith_mcp_bridge.rs b/src-tauri/src/bin/wraith_mcp_bridge.rs index d5f6ce8..7e978b6 100644 --- a/src-tauri/src/bin/wraith_mcp_bridge.rs +++ b/src-tauri/src/bin/wraith_mcp_bridge.rs @@ -224,6 +224,46 @@ fn handle_tools_list(id: Value) -> JsonRpcResponse { "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": "list_sessions", "description": "List all active Wraith sessions (SSH, RDP, PTY) with connection details", @@ -281,6 +321,14 @@ fn handle_tool_call(id: Value, port: u16, tool_name: &str, args: &Value) -> Json "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()), + "docker_ps" => call_wraith(port, "/mcp/docker/ps", args.clone()), + "docker_action" => call_wraith(port, "/mcp/docker/action", args.clone()), + "docker_exec" => call_wraith(port, "/mcp/docker/exec", args.clone()), + "service_status" => call_wraith(port, "/mcp/service/status", args.clone()), + "process_list" => call_wraith(port, "/mcp/process/list", args.clone()), + "git_status" => call_wraith(port, "/mcp/git/status", args.clone()), + "git_pull" => call_wraith(port, "/mcp/git/pull", args.clone()), + "git_log" => call_wraith(port, "/mcp/git/log", args.clone()), "terminal_screenshot" => { let result = call_wraith(port, "/mcp/screenshot", args.clone()); // Screenshot returns base64 PNG — wrap as image content for multimodal AI diff --git a/src-tauri/src/mcp/server.rs b/src-tauri/src/mcp/server.rs index c24d0d6..a2ab797 100644 --- a/src-tauri/src/mcp/server.rs +++ b/src-tauri/src/mcp/server.rs @@ -362,6 +362,76 @@ async fn tool_exec(handle: &std::sync::Arc>, Json(req): Json) -> Json> { + let session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) }; + match tool_exec(&session.handle, "docker ps -a --format '{{.Names}}|{{.Image}}|{{.Status}}|{{.Ports}}' 2>&1").await { Ok(o) => ok_response(o), Err(e) => err_response(e) } +} + +async fn handle_docker_action(AxumState(state): AxumState>, Json(req): Json) -> Json> { + let session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) }; + let cmd = match req.action.as_str() { + "start" => format!("docker start {} 2>&1", req.target), + "stop" => format!("docker stop {} 2>&1", req.target), + "restart" => format!("docker restart {} 2>&1", req.target), + "remove" => format!("docker rm -f {} 2>&1", req.target), + "logs" => format!("docker logs --tail 100 {} 2>&1", req.target), + "builder-prune" => "docker builder prune -f 2>&1".to_string(), + "system-prune" => "docker system prune -f 2>&1".to_string(), + _ => return err_response(format!("Unknown action: {}", req.action)), + }; + match tool_exec(&session.handle, &cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } +} + +async fn handle_docker_exec(AxumState(state): AxumState>, Json(req): Json) -> Json> { + let session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) }; + let cmd = format!("docker exec {} {} 2>&1", req.container, req.command); + match tool_exec(&session.handle, &cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } +} + +// ── Service/process handlers ───────────────────────────────────────────────── + +async fn handle_service_status(AxumState(state): AxumState>, Json(req): Json) -> Json> { + let session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) }; + match tool_exec(&session.handle, &format!("systemctl status {} --no-pager 2>&1 || service {} status 2>&1", req.target, req.target)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } +} + +async fn handle_process_list(AxumState(state): AxumState>, Json(req): Json) -> Json> { + let session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) }; + let filter = if req.target.is_empty() { "aux --sort=-%cpu | head -30".to_string() } else { format!("aux | grep -i {} | grep -v grep", req.target) }; + match tool_exec(&session.handle, &format!("ps {}", filter)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } +} + +// ── Git handlers ───────────────────────────────────────────────────────────── + +#[derive(Deserialize)] +struct GitRequest { session_id: String, path: String } + +async fn handle_git_status(AxumState(state): AxumState>, Json(req): Json) -> Json> { + let session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) }; + match tool_exec(&session.handle, &format!("cd {} && git status --short --branch 2>&1", req.path)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } +} + +async fn handle_git_pull(AxumState(state): AxumState>, Json(req): Json) -> Json> { + let session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) }; + match tool_exec(&session.handle, &format!("cd {} && git pull 2>&1", req.path)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } +} + +async fn handle_git_log(AxumState(state): AxumState>, Json(req): Json) -> Json> { + let session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) }; + match tool_exec(&session.handle, &format!("cd {} && git log --oneline -20 2>&1", req.path)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } +} + /// Start the MCP HTTP server and write the port to disk. pub async fn start_mcp_server( ssh: SshService, @@ -391,6 +461,14 @@ pub async fn start_mcp_server( .route("/mcp/tool/bandwidth", post(handle_tool_bandwidth)) .route("/mcp/tool/keygen", post(handle_tool_keygen)) .route("/mcp/tool/passgen", post(handle_tool_passgen)) + .route("/mcp/docker/ps", post(handle_docker_ps)) + .route("/mcp/docker/action", post(handle_docker_action)) + .route("/mcp/docker/exec", post(handle_docker_exec)) + .route("/mcp/service/status", post(handle_service_status)) + .route("/mcp/process/list", post(handle_process_list)) + .route("/mcp/git/status", post(handle_git_status)) + .route("/mcp/git/pull", post(handle_git_pull)) + .route("/mcp/git/log", post(handle_git_log)) .with_state(state); let listener = TcpListener::bind("127.0.0.1:0").await