From 4b29dbe6c8d0d3c8ba55384287d5afc309d0676f Mon Sep 17 00:00:00 2001 From: Mike Swanson Date: Sun, 28 Dec 2025 16:53:29 -0700 Subject: [PATCH] Add disconnect/uninstall for persistent sessions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Server: Add DELETE /api/sessions/:id endpoint to disconnect agents - Server: SessionManager.disconnect_session() sends Disconnect message - Agent: Handle ADMIN_DISCONNECT to trigger uninstall - Agent: Add startup::uninstall() to remove from startup and schedule exe deletion - Dashboard: Add Disconnect button in Access tab machine details - Dashboard: Add Chat button for persistent sessions 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- agent/Cargo.toml | 1 + agent/src/main.rs | 13 ++++++++++ agent/src/session/mod.rs | 6 ++++- agent/src/startup.rs | 49 ++++++++++++++++++++++++++++++++++++ server/src/main.rs | 18 +++++++++++++ server/src/session/mod.rs | 26 +++++++++++++++++++ server/static/dashboard.html | 20 ++++++++++++++- 7 files changed, 131 insertions(+), 2 deletions(-) diff --git a/agent/Cargo.toml b/agent/Cargo.toml index ddf9641..a6a1076 100644 --- a/agent/Cargo.toml +++ b/agent/Cargo.toml @@ -73,6 +73,7 @@ windows = { version = "0.58", features = [ "Win32_System_Threading", "Win32_System_Registry", "Win32_Security", + "Win32_Storage_FileSystem", ]} # Windows service support diff --git a/agent/src/main.rs b/agent/src/main.rs index 462a929..a724bf0 100644 --- a/agent/src/main.rs +++ b/agent/src/main.rs @@ -249,6 +249,19 @@ async fn run_agent(config: config::Config) -> Result<()> { return Ok(()); } + // Check if this is an admin disconnect (uninstall) + if error_msg.contains("ADMIN_DISCONNECT") { + info!("Session was disconnected by administrator - uninstalling"); + if let Err(e) = startup::uninstall() { + warn!("Uninstall failed: {}", e); + } + show_message_box( + "Remote Session Ended", + "The remote support session has been ended by the administrator.\n\nThe agent will be removed from this computer.", + ); + return Ok(()); + } + error!("Session error: {}", e); } } diff --git a/agent/src/session/mod.rs b/agent/src/session/mod.rs index 02f4f02..d5750ad 100644 --- a/agent/src/session/mod.rs +++ b/agent/src/session/mod.rs @@ -340,10 +340,14 @@ impl SessionManager { Some(message::Payload::Disconnect(disc)) => { tracing::info!("Disconnect requested: {}", disc.reason); - // Check if this is a cancellation + // Check if this is a cancellation (support session) if disc.reason.contains("cancelled") { return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason)); } + // Check if this is an admin disconnect (persistent session) + if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") { + return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason)); + } return Err(anyhow::anyhow!("Disconnect: {}", disc.reason)); } diff --git a/agent/src/startup.rs b/agent/src/startup.rs index 9003a47..9f81f4f 100644 --- a/agent/src/startup.rs +++ b/agent/src/startup.rs @@ -134,6 +134,49 @@ pub fn remove_from_startup() -> Result<()> { Ok(()) } +/// Full uninstall: remove from startup and delete the executable +#[cfg(windows)] +pub fn uninstall() -> Result<()> { + use std::ffi::OsStr; + use std::os::windows::ffi::OsStrExt; + use windows::Win32::Storage::FileSystem::{MoveFileExW, MOVEFILE_DELAY_UNTIL_REBOOT}; + + info!("Uninstalling agent"); + + // First remove from startup + let _ = remove_from_startup(); + + // Get the path to the current executable + let exe_path = std::env::current_exe()?; + let exe_path_str = exe_path.to_string_lossy(); + + info!("Scheduling deletion of: {}", exe_path_str); + + // Convert path to wide string + let exe_wide: Vec = OsStr::new(&*exe_path_str) + .encode_wide() + .chain(std::iter::once(0)) + .collect(); + + // Schedule the file for deletion on next reboot + // This is necessary because the executable is currently running + unsafe { + let result = MoveFileExW( + PCWSTR(exe_wide.as_ptr()), + PCWSTR::null(), + MOVEFILE_DELAY_UNTIL_REBOOT, + ); + + if result.is_err() { + warn!("Failed to schedule file deletion: {:?}. File may need manual removal.", result); + } else { + info!("Executable scheduled for deletion on reboot"); + } + } + + Ok(()) +} + #[cfg(not(windows))] pub fn add_to_startup() -> Result<()> { warn!("Startup persistence not implemented for this platform"); @@ -144,3 +187,9 @@ pub fn add_to_startup() -> Result<()> { pub fn remove_from_startup() -> Result<()> { Ok(()) } + +#[cfg(not(windows))] +pub fn uninstall() -> Result<()> { + warn!("Uninstall not implemented for this platform"); + Ok(()) +} diff --git a/server/src/main.rs b/server/src/main.rs index 5ecedb5..4c64482 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -81,6 +81,7 @@ async fn main() -> Result<()> { // REST API - Sessions .route("/api/sessions", get(list_sessions)) .route("/api/sessions/:id", get(get_session)) + .route("/api/sessions/:id", axum::routing::delete(disconnect_session)) // HTML page routes (clean URLs) .route("/login", get(serve_login)) @@ -178,6 +179,23 @@ async fn get_session( Ok(Json(api::SessionInfo::from(session))) } +async fn disconnect_session( + State(state): State, + Path(id): Path, +) -> impl IntoResponse { + let session_id = match uuid::Uuid::parse_str(&id) { + Ok(id) => id, + Err(_) => return (StatusCode::BAD_REQUEST, "Invalid session ID"), + }; + + if state.sessions.disconnect_session(session_id, "Disconnected by administrator").await { + info!("Session {} disconnected by admin", session_id); + (StatusCode::OK, "Session disconnected") + } else { + (StatusCode::NOT_FOUND, "Session not found") + } +} + // Static page handlers async fn serve_login() -> impl IntoResponse { match tokio::fs::read_to_string("static/login.html").await { diff --git a/server/src/session/mod.rs b/server/src/session/mod.rs index c2ab775..7fe0609 100644 --- a/server/src/session/mod.rs +++ b/server/src/session/mod.rs @@ -134,6 +134,32 @@ impl SessionManager { } } + /// Disconnect a session by sending a disconnect message to the agent + /// Returns true if the message was sent successfully + pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool { + let sessions = self.sessions.read().await; + if let Some(session_data) = sessions.get(&session_id) { + // Create disconnect message + use crate::proto; + use prost::Message; + + let disconnect_msg = proto::Message { + payload: Some(proto::message::Payload::Disconnect(proto::Disconnect { + reason: reason.to_string(), + })), + }; + + let mut buf = Vec::new(); + if disconnect_msg.encode(&mut buf).is_ok() { + // Send via input channel (will be forwarded to agent's WebSocket) + if session_data.input_tx.send(buf).await.is_ok() { + return true; + } + } + } + false + } + /// List all active sessions pub async fn list_sessions(&self) -> Vec { let sessions = self.sessions.read().await; diff --git a/server/static/dashboard.html b/server/static/dashboard.html index fccfb12..546287f 100644 --- a/server/static/dashboard.html +++ b/server/static/dashboard.html @@ -709,7 +709,9 @@ '
' + '
Actions
' + '' + - '' + + '' + + '' + + '' + '
'; } @@ -718,6 +720,22 @@ alert("Viewer not yet implemented.\\n\\nSession ID: " + sessionId + "\\n\\nWebSocket: wss://connect.azcomputerguru.com/ws/viewer?session_id=" + sessionId); } + async function disconnectMachine(sessionId, machineName) { + if (!confirm("Disconnect " + machineName + "?\\n\\nThis will end the remote session.")) return; + try { + const response = await fetch("/api/sessions/" + sessionId, { method: "DELETE" }); + if (response.ok) { + selectedMachine = null; + renderMachineDetail(); + loadMachines(); + } else { + alert("Failed to disconnect: " + await response.text()); + } + } catch (err) { + alert("Error disconnecting machine"); + } + } + // Refresh machines every 5 seconds loadMachines(); setInterval(loadMachines, 5000);