refactor: 5 code quality fixes — shared ssh exec, wraith_log!, idiomatic Clone, Clone derives, sync RDP commands
All checks were successful
Build & Sign Wraith / Build Windows + Sign (push) Successful in 3m46s

- Extract duplicate exec_on_session into ssh::exec module; remove 4 private copies from tools_commands, tools_commands_r2, docker_commands, mcp::server
- Replace eprintln! with wraith_log! in theme::mod and workspace::mod
- Replace .map(|entry| entry.clone()) with .map(|r| r.value().clone()) for DashMap refs in ssh::session, mcp::mod, sftp::mod
- Add #[derive(Clone)] to ThemeService and WorkspaceService
- Remove unnecessary async from rdp_send_mouse, rdp_send_key, rdp_send_clipboard, disconnect_rdp, list_rdp_sessions

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Vantz Stockwell 2026-03-29 16:56:55 -04:00
parent d4bfb3d5fd
commit c2afb6a50f
12 changed files with 96 additions and 126 deletions

View File

@ -3,6 +3,7 @@
use tauri::State; use tauri::State;
use serde::Serialize; use serde::Serialize;
use crate::AppState; use crate::AppState;
use crate::ssh::exec::exec_on_session;
use crate::utils::shell_escape; use crate::utils::shell_escape;
#[derive(Debug, Serialize)] #[derive(Debug, Serialize)]
@ -38,7 +39,7 @@ pub struct DockerVolume {
pub async fn docker_list_containers(session_id: String, all: Option<bool>, state: State<'_, AppState>) -> Result<Vec<DockerContainer>, String> { pub async fn docker_list_containers(session_id: String, all: Option<bool>, state: State<'_, AppState>) -> Result<Vec<DockerContainer>, String> {
let session = state.ssh.get_session(&session_id).ok_or("Session not found")?; let session = state.ssh.get_session(&session_id).ok_or("Session not found")?;
let flag = if all.unwrap_or(true) { "-a" } else { "" }; let flag = if all.unwrap_or(true) { "-a" } else { "" };
let output = exec(&session.handle, &format!("docker ps {} --format '{{{{.ID}}}}|{{{{.Names}}}}|{{{{.Image}}}}|{{{{.Status}}}}|{{{{.Ports}}}}|{{{{.CreatedAt}}}}' 2>&1", flag)).await?; let output = exec_on_session(&session.handle, &format!("docker ps {} --format '{{{{.ID}}}}|{{{{.Names}}}}|{{{{.Image}}}}|{{{{.Status}}}}|{{{{.Ports}}}}|{{{{.CreatedAt}}}}' 2>&1", flag)).await?;
Ok(output.lines().filter(|l| !l.is_empty() && !l.starts_with("CONTAINER")).map(|line| { Ok(output.lines().filter(|l| !l.is_empty() && !l.starts_with("CONTAINER")).map(|line| {
let p: Vec<&str> = line.splitn(6, '|').collect(); let p: Vec<&str> = line.splitn(6, '|').collect();
DockerContainer { DockerContainer {
@ -55,7 +56,7 @@ pub async fn docker_list_containers(session_id: String, all: Option<bool>, state
#[tauri::command] #[tauri::command]
pub async fn docker_list_images(session_id: String, state: State<'_, AppState>) -> Result<Vec<DockerImage>, String> { pub async fn docker_list_images(session_id: String, state: State<'_, AppState>) -> Result<Vec<DockerImage>, String> {
let session = state.ssh.get_session(&session_id).ok_or("Session not found")?; let session = state.ssh.get_session(&session_id).ok_or("Session not found")?;
let output = exec(&session.handle, "docker images --format '{{.ID}}|{{.Repository}}|{{.Tag}}|{{.Size}}|{{.CreatedAt}}' 2>&1").await?; let output = exec_on_session(&session.handle, "docker images --format '{{.ID}}|{{.Repository}}|{{.Tag}}|{{.Size}}|{{.CreatedAt}}' 2>&1").await?;
Ok(output.lines().filter(|l| !l.is_empty()).map(|line| { Ok(output.lines().filter(|l| !l.is_empty()).map(|line| {
let p: Vec<&str> = line.splitn(5, '|').collect(); let p: Vec<&str> = line.splitn(5, '|').collect();
DockerImage { DockerImage {
@ -71,7 +72,7 @@ pub async fn docker_list_images(session_id: String, state: State<'_, AppState>)
#[tauri::command] #[tauri::command]
pub async fn docker_list_volumes(session_id: String, state: State<'_, AppState>) -> Result<Vec<DockerVolume>, String> { pub async fn docker_list_volumes(session_id: String, state: State<'_, AppState>) -> Result<Vec<DockerVolume>, String> {
let session = state.ssh.get_session(&session_id).ok_or("Session not found")?; let session = state.ssh.get_session(&session_id).ok_or("Session not found")?;
let output = exec(&session.handle, "docker volume ls --format '{{.Name}}|{{.Driver}}|{{.Mountpoint}}' 2>&1").await?; let output = exec_on_session(&session.handle, "docker volume ls --format '{{.Name}}|{{.Driver}}|{{.Mountpoint}}' 2>&1").await?;
Ok(output.lines().filter(|l| !l.is_empty()).map(|line| { Ok(output.lines().filter(|l| !l.is_empty()).map(|line| {
let p: Vec<&str> = line.splitn(3, '|').collect(); let p: Vec<&str> = line.splitn(3, '|').collect();
DockerVolume { DockerVolume {
@ -99,19 +100,6 @@ pub async fn docker_action(session_id: String, action: String, target: String, s
"system-prune-all" => "docker system prune -a -f 2>&1".to_string(), "system-prune-all" => "docker system prune -a -f 2>&1".to_string(),
_ => return Err(format!("Unknown docker action: {}", action)), _ => return Err(format!("Unknown docker action: {}", action)),
}; };
exec(&session.handle, &cmd).await exec_on_session(&session.handle, &cmd).await
} }
async fn exec(handle: &std::sync::Arc<tokio::sync::Mutex<russh::client::Handle<crate::ssh::session::SshClient>>>, cmd: &str) -> Result<String, String> {
let mut channel = { let h = handle.lock().await; h.channel_open_session().await.map_err(|e| format!("Exec failed: {}", e))? };
channel.exec(true, cmd).await.map_err(|e| format!("Exec failed: {}", e))?;
let mut output = String::new();
loop {
match channel.wait().await {
Some(russh::ChannelMsg::Data { ref data }) => { if let Ok(t) = std::str::from_utf8(data.as_ref()) { output.push_str(t); } }
Some(russh::ChannelMsg::Eof) | Some(russh::ChannelMsg::Close) | None => break,
_ => {}
}
}
Ok(output)
}

View File

@ -45,7 +45,7 @@ pub async fn rdp_get_frame(
/// - 0x0100 = negative wheel direction /// - 0x0100 = negative wheel direction
/// - 0x0400 = horizontal wheel /// - 0x0400 = horizontal wheel
#[tauri::command] #[tauri::command]
pub async fn rdp_send_mouse( pub fn rdp_send_mouse(
session_id: String, session_id: String,
x: u16, x: u16,
y: u16, y: u16,
@ -63,7 +63,7 @@ pub async fn rdp_send_mouse(
/// ///
/// `pressed` is `true` for key-down, `false` for key-up. /// `pressed` is `true` for key-down, `false` for key-up.
#[tauri::command] #[tauri::command]
pub async fn rdp_send_key( pub fn rdp_send_key(
session_id: String, session_id: String,
scancode: u16, scancode: u16,
pressed: bool, pressed: bool,
@ -74,7 +74,7 @@ pub async fn rdp_send_key(
/// Send clipboard text to an RDP session by simulating keystrokes. /// Send clipboard text to an RDP session by simulating keystrokes.
#[tauri::command] #[tauri::command]
pub async fn rdp_send_clipboard( pub fn rdp_send_clipboard(
session_id: String, session_id: String,
text: String, text: String,
state: State<'_, AppState>, state: State<'_, AppState>,
@ -86,7 +86,7 @@ pub async fn rdp_send_clipboard(
/// ///
/// Sends a graceful shutdown to the RDP server and removes the session. /// Sends a graceful shutdown to the RDP server and removes the session.
#[tauri::command] #[tauri::command]
pub async fn disconnect_rdp( pub fn disconnect_rdp(
session_id: String, session_id: String,
state: State<'_, AppState>, state: State<'_, AppState>,
) -> Result<(), String> { ) -> Result<(), String> {
@ -95,7 +95,7 @@ pub async fn disconnect_rdp(
/// List all active RDP sessions (metadata only). /// List all active RDP sessions (metadata only).
#[tauri::command] #[tauri::command]
pub async fn list_rdp_sessions( pub fn list_rdp_sessions(
state: State<'_, AppState>, state: State<'_, AppState>,
) -> Result<Vec<RdpSessionInfo>, String> { ) -> Result<Vec<RdpSessionInfo>, String> {
Ok(state.rdp.list_sessions()) Ok(state.rdp.list_sessions())

View File

@ -4,6 +4,7 @@ use tauri::State;
use serde::Serialize; use serde::Serialize;
use crate::AppState; use crate::AppState;
use crate::ssh::exec::exec_on_session;
use crate::utils::shell_escape; use crate::utils::shell_escape;
// ── Ping ───────────────────────────────────────────────────────────────────── // ── Ping ─────────────────────────────────────────────────────────────────────
@ -185,32 +186,3 @@ pub fn tool_generate_password_inner(
Ok(password) Ok(password)
} }
// ── Helper ───────────────────────────────────────────────────────────────────
async fn exec_on_session(
handle: &std::sync::Arc<tokio::sync::Mutex<russh::client::Handle<crate::ssh::session::SshClient>>>,
cmd: &str,
) -> Result<String, String> {
let mut channel = {
let h = handle.lock().await;
h.channel_open_session().await.map_err(|e| format!("Exec channel failed: {}", e))?
};
channel.exec(true, cmd).await.map_err(|e| format!("Exec failed: {}", e))?;
let mut output = String::new();
loop {
match channel.wait().await {
Some(russh::ChannelMsg::Data { ref data }) => {
if let Ok(text) = std::str::from_utf8(data.as_ref()) {
output.push_str(text);
}
}
Some(russh::ChannelMsg::Eof) | Some(russh::ChannelMsg::Close) | None => break,
Some(russh::ChannelMsg::ExitStatus { .. }) => {}
_ => {}
}
}
Ok(output)
}

View File

@ -4,6 +4,7 @@ use tauri::State;
use serde::Serialize; use serde::Serialize;
use crate::AppState; use crate::AppState;
use crate::ssh::exec::exec_on_session;
use crate::utils::shell_escape; use crate::utils::shell_escape;
// ── DNS Lookup ─────────────────────────────────────────────────────────────── // ── DNS Lookup ───────────────────────────────────────────────────────────────
@ -181,27 +182,3 @@ fn to_ip(val: u32) -> String {
format!("{}.{}.{}.{}", val >> 24, (val >> 16) & 0xFF, (val >> 8) & 0xFF, val & 0xFF) format!("{}.{}.{}.{}", val >> 24, (val >> 16) & 0xFF, (val >> 8) & 0xFF, val & 0xFF)
} }
// ── Helper ───────────────────────────────────────────────────────────────────
async fn exec_on_session(
handle: &std::sync::Arc<tokio::sync::Mutex<russh::client::Handle<crate::ssh::session::SshClient>>>,
cmd: &str,
) -> Result<String, String> {
let mut channel = {
let h = handle.lock().await;
h.channel_open_session().await.map_err(|e| format!("Exec channel failed: {}", e))?
};
channel.exec(true, cmd).await.map_err(|e| format!("Exec failed: {}", e))?;
let mut output = String::new();
loop {
match channel.wait().await {
Some(russh::ChannelMsg::Data { ref data }) => {
if let Ok(text) = std::str::from_utf8(data.as_ref()) { output.push_str(text); }
}
Some(russh::ChannelMsg::Eof) | Some(russh::ChannelMsg::Close) | None => break,
Some(russh::ChannelMsg::ExitStatus { .. }) => {}
_ => {}
}
}
Ok(output)
}

View File

@ -36,7 +36,7 @@ impl ScrollbackRegistry {
/// Get the scrollback buffer for a session. /// Get the scrollback buffer for a session.
pub fn get(&self, session_id: &str) -> Option<Arc<ScrollbackBuffer>> { pub fn get(&self, session_id: &str) -> Option<Arc<ScrollbackBuffer>> {
self.buffers.get(session_id).map(|entry| entry.clone()) self.buffers.get(session_id).map(|r| r.value().clone())
} }
/// Remove a session's scrollback buffer. /// Remove a session's scrollback buffer.

View File

@ -19,6 +19,7 @@ use tokio::net::TcpListener;
use crate::mcp::ScrollbackRegistry; use crate::mcp::ScrollbackRegistry;
use crate::rdp::RdpService; use crate::rdp::RdpService;
use crate::sftp::SftpService; use crate::sftp::SftpService;
use crate::ssh::exec::exec_on_session;
use crate::ssh::session::SshService; use crate::ssh::session::SshService;
use crate::utils::shell_escape; use crate::utils::shell_escape;
@ -308,32 +309,32 @@ struct ToolPassgenRequest { length: Option<usize>, uppercase: Option<bool>, lowe
async fn handle_tool_ping(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> { async fn handle_tool_ping(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> {
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 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!("ping -c 4 {} 2>&1", shell_escape(&req.target))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("ping -c 4 {} 2>&1", shell_escape(&req.target))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_tool_traceroute(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> { async fn handle_tool_traceroute(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> {
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 session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) };
let t = shell_escape(&req.target); let t = shell_escape(&req.target);
match tool_exec(&session.handle, &format!("traceroute {} 2>&1 || tracert {} 2>&1", t, t)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("traceroute {} 2>&1 || tracert {} 2>&1", t, t)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_tool_dns(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolDnsRequest>) -> Json<McpResponse<String>> { async fn handle_tool_dns(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolDnsRequest>) -> Json<McpResponse<String>> {
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 session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) };
let rt = shell_escape(&req.record_type.unwrap_or_else(|| "A".to_string())); let rt = shell_escape(&req.record_type.unwrap_or_else(|| "A".to_string()));
let d = shell_escape(&req.domain); let d = shell_escape(&req.domain);
match tool_exec(&session.handle, &format!("dig {} {} +short 2>/dev/null || nslookup -type={} {} 2>/dev/null || host -t {} {} 2>/dev/null", d, rt, rt, d, rt, d)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("dig {} {} +short 2>/dev/null || nslookup -type={} {} 2>/dev/null || host -t {} {} 2>/dev/null", d, rt, rt, d, rt, d)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_tool_whois(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> { async fn handle_tool_whois(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> {
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 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!("whois {} 2>&1 | head -80", shell_escape(&req.target))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("whois {} 2>&1 | head -80", shell_escape(&req.target))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_tool_wol(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolWolRequest>) -> Json<McpResponse<String>> { async fn handle_tool_wol(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolWolRequest>) -> Json<McpResponse<String>> {
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 session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) };
let mac_clean = req.mac_address.replace([':', '-'], ""); let mac_clean = req.mac_address.replace([':', '-'], "");
let cmd = format!(r#"python3 -c "import socket;mac=bytes.fromhex({});pkt=b'\xff'*6+mac*16;s=socket.socket(socket.AF_INET,socket.SOCK_DGRAM);s.setsockopt(socket.SOL_SOCKET,socket.SO_BROADCAST,1);s.sendto(pkt,('255.255.255.255',9));s.close();print('WoL sent to {}')" 2>&1"#, shell_escape(&mac_clean), shell_escape(&req.mac_address)); let cmd = format!(r#"python3 -c "import socket;mac=bytes.fromhex({});pkt=b'\xff'*6+mac*16;s=socket.socket(socket.AF_INET,socket.SOCK_DGRAM);s.setsockopt(socket.SOL_SOCKET,socket.SO_BROADCAST,1);s.sendto(pkt,('255.255.255.255',9));s.close();print('WoL sent to {}')" 2>&1"#, shell_escape(&mac_clean), shell_escape(&req.mac_address));
match tool_exec(&session.handle, &cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_tool_scan_network(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolScanNetworkRequest>) -> Json<McpResponse<serde_json::Value>> { async fn handle_tool_scan_network(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolScanNetworkRequest>) -> Json<McpResponse<serde_json::Value>> {
@ -364,7 +365,7 @@ async fn handle_tool_subnet(_state: AxumState<Arc<McpServerState>>, Json(req): J
async fn handle_tool_bandwidth(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionOnly>) -> Json<McpResponse<String>> { async fn handle_tool_bandwidth(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionOnly>) -> Json<McpResponse<String>> {
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 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 = r#"if command -v speedtest-cli >/dev/null 2>&1; then speedtest-cli --simple 2>&1; elif command -v curl >/dev/null 2>&1; then curl -o /dev/null -w "Download: %{speed_download} bytes/sec\n" https://speed.cloudflare.com/__down?bytes=25000000 2>/dev/null; else echo "No speedtest tool found"; fi"#; let cmd = r#"if command -v speedtest-cli >/dev/null 2>&1; then speedtest-cli --simple 2>&1; elif command -v curl >/dev/null 2>&1; then curl -o /dev/null -w "Download: %{speed_download} bytes/sec\n" https://speed.cloudflare.com/__down?bytes=25000000 2>/dev/null; else echo "No speedtest tool found"; fi"#;
match tool_exec(&session.handle, cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_tool_keygen(_state: AxumState<Arc<McpServerState>>, Json(req): Json<ToolKeygenRequest>) -> Json<McpResponse<serde_json::Value>> { async fn handle_tool_keygen(_state: AxumState<Arc<McpServerState>>, Json(req): Json<ToolKeygenRequest>) -> Json<McpResponse<serde_json::Value>> {
@ -381,20 +382,6 @@ async fn handle_tool_passgen(_state: AxumState<Arc<McpServerState>>, Json(req):
} }
} }
async fn tool_exec(handle: &std::sync::Arc<tokio::sync::Mutex<russh::client::Handle<crate::ssh::session::SshClient>>>, cmd: &str) -> Result<String, String> {
let mut channel = { let h = handle.lock().await; h.channel_open_session().await.map_err(|e| format!("Exec failed: {}", e))? };
channel.exec(true, cmd).await.map_err(|e| format!("Exec failed: {}", e))?;
let mut output = String::new();
loop {
match channel.wait().await {
Some(russh::ChannelMsg::Data { ref data }) => { if let Ok(t) = std::str::from_utf8(data.as_ref()) { output.push_str(t); } }
Some(russh::ChannelMsg::Eof) | Some(russh::ChannelMsg::Close) | None => break,
_ => {}
}
}
Ok(output)
}
// ── Docker handlers ────────────────────────────────────────────────────────── // ── Docker handlers ──────────────────────────────────────────────────────────
#[derive(Deserialize)] #[derive(Deserialize)]
@ -408,7 +395,7 @@ struct DockerExecRequest { session_id: String, container: String, command: Strin
async fn handle_docker_ps(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<DockerListRequest>) -> Json<McpResponse<String>> { async fn handle_docker_ps(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<DockerListRequest>) -> Json<McpResponse<String>> {
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 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) } match exec_on_session(&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<Arc<McpServerState>>, Json(req): Json<DockerActionRequest>) -> Json<McpResponse<String>> { async fn handle_docker_action(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<DockerActionRequest>) -> Json<McpResponse<String>> {
@ -424,13 +411,13 @@ async fn handle_docker_action(AxumState(state): AxumState<Arc<McpServerState>>,
"system-prune" => "docker system prune -f 2>&1".to_string(), "system-prune" => "docker system prune -f 2>&1".to_string(),
_ => return err_response(format!("Unknown action: {}", req.action)), _ => 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) } match exec_on_session(&session.handle, &cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_docker_exec(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<DockerExecRequest>) -> Json<McpResponse<String>> { async fn handle_docker_exec(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<DockerExecRequest>) -> Json<McpResponse<String>> {
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 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", shell_escape(&req.container), shell_escape(&req.command)); let cmd = format!("docker exec {} {} 2>&1", shell_escape(&req.container), shell_escape(&req.command));
match tool_exec(&session.handle, &cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &cmd).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
// ── Service/process handlers ───────────────────────────────────────────────── // ── Service/process handlers ─────────────────────────────────────────────────
@ -438,13 +425,13 @@ async fn handle_docker_exec(AxumState(state): AxumState<Arc<McpServerState>>, Js
async fn handle_service_status(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> { async fn handle_service_status(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> {
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 session = match state.ssh.get_session(&req.session_id) { Some(s) => s, None => return err_response(format!("Session {} not found", req.session_id)) };
let t = shell_escape(&req.target); let t = shell_escape(&req.target);
match tool_exec(&session.handle, &format!("systemctl status {} --no-pager 2>&1 || service {} status 2>&1", t, t)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("systemctl status {} --no-pager 2>&1 || service {} status 2>&1", t, t)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_process_list(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> { async fn handle_process_list(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<ToolSessionTarget>) -> Json<McpResponse<String>> {
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 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", shell_escape(&req.target)) }; let filter = if req.target.is_empty() { "aux --sort=-%cpu | head -30".to_string() } else { format!("aux | grep -i {} | grep -v grep", shell_escape(&req.target)) };
match tool_exec(&session.handle, &format!("ps {}", filter)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("ps {}", filter)).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
// ── Git handlers ───────────────────────────────────────────────────────────── // ── Git handlers ─────────────────────────────────────────────────────────────
@ -454,17 +441,17 @@ struct GitRequest { session_id: String, path: String }
async fn handle_git_status(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<GitRequest>) -> Json<McpResponse<String>> { async fn handle_git_status(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<GitRequest>) -> Json<McpResponse<String>> {
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 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", shell_escape(&req.path))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("cd {} && git status --short --branch 2>&1", shell_escape(&req.path))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_git_pull(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<GitRequest>) -> Json<McpResponse<String>> { async fn handle_git_pull(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<GitRequest>) -> Json<McpResponse<String>> {
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 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", shell_escape(&req.path))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("cd {} && git pull 2>&1", shell_escape(&req.path))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
async fn handle_git_log(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<GitRequest>) -> Json<McpResponse<String>> { async fn handle_git_log(AxumState(state): AxumState<Arc<McpServerState>>, Json(req): Json<GitRequest>) -> Json<McpResponse<String>> {
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 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", shell_escape(&req.path))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) } match exec_on_session(&session.handle, &format!("cd {} && git log --oneline -20 2>&1", shell_escape(&req.path))).await { Ok(o) => ok_response(o), Err(e) => err_response(e) }
} }
// ── Session creation handlers ──────────────────────────────────────────────── // ── Session creation handlers ────────────────────────────────────────────────

View File

@ -5,7 +5,6 @@
//! provides all file operations needed by the frontend. //! provides all file operations needed by the frontend.
use std::sync::Arc; use std::sync::Arc;
use std::time::{Duration, UNIX_EPOCH};
use dashmap::DashMap; use dashmap::DashMap;
use log::{debug, info}; use log::{debug, info};
@ -35,9 +34,6 @@ pub struct FileEntry {
/// Format a Unix timestamp (seconds since epoch) as "Mon DD HH:MM". /// Format a Unix timestamp (seconds since epoch) as "Mon DD HH:MM".
fn format_mtime(unix_secs: u32) -> String { fn format_mtime(unix_secs: u32) -> String {
// Build a SystemTime from the raw epoch value.
let st = UNIX_EPOCH + Duration::from_secs(unix_secs as u64);
// Convert to seconds-since-epoch for manual formatting. We avoid pulling // Convert to seconds-since-epoch for manual formatting. We avoid pulling
// in chrono just for this; a simple manual decomposition is sufficient for // in chrono just for this; a simple manual decomposition is sufficient for
// the "Mar 17 14:30" display format expected by the frontend. // the "Mar 17 14:30" display format expected by the frontend.
@ -54,12 +50,10 @@ fn format_mtime(unix_secs: u32) -> String {
let era = if z >= 0 { z } else { z - 146_096 } / 146_097; let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
let doe = z - era * 146_097; let doe = z - era * 146_097;
let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365; let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365;
let y = yoe + era * 400;
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
let mp = (5 * doy + 2) / 153; let mp = (5 * doy + 2) / 153;
let d = doy - (153 * mp + 2) / 5 + 1; let d = doy - (153 * mp + 2) / 5 + 1;
let m = if mp < 10 { mp + 3 } else { mp - 9 }; let m = if mp < 10 { mp + 3 } else { mp - 9 };
let _y = if m <= 2 { y + 1 } else { y };
let month = match m { let month = match m {
1 => "Jan", 1 => "Jan",
@ -77,9 +71,6 @@ fn format_mtime(unix_secs: u32) -> String {
_ => "???", _ => "???",
}; };
// Suppress unused variable warning — st is only used as a sanity anchor.
let _ = st;
format!("{} {:2} {:02}:{:02}", month, d, hours, minutes) format!("{} {:2} {:02}:{:02}", month, d, hours, minutes)
} }
@ -319,7 +310,7 @@ impl SftpService {
) -> Result<Arc<TokioMutex<SftpSession>>, String> { ) -> Result<Arc<TokioMutex<SftpSession>>, String> {
self.clients self.clients
.get(session_id) .get(session_id)
.map(|r| r.clone()) .map(|r| r.value().clone())
.ok_or_else(|| format!("No SFTP client for session {}", session_id)) .ok_or_else(|| format!("No SFTP client for session {}", session_id))
} }
} }

51
src-tauri/src/ssh/exec.rs Normal file
View File

@ -0,0 +1,51 @@
//! Shared SSH exec-channel helper used by commands, MCP handlers, and tools.
//!
//! Opens a one-shot exec channel on an existing SSH handle, runs `cmd`, collects
//! all stdout/stderr, and returns it as a `String`. The caller is responsible
//! for ensuring the session is still alive.
use std::sync::Arc;
use tokio::sync::Mutex as TokioMutex;
use crate::ssh::session::SshClient;
/// Execute `cmd` on a separate exec channel and return all output as a `String`.
///
/// Locks the handle for only as long as it takes to open the channel, then
/// releases it before reading — this avoids holding the lock while waiting on
/// remote I/O.
pub async fn exec_on_session(
handle: &Arc<TokioMutex<russh::client::Handle<SshClient>>>,
cmd: &str,
) -> Result<String, String> {
let mut channel = {
let h = handle.lock().await;
h.channel_open_session()
.await
.map_err(|e| format!("Exec channel failed: {}", e))?
};
channel
.exec(true, cmd)
.await
.map_err(|e| format!("Exec failed: {}", e))?;
let mut output = String::new();
loop {
match channel.wait().await {
Some(russh::ChannelMsg::Data { ref data }) => {
if let Ok(text) = std::str::from_utf8(data.as_ref()) {
output.push_str(text);
}
}
Some(russh::ChannelMsg::Eof)
| Some(russh::ChannelMsg::Close)
| None => break,
Some(russh::ChannelMsg::ExitStatus { .. }) => {}
_ => {}
}
}
Ok(output)
}

View File

@ -2,3 +2,4 @@ pub mod session;
pub mod host_key; pub mod host_key;
pub mod cwd; pub mod cwd;
pub mod monitor; pub mod monitor;
pub mod exec;

View File

@ -258,7 +258,7 @@ impl SshService {
} }
pub fn get_session(&self, session_id: &str) -> Option<Arc<SshSession>> { pub fn get_session(&self, session_id: &str) -> Option<Arc<SshSession>> {
self.sessions.get(session_id).map(|entry| entry.clone()) self.sessions.get(session_id).map(|r| r.value().clone())
} }
pub fn list_sessions(&self) -> Vec<SessionInfo> { pub fn list_sessions(&self) -> Vec<SessionInfo> {
@ -405,22 +405,23 @@ fn extract_osc7_cwd(data: &[u8]) -> Option<String> {
} }
fn percent_decode(input: &str) -> String { fn percent_decode(input: &str) -> String {
let mut output = String::with_capacity(input.len()); let mut bytes: Vec<u8> = Vec::with_capacity(input.len());
let mut chars = input.chars(); let mut chars = input.chars();
while let Some(ch) = chars.next() { while let Some(ch) = chars.next() {
if ch == '%' { if ch == '%' {
let hex: String = chars.by_ref().take(2).collect(); let hex: String = chars.by_ref().take(2).collect();
if let Ok(byte) = u8::from_str_radix(&hex, 16) { if let Ok(byte) = u8::from_str_radix(&hex, 16) {
output.push(byte as char); bytes.push(byte);
} else { } else {
output.push('%'); bytes.extend_from_slice(b"%");
output.push_str(&hex); bytes.extend_from_slice(hex.as_bytes());
} }
} else { } else {
output.push(ch); let mut buf = [0u8; 4];
bytes.extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
} }
} }
output String::from_utf8_lossy(&bytes).into_owned()
} }
/// Resolve a private key string — if it looks like PEM content, return as-is. /// Resolve a private key string — if it looks like PEM content, return as-is.

View File

@ -59,6 +59,7 @@ struct BuiltinTheme {
// ── service ─────────────────────────────────────────────────────────────────── // ── service ───────────────────────────────────────────────────────────────────
#[derive(Clone)]
pub struct ThemeService { pub struct ThemeService {
db: Database, db: Database,
} }
@ -253,7 +254,7 @@ impl ThemeService {
t.bright_blue, t.bright_magenta, t.bright_cyan, t.bright_white, t.bright_blue, t.bright_magenta, t.bright_cyan, t.bright_white,
], ],
) { ) {
eprintln!("theme::seed_builtins: failed to seed '{}': {}", t.name, e); wraith_log!("theme::seed_builtins: failed to seed '{}': {}", t.name, e);
} }
} }
} }
@ -272,7 +273,7 @@ impl ThemeService {
) { ) {
Ok(s) => s, Ok(s) => s,
Err(e) => { Err(e) => {
eprintln!("theme::list: failed to prepare query: {}", e); wraith_log!("theme::list: failed to prepare query: {}", e);
return vec![]; return vec![];
} }
}; };
@ -280,12 +281,12 @@ impl ThemeService {
match stmt.query_map([], map_theme_row) { match stmt.query_map([], map_theme_row) {
Ok(rows) => rows Ok(rows) => rows
.filter_map(|r| { .filter_map(|r| {
r.map_err(|e| eprintln!("theme::list: row error: {}", e)) r.map_err(|e| wraith_log!("theme::list: row error: {}", e))
.ok() .ok()
}) })
.collect(), .collect(),
Err(e) => { Err(e) => {
eprintln!("theme::list: query failed: {}", e); wraith_log!("theme::list: query failed: {}", e);
vec![] vec![]
} }
} }

View File

@ -24,6 +24,7 @@ pub struct WorkspaceSnapshot {
const SNAPSHOT_KEY: &str = "workspace_snapshot"; const SNAPSHOT_KEY: &str = "workspace_snapshot";
const CLEAN_SHUTDOWN_KEY: &str = "clean_shutdown"; const CLEAN_SHUTDOWN_KEY: &str = "clean_shutdown";
#[derive(Clone)]
pub struct WorkspaceService { pub struct WorkspaceService {
settings: SettingsService, settings: SettingsService,
} }
@ -47,7 +48,7 @@ impl WorkspaceService {
pub fn load(&self) -> Option<WorkspaceSnapshot> { pub fn load(&self) -> Option<WorkspaceSnapshot> {
let json = self.settings.get(SNAPSHOT_KEY)?; let json = self.settings.get(SNAPSHOT_KEY)?;
serde_json::from_str(&json) serde_json::from_str(&json)
.map_err(|e| eprintln!("workspace::load: failed to deserialize snapshot: {e}")) .map_err(|e| wraith_log!("workspace::load: failed to deserialize snapshot: {e}"))
.ok() .ok()
} }