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

//! Installation and protocol handler registration
//!
//! Handles:
//! - Self-installation to Program Files (with UAC) or LocalAppData (fallback)
//! - Protocol handler registration (guruconnect://)
//! - UAC elevation with graceful fallback

use anyhow::{anyhow, Result};
use tracing::{info, warn, error};

#[cfg(windows)]
use windows::{
    core::PCWSTR,
    Win32::Foundation::HANDLE,
    Win32::Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION, TOKEN_QUERY},
    Win32::System::Threading::{GetCurrentProcess, OpenProcessToken},
    Win32::System::Registry::{
        RegCreateKeyExW, RegSetValueExW, RegCloseKey, HKEY, HKEY_CLASSES_ROOT,
        HKEY_CURRENT_USER, KEY_WRITE, REG_SZ, REG_OPTION_NON_VOLATILE,
    },
    Win32::UI::Shell::ShellExecuteW,
    Win32::UI::WindowsAndMessaging::SW_SHOWNORMAL,
};

#[cfg(windows)]
use std::ffi::OsStr;
#[cfg(windows)]
use std::os::windows::ffi::OsStrExt;

/// Install locations
pub const SYSTEM_INSTALL_PATH: &str = r"C:\Program Files\GuruConnect";
pub const USER_INSTALL_PATH: &str = r"GuruConnect"; // Relative to %LOCALAPPDATA%

/// Check if running with elevated privileges
#[cfg(windows)]
pub fn is_elevated() -> bool {
    unsafe {
        let mut token_handle = HANDLE::default();
        if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle).is_err() {
            return false;
        }

        let mut elevation = TOKEN_ELEVATION::default();
        let mut size = std::mem::size_of::<TOKEN_ELEVATION>() as u32;

        let result = GetTokenInformation(
            token_handle,
            TokenElevation,
            Some(&mut elevation as *mut _ as *mut _),
            size,
            &mut size,
        );

        let _ = windows::Win32::Foundation::CloseHandle(token_handle);

        result.is_ok() && elevation.TokenIsElevated != 0
    }
}

#[cfg(not(windows))]
pub fn is_elevated() -> bool {
    unsafe { libc::geteuid() == 0 }
}

/// Get the install path based on elevation status
pub fn get_install_path(elevated: bool) -> std::path::PathBuf {
    if elevated {
        std::path::PathBuf::from(SYSTEM_INSTALL_PATH)
    } else {
        let local_app_data = std::env::var("LOCALAPPDATA")
            .unwrap_or_else(|_| {
                let home = std::env::var("USERPROFILE").unwrap_or_else(|_| ".".to_string());
                format!(r"{}\AppData\Local", home)
            });
        std::path::PathBuf::from(local_app_data).join(USER_INSTALL_PATH)
    }
}

/// Get the executable path
pub fn get_exe_path(install_path: &std::path::Path) -> std::path::PathBuf {
    install_path.join("guruconnect.exe")
}

/// Attempt to elevate and re-run with install command
#[cfg(windows)]
pub fn try_elevate_and_install() -> Result<bool> {
    let exe_path = std::env::current_exe()?;
    let exe_path_wide: Vec<u16> = OsStr::new(exe_path.as_os_str())
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    let verb: Vec<u16> = OsStr::new("runas")
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    let params: Vec<u16> = OsStr::new("install --elevated")
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        let result = ShellExecuteW(
            None,
            PCWSTR(verb.as_ptr()),
            PCWSTR(exe_path_wide.as_ptr()),
            PCWSTR(params.as_ptr()),
            PCWSTR::null(),
            SW_SHOWNORMAL,
        );

        // ShellExecuteW returns > 32 on success
        if result.0 as usize > 32 {
            info!("UAC elevation requested");
            Ok(true)
        } else {
            warn!("UAC elevation denied or failed");
            Ok(false)
        }
    }
}

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

/// Register the guruconnect:// protocol handler
#[cfg(windows)]
pub fn register_protocol_handler(elevated: bool) -> Result<()> {
    let install_path = get_install_path(elevated);
    let exe_path = get_exe_path(&install_path);
    let exe_path_str = exe_path.to_string_lossy();

    // Command to execute: "C:\...\guruconnect.exe" "launch" "%1"
    let command = format!("\"{}\" launch \"%1\"", exe_path_str);

    // Choose registry root based on elevation
    let root_key = if elevated {
        HKEY_CLASSES_ROOT
    } else {
        // User-level registration under Software\Classes
        HKEY_CURRENT_USER
    };

    let base_path = if elevated {
        "guruconnect"
    } else {
        r"Software\Classes\guruconnect"
    };

    unsafe {
        // Create guruconnect key
        let mut protocol_key = HKEY::default();
        let key_path = to_wide(base_path);
        let result = RegCreateKeyExW(
            root_key,
            PCWSTR(key_path.as_ptr()),
            0,
            PCWSTR::null(),
            REG_OPTION_NON_VOLATILE,
            KEY_WRITE,
            None,
            &mut protocol_key,
            None,
        );
        if result.is_err() {
            return Err(anyhow!("Failed to create protocol key: {:?}", result));
        }

        // Set default value (protocol description)
        let description = to_wide("GuruConnect Protocol");
        let result = RegSetValueExW(
            protocol_key,
            PCWSTR::null(),
            0,
            REG_SZ,
            Some(&description_to_bytes(&description)),
        );
        if result.is_err() {
            let _ = RegCloseKey(protocol_key);
            return Err(anyhow!("Failed to set protocol description: {:?}", result));
        }

        // Set URL Protocol (empty string indicates this is a protocol handler)
        let url_protocol = to_wide("URL Protocol");
        let empty = to_wide("");
        let result = RegSetValueExW(
            protocol_key,
            PCWSTR(url_protocol.as_ptr()),
            0,
            REG_SZ,
            Some(&description_to_bytes(&empty)),
        );
        if result.is_err() {
            let _ = RegCloseKey(protocol_key);
            return Err(anyhow!("Failed to set URL Protocol: {:?}", result));
        }

        let _ = RegCloseKey(protocol_key);

        // Create shell\open\command key
        let command_path = if elevated {
            r"guruconnect\shell\open\command"
        } else {
            r"Software\Classes\guruconnect\shell\open\command"
        };
        let command_key_path = to_wide(command_path);
        let mut command_key = HKEY::default();
        let result = RegCreateKeyExW(
            root_key,
            PCWSTR(command_key_path.as_ptr()),
            0,
            PCWSTR::null(),
            REG_OPTION_NON_VOLATILE,
            KEY_WRITE,
            None,
            &mut command_key,
            None,
        );
        if result.is_err() {
            return Err(anyhow!("Failed to create command key: {:?}", result));
        }

        // Set the command
        let command_wide = to_wide(&command);
        let result = RegSetValueExW(
            command_key,
            PCWSTR::null(),
            0,
            REG_SZ,
            Some(&description_to_bytes(&command_wide)),
        );
        if result.is_err() {
            let _ = RegCloseKey(command_key);
            return Err(anyhow!("Failed to set command: {:?}", result));
        }

        let _ = RegCloseKey(command_key);
    }

    info!("Protocol handler registered: guruconnect://");
    Ok(())
}

#[cfg(not(windows))]
pub fn register_protocol_handler(_elevated: bool) -> Result<()> {
    warn!("Protocol handler registration not supported on this platform");
    Ok(())
}

/// Install the application
pub fn install(force_user_install: bool) -> Result<()> {
    let elevated = is_elevated();

    // If not elevated and not forcing user install, try to elevate
    if !elevated && !force_user_install {
        info!("Attempting UAC elevation for system-wide install...");
        match try_elevate_and_install() {
            Ok(true) => {
                // Elevation was requested, exit this instance
                // The elevated instance will continue the install
                info!("Elevated process started, exiting current instance");
                std::process::exit(0);
            }
            Ok(false) => {
                info!("UAC denied, falling back to user install");
            }
            Err(e) => {
                warn!("Elevation failed: {}, falling back to user install", e);
            }
        }
    }

    let install_path = get_install_path(elevated);
    let exe_path = get_exe_path(&install_path);

    info!("Installing to: {}", install_path.display());

    // Create install directory
    std::fs::create_dir_all(&install_path)?;

    // Copy ourselves to install location
    let current_exe = std::env::current_exe()?;
    if current_exe != exe_path {
        std::fs::copy(&current_exe, &exe_path)?;
        info!("Copied executable to: {}", exe_path.display());
    }

    // Register protocol handler
    register_protocol_handler(elevated)?;

    info!("Installation complete!");
    if elevated {
        info!("Installed system-wide to: {}", install_path.display());
    } else {
        info!("Installed for current user to: {}", install_path.display());
    }

    Ok(())
}

/// Check if the guruconnect:// protocol handler is registered
#[cfg(windows)]
pub fn is_protocol_handler_registered() -> bool {
    use windows::Win32::System::Registry::{
        RegOpenKeyExW, RegCloseKey, HKEY_CLASSES_ROOT, HKEY_CURRENT_USER, KEY_READ,
    };

    unsafe {
        // Check system-wide registration (HKCR\guruconnect)
        let mut key = HKEY::default();
        let key_path = to_wide("guruconnect");
        if RegOpenKeyExW(
            HKEY_CLASSES_ROOT,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_READ,
            &mut key,
        ).is_ok() {
            let _ = RegCloseKey(key);
            return true;
        }

        // Check user-level registration (HKCU\Software\Classes\guruconnect)
        let key_path = to_wide(r"Software\Classes\guruconnect");
        if RegOpenKeyExW(
            HKEY_CURRENT_USER,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_READ,
            &mut key,
        ).is_ok() {
            let _ = RegCloseKey(key);
            return true;
        }
    }

    false
}

#[cfg(not(windows))]
pub fn is_protocol_handler_registered() -> bool {
    // On non-Windows, assume not registered (or check ~/.local/share/applications)
    false
}

/// Parse a guruconnect:// URL and extract session parameters
pub fn parse_protocol_url(url_str: &str) -> Result<(String, String, Option<String>)> {
    // Expected formats:
    // guruconnect://view/SESSION_ID
    // guruconnect://view/SESSION_ID?token=API_KEY
    // guruconnect://connect/SESSION_ID?server=wss://...&token=API_KEY
    //
    // Note: In URL parsing, "view" becomes the host, SESSION_ID is the path

    let url = url::Url::parse(url_str)
        .map_err(|e| anyhow!("Invalid URL: {}", e))?;

    if url.scheme() != "guruconnect" {
        return Err(anyhow!("Invalid scheme: expected guruconnect://"));
    }

    // The "action" (view/connect) is parsed as the host
    let action = url.host_str()
        .ok_or_else(|| anyhow!("Missing action in URL"))?;

    // The session ID is the first path segment
    let path = url.path().trim_start_matches('/');
    info!("URL path: '{}', host: '{:?}'", path, url.host_str());
    let session_id = if path.is_empty() {
        return Err(anyhow!("Invalid URL: Missing session ID (path was empty, full URL: {})", url_str));
    } else {
        path.split('/').next().unwrap_or("").to_string()
    };

    if session_id.is_empty() {
        return Err(anyhow!("Missing session ID"));
    }

    // Extract query parameters
    let mut server = None;
    let mut token = None;

    for (key, value) in url.query_pairs() {
        match key.as_ref() {
            "server" => server = Some(value.to_string()),
            "token" | "api_key" => token = Some(value.to_string()),
            _ => {}
        }
    }

    // Default server if not specified
    let server = server.unwrap_or_else(|| "wss://connect.azcomputerguru.com/ws/viewer".to_string());

    match action {
        "view" | "connect" => Ok((server, session_id, token)),
        _ => Err(anyhow!("Unknown action: {}", action)),
    }
}

// Helper functions for Windows registry operations
#[cfg(windows)]
fn to_wide(s: &str) -> Vec<u16> {
    OsStr::new(s)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect()
}

#[cfg(windows)]
fn description_to_bytes(wide: &[u16]) -> Vec<u8> {
    wide.iter()
        .flat_map(|w| w.to_le_bytes())
        .collect()
}