     1→//! WebSocket client for server communication
     2→//!
     3→//! Handles the WebSocket connection lifecycle including:
     4→//! - Connection establishment
     5→//! - Authentication handshake
     6→//! - Message sending/receiving
     7→//! - Heartbeat maintenance
     8→//! - Command handling
     9→
    10→use std::sync::Arc;
    11→use std::time::Duration;
    12→
    13→use anyhow::{Context, Result};
    14→use futures_util::{SinkExt, StreamExt};
    15→use tokio::sync::mpsc;
    16→use tokio::time::{interval, timeout};
    17→use tokio_tungstenite::{connect_async, tungstenite::Message};
    18→use tracing::{debug, error, info, warn};
    19→
    20→use super::{AgentMessage, AuthPayload, CommandPayload, ServerMessage, UpdatePayload, UpdateResultPayload, UpdateStatus};
    21→use crate::metrics::NetworkState;
    22→use crate::updater::{AgentUpdater, UpdaterConfig};
    23→use crate::AppState;
    24→
    25→/// WebSocket client for communicating with the GuruRMM server
    26→pub struct WebSocketClient;
    27→
    28→impl WebSocketClient {
    29→    /// Connect to the server and run the message loop
    30→    ///
    31→    /// This function will return when the connection is closed or an error occurs.
    32→    /// The caller should handle reconnection logic.
    33→    pub async fn connect_and_run(state: Arc<AppState>) -> Result<()> {
    34→        let url = &state.config.server.url;
    35→
    36→        // Connect to WebSocket server
    37→        info!("Connecting to {}", url);
    38→        let (ws_stream, response) = connect_async(url)
    39→            .await
    40→            .context("Failed to connect to WebSocket server")?;
    41→
    42→        info!(
    43→            "WebSocket connected (HTTP status: {})",
    44→            response.status()
    45→        );
    46→
    47→        let (mut write, mut read) = ws_stream.split();
    48→
    49→        // Check for pending update (from previous update attempt)
    50→        let updater_config = UpdaterConfig::default();
    51→        let pending_update = AgentUpdater::load_pending_update(&updater_config).await;
    52→
    53→        // If we have pending update info, we just restarted after an update
    54→        let (previous_version, pending_update_id) = if let Some(ref info) = pending_update {
    55→            info!(
    56→                "Found pending update info: {} -> {} (id: {})",
    57→                info.old_version, info.target_version, info.update_id
    58→            );
    59→            (Some(info.old_version.clone()), Some(info.update_id))
    60→        } else {
    61→            (None, None)
    62→        };
    63→
    64→        // Send authentication message
    65→        let auth_msg = AgentMessage::Auth(AuthPayload {
    66→            api_key: state.config.server.api_key.clone(),
    67→            device_id: crate::device_id::get_device_id(),
    68→            hostname: state.config.get_hostname(),
    69→            os_type: std::env::consts::OS.to_string(),
    70→            os_version: sysinfo::System::os_version().unwrap_or_else(|| "unknown".to_string()),
    71→            agent_version: env!("CARGO_PKG_VERSION").to_string(),
    72→            architecture: Self::get_architecture().to_string(),
    73→            previous_version,
    74→            pending_update_id,
    75→        });
    76→
    77→        let auth_json = serde_json::to_string(&auth_msg)?;
    78→        write.send(Message::Text(auth_json)).await?;
    79→        debug!("Sent authentication message");
    80→
    81→        // Wait for auth response with timeout
    82→        let auth_response = timeout(Duration::from_secs(10), read.next())
    83→            .await
    84→            .context("Authentication timeout")?
    85→            .ok_or_else(|| anyhow::anyhow!("Connection closed before auth response"))?
    86→            .context("Failed to receive auth response")?;
    87→
    88→        // Parse auth response
    89→        if let Message::Text(text) = auth_response {
    90→            let server_msg: ServerMessage =
    91→                serde_json::from_str(&text).context("Failed to parse auth response")?;
    92→
    93→            match server_msg {
    94→                ServerMessage::AuthAck(ack) => {
    95→                    if ack.success {
    96→                        info!("Authentication successful, agent_id: {:?}", ack.agent_id);
    97→                        *state.connected.write().await = true;
    98→
    99→                        // Send initial network state immediately after auth
   100→                        let network_state = NetworkState::collect();
   101→                        info!(
   102→                            "Sending initial network state ({} interfaces)",
   103→                            network_state.interfaces.len()
   104→                        );
   105→                        let network_msg = AgentMessage::NetworkState(network_state);
   106→                        let network_json = serde_json::to_string(&network_msg)?;
   107→                        write.send(Message::Text(network_json)).await?;
   108→                    } else {
   109→                        error!("Authentication failed: {:?}", ack.error);
   110→                        return Err(anyhow::anyhow!(
   111→                            "Authentication failed: {}",
   112→                            ack.error.unwrap_or_else(|| "Unknown error".to_string())
   113→                        ));
   114→                    }
   115→                }
   116→                ServerMessage::Error { code, message } => {
   117→                    error!("Server error during auth: {} - {}", code, message);
   118→                    return Err(anyhow::anyhow!("Server error: {} - {}", code, message));
   119→                }
   120→                _ => {
   121→                    warn!("Unexpected message during auth: {:?}", server_msg);
   122→                }
   123→            }
   124→        }
   125→
   126→        // Create channel for outgoing messages
   127→        let (tx, mut rx) = mpsc::channel::<AgentMessage>(100);
   128→
   129→        // Spawn metrics sender task
   130→        let metrics_tx = tx.clone();
   131→        let metrics_state = Arc::clone(&state);
   132→        let metrics_interval = state.config.metrics.interval_seconds;
   133→
   134→        let metrics_task = tokio::spawn(async move {
   135→            let mut timer = interval(Duration::from_secs(metrics_interval));
   136→
   137→            loop {
   138→                timer.tick().await;
   139→
   140→                let metrics = metrics_state.metrics_collector.collect().await;
   141→                if metrics_tx.send(AgentMessage::Metrics(metrics)).await.is_err() {
   142→                    debug!("Metrics channel closed");
   143→                    break;
   144→                }
   145→            }
   146→        });
   147→
   148→        // Spawn network state monitor task (checks for changes every 30 seconds)
   149→        let network_tx = tx.clone();
   150→        let network_task = tokio::spawn(async move {
   151→            // Check for network changes every 30 seconds
   152→            let mut timer = interval(Duration::from_secs(30));
   153→            let mut last_state = NetworkState::collect();
   154→
   155→            loop {
   156→                timer.tick().await;
   157→
   158→                let current_state = NetworkState::collect();
   159→                if current_state.has_changed(&last_state) {
   160→                    info!(
   161→                        "Network state changed (hash: {} -> {}), sending update",
   162→                        last_state.state_hash, current_state.state_hash
   163→                    );
   164→
   165→                    // Log the changes for debugging
   166→                    for iface in &current_state.interfaces {
   167→                        debug!(
   168→                            "  Interface {}: IPv4={:?}",
   169→                            iface.name, iface.ipv4_addresses
   170→                        );
   171→                    }
   172→
   173→                    if network_tx
   174→                        .send(AgentMessage::NetworkState(current_state.clone()))
   175→                        .await
   176→                        .is_err()
   177→                    {
   178→                        debug!("Network channel closed");
   179→                        break;
   180→                    }
   181→                    last_state = current_state;
   182→                }
   183→            }
   184→        });
   185→
   186→        // Spawn heartbeat task
   187→        let heartbeat_tx = tx.clone();
   188→        let heartbeat_task = tokio::spawn(async move {
   189→            let mut timer = interval(Duration::from_secs(30));
   190→
   191→            loop {
   192→                timer.tick().await;
   193→
   194→                if heartbeat_tx.send(AgentMessage::Heartbeat).await.is_err() {
   195→                    debug!("Heartbeat channel closed");
   196→                    break;
   197→                }
   198→            }
   199→        });
   200→
   201→        // Main message loop
   202→        let result: Result<()> = loop {
   203→            tokio::select! {
   204→                // Handle outgoing messages
   205→                Some(msg) = rx.recv() => {
   206→                    let json = serde_json::to_string(&msg)?;
   207→                    if let Err(e) = write.send(Message::Text(json)).await {
   208→                        break Err(e.into());
   209→                    }
   210→
   211→                    match &msg {
   212→                        AgentMessage::Metrics(m) => {
   213→                            debug!("Sent metrics: CPU={:.1}%", m.cpu_percent);
   214→                        }
   215→                        AgentMessage::NetworkState(n) => {
   216→                            debug!("Sent network state: {} interfaces, hash={}",
   217→                                n.interfaces.len(), n.state_hash);
   218→                        }
   219→                        AgentMessage::Heartbeat => {
   220→                            debug!("Sent heartbeat");
   221→                        }
   222→                        _ => {
   223→                            debug!("Sent message: {:?}", std::mem::discriminant(&msg));
   224→                        }
   225→                    }
   226→                }
   227→
   228→                // Handle incoming messages
   229→                Some(msg_result) = read.next() => {
   230→                    match msg_result {
   231→                        Ok(Message::Text(text)) => {
   232→                            if let Err(e) = Self::handle_server_message(&text, &tx).await {
   233→                                error!("Error handling message: {}", e);
   234→                            }
   235→                        }
   236→                        Ok(Message::Ping(data)) => {
   237→                            if let Err(e) = write.send(Message::Pong(data)).await {
   238→                                break Err(e.into());
   239→                            }
   240→                        }
   241→                        Ok(Message::Pong(_)) => {
   242→                            debug!("Received pong");
   243→                        }
   244→                        Ok(Message::Close(frame)) => {
   245→                            info!("Server closed connection: {:?}", frame);
   246→                            break Ok(());
   247→                        }
   248→                        Ok(Message::Binary(_)) => {
   249→                            warn!("Received unexpected binary message");
   250→                        }
   251→                        Ok(Message::Frame(_)) => {
   252→                            // Raw frame, usually not seen
   253→                        }
   254→                        Err(e) => {
   255→                            error!("WebSocket error: {}", e);
   256→                            break Err(e.into());
   257→                        }
   258→                    }
   259→                }
   260→
   261→                // Connection timeout (no activity)
   262→                _ = tokio::time::sleep(Duration::from_secs(90)) => {
   263→                    warn!("Connection timeout, no activity for 90 seconds");
   264→                    break Err(anyhow::anyhow!("Connection timeout"));
   265→                }
   266→            }
   267→        };
   268→
   269→        // Cleanup
   270→        metrics_task.abort();
   271→        network_task.abort();
   272→        heartbeat_task.abort();
   273→        *state.connected.write().await = false;
   274→
   275→        result
   276→    }
   277→
   278→    /// Handle a message received from the server
   279→    async fn handle_server_message(
   280→        text: &str,
   281→        tx: &mpsc::Sender<AgentMessage>,
   282→    ) -> Result<()> {
   283→        let msg: ServerMessage =
   284→            serde_json::from_str(text).context("Failed to parse server message")?;
   285→
   286→        match msg {
   287→            ServerMessage::Command(cmd) => {
   288→                info!("Received command: {:?} (id: {})", cmd.command_type, cmd.id);
   289→                Self::execute_command(cmd, tx.clone()).await;
   290→            }
   291→            ServerMessage::ConfigUpdate(update) => {
   292→                info!("Received config update: {:?}", update);
   293→                // Config updates will be handled in a future phase
   294→            }
   295→            ServerMessage::Ack { message_id } => {
   296→                debug!("Received ack for message: {:?}", message_id);
   297→            }
   298→            ServerMessage::AuthAck(_) => {
   299→                // Already handled during initial auth
   300→            }
   301→            ServerMessage::Error { code, message } => {
   302→                error!("Server error: {} - {}", code, message);
   303→            }
   304→            ServerMessage::Update(payload) => {
   305→                info!(
   306→                    "Received update command: {} -> {} (id: {})",
   307→                    env!("CARGO_PKG_VERSION"),
   308→                    payload.target_version,
   309→                    payload.update_id
   310→                );
   311→                Self::handle_update(payload, tx.clone()).await;
   312→            }
   313→            ServerMessage::TrayPolicyUpdate(policy) => {
   314→                info!("Received tray policy update: enabled={}", policy.enabled);
   315→                // Tray policy updates are handled via IPC broadcast
   316→                // The IPC server will pick this up and notify connected tray apps
   317→            }
   318→        }
   319→
   320→        Ok(())
   321→    }
   322→
   323→    /// Handle an update command from the server
   324→    async fn handle_update(payload: UpdatePayload, tx: mpsc::Sender<AgentMessage>) {
   325→        // Send starting status
   326→        let starting_result = UpdateResultPayload {
   327→            update_id: payload.update_id,
   328→            status: UpdateStatus::Starting,
   329→            old_version: env!("CARGO_PKG_VERSION").to_string(),
   330→            new_version: None,
   331→            error: None,
   332→        };
   333→        let _ = tx.send(AgentMessage::UpdateResult(starting_result)).await;
   334→
   335→        // Spawn update in background (it will restart the service)
   336→        tokio::spawn(async move {
   337→            let config = UpdaterConfig::default();
   338→            let updater = AgentUpdater::new(config);
   339→            let result = updater.perform_update(payload).await;
   340→
   341→            // If we reach here, the update failed (successful update restarts the process)
   342→            let _ = tx.send(AgentMessage::UpdateResult(result)).await;
   343→        });
   344→    }
   345→
   346→    /// Get the current architecture
   347→    fn get_architecture() -> &'static str {
   348→        #[cfg(target_arch = "x86_64")]
   349→        { "amd64" }
   350→        #[cfg(target_arch = "aarch64")]
   351→        { "arm64" }
   352→        #[cfg(target_arch = "x86")]
   353→        { "386" }
   354→        #[cfg(target_arch = "arm")]
   355→        { "arm" }
   356→        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64", target_arch = "x86", target_arch = "arm")))]
   357→        { "unknown" }
   358→    }
   359→
   360→    /// Execute a command received from the server
   361→    async fn execute_command(cmd: CommandPayload, tx: mpsc::Sender<AgentMessage>) {
   362→        let command_id = cmd.id;
   363→
   364→        // Spawn command execution in background
   365→        tokio::spawn(async move {
   366→            let start = std::time::Instant::now();
   367→
   368→            let result = Self::run_command(&cmd).await;
   369→            let duration_ms = start.elapsed().as_millis() as u64;
   370→
   371→            let (exit_code, stdout, stderr) = match result {
   372→                Ok((code, out, err)) => (code, out, err),
   373→                Err(e) => (-1, String::new(), format!("Execution error: {}", e)),
   374→            };
   375→
   376→            let result_msg = AgentMessage::CommandResult(super::CommandResultPayload {
   377→                command_id,
   378→                exit_code,
   379→                stdout,
   380→                stderr,
   381→                duration_ms,
   382→            });
   383→
   384→            if tx.send(result_msg).await.is_err() {
   385→                error!("Failed to send command result");
   386→            }
   387→        });
   388→    }
   389→
   390→    /// Run a command and capture output
   391→    async fn run_command(cmd: &CommandPayload) -> Result<(i32, String, String)> {
   392→        use tokio::process::Command;
   393→
   394→        let timeout_secs = cmd.timeout_seconds.unwrap_or(300); // 5 minute default
   395→
   396→        let mut command = match &cmd.command_type {
   397→            super::CommandType::Shell => {
   398→                #[cfg(windows)]
   399→                {
   400→                    let mut c = Command::new("cmd");
   401→                    c.args(["/C", &cmd.command]);
   402→                    c
   403→                }
   404→                #[cfg(unix)]
   405→                {
   406→                    let mut c = Command::new("sh");
   407→                    c.args(["-c", &cmd.command]);
   408→                    c
   409→                }
   410→            }
   411→            super::CommandType::PowerShell => {
   412→                let mut c = Command::new("powershell");
   413→                c.args(["-NoProfile", "-NonInteractive", "-Command", &cmd.command]);
   414→                c
   415→            }
   416→            super::CommandType::Python => {
   417→                let mut c = Command::new("python");
   418→                c.args(["-c", &cmd.command]);
   419→                c
   420→            }
   421→            super::CommandType::Script { interpreter } => {
   422→                let mut c = Command::new(interpreter);
   423→                c.args(["-c", &cmd.command]);
   424→                c
   425→            }
   426→        };
   427→
   428→        // Capture output
   429→        command.stdout(std::process::Stdio::piped());
   430→        command.stderr(std::process::Stdio::piped());
   431→
   432→        // Execute with timeout
   433→        let output = timeout(Duration::from_secs(timeout_secs), command.output())
   434→            .await
   435→            .context("Command timeout")?
   436→            .context("Failed to execute command")?;
   437→
   438→        let exit_code = output.status.code().unwrap_or(-1);
   439→        let stdout = String::from_utf8_lossy(&output.stdout).to_string();
   440→        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
   441→
   442→        Ok((exit_code, stdout, stderr))
   443→    }
   444→}
   445→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! WebSocket client for server communication
//!
//! Handles the WebSocket connection lifecycle including:
//! - Connection establishment
//! - Authentication handshake
//! - Message sending/receiving
//! - Heartbeat maintenance
//! - Command handling

use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use futures_util::{SinkExt, StreamExt};
use tokio::sync::mpsc;
use tokio::time::{interval, timeout};
use tokio_tungstenite::{connect_async, tungstenite::Message};
use tracing::{debug, error, info, warn};

use super::{AgentMessage, AuthPayload, CommandPayload, ServerMessage, UpdatePayload, UpdateResultPayload, UpdateStatus};
use crate::metrics::NetworkState;
use crate::updater::{AgentUpdater, UpdaterConfig};
use crate::AppState;

/// WebSocket client for communicating with the GuruRMM server
pub struct WebSocketClient;

impl WebSocketClient {
    /// Connect to the server and run the message loop
    ///
    /// This function will return when the connection is closed or an error occurs.
    /// The caller should handle reconnection logic.
    pub async fn connect_and_run(state: Arc<AppState>) -> Result<()> {
        let url = &state.config.server.url;

        // Connect to WebSocket server
        info!("Connecting to {}", url);
        let (ws_stream, response) = connect_async(url)
            .await
            .context("Failed to connect to WebSocket server")?;

        info!(
            "WebSocket connected (HTTP status: {})",
            response.status()
        );

        let (mut write, mut read) = ws_stream.split();

        // Check for pending update (from previous update attempt)
        let updater_config = UpdaterConfig::default();
        let pending_update = AgentUpdater::load_pending_update(&updater_config).await;

        // If we have pending update info, we just restarted after an update
        let (previous_version, pending_update_id) = if let Some(ref info) = pending_update {
            info!(
                "Found pending update info: {} -> {} (id: {})",
                info.old_version, info.target_version, info.update_id
            );
            (Some(info.old_version.clone()), Some(info.update_id))
        } else {
            (None, None)
        };

        // Send authentication message
        let auth_msg = AgentMessage::Auth(AuthPayload {
            api_key: state.config.server.api_key.clone(),
            device_id: crate::device_id::get_device_id(),
            hostname: state.config.get_hostname(),
            os_type: std::env::consts::OS.to_string(),
            os_version: sysinfo::System::os_version().unwrap_or_else(|| "unknown".to_string()),
            agent_version: env!("CARGO_PKG_VERSION").to_string(),
            architecture: Self::get_architecture().to_string(),
            previous_version,
            pending_update_id,
        });

        let auth_json = serde_json::to_string(&auth_msg)?;
        write.send(Message::Text(auth_json)).await?;
        debug!("Sent authentication message");

        // Wait for auth response with timeout
        let auth_response = timeout(Duration::from_secs(10), read.next())
            .await
            .context("Authentication timeout")?
            .ok_or_else(|| anyhow::anyhow!("Connection closed before auth response"))?
            .context("Failed to receive auth response")?;

        // Parse auth response
        if let Message::Text(text) = auth_response {
            let server_msg: ServerMessage =
                serde_json::from_str(&text).context("Failed to parse auth response")?;

            match server_msg {
                ServerMessage::AuthAck(ack) => {
                    if ack.success {
                        info!("Authentication successful, agent_id: {:?}", ack.agent_id);
                        *state.connected.write().await = true;

                        // Send initial network state immediately after auth
                        let network_state = NetworkState::collect();
                        info!(
                            "Sending initial network state ({} interfaces)",
                            network_state.interfaces.len()
                        );
                        let network_msg = AgentMessage::NetworkState(network_state);
                        let network_json = serde_json::to_string(&network_msg)?;
                        write.send(Message::Text(network_json)).await?;
                    } else {
                        error!("Authentication failed: {:?}", ack.error);
                        return Err(anyhow::anyhow!(
                            "Authentication failed: {}",
                            ack.error.unwrap_or_else(|| "Unknown error".to_string())
                        ));
                    }
                }
                ServerMessage::Error { code, message } => {
                    error!("Server error during auth: {} - {}", code, message);
                    return Err(anyhow::anyhow!("Server error: {} - {}", code, message));
                }
                _ => {
                    warn!("Unexpected message during auth: {:?}", server_msg);
                }
            }
        }

        // Create channel for outgoing messages
        let (tx, mut rx) = mpsc::channel::<AgentMessage>(100);

        // Spawn metrics sender task
        let metrics_tx = tx.clone();
        let metrics_state = Arc::clone(&state);
        let metrics_interval = state.config.metrics.interval_seconds;

        let metrics_task = tokio::spawn(async move {
            let mut timer = interval(Duration::from_secs(metrics_interval));

            loop {
                timer.tick().await;

                let metrics = metrics_state.metrics_collector.collect().await;
                if metrics_tx.send(AgentMessage::Metrics(metrics)).await.is_err() {
                    debug!("Metrics channel closed");
                    break;
                }
            }
        });

        // Spawn network state monitor task (checks for changes every 30 seconds)
        let network_tx = tx.clone();
        let network_task = tokio::spawn(async move {
            // Check for network changes every 30 seconds
            let mut timer = interval(Duration::from_secs(30));
            let mut last_state = NetworkState::collect();

            loop {
                timer.tick().await;

                let current_state = NetworkState::collect();
                if current_state.has_changed(&last_state) {
                    info!(
                        "Network state changed (hash: {} -> {}), sending update",
                        last_state.state_hash, current_state.state_hash
                    );

                    // Log the changes for debugging
                    for iface in &current_state.interfaces {
                        debug!(
                            "  Interface {}: IPv4={:?}",
                            iface.name, iface.ipv4_addresses
                        );
                    }

                    if network_tx
                        .send(AgentMessage::NetworkState(current_state.clone()))
                        .await
                        .is_err()
                    {
                        debug!("Network channel closed");
                        break;
                    }
                    last_state = current_state;
                }
            }
        });

        // Spawn heartbeat task
        let heartbeat_tx = tx.clone();
        let heartbeat_task = tokio::spawn(async move {
            let mut timer = interval(Duration::from_secs(30));

            loop {
                timer.tick().await;

                if heartbeat_tx.send(AgentMessage::Heartbeat).await.is_err() {
                    debug!("Heartbeat channel closed");
                    break;
                }
            }
        });

        // Main message loop
        let result: Result<()> = loop {
            tokio::select! {
                // Handle outgoing messages
                Some(msg) = rx.recv() => {
                    let json = serde_json::to_string(&msg)?;
                    if let Err(e) = write.send(Message::Text(json)).await {
                        break Err(e.into());
                    }

                    match &msg {
                        AgentMessage::Metrics(m) => {
                            debug!("Sent metrics: CPU={:.1}%", m.cpu_percent);
                        }
                        AgentMessage::NetworkState(n) => {
                            debug!("Sent network state: {} interfaces, hash={}",
                                n.interfaces.len(), n.state_hash);
                        }
                        AgentMessage::Heartbeat => {
                            debug!("Sent heartbeat");
                        }
                        _ => {
                            debug!("Sent message: {:?}", std::mem::discriminant(&msg));
                        }
                    }
                }

                // Handle incoming messages
                Some(msg_result) = read.next() => {
                    match msg_result {
                        Ok(Message::Text(text)) => {
                            if let Err(e) = Self::handle_server_message(&text, &tx).await {
                                error!("Error handling message: {}", e);
                            }
                        }
                        Ok(Message::Ping(data)) => {
                            if let Err(e) = write.send(Message::Pong(data)).await {
                                break Err(e.into());
                            }
                        }
                        Ok(Message::Pong(_)) => {
                            debug!("Received pong");
                        }
                        Ok(Message::Close(frame)) => {
                            info!("Server closed connection: {:?}", frame);
                            break Ok(());
                        }
                        Ok(Message::Binary(_)) => {
                            warn!("Received unexpected binary message");
                        }
                        Ok(Message::Frame(_)) => {
                            // Raw frame, usually not seen
                        }
                        Err(e) => {
                            error!("WebSocket error: {}", e);
                            break Err(e.into());
                        }
                    }
                }

                // Connection timeout (no activity)
                _ = tokio::time::sleep(Duration::from_secs(90)) => {
                    warn!("Connection timeout, no activity for 90 seconds");
                    break Err(anyhow::anyhow!("Connection timeout"));
                }
            }
        };

        // Cleanup
        metrics_task.abort();
        network_task.abort();
        heartbeat_task.abort();
        *state.connected.write().await = false;

        result
    }

    /// Handle a message received from the server
    async fn handle_server_message(
        text: &str,
        tx: &mpsc::Sender<AgentMessage>,
    ) -> Result<()> {
        let msg: ServerMessage =
            serde_json::from_str(text).context("Failed to parse server message")?;

        match msg {
            ServerMessage::Command(cmd) => {
                info!("Received command: {:?} (id: {})", cmd.command_type, cmd.id);
                Self::execute_command(cmd, tx.clone()).await;
            }
            ServerMessage::ConfigUpdate(update) => {
                info!("Received config update: {:?}", update);
                // Config updates will be handled in a future phase
            }
            ServerMessage::Ack { message_id } => {
                debug!("Received ack for message: {:?}", message_id);
            }
            ServerMessage::AuthAck(_) => {
                // Already handled during initial auth
            }
            ServerMessage::Error { code, message } => {
                error!("Server error: {} - {}", code, message);
            }
            ServerMessage::Update(payload) => {
                info!(
                    "Received update command: {} -> {} (id: {})",
                    env!("CARGO_PKG_VERSION"),
                    payload.target_version,
                    payload.update_id
                );
                Self::handle_update(payload, tx.clone()).await;
            }
            ServerMessage::TrayPolicyUpdate(policy) => {
                info!("Received tray policy update: enabled={}", policy.enabled);
                // Tray policy updates are handled via IPC broadcast
                // The IPC server will pick this up and notify connected tray apps
            }
        }

        Ok(())
    }

    /// Handle an update command from the server
    async fn handle_update(payload: UpdatePayload, tx: mpsc::Sender<AgentMessage>) {
        // Send starting status
        let starting_result = UpdateResultPayload {
            update_id: payload.update_id,
            status: UpdateStatus::Starting,
            old_version: env!("CARGO_PKG_VERSION").to_string(),
            new_version: None,
            error: None,
        };
        let _ = tx.send(AgentMessage::UpdateResult(starting_result)).await;

        // Spawn update in background (it will restart the service)
        tokio::spawn(async move {
            let config = UpdaterConfig::default();
            let updater = AgentUpdater::new(config);
            let result = updater.perform_update(payload).await;

            // If we reach here, the update failed (successful update restarts the process)
            let _ = tx.send(AgentMessage::UpdateResult(result)).await;
        });
    }

    /// Get the current architecture
    fn get_architecture() -> &'static str {
        #[cfg(target_arch = "x86_64")]
        { "amd64" }
        #[cfg(target_arch = "aarch64")]
        { "arm64" }
        #[cfg(target_arch = "x86")]
        { "386" }
        #[cfg(target_arch = "arm")]
        { "arm" }
        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64", target_arch = "x86", target_arch = "arm")))]
        { "unknown" }
    }

    /// Execute a command received from the server
    async fn execute_command(cmd: CommandPayload, tx: mpsc::Sender<AgentMessage>) {
        let command_id = cmd.id;

        // Spawn command execution in background
        tokio::spawn(async move {
            let start = std::time::Instant::now();

            let result = Self::run_command(&cmd).await;
            let duration_ms = start.elapsed().as_millis() as u64;

            let (exit_code, stdout, stderr) = match result {
                Ok((code, out, err)) => (code, out, err),
                Err(e) => (-1, String::new(), format!("Execution error: {}", e)),
            };

            let result_msg = AgentMessage::CommandResult(super::CommandResultPayload {
                command_id,
                exit_code,
                stdout,
                stderr,
                duration_ms,
            });

            if tx.send(result_msg).await.is_err() {
                error!("Failed to send command result");
            }
        });
    }

    /// Run a command and capture output
    async fn run_command(cmd: &CommandPayload) -> Result<(i32, String, String)> {
        use tokio::process::Command;

        let timeout_secs = cmd.timeout_seconds.unwrap_or(300); // 5 minute default

        let mut command = match &cmd.command_type {
            super::CommandType::Shell => {
                #[cfg(windows)]
                {
                    let mut c = Command::new("cmd");
                    c.args(["/C", &cmd.command]);
                    c
                }
                #[cfg(unix)]
                {
                    let mut c = Command::new("sh");
                    c.args(["-c", &cmd.command]);
                    c
                }
            }
            super::CommandType::PowerShell => {
                let mut c = Command::new("powershell");
                c.args(["-NoProfile", "-NonInteractive", "-Command", &cmd.command]);
                c
            }
            super::CommandType::Python => {
                let mut c = Command::new("python");
                c.args(["-c", &cmd.command]);
                c
            }
            super::CommandType::Script { interpreter } => {
                let mut c = Command::new(interpreter);
                c.args(["-c", &cmd.command]);
                c
            }
        };

        // Capture output
        command.stdout(std::process::Stdio::piped());
        command.stderr(std::process::Stdio::piped());

        // Execute with timeout
        let output = timeout(Duration::from_secs(timeout_secs), command.output())
            .await
            .context("Command timeout")?
            .context("Failed to execute command")?;

        let exit_code = output.status.code().unwrap_or(-1);
        let stdout = String::from_utf8_lossy(&output.stdout).to_string();
        let stderr = String::from_utf8_lossy(&output.stderr).to_string();

        Ok((exit_code, stdout, stderr))
    }
}


Now I have a comprehensive understanding of the codebase. Let me examine one more file to understand the agent's main AppState structure:


     1→//! GuruRMM Agent - Cross-platform Remote Monitoring and Management Agent
     2→//!
     3→//! This agent connects to the GuruRMM server, reports system metrics,
     4→//! monitors services (watchdog), and executes remote commands.
     5→
     6→mod config;
     7→mod device_id;
     8→mod ipc;
     9→mod metrics;
    10→mod service;
    11→mod transport;
    12→mod tray_launcher;
    13→mod updater;
    14→
    15→use anyhow::{Context, Result};
    16→use clap::{Parser, Subcommand};
    17→use std::path::PathBuf;
    18→use std::sync::Arc;
    19→use tokio::sync::RwLock;
    20→use tracing::{error, info, warn};
    21→
    22→use crate::config::AgentConfig;
    23→use crate::metrics::MetricsCollector;
    24→use crate::transport::WebSocketClient;
    25→
    26→/// GuruRMM Agent - Remote Monitoring and Management
    27→#[derive(Parser)]
    28→#[command(name = "gururmm-agent")]
    29→#[command(author, version, about, long_about = None)]
    30→struct Cli {
    31→    /// Path to configuration file
    32→    #[arg(short, long, default_value = "agent.toml")]
    33→    config: PathBuf,
    34→
    35→    /// Subcommand to run
    36→    #[command(subcommand)]
    37→    command: Option<Commands>,
    38→}
    39→
    40→#[derive(Subcommand)]
    41→enum Commands {
    42→    /// Run the agent (default)
    43→    Run,
    44→
    45→    /// Install as a system service
    46→    Install {
    47→        /// Server WebSocket URL (e.g., wss://rmm-api.example.com/ws)
    48→        #[arg(long)]
    49→        server_url: Option<String>,
    50→
    51→        /// API key for authentication
    52→        #[arg(long)]
    53→        api_key: Option<String>,
    54→
    55→        /// Skip legacy service detection and cleanup
    56→        #[arg(long, default_value = "false")]
    57→        skip_legacy_check: bool,
    58→    },
    59→
    60→    /// Uninstall the system service
    61→    Uninstall,
    62→
    63→    /// Start the installed service
    64→    Start,
    65→
    66→    /// Stop the installed service
    67→    Stop,
    68→
    69→    /// Show agent status
    70→    Status,
    71→
    72→    /// Generate a sample configuration file
    73→    GenerateConfig {
    74→        /// Output path for config file
    75→        #[arg(short, long, default_value = "agent.toml")]
    76→        output: PathBuf,
    77→    },
    78→
    79→    /// Run as Windows service (called by SCM, not for manual use)
    80→    #[command(hide = true)]
    81→    Service,
    82→}
    83→
    84→/// Shared application state
    85→pub struct AppState {
    86→    pub config: AgentConfig,
    87→    pub metrics_collector: MetricsCollector,
    88→    pub connected: RwLock<bool>,
    89→    pub last_checkin: RwLock<Option<chrono::DateTime<chrono::Utc>>>,
    90→    pub tray_policy: RwLock<ipc::TrayPolicy>,
    91→}
    92→
    93→#[tokio::main]
    94→async fn main() -> Result<()> {
    95→    // Initialize logging
    96→    tracing_subscriber::fmt()
    97→        .with_env_filter(
    98→            tracing_subscriber::EnvFilter::from_default_env()
    99→                .add_directive("gururmm_agent=info".parse()?)
   100→                .add_directive("info".parse()?),
   101→        )
   102→        .init();
   103→
   104→    let cli = Cli::parse();
   105→
   106→    match cli.command.unwrap_or(Commands::Run) {
   107→        Commands::Run => run_agent(cli.config).await,
   108→        Commands::Install { server_url, api_key, skip_legacy_check } => {
   109→            install_service(server_url, api_key, skip_legacy_check).await
   110→        }
   111→        Commands::Uninstall => uninstall_service().await,
   112→        Commands::Start => start_service().await,
   113→        Commands::Stop => stop_service().await,
   114→        Commands::Status => show_status(cli.config).await,
   115→        Commands::GenerateConfig { output } => generate_config(output).await,
   116→        Commands::Service => run_as_windows_service(),
   117→    }
   118→}
   119→
   120→/// Run as a Windows service (called by SCM)
   121→fn run_as_windows_service() -> Result<()> {
   122→    #[cfg(windows)]
   123→    {
   124→        service::windows::run_as_service()
   125→    }
   126→
   127→    #[cfg(not(windows))]
   128→    {
   129→        anyhow::bail!("Windows service mode is only available on Windows");
   130→    }
   131→}
   132→
   133→/// Main agent runtime loop
   134→async fn run_agent(config_path: PathBuf) -> Result<()> {
   135→    info!("GuruRMM Agent starting...");
   136→
   137→    // Load configuration
   138→    let config = AgentConfig::load(&config_path)?;
   139→    info!("Loaded configuration from {:?}", config_path);
   140→    info!("Server URL: {}", config.server.url);
   141→
   142→    // Initialize metrics collector
   143→    let metrics_collector = MetricsCollector::new();
   144→    info!("Metrics collector initialized");
   145→
   146→    // Create channels for IPC commands
   147→    let (force_checkin_tx, mut force_checkin_rx) = tokio::sync::mpsc::channel::<()>(8);
   148→    let (stop_agent_tx, mut stop_agent_rx) = tokio::sync::mpsc::channel::<()>(1);
   149→
   150→    // Create shared state
   151→    let state = Arc::new(AppState {
   152→        config: config.clone(),
   153→        metrics_collector,
   154→        connected: RwLock::new(false),
   155→        last_checkin: RwLock::new(None),
   156→        tray_policy: RwLock::new(ipc::TrayPolicy::default_permissive()),
   157→    });
   158→
   159→    // Create IPC state for tray communication
   160→    let device_id = device_id::get_device_id();
   161→    let hostname = hostname::get()
   162→        .map(|h| h.to_string_lossy().to_string())
   163→        .unwrap_or_else(|_| "unknown".to_string());
   164→
   165→    let initial_status = ipc::AgentStatus {
   166→        connected: false,
   167→        last_checkin: None,
   168→        server_url: config.server.url.clone(),
   169→        agent_version: env!("CARGO_PKG_VERSION").to_string(),
   170→        device_id,
   171→        hostname,
   172→        error: None,
   173→    };
   174→
   175→    let ipc_state = Arc::new(ipc::IpcState::new(
   176→        initial_status,
   177→        ipc::TrayPolicy::default_permissive(),
   178→        force_checkin_tx,
   179→        stop_agent_tx,
   180→    ));
   181→
   182→    // Start IPC server for tray application
   183→    let ipc_server_state = Arc::clone(&ipc_state);
   184→    let ipc_handle = tokio::spawn(async move {
   185→        if let Err(e) = ipc::server::run_ipc_server(ipc_server_state).await {
   186→            error!("IPC server error: {}", e);
   187→        }
   188→    });
   189→    info!("IPC server started");
   190→
   191→    // Start tray launcher (Windows only for now)
   192→    #[cfg(windows)]
   193→    let tray_launcher = {
   194→        let tray_exe_path = std::path::PathBuf::from(
   195→            format!("{}\\gururmm-tray.exe", service::windows::INSTALL_DIR)
   196→        );
   197→        let launcher = Arc::new(tray_launcher::TrayLauncher::new(tray_exe_path));
   198→
   199→        // Launch tray in existing sessions
   200→        let launcher_init = Arc::clone(&launcher);
   201→        tokio::spawn(async move {
   202→            // Small delay to let IPC server fully start
   203→            tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
   204→            launcher_init.ensure_running().await;
   205→        });
   206→
   207→        // Periodic check for new sessions (every 30 seconds)
   208→        let launcher_periodic = Arc::clone(&launcher);
   209→        tokio::spawn(async move {
   210→            let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(30));
   211→            loop {
   212→                interval.tick().await;
   213→                launcher_periodic.ensure_running().await;
   214→            }
   215→        });
   216→
   217→        info!("Tray launcher initialized");
   218→        Some(launcher)
   219→    };
   220→
   221→    #[cfg(not(windows))]
   222→    let tray_launcher: Option<Arc<tray_launcher::TrayLauncher>> = None;
   223→
   224→    // Start the WebSocket client with auto-reconnect
   225→    let ws_state = Arc::clone(&state);
   226→    let ws_handle = tokio::spawn(async move {
   227→        loop {
   228→            info!("Connecting to server...");
   229→            match WebSocketClient::connect_and_run(Arc::clone(&ws_state)).await {
   230→                Ok(_) => {
   231→                    warn!("WebSocket connection closed normally, reconnecting...");
   232→                }
   233→                Err(e) => {
   234→                    error!("WebSocket error: {}, reconnecting in 10 seconds...", e);
   235→                }
   236→            }
   237→
   238→            // Mark as disconnected
   239→            *ws_state.connected.write().await = false;
   240→
   241→            // Wait before reconnecting
   242→            tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
   243→        }
   244→    });
   245→
   246→    // Start metrics collection loop
   247→    let metrics_state = Arc::clone(&state);
   248→    let metrics_handle = tokio::spawn(async move {
   249→        let interval = metrics_state.config.metrics.interval_seconds;
   250→        let mut interval_timer = tokio::time::interval(tokio::time::Duration::from_secs(interval));
   251→
   252→        loop {
   253→            interval_timer.tick().await;
   254→
   255→            // Collect metrics (they'll be sent via WebSocket if connected)
   256→            let metrics = metrics_state.metrics_collector.collect().await;
   257→            if *metrics_state.connected.read().await {
   258→                info!(
   259→                    "Metrics: CPU={:.1}%, Mem={:.1}%, Disk={:.1}%",
   260→                    metrics.cpu_percent, metrics.memory_percent, metrics.disk_percent
   261→                );
   262→            }
   263→        }
   264→    });
   265→
   266→    // Task to update IPC status periodically
   267→    let ipc_update_state = Arc::clone(&state);
   268→    let ipc_update_ipc = Arc::clone(&ipc_state);
   269→    let ipc_update_handle = tokio::spawn(async move {
   270→        let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(5));
   271→        loop {
   272→            interval.tick().await;
   273→
   274→            let connected = *ipc_update_state.connected.read().await;
   275→            let last_checkin = ipc_update_state.last_checkin.read().await.map(|dt| dt.to_rfc3339());
   276→
   277→            let status = ipc::AgentStatus {
   278→                connected,
   279→                last_checkin,
   280→                server_url: ipc_update_state.config.server.url.clone(),
   281→                agent_version: env!("CARGO_PKG_VERSION").to_string(),
   282→                device_id: device_id::get_device_id(),
   283→                hostname: hostname::get()
   284→                    .map(|h| h.to_string_lossy().to_string())
   285→                    .unwrap_or_else(|_| "unknown".to_string()),
   286→                error: if connected { None } else { Some("Disconnected".to_string()) },
   287→            };
   288→
   289→            ipc_update_ipc.update_status(status).await;
   290→        }
   291→    });
   292→
   293→    // Wait for shutdown signal or stop request
   294→    tokio::select! {
   295→        _ = tokio::signal::ctrl_c() => {
   296→            info!("Received shutdown signal");
   297→        }
   298→        _ = stop_agent_rx.recv() => {
   299→            info!("Received stop request from tray");
   300→        }
   301→        _ = ws_handle => {
   302→            error!("WebSocket task ended unexpectedly");
   303→        }
   304→        _ = metrics_handle => {
   305→            error!("Metrics task ended unexpectedly");
   306→        }
   307→        _ = ipc_handle => {
   308→            error!("IPC server ended unexpectedly");
   309→        }
   310→        _ = ipc_update_handle => {
   311→            error!("IPC update task ended unexpectedly");
   312→        }
   313→    }
   314→
   315→    // Terminate tray processes on shutdown
   316→    #[cfg(windows)]
   317→    if let Some(launcher) = tray_launcher {
   318→        info!("Terminating tray processes...");
   319→        launcher.terminate_all().await;
   320→    }
   321→
   322→    #[cfg(not(windows))]
   323→    let _ = tray_launcher; // Suppress unused warning
   324→
   325→    info!("GuruRMM Agent shutting down");
   326→    Ok(())
   327→}
   328→
   329→/// Install the agent as a system service
   330→async fn install_service(
   331→    server_url: Option<String>,
   332→    api_key: Option<String>,
   333→    skip_legacy_check: bool,
   334→) -> Result<()> {
   335→    #[cfg(windows)]
   336→    {
   337→        service::windows::install(server_url, api_key, skip_legacy_check)
   338→    }
   339→
   340→    #[cfg(target_os = "linux")]
   341→    {
   342→        install_systemd_service(server_url, api_key, skip_legacy_check).await
   343→    }
   344→
   345→    #[cfg(target_os = "macos")]
   346→    {
   347→        let _ = (server_url, api_key, skip_legacy_check); // Suppress unused warnings
   348→        info!("Installing GuruRMM Agent as launchd service...");
   349→        todo!("macOS launchd service installation not yet implemented");
   350→    }
   351→}
   352→
   353→/// Legacy service names to check for and clean up (Linux)
   354→#[cfg(target_os = "linux")]
   355→const LINUX_LEGACY_SERVICE_NAMES: &[&str] = &[
   356→    "gururmm",           // Old name without -agent suffix
   357→    "guru-rmm-agent",    // Alternative naming
   358→    "GuruRMM-Agent",     // Case variant
   359→];
   360→
   361→/// Clean up legacy Linux service installations
   362→#[cfg(target_os = "linux")]
   363→fn cleanup_legacy_linux_services() -> Result<()> {
   364→    use std::process::Command;
   365→
   366→    info!("Checking for legacy service installations...");
   367→
   368→    for legacy_name in LINUX_LEGACY_SERVICE_NAMES {
   369→        // Check if service exists
   370→        let status = Command::new("systemctl")
   371→            .args(["status", legacy_name])
   372→            .output();
   373→
   374→        if let Ok(output) = status {
   375→            if output.status.success() || String::from_utf8_lossy(&output.stderr).contains("Loaded:") {
   376→                info!("Found legacy service '{}', removing...", legacy_name);
   377→
   378→                // Stop the service
   379→                let _ = Command::new("systemctl")
   380→                    .args(["stop", legacy_name])
   381→                    .status();
   382→
   383→                // Disable the service
   384→                let _ = Command::new("systemctl")
   385→                    .args(["disable", legacy_name])
   386→                    .status();
   387→
   388→                // Remove unit file
   389→                let unit_file = format!("/etc/systemd/system/{}.service", legacy_name);
   390→                if std::path::Path::new(&unit_file).exists() {
   391→                    info!("Removing legacy unit file: {}", unit_file);
   392→                    let _ = std::fs::remove_file(&unit_file);
   393→                }
   394→            }
   395→        }
   396→    }
   397→
   398→    // Check for legacy binaries in common locations
   399→    let legacy_binary_locations = [
   400→        "/usr/local/bin/gururmm",
   401→        "/usr/bin/gururmm",
   402→        "/opt/gururmm/gururmm",
   403→        "/opt/gururmm/agent",
   404→    ];
   405→
   406→    for legacy_path in legacy_binary_locations {
   407→        if std::path::Path::new(legacy_path).exists() {
   408→            info!("Found legacy binary at '{}', removing...", legacy_path);
   409→            let _ = std::fs::remove_file(legacy_path);
   410→        }
   411→    }
   412→
   413→    // Reload systemd to pick up removed unit files
   414→    let _ = Command::new("systemctl")
   415→        .args(["daemon-reload"])
   416→        .status();
   417→
   418→    Ok(())
   419→}
   420→
   421→/// Install as a systemd service (Linux)
   422→#[cfg(target_os = "linux")]
   423→async fn install_systemd_service(
   424→    server_url: Option<String>,
   425→    api_key: Option<String>,
   426→    skip_legacy_check: bool,
   427→) -> Result<()> {
   428→    use std::process::Command;
   429→
   430→    const SERVICE_NAME: &str = "gururmm-agent";
   431→    const INSTALL_DIR: &str = "/usr/local/bin";
   432→    const CONFIG_DIR: &str = "/etc/gururmm";
   433→    const SYSTEMD_DIR: &str = "/etc/systemd/system";
   434→
   435→    info!("Installing GuruRMM Agent as systemd service...");
   436→
   437→    // Check if running as root
   438→    if !nix::unistd::geteuid().is_root() {
   439→        anyhow::bail!("Installation requires root privileges. Please run with sudo.");
   440→    }
   441→
   442→    // Clean up legacy installations unless skipped
   443→    if !skip_legacy_check {
   444→        if let Err(e) = cleanup_legacy_linux_services() {
   445→            warn!("Legacy cleanup warning: {}", e);
   446→        }
   447→    }
   448→
   449→    // Get the current executable path
   450→    let current_exe = std::env::current_exe()
   451→        .context("Failed to get current executable path")?;
   452→
   453→    let binary_dest = format!("{}/{}", INSTALL_DIR, SERVICE_NAME);
   454→    let config_dest = format!("{}/agent.toml", CONFIG_DIR);
   455→    let unit_file = format!("{}/{}.service", SYSTEMD_DIR, SERVICE_NAME);
   456→
   457→    // Create config directory
   458→    info!("Creating config directory: {}", CONFIG_DIR);
   459→    std::fs::create_dir_all(CONFIG_DIR)
   460→        .context("Failed to create config directory")?;
   461→
   462→    // Copy binary
   463→    info!("Copying binary to: {}", binary_dest);
   464→    std::fs::copy(&current_exe, &binary_dest)
   465→        .context("Failed to copy binary")?;
   466→
   467→    // Make binary executable
   468→    Command::new("chmod")
   469→        .args(["+x", &binary_dest])
   470→        .status()
   471→        .context("Failed to set binary permissions")?;
   472→
   473→    // Handle configuration
   474→    let config_needs_manual_edit;
   475→    if !std::path::Path::new(&config_dest).exists() {
   476→        info!("Creating config: {}", config_dest);
   477→
   478→        // Start with sample config
   479→        let mut config = crate::config::AgentConfig::sample();
   480→
   481→        // Apply provided values
   482→        if let Some(url) = &server_url {
   483→            config.server.url = url.clone();
   484→        }
   485→        if let Some(key) = &api_key {
   486→            config.server.api_key = key.clone();
   487→        }
   488→
   489→        let toml_str = toml::to_string_pretty(&config)?;
   490→        std::fs::write(&config_dest, toml_str)
   491→            .context("Failed to write config file")?;
   492→
   493→        // Set restrictive permissions on config (contains API key)
   494→        Command::new("chmod")
   495→            .args(["600", &config_dest])
   496→            .status()
   497→            .context("Failed to set config permissions")?;
   498→
   499→        config_needs_manual_edit = server_url.is_none() || api_key.is_none();
   500→    } else {
   501→        info!("Config already exists: {}", config_dest);
   502→        config_needs_manual_edit = false;
   503→
   504→        // If server_url or api_key provided, update existing config
   505→        if server_url.is_some() || api_key.is_some() {
   506→            info!("Updating existing configuration...");
   507→            let config_content = std::fs::read_to_string(&config_dest)?;
   508→            let mut config: crate::config::AgentConfig = toml::from_str(&config_content)
   509→                .context("Failed to parse existing config")?;
   510→
   511→            if let Some(url) = &server_url {
   512→                config.server.url = url.clone();
   513→            }
   514→            if let Some(key) = &api_key {
   515→                config.server.api_key = key.clone();
   516→            }
   517→
   518→            let toml_str = toml::to_string_pretty(&config)?;
   519→            std::fs::write(&config_dest, toml_str)
   520→                .context("Failed to update config file")?;
   521→        }
   522→    }
   523→
   524→    // Create systemd unit file
   525→    let unit_content = format!(r#"[Unit]
   526→Description=GuruRMM Agent - Remote Monitoring and Management
   527→Documentation=https://github.com/azcomputerguru/gururmm
   528→After=network-online.target
   529→Wants=network-online.target
   530→
   531→[Service]
   532→Type=simple
   533→ExecStart={binary} --config {config} run
   534→Restart=always
   535→RestartSec=10
   536→StandardOutput=journal
   537→StandardError=journal
   538→SyslogIdentifier={service}
   539→
   540→# Security hardening
   541→NoNewPrivileges=true
   542→ProtectSystem=strict
   543→ProtectHome=read-only
   544→PrivateTmp=true
   545→ReadWritePaths=/var/log
   546→
   547→[Install]
   548→WantedBy=multi-user.target
   549→"#,
   550→        binary = binary_dest,
   551→        config = config_dest,
   552→        service = SERVICE_NAME
   553→    );
   554→
   555→    info!("Creating systemd unit file: {}", unit_file);
   556→    std::fs::write(&unit_file, unit_content)
   557→        .context("Failed to write systemd unit file")?;
   558→
   559→    // Reload systemd daemon
   560→    info!("Reloading systemd daemon...");
   561→    let status = Command::new("systemctl")
   562→        .args(["daemon-reload"])
   563→        .status()
   564→        .context("Failed to reload systemd")?;
   565→
   566→    if !status.success() {
   567→        anyhow::bail!("systemctl daemon-reload failed");
   568→    }
   569→
   570→    // Enable the service
   571→    info!("Enabling service...");
   572→    let status = Command::new("systemctl")
   573→        .args(["enable", SERVICE_NAME])
   574→        .status()
   575→        .context("Failed to enable service")?;
   576→
   577→    if !status.success() {
   578→        anyhow::bail!("systemctl enable failed");
   579→    }
   580→
   581→    println!("\n✓ GuruRMM Agent installed successfully!");
   582→    println!("\nInstalled files:");
   583→    println!("  Binary: {}", binary_dest);
   584→    println!("  Config: {}", config_dest);
   585→    println!("  Service: {}", unit_file);
   586→
   587→    if config_needs_manual_edit {
   588→        println!("\n⚠  IMPORTANT: Edit {} with your server URL and API key!", config_dest);
   589→        println!("\nNext steps:");
   590→        println!("  1. Edit {} with your server URL and API key", config_dest);
   591→        println!("  2. Start the service: sudo systemctl start {}", SERVICE_NAME);
   592→    } else {
   593→        println!("\nStarting service...");
   594→        let status = Command::new("systemctl")
   595→            .args(["start", SERVICE_NAME])
   596→            .status();
   597→
   598→        if status.is_ok() && status.unwrap().success() {
   599→            println!("✓ Service started successfully!");
   600→        } else {
   601→            println!("⚠  Failed to start service. Check logs: sudo journalctl -u {} -f", SERVICE_NAME);
   602→        }
   603→    }
   604→
   605→    println!("\nUseful commands:");
   606→    println!("  Status: sudo systemctl status {}", SERVICE_NAME);
   607→    println!("  Logs:   sudo journalctl -u {} -f", SERVICE_NAME);
   608→    println!("  Stop:   sudo systemctl stop {}", SERVICE_NAME);
   609→    println!("  Start:  sudo systemctl start {}", SERVICE_NAME);
   610→
   611→    Ok(())
   612→}
   613→
   614→/// Uninstall the system service
   615→async fn uninstall_service() -> Result<()> {
   616→    #[cfg(windows)]
   617→    {
   618→        service::windows::uninstall()
   619→    }
   620→
   621→    #[cfg(target_os = "linux")]
   622→    {
   623→        uninstall_systemd_service().await
   624→    }
   625→
   626→    #[cfg(target_os = "macos")]
   627→    {
   628→        todo!("macOS service uninstallation not yet implemented");
   629→    }
   630→}
   631→
   632→/// Uninstall systemd service (Linux)
   633→#[cfg(target_os = "linux")]
   634→async fn uninstall_systemd_service() -> Result<()> {
   635→    use std::process::Command;
   636→
   637→    const SERVICE_NAME: &str = "gururmm-agent";
   638→    const INSTALL_DIR: &str = "/usr/local/bin";
   639→    const CONFIG_DIR: &str = "/etc/gururmm";
   640→    const SYSTEMD_DIR: &str = "/etc/systemd/system";
   641→
   642→    info!("Uninstalling GuruRMM Agent...");
   643→
   644→    if !nix::unistd::geteuid().is_root() {
   645→        anyhow::bail!("Uninstallation requires root privileges. Please run with sudo.");
   646→    }
   647→
   648→    let binary_path = format!("{}/{}", INSTALL_DIR, SERVICE_NAME);
   649→    let unit_file = format!("{}/{}.service", SYSTEMD_DIR, SERVICE_NAME);
   650→
   651→    // Stop the service if running
   652→    info!("Stopping service...");
   653→    let _ = Command::new("systemctl")
   654→        .args(["stop", SERVICE_NAME])
   655→        .status();
   656→
   657→    // Disable the service
   658→    info!("Disabling service...");
   659→    let _ = Command::new("systemctl")
   660→        .args(["disable", SERVICE_NAME])
   661→        .status();
   662→
   663→    // Remove unit file
   664→    if std::path::Path::new(&unit_file).exists() {
   665→        info!("Removing unit file: {}", unit_file);
   666→        std::fs::remove_file(&unit_file)?;
   667→    }
   668→
   669→    // Remove binary
   670→    if std::path::Path::new(&binary_path).exists() {
   671→        info!("Removing binary: {}", binary_path);
   672→        std::fs::remove_file(&binary_path)?;
   673→    }
   674→
   675→    // Reload systemd
   676→    let _ = Command::new("systemctl")
   677→        .args(["daemon-reload"])
   678→        .status();
   679→
   680→    println!("\n✓ GuruRMM Agent uninstalled successfully!");
   681→    println!("\nNote: Config directory {} was preserved.", CONFIG_DIR);
   682→    println!("Remove it manually if no longer needed: sudo rm -rf {}", CONFIG_DIR);
   683→
   684→    Ok(())
   685→}
   686→
   687→/// Start the installed service
   688→async fn start_service() -> Result<()> {
   689→    #[cfg(windows)]
   690→    {
   691→        service::windows::start()
   692→    }
   693→
   694→    #[cfg(target_os = "linux")]
   695→    {
   696→        use std::process::Command;
   697→
   698→        info!("Starting GuruRMM Agent service...");
   699→
   700→        let status = Command::new("systemctl")
   701→            .args(["start", "gururmm-agent"])
   702→            .status()
   703→            .context("Failed to start service")?;
   704→
   705→        if status.success() {
   706→            println!("** Service started successfully");
   707→            println!("Check status: sudo systemctl status gururmm-agent");
   708→        } else {
   709→            anyhow::bail!("Failed to start service. Check: sudo journalctl -u gururmm-agent -n 50");
   710→        }
   711→
   712→        Ok(())
   713→    }
   714→
   715→    #[cfg(target_os = "macos")]
   716→    {
   717→        todo!("macOS service start not yet implemented");
   718→    }
   719→}
   720→
   721→/// Stop the installed service
   722→async fn stop_service() -> Result<()> {
   723→    #[cfg(windows)]
   724→    {
   725→        service::windows::stop()
   726→    }
   727→
   728→    #[cfg(target_os = "linux")]
   729→    {
   730→        use std::process::Command;
   731→
   732→        info!("Stopping GuruRMM Agent service...");
   733→
   734→        let status = Command::new("systemctl")
   735→            .args(["stop", "gururmm-agent"])
   736→            .status()
   737→            .context("Failed to stop service")?;
   738→
   739→        if status.success() {
   740→            println!("** Service stopped successfully");
   741→        } else {
   742→            anyhow::bail!("Failed to stop service");
   743→        }
   744→
   745→        Ok(())
   746→    }
   747→
   748→    #[cfg(target_os = "macos")]
   749→    {
   750→        todo!("macOS service stop not yet implemented");
   751→    }
   752→}
   753→
   754→/// Show agent status
   755→async fn show_status(config_path: PathBuf) -> Result<()> {
   756→    // On Windows, show service status
   757→    #[cfg(windows)]
   758→    {
   759→        service::windows::status()?;
   760→        println!();
   761→    }
   762→
   763→    // Try to load config for additional info
   764→    match AgentConfig::load(&config_path) {
   765→        Ok(config) => {
   766→            println!("Configuration");
   767→            println!("=============");
   768→            println!("Config file: {:?}", config_path);
   769→            println!("Server URL: {}", config.server.url);
   770→            println!("Metrics interval: {} seconds", config.metrics.interval_seconds);
   771→            println!("Watchdog enabled: {}", config.watchdog.enabled);
   772→
   773→            // Collect current metrics
   774→            let collector = MetricsCollector::new();
   775→            let metrics = collector.collect().await;
   776→
   777→            println!("\nCurrent System Metrics:");
   778→            println!("  CPU Usage: {:.1}%", metrics.cpu_percent);
   779→            println!("  Memory Usage: {:.1}%", metrics.memory_percent);
   780→            println!(
   781→                "  Memory Used: {:.2} GB",
   782→                metrics.memory_used_bytes as f64 / 1_073_741_824.0
   783→            );
   784→            println!("  Disk Usage: {:.1}%", metrics.disk_percent);
   785→            println!(
   786→                "  Disk Used: {:.2} GB",
   787→                metrics.disk_used_bytes as f64 / 1_073_741_824.0
   788→            );
   789→        }
   790→        Err(_) => {
   791→            println!("\nConfig file {:?} not found or invalid.", config_path);
   792→            #[cfg(windows)]
   793→            println!("Service config location: {}\\agent.toml", service::windows::CONFIG_DIR);
   794→        }
   795→    }
   796→
   797→    Ok(())
   798→}
   799→
   800→/// Generate a sample configuration file
   801→async fn generate_config(output: PathBuf) -> Result<()> {
   802→    let sample_config = AgentConfig::sample();
   803→    let toml_str = toml::to_string_pretty(&sample_config)?;
   804→
   805→    std::fs::write(&output, toml_str)?;
   806→    println!("Sample configuration written to {:?}", output);
   807→    println!("\nEdit this file with your server URL and API key, then run:");
   808→    println!("  gururmm-agent --config {:?} run", output);
   809→
   810→    Ok(())
   811→}
   812→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! GuruRMM Agent - Cross-platform Remote Monitoring and Management Agent
//!
//! This agent connects to the GuruRMM server, reports system metrics,
//! monitors services (watchdog), and executes remote commands.

mod config;
mod device_id;
mod ipc;
mod metrics;
mod service;
mod transport;
mod tray_launcher;
mod updater;

use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{error, info, warn};

use crate::config::AgentConfig;
use crate::metrics::MetricsCollector;
use crate::transport::WebSocketClient;

/// GuruRMM Agent - Remote Monitoring and Management
#[derive(Parser)]
#[command(name = "gururmm-agent")]
#[command(author, version, about, long_about = None)]
struct Cli {
    /// Path to configuration file
    #[arg(short, long, default_value = "agent.toml")]
    config: PathBuf,

    /// Subcommand to run
    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Subcommand)]
enum Commands {
    /// Run the agent (default)
    Run,

    /// Install as a system service
    Install {
        /// Server WebSocket URL (e.g., wss://rmm-api.example.com/ws)
        #[arg(long)]
        server_url: Option<String>,

        /// API key for authentication
        #[arg(long)]
        api_key: Option<String>,

        /// Skip legacy service detection and cleanup
        #[arg(long, default_value = "false")]
        skip_legacy_check: bool,
    },

    /// Uninstall the system service
    Uninstall,

    /// Start the installed service
    Start,

    /// Stop the installed service
    Stop,

    /// Show agent status
    Status,

    /// Generate a sample configuration file
    GenerateConfig {
        /// Output path for config file
        #[arg(short, long, default_value = "agent.toml")]
        output: PathBuf,
    },

    /// Run as Windows service (called by SCM, not for manual use)
    #[command(hide = true)]
    Service,
}

/// Shared application state
pub struct AppState {
    pub config: AgentConfig,
    pub metrics_collector: MetricsCollector,
    pub connected: RwLock<bool>,
    pub last_checkin: RwLock<Option<chrono::DateTime<chrono::Utc>>>,
    pub tray_policy: RwLock<ipc::TrayPolicy>,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive("gururmm_agent=info".parse()?)
                .add_directive("info".parse()?),
        )
        .init();

    let cli = Cli::parse();

    match cli.command.unwrap_or(Commands::Run) {
        Commands::Run => run_agent(cli.config).await,
        Commands::Install { server_url, api_key, skip_legacy_check } => {
            install_service(server_url, api_key, skip_legacy_check).await
        }
        Commands::Uninstall => uninstall_service().await,
        Commands::Start => start_service().await,
        Commands::Stop => stop_service().await,
        Commands::Status => show_status(cli.config).await,
        Commands::GenerateConfig { output } => generate_config(output).await,
        Commands::Service => run_as_windows_service(),
    }
}

/// Run as a Windows service (called by SCM)
fn run_as_windows_service() -> Result<()> {
    #[cfg(windows)]
    {
        service::windows::run_as_service()
    }

    #[cfg(not(windows))]
    {
        anyhow::bail!("Windows service mode is only available on Windows");
    }
}

/// Main agent runtime loop
async fn run_agent(config_path: PathBuf) -> Result<()> {
    info!("GuruRMM Agent starting...");

    // Load configuration
    let config = AgentConfig::load(&config_path)?;
    info!("Loaded configuration from {:?}", config_path);
    info!("Server URL: {}", config.server.url);

    // Initialize metrics collector
    let metrics_collector = MetricsCollector::new();
    info!("Metrics collector initialized");

    // Create channels for IPC commands
    let (force_checkin_tx, mut force_checkin_rx) = tokio::sync::mpsc::channel::<()>(8);
    let (stop_agent_tx, mut stop_agent_rx) = tokio::sync::mpsc::channel::<()>(1);

    // Create shared state
    let state = Arc::new(AppState {
        config: config.clone(),
        metrics_collector,
        connected: RwLock::new(false),
        last_checkin: RwLock::new(None),
        tray_policy: RwLock::new(ipc::TrayPolicy::default_permissive()),
    });

    // Create IPC state for tray communication
    let device_id = device_id::get_device_id();
    let hostname = hostname::get()
        .map(|h| h.to_string_lossy().to_string())
        .unwrap_or_else(|_| "unknown".to_string());

    let initial_status = ipc::AgentStatus {
        connected: false,
        last_checkin: None,
        server_url: config.server.url.clone(),
        agent_version: env!("CARGO_PKG_VERSION").to_string(),
        device_id,
        hostname,
        error: None,
    };

    let ipc_state = Arc::new(ipc::IpcState::new(
        initial_status,
        ipc::TrayPolicy::default_permissive(),
        force_checkin_tx,
        stop_agent_tx,
    ));

    // Start IPC server for tray application
    let ipc_server_state = Arc::clone(&ipc_state);
    let ipc_handle = tokio::spawn(async move {
        if let Err(e) = ipc::server::run_ipc_server(ipc_server_state).await {
            error!("IPC server error: {}", e);
        }
    });
    info!("IPC server started");

    // Start tray launcher (Windows only for now)
    #[cfg(windows)]
    let tray_launcher = {
        let tray_exe_path = std::path::PathBuf::from(
            format!("{}\\gururmm-tray.exe", service::windows::INSTALL_DIR)
        );
        let launcher = Arc::new(tray_launcher::TrayLauncher::new(tray_exe_path));

        // Launch tray in existing sessions
        let launcher_init = Arc::clone(&launcher);
        tokio::spawn(async move {
            // Small delay to let IPC server fully start
            tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
            launcher_init.ensure_running().await;
        });

        // Periodic check for new sessions (every 30 seconds)
        let launcher_periodic = Arc::clone(&launcher);
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(30));
            loop {
                interval.tick().await;
                launcher_periodic.ensure_running().await;
            }
        });

        info!("Tray launcher initialized");
        Some(launcher)
    };

    #[cfg(not(windows))]
    let tray_launcher: Option<Arc<tray_launcher::TrayLauncher>> = None;

    // Start the WebSocket client with auto-reconnect
    let ws_state = Arc::clone(&state);
    let ws_handle = tokio::spawn(async move {
        loop {
            info!("Connecting to server...");
            match WebSocketClient::connect_and_run(Arc::clone(&ws_state)).await {
                Ok(_) => {
                    warn!("WebSocket connection closed normally, reconnecting...");
                }
                Err(e) => {
                    error!("WebSocket error: {}, reconnecting in 10 seconds...", e);
                }
            }

            // Mark as disconnected
            *ws_state.connected.write().await = false;

            // Wait before reconnecting
            tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
        }
    });

    // Start metrics collection loop
    let metrics_state = Arc::clone(&state);
    let metrics_handle = tokio::spawn(async move {
        let interval = metrics_state.config.metrics.interval_seconds;
        let mut interval_timer = tokio::time::interval(tokio::time::Duration::from_secs(interval));

        loop {
            interval_timer.tick().await;

            // Collect metrics (they'll be sent via WebSocket if connected)
            let metrics = metrics_state.metrics_collector.collect().await;
            if *metrics_state.connected.read().await {
                info!(
                    "Metrics: CPU={:.1}%, Mem={:.1}%, Disk={:.1}%",
                    metrics.cpu_percent, metrics.memory_percent, metrics.disk_percent
                );
            }
        }
    });

    // Task to update IPC status periodically
    let ipc_update_state = Arc::clone(&state);
    let ipc_update_ipc = Arc::clone(&ipc_state);
    let ipc_update_handle = tokio::spawn(async move {
        let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(5));
        loop {
            interval.tick().await;

            let connected = *ipc_update_state.connected.read().await;
            let last_checkin = ipc_update_state.last_checkin.read().await.map(|dt| dt.to_rfc3339());

            let status = ipc::AgentStatus {
                connected,
                last_checkin,
                server_url: ipc_update_state.config.server.url.clone(),
                agent_version: env!("CARGO_PKG_VERSION").to_string(),
                device_id: device_id::get_device_id(),
                hostname: hostname::get()
                    .map(|h| h.to_string_lossy().to_string())
                    .unwrap_or_else(|_| "unknown".to_string()),
                error: if connected { None } else { Some("Disconnected".to_string()) },
            };

            ipc_update_ipc.update_status(status).await;
        }
    });

    // Wait for shutdown signal or stop request
    tokio::select! {
        _ = tokio::signal::ctrl_c() => {
            info!("Received shutdown signal");
        }
        _ = stop_agent_rx.recv() => {
            info!("Received stop request from tray");
        }
        _ = ws_handle => {
            error!("WebSocket task ended unexpectedly");
        }
        _ = metrics_handle => {
            error!("Metrics task ended unexpectedly");
        }
        _ = ipc_handle => {
            error!("IPC server ended unexpectedly");
        }
        _ = ipc_update_handle => {
            error!("IPC update task ended unexpectedly");
        }
    }

    // Terminate tray processes on shutdown
    #[cfg(windows)]
    if let Some(launcher) = tray_launcher {
        info!("Terminating tray processes...");
        launcher.terminate_all().await;
    }

    #[cfg(not(windows))]
    let _ = tray_launcher; // Suppress unused warning

    info!("GuruRMM Agent shutting down");
    Ok(())
}

/// Install the agent as a system service
async fn install_service(
    server_url: Option<String>,
    api_key: Option<String>,
    skip_legacy_check: bool,
) -> Result<()> {
    #[cfg(windows)]
    {
        service::windows::install(server_url, api_key, skip_legacy_check)
    }

    #[cfg(target_os = "linux")]
    {
        install_systemd_service(server_url, api_key, skip_legacy_check).await
    }

    #[cfg(target_os = "macos")]
    {
        let _ = (server_url, api_key, skip_legacy_check); // Suppress unused warnings
        info!("Installing GuruRMM Agent as launchd service...");
        todo!("macOS launchd service installation not yet implemented");
    }
}

/// Legacy service names to check for and clean up (Linux)
#[cfg(target_os = "linux")]
const LINUX_LEGACY_SERVICE_NAMES: &[&str] = &[
    "gururmm",           // Old name without -agent suffix
    "guru-rmm-agent",    // Alternative naming
    "GuruRMM-Agent",     // Case variant
];

/// Clean up legacy Linux service installations
#[cfg(target_os = "linux")]
fn cleanup_legacy_linux_services() -> Result<()> {
    use std::process::Command;

    info!("Checking for legacy service installations...");

    for legacy_name in LINUX_LEGACY_SERVICE_NAMES {
        // Check if service exists
        let status = Command::new("systemctl")
            .args(["status", legacy_name])
            .output();

        if let Ok(output) = status {
            if output.status.success() || String::from_utf8_lossy(&output.stderr).contains("Loaded:") {
                info!("Found legacy service '{}', removing...", legacy_name);

                // Stop the service
                let _ = Command::new("systemctl")
                    .args(["stop", legacy_name])
                    .status();

                // Disable the service
                let _ = Command::new("systemctl")
                    .args(["disable", legacy_name])
                    .status();

                // Remove unit file
                let unit_file = format!("/etc/systemd/system/{}.service", legacy_name);
                if std::path::Path::new(&unit_file).exists() {
                    info!("Removing legacy unit file: {}", unit_file);
                    let _ = std::fs::remove_file(&unit_file);
                }
            }
        }
    }

    // Check for legacy binaries in common locations
    let legacy_binary_locations = [
        "/usr/local/bin/gururmm",
        "/usr/bin/gururmm",
        "/opt/gururmm/gururmm",
        "/opt/gururmm/agent",
    ];

    for legacy_path in legacy_binary_locations {
        if std::path::Path::new(legacy_path).exists() {
            info!("Found legacy binary at '{}', removing...", legacy_path);
            let _ = std::fs::remove_file(legacy_path);
        }
    }

    // Reload systemd to pick up removed unit files
    let _ = Command::new("systemctl")
        .args(["daemon-reload"])
        .status();

    Ok(())
}

/// Install as a systemd service (Linux)
#[cfg(target_os = "linux")]
async fn install_systemd_service(
    server_url: Option<String>,
    api_key: Option<String>,
    skip_legacy_check: bool,
) -> Result<()> {
    use std::process::Command;

    const SERVICE_NAME: &str = "gururmm-agent";
    const INSTALL_DIR: &str = "/usr/local/bin";
    const CONFIG_DIR: &str = "/etc/gururmm";
    const SYSTEMD_DIR: &str = "/etc/systemd/system";

    info!("Installing GuruRMM Agent as systemd service...");

    // Check if running as root
    if !nix::unistd::geteuid().is_root() {
        anyhow::bail!("Installation requires root privileges. Please run with sudo.");
    }

    // Clean up legacy installations unless skipped
    if !skip_legacy_check {
        if let Err(e) = cleanup_legacy_linux_services() {
            warn!("Legacy cleanup warning: {}", e);
        }
    }

    // Get the current executable path
    let current_exe = std::env::current_exe()
        .context("Failed to get current executable path")?;

    let binary_dest = format!("{}/{}", INSTALL_DIR, SERVICE_NAME);
    let config_dest = format!("{}/agent.toml", CONFIG_DIR);
    let unit_file = format!("{}/{}.service", SYSTEMD_DIR, SERVICE_NAME);

    // Create config directory
    info!("Creating config directory: {}", CONFIG_DIR);
    std::fs::create_dir_all(CONFIG_DIR)
        .context("Failed to create config directory")?;

    // Copy binary
    info!("Copying binary to: {}", binary_dest);
    std::fs::copy(&current_exe, &binary_dest)
        .context("Failed to copy binary")?;

    // Make binary executable
    Command::new("chmod")
        .args(["+x", &binary_dest])
        .status()
        .context("Failed to set binary permissions")?;

    // Handle configuration
    let config_needs_manual_edit;
    if !std::path::Path::new(&config_dest).exists() {
        info!("Creating config: {}", config_dest);

        // Start with sample config
        let mut config = crate::config::AgentConfig::sample();

        // Apply provided values
        if let Some(url) = &server_url {
            config.server.url = url.clone();
        }
        if let Some(key) = &api_key {
            config.server.api_key = key.clone();
        }

        let toml_str = toml::to_string_pretty(&config)?;
        std::fs::write(&config_dest, toml_str)
            .context("Failed to write config file")?;

        // Set restrictive permissions on config (contains API key)
        Command::new("chmod")
            .args(["600", &config_dest])
            .status()
            .context("Failed to set config permissions")?;

        config_needs_manual_edit = server_url.is_none() || api_key.is_none();
    } else {
        info!("Config already exists: {}", config_dest);
        config_needs_manual_edit = false;

        // If server_url or api_key provided, update existing config
        if server_url.is_some() || api_key.is_some() {
            info!("Updating existing configuration...");
            let config_content = std::fs::read_to_string(&config_dest)?;
            let mut config: crate::config::AgentConfig = toml::from_str(&config_content)
                .context("Failed to parse existing config")?;

            if let Some(url) = &server_url {
                config.server.url = url.clone();
            }
            if let Some(key) = &api_key {
                config.server.api_key = key.clone();
            }

            let toml_str = toml::to_string_pretty(&config)?;
            std::fs::write(&config_dest, toml_str)
                .context("Failed to update config file")?;
        }
    }

    // Create systemd unit file
    let unit_content = format!(r#"[Unit]
Description=GuruRMM Agent - Remote Monitoring and Management
Documentation=https://github.com/azcomputerguru/gururmm
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
ExecStart={binary} --config {config} run
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal
SyslogIdentifier={service}

# Security hardening
NoNewPrivileges=true
ProtectSystem=strict
ProtectHome=read-only
PrivateTmp=true
ReadWritePaths=/var/log

[Install]
WantedBy=multi-user.target
"#,
        binary = binary_dest,
        config = config_dest,
        service = SERVICE_NAME
    );

    info!("Creating systemd unit file: {}", unit_file);
    std::fs::write(&unit_file, unit_content)
        .context("Failed to write systemd unit file")?;

    // Reload systemd daemon
    info!("Reloading systemd daemon...");
    let status = Command::new("systemctl")
        .args(["daemon-reload"])
        .status()
        .context("Failed to reload systemd")?;

    if !status.success() {
        anyhow::bail!("systemctl daemon-reload failed");
    }

    // Enable the service
    info!("Enabling service...");
    let status = Command::new("systemctl")
        .args(["enable", SERVICE_NAME])
        .status()
        .context("Failed to enable service")?;

    if !status.success() {
        anyhow::bail!("systemctl enable failed");
    }

    println!("\n✓ GuruRMM Agent installed successfully!");
    println!("\nInstalled files:");
    println!("  Binary: {}", binary_dest);
    println!("  Config: {}", config_dest);
    println!("  Service: {}", unit_file);

    if config_needs_manual_edit {
        println!("\n⚠  IMPORTANT: Edit {} with your server URL and API key!", config_dest);
        println!("\nNext steps:");
        println!("  1. Edit {} with your server URL and API key", config_dest);
        println!("  2. Start the service: sudo systemctl start {}", SERVICE_NAME);
    } else {
        println!("\nStarting service...");
        let status = Command::new("systemctl")
            .args(["start", SERVICE_NAME])
            .status();

        if status.is_ok() && status.unwrap().success() {
            println!("✓ Service started successfully!");
        } else {
            println!("⚠  Failed to start service. Check logs: sudo journalctl -u {} -f", SERVICE_NAME);
        }
    }

    println!("\nUseful commands:");
    println!("  Status: sudo systemctl status {}", SERVICE_NAME);
    println!("  Logs:   sudo journalctl -u {} -f", SERVICE_NAME);
    println!("  Stop:   sudo systemctl stop {}", SERVICE_NAME);
    println!("  Start:  sudo systemctl start {}", SERVICE_NAME);

    Ok(())
}

/// Uninstall the system service
async fn uninstall_service() -> Result<()> {
    #[cfg(windows)]
    {
        service::windows::uninstall()
    }

    #[cfg(target_os = "linux")]
    {
        uninstall_systemd_service().await
    }

    #[cfg(target_os = "macos")]
    {
        todo!("macOS service uninstallation not yet implemented");
    }
}

/// Uninstall systemd service (Linux)
#[cfg(target_os = "linux")]
async fn uninstall_systemd_service() -> Result<()> {
    use std::process::Command;

    const SERVICE_NAME: &str = "gururmm-agent";
    const INSTALL_DIR: &str = "/usr/local/bin";
    const CONFIG_DIR: &str = "/etc/gururmm";
    const SYSTEMD_DIR: &str = "/etc/systemd/system";

    info!("Uninstalling GuruRMM Agent...");

    if !nix::unistd::geteuid().is_root() {
        anyhow::bail!("Uninstallation requires root privileges. Please run with sudo.");
    }

    let binary_path = format!("{}/{}", INSTALL_DIR, SERVICE_NAME);
    let unit_file = format!("{}/{}.service", SYSTEMD_DIR, SERVICE_NAME);

    // Stop the service if running
    info!("Stopping service...");
    let _ = Command::new("systemctl")
        .args(["stop", SERVICE_NAME])
        .status();

    // Disable the service
    info!("Disabling service...");
    let _ = Command::new("systemctl")
        .args(["disable", SERVICE_NAME])
        .status();

    // Remove unit file
    if std::path::Path::new(&unit_file).exists() {
        info!("Removing unit file: {}", unit_file);
        std::fs::remove_file(&unit_file)?;
    }

    // Remove binary
    if std::path::Path::new(&binary_path).exists() {
        info!("Removing binary: {}", binary_path);
        std::fs::remove_file(&binary_path)?;
    }

    // Reload systemd
    let _ = Command::new("systemctl")
        .args(["daemon-reload"])
        .status();

    println!("\n✓ GuruRMM Agent uninstalled successfully!");
    println!("\nNote: Config directory {} was preserved.", CONFIG_DIR);
    println!("Remove it manually if no longer needed: sudo rm -rf {}", CONFIG_DIR);

    Ok(())
}

/// Start the installed service
async fn start_service() -> Result<()> {
    #[cfg(windows)]
    {
        service::windows::start()
    }

    #[cfg(target_os = "linux")]
    {
        use std::process::Command;

        info!("Starting GuruRMM Agent service...");

        let status = Command::new("systemctl")
            .args(["start", "gururmm-agent"])
            .status()
            .context("Failed to start service")?;

        if status.success() {
            println!("** Service started successfully");
            println!("Check status: sudo systemctl status gururmm-agent");
        } else {
            anyhow::bail!("Failed to start service. Check: sudo journalctl -u gururmm-agent -n 50");
        }

        Ok(())
    }

    #[cfg(target_os = "macos")]
    {
        todo!("macOS service start not yet implemented");
    }
}

/// Stop the installed service
async fn stop_service() -> Result<()> {
    #[cfg(windows)]
    {
        service::windows::stop()
    }

    #[cfg(target_os = "linux")]
    {
        use std::process::Command;

        info!("Stopping GuruRMM Agent service...");

        let status = Command::new("systemctl")
            .args(["stop", "gururmm-agent"])
            .status()
            .context("Failed to stop service")?;

        if status.success() {
            println!("** Service stopped successfully");
        } else {
            anyhow::bail!("Failed to stop service");
        }

        Ok(())
    }

    #[cfg(target_os = "macos")]
    {
        todo!("macOS service stop not yet implemented");
    }
}

/// Show agent status
async fn show_status(config_path: PathBuf) -> Result<()> {
    // On Windows, show service status
    #[cfg(windows)]
    {
        service::windows::status()?;
        println!();
    }

    // Try to load config for additional info
    match AgentConfig::load(&config_path) {
        Ok(config) => {
            println!("Configuration");
            println!("=============");
            println!("Config file: {:?}", config_path);
            println!("Server URL: {}", config.server.url);
            println!("Metrics interval: {} seconds", config.metrics.interval_seconds);
            println!("Watchdog enabled: {}", config.watchdog.enabled);

            // Collect current metrics
            let collector = MetricsCollector::new();
            let metrics = collector.collect().await;

            println!("\nCurrent System Metrics:");
            println!("  CPU Usage: {:.1}%", metrics.cpu_percent);
            println!("  Memory Usage: {:.1}%", metrics.memory_percent);
            println!(
                "  Memory Used: {:.2} GB",
                metrics.memory_used_bytes as f64 / 1_073_741_824.0
            );
            println!("  Disk Usage: {:.1}%", metrics.disk_percent);
            println!(
                "  Disk Used: {:.2} GB",
                metrics.disk_used_bytes as f64 / 1_073_741_824.0
            );
        }
        Err(_) => {
            println!("\nConfig file {:?} not found or invalid.", config_path);
            #[cfg(windows)]
            println!("Service config location: {}\\agent.toml", service::windows::CONFIG_DIR);
        }
    }

    Ok(())
}

/// Generate a sample configuration file
async fn generate_config(output: PathBuf) -> Result<()> {
    let sample_config = AgentConfig::sample();
    let toml_str = toml::to_string_pretty(&sample_config)?;

    std::fs::write(&output, toml_str)?;
    println!("Sample configuration written to {:?}", output);
    println!("\nEdit this file with your server URL and API key, then run:");
    println!("  gururmm-agent --config {:?} run", output);

    Ok(())
}