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

//! GuruConnect - Remote Desktop Agent and Viewer
//!
//! Single binary for both agent (receiving connections) and viewer (initiating connections).
//!
//! Usage:
//!   guruconnect agent              - Run as background agent
//!   guruconnect view <session_id>  - View a remote session
//!   guruconnect install            - Install and register protocol handler
//!   guruconnect launch <url>       - Handle guruconnect:// URL
//!   guruconnect [support_code]     - Legacy: run agent with support code

// Hide console window by default on Windows (release builds)
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]

mod capture;
mod chat;
mod config;
mod encoder;
mod input;
mod install;
mod sas_client;
mod session;
mod startup;
mod transport;
mod tray;
mod update;
mod viewer;

pub mod proto {
    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
}

/// Build information embedded at compile time
pub mod build_info {
    /// Cargo package version (from Cargo.toml)
    pub const VERSION: &str = env!("CARGO_PKG_VERSION");

    /// Git commit hash (short, 8 chars)
    pub const GIT_HASH: &str = env!("GIT_HASH");

    /// Git commit hash (full)
    pub const GIT_HASH_FULL: &str = env!("GIT_HASH_FULL");

    /// Git branch name
    pub const GIT_BRANCH: &str = env!("GIT_BRANCH");

    /// Git dirty state ("clean" or "dirty")
    pub const GIT_DIRTY: &str = env!("GIT_DIRTY");

    /// Git commit date
    pub const GIT_COMMIT_DATE: &str = env!("GIT_COMMIT_DATE");

    /// Build timestamp (UTC)
    pub const BUILD_TIMESTAMP: &str = env!("BUILD_TIMESTAMP");

    /// Build profile (debug/release)
    pub const BUILD_PROFILE: &str = env!("BUILD_PROFILE");

    /// Target triple (e.g., x86_64-pc-windows-msvc)
    pub const BUILD_TARGET: &str = env!("BUILD_TARGET");

    /// Short version string for display (version + git hash)
    pub fn short_version() -> String {
        if GIT_DIRTY == "dirty" {
            format!("{}-{}-dirty", VERSION, GIT_HASH)
        } else {
            format!("{}-{}", VERSION, GIT_HASH)
        }
    }

    /// Full version string with all details
    pub fn full_version() -> String {
        format!(
            "GuruConnect v{}\n\
             Git:     {} ({})\n\
             Branch:  {}\n\
             Commit:  {}\n\
             Built:   {}\n\
             Profile: {}\n\
             Target:  {}",
            VERSION,
            GIT_HASH,
            GIT_DIRTY,
            GIT_BRANCH,
            GIT_COMMIT_DATE,
            BUILD_TIMESTAMP,
            BUILD_PROFILE,
            BUILD_TARGET
        )
    }
}

use anyhow::Result;
use clap::{Parser, Subcommand};
use tracing::{info, error, warn, Level};
use tracing_subscriber::FmtSubscriber;

#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MB_OK, MB_ICONINFORMATION, MB_ICONERROR};
#[cfg(windows)]
use windows::core::PCWSTR;
#[cfg(windows)]
use windows::Win32::System::Console::{AllocConsole, GetConsoleWindow};
#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{ShowWindow, SW_SHOW};

/// GuruConnect Remote Desktop
#[derive(Parser)]
#[command(name = "guruconnect")]
#[command(version = concat!(env!("CARGO_PKG_VERSION"), "-", env!("GIT_HASH")), about = "Remote desktop agent and viewer")]
struct Cli {
    #[command(subcommand)]
    command: Option<Commands>,

    /// Support code for legacy mode (runs agent with code)
    #[arg(value_name = "SUPPORT_CODE")]
    support_code: Option<String>,

    /// Enable verbose logging
    #[arg(short, long, global = true)]
    verbose: bool,

    /// Internal flag: set after auto-update to trigger cleanup
    #[arg(long, hide = true)]
    post_update: bool,
}

#[derive(Subcommand)]
enum Commands {
    /// Run as background agent (receive remote connections)
    Agent {
        /// Support code for one-time session
        #[arg(short, long)]
        code: Option<String>,
    },

    /// View a remote session (connect to an agent)
    View {
        /// Session ID to connect to
        session_id: String,

        /// Server URL
        #[arg(short, long, default_value = "wss://connect.azcomputerguru.com/ws/viewer")]
        server: String,

        /// API key for authentication
        #[arg(short, long, default_value = "")]
        api_key: String,
    },

    /// Install GuruConnect and register protocol handler
    Install {
        /// Skip UAC elevation, install for current user only
        #[arg(long)]
        user_only: bool,

        /// Called internally when running elevated
        #[arg(long, hide = true)]
        elevated: bool,
    },

    /// Uninstall GuruConnect
    Uninstall,

    /// Handle a guruconnect:// protocol URL
    Launch {
        /// The guruconnect:// URL to handle
        url: String,
    },

    /// Show detailed version and build information
    #[command(name = "version-info")]
    VersionInfo,
}

fn main() -> Result<()> {
    let cli = Cli::parse();

    // Initialize logging
    let level = if cli.verbose { Level::DEBUG } else { Level::INFO };
    FmtSubscriber::builder()
        .with_max_level(level)
        .with_target(true)
        .with_thread_ids(true)
        .init();

    info!("GuruConnect {} ({})", build_info::short_version(), build_info::BUILD_TARGET);
    info!("Built: {} | Commit: {}", build_info::BUILD_TIMESTAMP, build_info::GIT_COMMIT_DATE);

    // Handle post-update cleanup
    if cli.post_update {
        info!("Post-update mode: cleaning up old executable");
        update::cleanup_post_update();
    }

    match cli.command {
        Some(Commands::Agent { code }) => {
            run_agent_mode(code)
        }
        Some(Commands::View { session_id, server, api_key }) => {
            run_viewer_mode(&server, &session_id, &api_key)
        }
        Some(Commands::Install { user_only, elevated }) => {
            run_install(user_only || elevated)
        }
        Some(Commands::Uninstall) => {
            run_uninstall()
        }
        Some(Commands::Launch { url }) => {
            run_launch(&url)
        }
        Some(Commands::VersionInfo) => {
            // Show detailed version info (allocate console on Windows for visibility)
            #[cfg(windows)]
            show_debug_console();
            println!("{}", build_info::full_version());
            Ok(())
        }
        None => {
            // No subcommand - detect mode from filename or embedded config
            // Legacy: if support_code arg provided, use that
            if let Some(code) = cli.support_code {
                return run_agent_mode(Some(code));
            }

            // Detect run mode from filename
            use config::RunMode;
            match config::Config::detect_run_mode() {
                RunMode::Viewer => {
                    // Filename indicates viewer-only (e.g., "GuruConnect-Viewer.exe")
                    info!("Viewer mode detected from filename");
                    if !install::is_protocol_handler_registered() {
                        info!("Installing protocol handler for viewer");
                        run_install(false)
                    } else {
                        info!("Viewer already installed, nothing to do");
                        show_message_box("GuruConnect Viewer", "GuruConnect viewer is installed.\n\nUse guruconnect:// links to connect to remote sessions.");
                        Ok(())
                    }
                }
                RunMode::TempSupport(code) => {
                    // Filename contains support code (e.g., "GuruConnect-123456.exe")
                    info!("Temp support session detected from filename: {}", code);
                    run_agent_mode(Some(code))
                }
                RunMode::PermanentAgent => {
                    // Embedded config found - run as permanent agent
                    info!("Permanent agent mode detected (embedded config)");
                    if !install::is_protocol_handler_registered() {
                        // First run - install then run as agent
                        info!("First run - installing agent");
                        if let Err(e) = install::install(false) {
                            warn!("Installation failed: {}", e);
                        }
                    }
                    run_agent_mode(None)
                }
                RunMode::Default => {
                    // No special mode detected - use legacy logic
                    if !install::is_protocol_handler_registered() {
                        // Protocol handler not registered - user likely downloaded from web
                        info!("Protocol handler not registered, running installer");
                        run_install(false)
                    } else if config::Config::has_agent_config() {
                        // Has agent config - run as agent
                        info!("Agent config found, running as agent");
                        run_agent_mode(None)
                    } else {
                        // Viewer-only installation - just exit silently
                        info!("Viewer-only installation, exiting");
                        Ok(())
                    }
                }
            }
        }
    }
}

/// Run in agent mode (receive remote connections)
fn run_agent_mode(support_code: Option<String>) -> Result<()> {
    info!("Running in agent mode");

    // Check elevation status
    if install::is_elevated() {
        info!("Running with elevated (administrator) privileges");
    } else {
        info!("Running with standard user privileges");
    }

    // Load configuration
    let mut config = config::Config::load()?;

    // Set support code if provided
    if let Some(code) = support_code {
        info!("Support code: {}", code);
        config.support_code = Some(code);
    }

    info!("Server: {}", config.server_url);
    if let Some(ref company) = config.company {
        info!("Company: {}", company);
    }
    if let Some(ref site) = config.site {
        info!("Site: {}", site);
    }

    // Run the agent
    let rt = tokio::runtime::Runtime::new()?;
    rt.block_on(run_agent(config))
}

/// Run in viewer mode (connect to remote session)
fn run_viewer_mode(server: &str, session_id: &str, api_key: &str) -> Result<()> {
    info!("Running in viewer mode");
    info!("Connecting to session: {}", session_id);

    let rt = tokio::runtime::Runtime::new()?;
    rt.block_on(viewer::run(server, session_id, api_key))
}

/// Handle guruconnect:// URL launch
fn run_launch(url: &str) -> Result<()> {
    info!("Handling protocol URL: {}", url);

    match install::parse_protocol_url(url) {
        Ok((server, session_id, token)) => {
            let api_key = token.unwrap_or_default();
            run_viewer_mode(&server, &session_id, &api_key)
        }
        Err(e) => {
            error!("Failed to parse URL: {}", e);
            show_error_box("GuruConnect", &format!("Invalid URL: {}", e));
            Err(e)
        }
    }
}

/// Install GuruConnect
fn run_install(force_user_install: bool) -> Result<()> {
    info!("Installing GuruConnect...");

    match install::install(force_user_install) {
        Ok(()) => {
            show_message_box("GuruConnect", "Installation complete!\n\nYou can now use guruconnect:// links.");
            Ok(())
        }
        Err(e) => {
            error!("Installation failed: {}", e);
            show_error_box("GuruConnect", &format!("Installation failed: {}", e));
            Err(e)
        }
    }
}

/// Uninstall GuruConnect
fn run_uninstall() -> Result<()> {
    info!("Uninstalling GuruConnect...");

    // Remove from startup
    if let Err(e) = startup::remove_from_startup() {
        warn!("Failed to remove from startup: {}", e);
    }

    // TODO: Remove registry keys for protocol handler
    // TODO: Remove install directory

    show_message_box("GuruConnect", "Uninstall complete.");
    Ok(())
}

/// Show a message box (Windows only)
#[cfg(windows)]
fn show_message_box(title: &str, message: &str) {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    let title_wide: Vec<u16> = OsStr::new(title)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let message_wide: Vec<u16> = OsStr::new(message)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        MessageBoxW(
            None,
            PCWSTR(message_wide.as_ptr()),
            PCWSTR(title_wide.as_ptr()),
            MB_OK | MB_ICONINFORMATION,
        );
    }
}

#[cfg(not(windows))]
fn show_message_box(_title: &str, message: &str) {
    println!("{}", message);
}

/// Show an error message box (Windows only)
#[cfg(windows)]
fn show_error_box(title: &str, message: &str) {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    let title_wide: Vec<u16> = OsStr::new(title)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let message_wide: Vec<u16> = OsStr::new(message)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        MessageBoxW(
            None,
            PCWSTR(message_wide.as_ptr()),
            PCWSTR(title_wide.as_ptr()),
            MB_OK | MB_ICONERROR,
        );
    }
}

#[cfg(not(windows))]
fn show_error_box(_title: &str, message: &str) {
    eprintln!("ERROR: {}", message);
}

/// Show debug console window (Windows only)
#[cfg(windows)]
#[allow(dead_code)]
fn show_debug_console() {
    unsafe {
        let hwnd = GetConsoleWindow();
        if hwnd.0 == std::ptr::null_mut() {
            let _ = AllocConsole();
        } else {
            let _ = ShowWindow(hwnd, SW_SHOW);
        }
    }
}

#[cfg(not(windows))]
#[allow(dead_code)]
fn show_debug_console() {}

/// Clean up before exiting
fn cleanup_on_exit() {
    info!("Cleaning up before exit");
    if let Err(e) = startup::remove_from_startup() {
        warn!("Failed to remove from startup: {}", e);
    }
}

/// Run the agent main loop
async fn run_agent(config: config::Config) -> Result<()> {
    let elevated = install::is_elevated();
    let mut session = session::SessionManager::new(config.clone(), elevated);
    let is_support_session = config.support_code.is_some();
    let hostname = config.hostname();

    // Add to startup
    if let Err(e) = startup::add_to_startup() {
        warn!("Failed to add to startup: {}", e);
    }

    // Create tray icon
    let tray = match tray::TrayController::new(&hostname, config.support_code.as_deref(), is_support_session) {
        Ok(t) => {
            info!("Tray icon created");
            Some(t)
        }
        Err(e) => {
            warn!("Failed to create tray icon: {}", e);
            None
        }
    };

    // Create chat controller
    let chat_ctrl = chat::ChatController::new();

    // Connect to server and run main loop
    loop {
        info!("Connecting to server...");

        if is_support_session {
            if let Some(ref t) = tray {
                if t.exit_requested() {
                    info!("Exit requested by user");
                    cleanup_on_exit();
                    return Ok(());
                }
            }
        }

        match session.connect().await {
            Ok(_) => {
                info!("Connected to server");

                if let Some(ref t) = tray {
                    t.update_status("Status: Connected");
                }

                if let Err(e) = session.run_with_tray(tray.as_ref(), chat_ctrl.as_ref()).await {
                    let error_msg = e.to_string();

                    if error_msg.contains("USER_EXIT") {
                        info!("Session ended by user");
                        cleanup_on_exit();
                        return Ok(());
                    }

                    if error_msg.contains("SESSION_CANCELLED") {
                        info!("Session was cancelled by technician");
                        cleanup_on_exit();
                        show_message_box("Support Session Ended", "The support session was cancelled.");
                        return Ok(());
                    }

                    if error_msg.contains("ADMIN_DISCONNECT") {
                        info!("Session disconnected by administrator - uninstalling");
                        if let Err(e) = startup::uninstall() {
                            warn!("Uninstall failed: {}", e);
                        }
                        show_message_box("Remote Session Ended", "The session was ended by the administrator.");
                        return Ok(());
                    }

                    if error_msg.contains("ADMIN_UNINSTALL") {
                        info!("Uninstall command received from server - uninstalling");
                        if let Err(e) = startup::uninstall() {
                            warn!("Uninstall failed: {}", e);
                        }
                        show_message_box("GuruConnect Removed", "This computer has been removed from remote management.");
                        return Ok(());
                    }

                    if error_msg.contains("ADMIN_RESTART") {
                        info!("Restart command received - will reconnect");
                        // Don't exit, just let the loop continue to reconnect
                    } else {
                        error!("Session error: {}", e);
                    }
                }
            }
            Err(e) => {
                let error_msg = e.to_string();

                if error_msg.contains("cancelled") {
                    info!("Support code was cancelled");
                    cleanup_on_exit();
                    show_message_box("Support Session Cancelled", "This support session has been cancelled.");
                    return Ok(());
                }

                error!("Connection failed: {}", e);
            }
        }

        if is_support_session {
            info!("Support session ended, not reconnecting");
            cleanup_on_exit();
            return Ok(());
        }

        info!("Reconnecting in 5 seconds...");
        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
    }
}