claudetools/projects/msp-tools/guru-connect/agent/src/startup.rs

//! 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<u16> = OsStr::new(STARTUP_KEY)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_name: Vec<u16> = OsStr::new(STARTUP_VALUE_NAME)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_data: Vec<u16> = 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<u16> = OsStr::new(STARTUP_KEY)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_name: Vec<u16> = 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(())
}

/// 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<u16> = 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(())
}

/// Install the SAS service if the binary is available
/// This allows the agent to send Ctrl+Alt+Del even without SYSTEM privileges
#[cfg(windows)]
pub fn install_sas_service() -> Result<()> {
    info!("Attempting to install SAS service...");

    // Check if the SAS service binary exists alongside the agent
    let exe_path = std::env::current_exe()?;
    let exe_dir = exe_path.parent().ok_or_else(|| anyhow::anyhow!("No parent directory"))?;
    let sas_binary = exe_dir.join("guruconnect-sas-service.exe");

    if !sas_binary.exists() {
        // Also check in Program Files
        let program_files = std::path::PathBuf::from(r"C:\Program Files\GuruConnect\guruconnect-sas-service.exe");
        if !program_files.exists() {
            warn!("SAS service binary not found");
            return Ok(());
        }
    }

    // Run the install command
    let sas_path = if sas_binary.exists() {
        sas_binary
    } else {
        std::path::PathBuf::from(r"C:\Program Files\GuruConnect\guruconnect-sas-service.exe")
    };

    let output = std::process::Command::new(&sas_path)
        .arg("install")
        .output();

    match output {
        Ok(result) => {
            if result.status.success() {
                info!("SAS service installed successfully");
            } else {
                let stderr = String::from_utf8_lossy(&result.stderr);
                warn!("SAS service install failed: {}", stderr);
            }
        }
        Err(e) => {
            warn!("Failed to run SAS service installer: {}", e);
        }
    }

    Ok(())
}

/// Uninstall the SAS service
#[cfg(windows)]
pub fn uninstall_sas_service() -> Result<()> {
    info!("Attempting to uninstall SAS service...");

    // Try to find and run the uninstall command
    let paths = [
        std::env::current_exe().ok().and_then(|p| p.parent().map(|d| d.join("guruconnect-sas-service.exe"))),
        Some(std::path::PathBuf::from(r"C:\Program Files\GuruConnect\guruconnect-sas-service.exe")),
    ];

    for path_opt in paths.iter() {
        if let Some(ref path) = path_opt {
            if path.exists() {
                let output = std::process::Command::new(path)
                    .arg("uninstall")
                    .output();

                if let Ok(result) = output {
                    if result.status.success() {
                        info!("SAS service uninstalled successfully");
                        return Ok(());
                    }
                }
            }
        }
    }

    warn!("SAS service binary not found for uninstall");
    Ok(())
}

/// Check if the SAS service is installed and running
#[cfg(windows)]
pub fn check_sas_service() -> bool {
    use crate::sas_client;
    sas_client::is_service_available()
}

#[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(())
}

#[cfg(not(windows))]
pub fn uninstall() -> Result<()> {
    warn!("Uninstall not implemented for this platform");
    Ok(())
}

#[cfg(not(windows))]
pub fn install_sas_service() -> Result<()> {
    warn!("SAS service only available on Windows");
    Ok(())
}

#[cfg(not(windows))]
pub fn uninstall_sas_service() -> Result<()> {
    Ok(())
}

#[cfg(not(windows))]
pub fn check_sas_service() -> bool {
    false
}