From 0dcbae69a08ab261742b8788cbc91b24d81260cb Mon Sep 17 00:00:00 2001 From: Mike Swanson Date: Sun, 28 Dec 2025 16:15:24 -0700 Subject: [PATCH] Add startup persistence for support sessions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Added startup.rs module for Windows registry operations - Agent adds itself to HKCU\...\Run on session start - Agent removes itself from startup on session end - Works with user-level registry (no admin required) - Added Win32_System_Registry feature to windows crate 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- agent/Cargo.toml | 1 + agent/src/main.rs | 23 +++++++ agent/src/startup.rs | 146 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 170 insertions(+) create mode 100644 agent/src/startup.rs diff --git a/agent/Cargo.toml b/agent/Cargo.toml index 60c0f3f..ddf9641 100644 --- a/agent/Cargo.toml +++ b/agent/Cargo.toml @@ -71,6 +71,7 @@ windows = { version = "0.58", features = [ "Win32_UI_WindowsAndMessaging", "Win32_System_LibraryLoader", "Win32_System_Threading", + "Win32_System_Registry", "Win32_Security", ]} diff --git a/agent/src/main.rs b/agent/src/main.rs index d27a00f..5cf03f2 100644 --- a/agent/src/main.rs +++ b/agent/src/main.rs @@ -13,6 +13,7 @@ mod config; mod encoder; mod input; mod session; +mod startup; mod transport; mod tray; @@ -162,12 +163,29 @@ async fn main() -> Result<()> { Ok(()) } +/// Clean up before exiting (remove from startup, etc.) +fn cleanup_on_exit(is_support_session: bool) { + if is_support_session { + info!("Cleaning up before exit"); + if let Err(e) = startup::remove_from_startup() { + warn!("Failed to remove from startup: {}", e); + } + } +} + async fn run_agent(config: config::Config) -> Result<()> { // Create session manager let mut session = session::SessionManager::new(config.clone()); let is_support_session = config.support_code.is_some(); let hostname = config.hostname(); + // Add to startup so we reconnect after reboot + if is_support_session { + if let Err(e) = startup::add_to_startup() { + warn!("Failed to add to startup: {}. Agent won't persist through reboot.", e); + } + } + // Create tray icon let tray = match tray::TrayController::new(&hostname, config.support_code.as_deref()) { Ok(t) => { @@ -208,12 +226,14 @@ async fn run_agent(config: config::Config) -> Result<()> { // Check if this is a user-initiated exit if error_msg.contains("USER_EXIT") { info!("Session ended by user"); + cleanup_on_exit(is_support_session); return Ok(()); } // Check if this is a cancellation if error_msg.contains("SESSION_CANCELLED") { info!("Session was cancelled by technician"); + cleanup_on_exit(is_support_session); show_message_box( "Support Session Ended", "The support session was cancelled by the technician.\n\nThis window will close automatically.", @@ -231,6 +251,7 @@ async fn run_agent(config: config::Config) -> Result<()> { // Check if connection was rejected due to cancelled code if error_msg.contains("cancelled") { info!("Support code was cancelled before connection"); + cleanup_on_exit(is_support_session); show_message_box( "Support Session Cancelled", "This support session has been cancelled.\n\nPlease contact your technician for a new support code.", @@ -246,6 +267,7 @@ async fn run_agent(config: config::Config) -> Result<()> { // For support sessions, don't reconnect if something goes wrong if is_support_session { info!("Support session ended, not reconnecting"); + cleanup_on_exit(is_support_session); return Ok(()); } @@ -253,6 +275,7 @@ async fn run_agent(config: config::Config) -> Result<()> { if let Some(ref t) = tray { if t.exit_requested() { info!("Exit requested by user"); + cleanup_on_exit(is_support_session); return Ok(()); } } diff --git a/agent/src/startup.rs b/agent/src/startup.rs new file mode 100644 index 0000000..9003a47 --- /dev/null +++ b/agent/src/startup.rs @@ -0,0 +1,146 @@ +//! Startup persistence for the agent +//! +//! Handles adding/removing the agent from Windows startup. + +use anyhow::Result; +use tracing::{info, warn, error}; + +#[cfg(windows)] +use windows::Win32::System::Registry::{ + RegOpenKeyExW, RegSetValueExW, RegDeleteValueW, RegCloseKey, + HKEY_CURRENT_USER, KEY_WRITE, REG_SZ, +}; +#[cfg(windows)] +use windows::core::PCWSTR; + +const STARTUP_KEY: &str = r"Software\Microsoft\Windows\CurrentVersion\Run"; +const STARTUP_VALUE_NAME: &str = "GuruConnect"; + +/// Add the current executable to Windows startup +#[cfg(windows)] +pub fn add_to_startup() -> Result<()> { + use std::ffi::OsStr; + use std::os::windows::ffi::OsStrExt; + + // Get the path to the current executable + let exe_path = std::env::current_exe()?; + let exe_path_str = exe_path.to_string_lossy(); + + info!("Adding to startup: {}", exe_path_str); + + // Convert strings to wide strings + let key_path: Vec = OsStr::new(STARTUP_KEY) + .encode_wide() + .chain(std::iter::once(0)) + .collect(); + let value_name: Vec = OsStr::new(STARTUP_VALUE_NAME) + .encode_wide() + .chain(std::iter::once(0)) + .collect(); + let value_data: Vec = OsStr::new(&*exe_path_str) + .encode_wide() + .chain(std::iter::once(0)) + .collect(); + + unsafe { + let mut hkey = windows::Win32::Foundation::HANDLE::default(); + + // Open the Run key + let result = RegOpenKeyExW( + HKEY_CURRENT_USER, + PCWSTR(key_path.as_ptr()), + 0, + KEY_WRITE, + &mut hkey as *mut _ as *mut _, + ); + + if result.is_err() { + anyhow::bail!("Failed to open registry key: {:?}", result); + } + + let hkey_raw = std::mem::transmute::<_, windows::Win32::System::Registry::HKEY>(hkey); + + // Set the value + let data_bytes = std::slice::from_raw_parts( + value_data.as_ptr() as *const u8, + value_data.len() * 2, + ); + + let set_result = RegSetValueExW( + hkey_raw, + PCWSTR(value_name.as_ptr()), + 0, + REG_SZ, + Some(data_bytes), + ); + + let _ = RegCloseKey(hkey_raw); + + if set_result.is_err() { + anyhow::bail!("Failed to set registry value: {:?}", set_result); + } + } + + info!("Successfully added to startup"); + Ok(()) +} + +/// Remove the agent from Windows startup +#[cfg(windows)] +pub fn remove_from_startup() -> Result<()> { + use std::ffi::OsStr; + use std::os::windows::ffi::OsStrExt; + + info!("Removing from startup"); + + let key_path: Vec = OsStr::new(STARTUP_KEY) + .encode_wide() + .chain(std::iter::once(0)) + .collect(); + let value_name: Vec = OsStr::new(STARTUP_VALUE_NAME) + .encode_wide() + .chain(std::iter::once(0)) + .collect(); + + unsafe { + let mut hkey = windows::Win32::Foundation::HANDLE::default(); + + let result = RegOpenKeyExW( + HKEY_CURRENT_USER, + PCWSTR(key_path.as_ptr()), + 0, + KEY_WRITE, + &mut hkey as *mut _ as *mut _, + ); + + if result.is_err() { + warn!("Failed to open registry key for removal: {:?}", result); + return Ok(()); // Not an error if key doesn't exist + } + + let hkey_raw = std::mem::transmute::<_, windows::Win32::System::Registry::HKEY>(hkey); + + let delete_result = RegDeleteValueW(hkey_raw, PCWSTR(value_name.as_ptr())); + + let _ = RegCloseKey(hkey_raw); + + if delete_result.is_err() { + warn!("Registry value may not exist: {:?}", delete_result); + } else { + info!("Successfully removed from startup"); + } + } + + Ok(()) +} + +#[cfg(not(windows))] +pub fn add_to_startup() -> Result<()> { + warn!("Startup persistence not implemented for this platform"); + Ok(()) +} + +#[cfg(not(windows))] +pub fn remove_from_startup() -> Result<()> { + Ok(()) +}