claudetools/projects/msp-tools/guru-rmm/agent/src/transport/websocket.rs

//! 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 once_cell::sync::Lazy;
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, TunnelDataPayload, UpdatePayload, UpdateResultPayload, UpdateStatus};
use crate::claude::{ClaudeExecutor, ClaudeTaskCommand};
use crate::metrics::NetworkState;
use crate::tunnel::TunnelManager;
use crate::updater::{AgentUpdater, UpdaterConfig};
use crate::AppState;

/// Global Claude executor for handling Claude Code tasks
static CLAUDE_EXECUTOR: Lazy<ClaudeExecutor> = Lazy::new(|| ClaudeExecutor::new());

/// 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;
                }
            }
        });

        // Create tunnel manager for mode switching
        let mut tunnel_manager = TunnelManager::new();

        // 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");
                        }
                        AgentMessage::TunnelReady { session_id } => {
                            info!("Sent TunnelReady for session: {}", session_id);
                        }
                        AgentMessage::TunnelData { channel_id, .. } => {
                            debug!("Sent TunnelData on channel: {}", channel_id);
                        }
                        AgentMessage::TunnelError { channel_id, error } => {
                            warn!("Sent TunnelError on channel {}: {}", channel_id, error);
                        }
                        _ => {
                            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, &mut tunnel_manager).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;

        // Force close tunnel if active
        tunnel_manager.force_close();

        result
    }

    /// Handle a message received from the server
    async fn handle_server_message(
        text: &str,
        tx: &mpsc::Sender<AgentMessage>,
        tunnel_manager: &mut TunnelManager,
    ) -> 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::TunnelOpen { session_id, tech_id } => {
                info!(
                    "Received tunnel open request: session={}, tech={}",
                    session_id, tech_id
                );
                Self::handle_tunnel_open(session_id, tech_id, tunnel_manager, tx.clone()).await;
            }
            ServerMessage::TunnelClose { session_id } => {
                info!("Received tunnel close request: session={}", session_id);
                Self::handle_tunnel_close(session_id, tunnel_manager, tx.clone()).await;
            }
            ServerMessage::TunnelData { channel_id, data } => {
                debug!("Received tunnel data on channel: {}", channel_id);
                Self::handle_tunnel_data(channel_id, data, tunnel_manager, tx.clone()).await;
            }
        }

        Ok(())
    }

    /// Handle tunnel open request
    async fn handle_tunnel_open(
        session_id: String,
        tech_id: uuid::Uuid,
        tunnel_manager: &mut TunnelManager,
        tx: mpsc::Sender<AgentMessage>,
    ) {
        match tunnel_manager.open_tunnel(session_id.clone(), tech_id) {
            Ok(_) => {
                info!("Tunnel opened successfully: {}", session_id);
                // Send TunnelReady confirmation
                let ready_msg = AgentMessage::TunnelReady {
                    session_id: session_id.clone(),
                };
                if let Err(e) = tx.send(ready_msg).await {
                    error!("Failed to send TunnelReady message: {}", e);
                }
            }
            Err(e) => {
                error!("Failed to open tunnel: {}", e);
                // Send error back to server
                let error_msg = AgentMessage::TunnelError {
                    channel_id: "system".to_string(),
                    error: format!("Failed to open tunnel: {}", e),
                };
                let _ = tx.send(error_msg).await;
            }
        }
    }

    /// Handle tunnel close request
    async fn handle_tunnel_close(
        session_id: String,
        tunnel_manager: &mut TunnelManager,
        tx: mpsc::Sender<AgentMessage>,
    ) {
        match tunnel_manager.close_tunnel(&session_id) {
            Ok(_) => {
                info!("Tunnel closed successfully: {}", session_id);
            }
            Err(e) => {
                warn!("Error closing tunnel: {}", e);
                // Send error back to server
                let error_msg = AgentMessage::TunnelError {
                    channel_id: "system".to_string(),
                    error: format!("Failed to close tunnel: {}", e),
                };
                let _ = tx.send(error_msg).await;
            }
        }
    }

    /// Handle tunnel data (Phase 1: Terminal commands only)
    async fn handle_tunnel_data(
        channel_id: String,
        data: TunnelDataPayload,
        _tunnel_manager: &TunnelManager,
        tx: mpsc::Sender<AgentMessage>,
    ) {
        match data {
            TunnelDataPayload::Terminal { command } => {
                info!("Terminal command on channel {}: {}", channel_id, command);
                // Phase 1: Just log and respond with placeholder
                // Phase 2 will implement actual command execution
                let response = AgentMessage::TunnelData {
                    channel_id,
                    data: TunnelDataPayload::TerminalOutput {
                        stdout: String::new(),
                        stderr: "Terminal execution not yet implemented (Phase 2)".to_string(),
                        exit_code: Some(-1),
                    },
                };
                let _ = tx.send(response).await;
            }
            TunnelDataPayload::TerminalOutput { .. } => {
                // This shouldn't be sent to the agent, it's agent → server only
                warn!("Received TerminalOutput on agent (unexpected)");
            }
        }
    }

    /// 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

        match &cmd.command_type {
            super::CommandType::ClaudeTask {
                task,
                working_directory,
                context_files,
            } => {
                // Handle Claude Code task
                info!("Executing Claude Code task: {}", task);

                let claude_cmd = ClaudeTaskCommand {
                    task: task.clone(),
                    working_directory: working_directory.clone(),
                    timeout: Some(timeout_secs),
                    context_files: context_files.clone(),
                };

                match CLAUDE_EXECUTOR.execute_task(claude_cmd).await {
                    Ok(result) => {
                        let exit_code = match result.status {
                            crate::claude::TaskStatus::Completed => 0,
                            crate::claude::TaskStatus::Failed => 1,
                            crate::claude::TaskStatus::Timeout => 124,
                        };

                        let stdout = result.output.unwrap_or_default();
                        let stderr = result.error.unwrap_or_default();

                        Ok((exit_code, stdout, stderr))
                    }
                    Err(e) => {
                        error!("Claude task execution error: {}", e);
                        Ok((-1, String::new(), e))
                    }
                }
            }
            _ => {
                // Handle regular commands
                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
                    }
                    super::CommandType::ClaudeTask { .. } => {
                        unreachable!("ClaudeTask already handled above")
                    }
                };

                // 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))
            }
        }
    }
}