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

//! Auto-update module for GuruConnect agent
//!
//! Handles checking for updates, downloading new versions, and performing
//! in-place binary replacement with restart.

use anyhow::{anyhow, Result};
use sha2::{Sha256, Digest};
use std::path::PathBuf;
use tracing::{info, warn, error};

use crate::build_info;

/// Version information from the server
#[derive(Debug, Clone, serde::Deserialize)]
pub struct VersionInfo {
    pub latest_version: String,
    pub download_url: String,
    pub checksum_sha256: String,
    pub is_mandatory: bool,
    pub release_notes: Option<String>,
}

/// Update state tracking
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum UpdateState {
    Idle,
    Checking,
    Downloading,
    Verifying,
    Installing,
    Restarting,
    Failed,
}

/// Check if an update is available
pub async fn check_for_update(server_base_url: &str) -> Result<Option<VersionInfo>> {
    let url = format!("{}/api/version", server_base_url.trim_end_matches('/'));
    info!("Checking for updates at {}", url);

    let client = reqwest::Client::builder()
        .danger_accept_invalid_certs(true)  // For self-signed certs in dev
        .build()?;

    let response = client
        .get(&url)
        .timeout(std::time::Duration::from_secs(30))
        .send()
        .await?;

    if response.status() == reqwest::StatusCode::NOT_FOUND {
        info!("No stable release available on server");
        return Ok(None);
    }

    if !response.status().is_success() {
        return Err(anyhow!("Version check failed: HTTP {}", response.status()));
    }

    let version_info: VersionInfo = response.json().await?;

    // Compare versions
    let current = build_info::VERSION;
    if is_newer_version(&version_info.latest_version, current) {
        info!(
            "Update available: {} -> {} (mandatory: {})",
            current, version_info.latest_version, version_info.is_mandatory
        );
        Ok(Some(version_info))
    } else {
        info!("Already running latest version: {}", current);
        Ok(None)
    }
}

/// Simple semantic version comparison
/// Returns true if `available` is newer than `current`
fn is_newer_version(available: &str, current: &str) -> bool {
    // Strip any git hash suffix (e.g., "0.1.0-abc123" -> "0.1.0")
    let available_clean = available.split('-').next().unwrap_or(available);
    let current_clean = current.split('-').next().unwrap_or(current);

    let parse_version = |s: &str| -> Vec<u32> {
        s.split('.')
            .filter_map(|p| p.parse().ok())
            .collect()
    };

    let av = parse_version(available_clean);
    let cv = parse_version(current_clean);

    // Compare component by component
    for i in 0..av.len().max(cv.len()) {
        let a = av.get(i).copied().unwrap_or(0);
        let c = cv.get(i).copied().unwrap_or(0);
        if a > c {
            return true;
        }
        if a < c {
            return false;
        }
    }
    false
}

/// Download update to temporary file
pub async fn download_update(version_info: &VersionInfo) -> Result<PathBuf> {
    info!("Downloading update from {}", version_info.download_url);

    let client = reqwest::Client::builder()
        .danger_accept_invalid_certs(true)
        .build()?;

    let response = client
        .get(&version_info.download_url)
        .timeout(std::time::Duration::from_secs(300))  // 5 minutes for large files
        .send()
        .await?;

    if !response.status().is_success() {
        return Err(anyhow!("Download failed: HTTP {}", response.status()));
    }

    // Get temp directory
    let temp_dir = std::env::temp_dir();
    let temp_path = temp_dir.join("guruconnect-update.exe");

    // Download to file
    let bytes = response.bytes().await?;
    std::fs::write(&temp_path, &bytes)?;

    info!("Downloaded {} bytes to {:?}", bytes.len(), temp_path);
    Ok(temp_path)
}

/// Verify downloaded file checksum
pub fn verify_checksum(file_path: &PathBuf, expected_sha256: &str) -> Result<bool> {
    info!("Verifying checksum...");

    let contents = std::fs::read(file_path)?;
    let mut hasher = Sha256::new();
    hasher.update(&contents);
    let result = hasher.finalize();
    let computed = format!("{:x}", result);

    let matches = computed.eq_ignore_ascii_case(expected_sha256);

    if matches {
        info!("Checksum verified: {}", computed);
    } else {
        error!("Checksum mismatch! Expected: {}, Got: {}", expected_sha256, computed);
    }

    Ok(matches)
}

/// Perform the actual update installation
/// This renames the current executable and copies the new one in place
pub fn install_update(temp_path: &PathBuf) -> Result<PathBuf> {
    info!("Installing update...");

    // Get current executable path
    let current_exe = std::env::current_exe()?;
    let exe_dir = current_exe.parent()
        .ok_or_else(|| anyhow!("Cannot get executable directory"))?;

    // Create paths for backup and new executable
    let backup_path = exe_dir.join("guruconnect.exe.old");

    // Delete any existing backup
    if backup_path.exists() {
        if let Err(e) = std::fs::remove_file(&backup_path) {
            warn!("Could not remove old backup: {}", e);
        }
    }

    // Rename current executable to .old (this works even while running)
    info!("Renaming current exe to backup: {:?}", backup_path);
    std::fs::rename(&current_exe, &backup_path)?;

    // Copy new executable to original location
    info!("Copying new exe to: {:?}", current_exe);
    std::fs::copy(temp_path, &current_exe)?;

    // Clean up temp file
    let _ = std::fs::remove_file(temp_path);

    info!("Update installed successfully");
    Ok(current_exe)
}

/// Spawn new process and exit current one
pub fn restart_with_new_version(exe_path: &PathBuf, args: &[String]) -> Result<()> {
    info!("Restarting with new version...");

    // Build command with --post-update flag
    let mut cmd_args = vec!["--post-update".to_string()];
    cmd_args.extend(args.iter().cloned());

    #[cfg(windows)]
    {
        use std::os::windows::process::CommandExt;
        const CREATE_NEW_PROCESS_GROUP: u32 = 0x00000200;
        const DETACHED_PROCESS: u32 = 0x00000008;

        std::process::Command::new(exe_path)
            .args(&cmd_args)
            .creation_flags(CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS)
            .spawn()?;
    }

    #[cfg(not(windows))]
    {
        std::process::Command::new(exe_path)
            .args(&cmd_args)
            .spawn()?;
    }

    info!("New process spawned, exiting current process");
    Ok(())
}

/// Clean up old executable after successful update
pub fn cleanup_post_update() {
    let current_exe = match std::env::current_exe() {
        Ok(p) => p,
        Err(e) => {
            warn!("Could not get current exe path for cleanup: {}", e);
            return;
        }
    };

    let exe_dir = match current_exe.parent() {
        Some(d) => d,
        None => {
            warn!("Could not get executable directory for cleanup");
            return;
        }
    };

    let backup_path = exe_dir.join("guruconnect.exe.old");

    if backup_path.exists() {
        info!("Cleaning up old executable: {:?}", backup_path);
        match std::fs::remove_file(&backup_path) {
            Ok(_) => info!("Old executable removed successfully"),
            Err(e) => {
                warn!("Could not remove old executable (may be in use): {}", e);
                // On Windows, we might need to schedule deletion on reboot
                #[cfg(windows)]
                schedule_delete_on_reboot(&backup_path);
            }
        }
    }
}

/// Schedule file deletion on reboot (Windows)
#[cfg(windows)]
fn schedule_delete_on_reboot(path: &PathBuf) {
    use std::os::windows::ffi::OsStrExt;
    use windows::Win32::Storage::FileSystem::{MoveFileExW, MOVEFILE_DELAY_UNTIL_REBOOT};
    use windows::core::PCWSTR;

    let path_wide: Vec<u16> = path.as_os_str()
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        let result = MoveFileExW(
            PCWSTR(path_wide.as_ptr()),
            PCWSTR::null(),
            MOVEFILE_DELAY_UNTIL_REBOOT,
        );
        if result.is_ok() {
            info!("Scheduled {:?} for deletion on reboot", path);
        } else {
            warn!("Failed to schedule {:?} for deletion on reboot", path);
        }
    }
}

/// Perform complete update process
pub async fn perform_update(version_info: &VersionInfo) -> Result<()> {
    // Download
    let temp_path = download_update(version_info).await?;

    // Verify
    if !verify_checksum(&temp_path, &version_info.checksum_sha256)? {
        let _ = std::fs::remove_file(&temp_path);
        return Err(anyhow!("Update verification failed: checksum mismatch"));
    }

    // Install
    let exe_path = install_update(&temp_path)?;

    // Restart
    // Get current args (without the current executable name)
    let args: Vec<String> = std::env::args().skip(1).collect();
    restart_with_new_version(&exe_path, &args)?;

    // Exit current process
    std::process::exit(0);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_version_comparison() {
        assert!(is_newer_version("0.2.0", "0.1.0"));
        assert!(is_newer_version("1.0.0", "0.9.9"));
        assert!(is_newer_version("0.1.1", "0.1.0"));
        assert!(!is_newer_version("0.1.0", "0.1.0"));
        assert!(!is_newer_version("0.1.0", "0.2.0"));
        assert!(is_newer_version("0.2.0-abc123", "0.1.0-def456"));
    }
}