Explore the GuruRMM server to understand the API and WebSocket architecture. Look at:

1. server/src/api/mod.rs - API route structure
2. server/src/ws/mod.rs - WebSocket message handling
3. How server sends configuration/commands to agents
4. Any existing alert or notification patterns

I need to understand the server architecture to design policy delivery and alert threshold evaluation. Report:
- API route patterns and how new routes would be added
- WebSocket message flow to agents
- Any existing alerting or threshold checking logic


I'll explore the GuruRMM server architecture to understand the API and WebSocket implementation. Let me start by examining the key files.


Bash

find /c/Users/MikeSwanson/claude-projects/gururmm -type f -name "*.rs" | head -20

Find all Rust files in GuruRMM directory


Now let me read the main API and WebSocket modules:


     1→//! REST API routes
     2→//!
     3→//! Provides endpoints for:
     4→//! - Agent management (registration, listing, deletion)
     5→//! - Client and site management
     6→//! - Metrics retrieval
     7→//! - Command execution
     8→//! - User authentication
     9→
    10→pub mod agents;
    11→pub mod auth;
    12→pub mod clients;
    13→pub mod commands;
    14→pub mod metrics;
    15→pub mod sites;
    16→pub mod sso;
    17→pub mod users;
    18→
    19→use axum::{
    20→    routing::{delete, get, post, put},
    21→    Router,
    22→};
    23→
    24→use crate::AppState;
    25→
    26→/// Build all API routes
    27→pub fn routes() -> Router<AppState> {
    28→    Router::new()
    29→        // Authentication
    30→        .route("/auth/login", post(auth::login))
    31→        .route("/auth/register", post(auth::register))
    32→        .route("/auth/me", get(auth::me))
    33→        // SSO
    34→        .route("/auth/sso/entra/config", get(sso::entra_config))
    35→        .route("/auth/sso/entra/callback", post(sso::entra_callback))
    36→        // Users (admin)
    37→        .route("/users", get(users::list_users))
    38→        .route("/users/:id", get(users::get_user))
    39→        .route("/users/:id", delete(users::delete_user))
    40→        .route("/users/:id/role", put(users::update_user_role))
    41→        // Clients
    42→        .route("/clients", get(clients::list_clients))
    43→        .route("/clients", post(clients::create_client))
    44→        .route("/clients/:id", get(clients::get_client))
    45→        .route("/clients/:id", put(clients::update_client))
    46→        .route("/clients/:id", delete(clients::delete_client))
    47→        .route("/clients/:id/sites", get(sites::list_sites_by_client))
    48→        // Sites
    49→        .route("/sites", get(sites::list_sites))
    50→        .route("/sites", post(sites::create_site))
    51→        .route("/sites/:id", get(sites::get_site))
    52→        .route("/sites/:id", put(sites::update_site))
    53→        .route("/sites/:id", delete(sites::delete_site))
    54→        .route("/sites/:id/regenerate-key", post(sites::regenerate_api_key))
    55→        // Agents
    56→        .route("/agents", get(agents::list_agents_with_details))
    57→        .route("/agents", post(agents::register_agent))
    58→        .route("/agents/stats", get(agents::get_stats))
    59→        .route("/agents/unassigned", get(agents::list_unassigned_agents))
    60→        .route("/agents/:id", get(agents::get_agent))
    61→        .route("/agents/:id", delete(agents::delete_agent))
    62→        .route("/agents/:id/move", post(agents::move_agent))
    63→        .route("/agents/:id/state", get(agents::get_agent_state))
    64→        .route("/agents/:id/update", post(agents::trigger_update))
    65→        // Metrics
    66→        .route("/agents/:id/metrics", get(metrics::get_agent_metrics))
    67→        .route("/metrics/summary", get(metrics::get_summary))
    68→        // Commands
    69→        .route("/agents/:id/command", post(commands::send_command))
    70→        .route("/commands", get(commands::list_commands))
    71→        .route("/commands/:id", get(commands::get_command))
    72→        // Legacy Agent (PowerShell for 2008 R2)
    73→        .route("/agent/register-legacy", post(agents::register_legacy))
    74→        .route("/agent/heartbeat", post(agents::heartbeat))
    75→        .route("/agent/command-result", post(agents::command_result))
    76→}
    77→

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

//! REST API routes
//!
//! Provides endpoints for:
//! - Agent management (registration, listing, deletion)
//! - Client and site management
//! - Metrics retrieval
//! - Command execution
//! - User authentication

pub mod agents;
pub mod auth;
pub mod clients;
pub mod commands;
pub mod metrics;
pub mod sites;
pub mod sso;
pub mod users;

use axum::{
    routing::{delete, get, post, put},
    Router,
};

use crate::AppState;

/// Build all API routes
pub fn routes() -> Router<AppState> {
    Router::new()
        // Authentication
        .route("/auth/login", post(auth::login))
        .route("/auth/register", post(auth::register))
        .route("/auth/me", get(auth::me))
        // SSO
        .route("/auth/sso/entra/config", get(sso::entra_config))
        .route("/auth/sso/entra/callback", post(sso::entra_callback))
        // Users (admin)
        .route("/users", get(users::list_users))
        .route("/users/:id", get(users::get_user))
        .route("/users/:id", delete(users::delete_user))
        .route("/users/:id/role", put(users::update_user_role))
        // Clients
        .route("/clients", get(clients::list_clients))
        .route("/clients", post(clients::create_client))
        .route("/clients/:id", get(clients::get_client))
        .route("/clients/:id", put(clients::update_client))
        .route("/clients/:id", delete(clients::delete_client))
        .route("/clients/:id/sites", get(sites::list_sites_by_client))
        // Sites
        .route("/sites", get(sites::list_sites))
        .route("/sites", post(sites::create_site))
        .route("/sites/:id", get(sites::get_site))
        .route("/sites/:id", put(sites::update_site))
        .route("/sites/:id", delete(sites::delete_site))
        .route("/sites/:id/regenerate-key", post(sites::regenerate_api_key))
        // Agents
        .route("/agents", get(agents::list_agents_with_details))
        .route("/agents", post(agents::register_agent))
        .route("/agents/stats", get(agents::get_stats))
        .route("/agents/unassigned", get(agents::list_unassigned_agents))
        .route("/agents/:id", get(agents::get_agent))
        .route("/agents/:id", delete(agents::delete_agent))
        .route("/agents/:id/move", post(agents::move_agent))
        .route("/agents/:id/state", get(agents::get_agent_state))
        .route("/agents/:id/update", post(agents::trigger_update))
        // Metrics
        .route("/agents/:id/metrics", get(metrics::get_agent_metrics))
        .route("/metrics/summary", get(metrics::get_summary))
        // Commands
        .route("/agents/:id/command", post(commands::send_command))
        .route("/commands", get(commands::list_commands))
        .route("/commands/:id", get(commands::get_command))
        // Legacy Agent (PowerShell for 2008 R2)
        .route("/agent/register-legacy", post(agents::register_legacy))
        .route("/agent/heartbeat", post(agents::heartbeat))
        .route("/agent/command-result", post(agents::command_result))
}


     1→//! WebSocket handler for agent connections
     2→//!
     3→//! Handles real-time communication with agents including:
     4→//! - Authentication handshake
     5→//! - Metrics ingestion
     6→//! - Command dispatching
     7→//! - Watchdog event handling
     8→
     9→use std::collections::HashMap;
    10→use std::sync::Arc;
    11→
    12→use axum::{
    13→    extract::{
    14→        ws::{Message, WebSocket},
    15→        State, WebSocketUpgrade,
    16→    },
    17→    response::Response,
    18→};
    19→use futures_util::{SinkExt, StreamExt};
    20→use serde::{Deserialize, Serialize};
    21→use tokio::sync::{mpsc, RwLock};
    22→use tracing::{debug, error, info, warn};
    23→use uuid::Uuid;
    24→
    25→use crate::db;
    26→use crate::AppState;
    27→
    28→/// Connected agents manager
    29→pub struct AgentConnections {
    30→    /// Map of agent ID to sender channel
    31→    connections: HashMap<Uuid, mpsc::Sender<ServerMessage>>,
    32→}
    33→
    34→impl AgentConnections {
    35→    pub fn new() -> Self {
    36→        Self {
    37→            connections: HashMap::new(),
    38→        }
    39→    }
    40→
    41→    /// Add a new agent connection
    42→    pub fn add(&mut self, agent_id: Uuid, tx: mpsc::Sender<ServerMessage>) {
    43→        self.connections.insert(agent_id, tx);
    44→    }
    45→
    46→    /// Remove an agent connection
    47→    pub fn remove(&mut self, agent_id: &Uuid) {
    48→        self.connections.remove(agent_id);
    49→    }
    50→
    51→    /// Send a message to a specific agent
    52→    pub async fn send_to(&self, agent_id: &Uuid, msg: ServerMessage) -> bool {
    53→        if let Some(tx) = self.connections.get(agent_id) {
    54→            tx.send(msg).await.is_ok()
    55→        } else {
    56→            false
    57→        }
    58→    }
    59→
    60→    /// Check if an agent is connected
    61→    pub fn is_connected(&self, agent_id: &Uuid) -> bool {
    62→        self.connections.contains_key(agent_id)
    63→    }
    64→
    65→    /// Get count of connected agents
    66→    pub fn count(&self) -> usize {
    67→        self.connections.len()
    68→    }
    69→}
    70→
    71→impl Default for AgentConnections {
    72→    fn default() -> Self {
    73→        Self::new()
    74→    }
    75→}
    76→
    77→/// Messages from agent to server
    78→#[derive(Debug, Clone, Serialize, Deserialize)]
    79→#[serde(tag = "type", content = "payload")]
    80→#[serde(rename_all = "snake_case")]
    81→pub enum AgentMessage {
    82→    Auth(AuthPayload),
    83→    Metrics(MetricsPayload),
    84→    NetworkState(NetworkStatePayload),
    85→    CommandResult(CommandResultPayload),
    86→    WatchdogEvent(WatchdogEventPayload),
    87→    UpdateResult(UpdateResultPayload),
    88→    Heartbeat,
    89→}
    90→
    91→/// Messages from server to agent
    92→#[derive(Debug, Clone, Serialize, Deserialize)]
    93→#[serde(tag = "type", content = "payload")]
    94→#[serde(rename_all = "snake_case")]
    95→pub enum ServerMessage {
    96→    AuthAck(AuthAckPayload),
    97→    Command(CommandPayload),
    98→    ConfigUpdate(serde_json::Value),
    99→    Update(UpdatePayload),
   100→    Ack { message_id: Option<String> },
   101→    Error { code: String, message: String },
   102→}
   103→
   104→#[derive(Debug, Clone, Serialize, Deserialize)]
   105→pub struct AuthPayload {
   106→    pub api_key: String,
   107→    /// Hardware-derived device ID (for site-based registration)
   108→    #[serde(default)]
   109→    pub device_id: Option<String>,
   110→    pub hostname: String,
   111→    pub os_type: String,
   112→    pub os_version: String,
   113→    pub agent_version: String,
   114→    /// Architecture (amd64, arm64, etc.)
   115→    #[serde(default = "default_arch")]
   116→    pub architecture: String,
   117→    /// Previous version if reconnecting after update
   118→    #[serde(default)]
   119→    pub previous_version: Option<String>,
   120→    /// Update ID if reconnecting after update
   121→    #[serde(default)]
   122→    pub pending_update_id: Option<Uuid>,
   123→}
   124→
   125→fn default_arch() -> String {
   126→    "amd64".to_string()
   127→}
   128→
   129→#[derive(Debug, Clone, Serialize, Deserialize)]
   130→pub struct AuthAckPayload {
   131→    pub success: bool,
   132→    pub agent_id: Option<Uuid>,
   133→    pub error: Option<String>,
   134→}
   135→
   136→#[derive(Debug, Clone, Serialize, Deserialize)]
   137→pub struct MetricsPayload {
   138→    pub timestamp: chrono::DateTime<chrono::Utc>,
   139→    pub cpu_percent: f32,
   140→    pub memory_percent: f32,
   141→    pub memory_used_bytes: u64,
   142→    pub memory_total_bytes: u64,
   143→    pub disk_percent: f32,
   144→    pub disk_used_bytes: u64,
   145→    pub disk_total_bytes: u64,
   146→    pub network_rx_bytes: u64,
   147→    pub network_tx_bytes: u64,
   148→    pub os_type: String,
   149→    pub os_version: String,
   150→    pub hostname: String,
   151→    // Extended metrics (optional for backwards compatibility)
   152→    #[serde(default)]
   153→    pub uptime_seconds: Option<u64>,
   154→    #[serde(default)]
   155→    pub boot_time: Option<i64>,
   156→    #[serde(default)]
   157→    pub logged_in_user: Option<String>,
   158→    #[serde(default)]
   159→    pub user_idle_seconds: Option<u64>,
   160→    #[serde(default)]
   161→    pub public_ip: Option<String>,
   162→    // Temperature metrics (optional for backwards compatibility)
   163→    #[serde(default)]
   164→    pub cpu_temp_celsius: Option<f32>,
   165→    #[serde(default)]
   166→    pub gpu_temp_celsius: Option<f32>,
   167→    #[serde(default)]
   168→    pub temperatures: Option<Vec<TemperatureSensor>>,
   169→}
   170→
   171→/// Temperature sensor reading from agent
   172→#[derive(Debug, Clone, Serialize, Deserialize)]
   173→pub struct TemperatureSensor {
   174→    pub label: String,
   175→    pub temp_celsius: f32,
   176→    #[serde(default)]
   177→    pub critical_celsius: Option<f32>,
   178→}
   179→
   180→#[derive(Debug, Clone, Serialize, Deserialize)]
   181→pub struct CommandResultPayload {
   182→    pub command_id: Uuid,
   183→    pub exit_code: i32,
   184→    pub stdout: String,
   185→    pub stderr: String,
   186→    pub duration_ms: u64,
   187→}
   188→
   189→#[derive(Debug, Clone, Serialize, Deserialize)]
   190→pub struct WatchdogEventPayload {
   191→    pub name: String,
   192→    pub event: String,
   193→    pub details: Option<String>,
   194→}
   195→
   196→/// Network interface information
   197→#[derive(Debug, Clone, Serialize, Deserialize)]
   198→pub struct NetworkInterface {
   199→    pub name: String,
   200→    pub mac_address: Option<String>,
   201→    pub ipv4_addresses: Vec<String>,
   202→    pub ipv6_addresses: Vec<String>,
   203→}
   204→
   205→/// Network state payload from agent
   206→#[derive(Debug, Clone, Serialize, Deserialize)]
   207→pub struct NetworkStatePayload {
   208→    pub timestamp: chrono::DateTime<chrono::Utc>,
   209→    pub interfaces: Vec<NetworkInterface>,
   210→    pub state_hash: String,
   211→}
   212→
   213→#[derive(Debug, Clone, Serialize, Deserialize)]
   214→pub struct CommandPayload {
   215→    pub id: Uuid,
   216→    pub command_type: String,
   217→    pub command: String,
   218→    pub timeout_seconds: Option<u64>,
   219→    pub elevated: bool,
   220→}
   221→
   222→/// Update command payload from server to agent
   223→#[derive(Debug, Clone, Serialize, Deserialize)]
   224→pub struct UpdatePayload {
   225→    /// Unique update ID for tracking
   226→    pub update_id: Uuid,
   227→    /// Target version to update to
   228→    pub target_version: String,
   229→    /// Download URL for the new binary
   230→    pub download_url: String,
   231→    /// SHA256 checksum of the binary
   232→    pub checksum_sha256: String,
   233→    /// Whether to force update (skip version check)
   234→    #[serde(default)]
   235→    pub force: bool,
   236→}
   237→
   238→/// Update result payload from agent to server
   239→#[derive(Debug, Clone, Serialize, Deserialize)]
   240→pub struct UpdateResultPayload {
   241→    /// Update ID (from the server)
   242→    pub update_id: Uuid,
   243→    /// Update status
   244→    pub status: String,
   245→    /// Old version before update
   246→    pub old_version: String,
   247→    /// New version after update (if successful)
   248→    pub new_version: Option<String>,
   249→    /// Error message if failed
   250→    pub error: Option<String>,
   251→}
   252→
   253→/// Result of successful agent authentication
   254→struct AuthResult {
   255→    agent_id: Uuid,
   256→    agent_version: String,
   257→    os_type: String,
   258→    architecture: String,
   259→    previous_version: Option<String>,
   260→    pending_update_id: Option<Uuid>,
   261→}
   262→
   263→/// WebSocket upgrade handler
   264→pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> Response {
   265→    ws.on_upgrade(|socket| handle_socket(socket, state))
   266→}
   267→
   268→/// Handle a WebSocket connection
   269→async fn handle_socket(socket: WebSocket, state: AppState) {
   270→    let (mut sender, mut receiver) = socket.split();
   271→
   272→    // Create channel for outgoing messages
   273→    let (tx, mut rx) = mpsc::channel::<ServerMessage>(100);
   274→
   275→    // Wait for authentication message
   276→    let auth_result = match authenticate(&mut receiver, &mut sender, &state).await {
   277→        Ok(result) => {
   278→            info!("Agent authenticated: {}", result.agent_id);
   279→
   280→            // Send auth success
   281→            let ack = ServerMessage::AuthAck(AuthAckPayload {
   282→                success: true,
   283→                agent_id: Some(result.agent_id),
   284→                error: None,
   285→            });
   286→
   287→            if let Ok(json) = serde_json::to_string(&ack) {
   288→                let _ = sender.send(Message::Text(json)).await;
   289→            }
   290→
   291→            // Register connection
   292→            state.agents.write().await.add(result.agent_id, tx.clone());
   293→
   294→            // Update agent status
   295→            let _ = db::update_agent_status(&state.db, result.agent_id, "online").await;
   296→
   297→            // Check if this is a post-update reconnection
   298→            if let Some(prev_version) = &result.previous_version {
   299→                if prev_version != &result.agent_version {
   300→                    info!(
   301→                        "Agent {} reconnected after update: {} -> {}",
   302→                        result.agent_id, prev_version, result.agent_version
   303→                    );
   304→                    // Mark update as completed
   305→                    let _ = db::complete_update_by_agent(
   306→                        &state.db,
   307→                        result.agent_id,
   308→                        result.pending_update_id,
   309→                        prev_version,
   310→                        &result.agent_version,
   311→                    ).await;
   312→                }
   313→            }
   314→
   315→            // Check if agent needs update (auto-update enabled)
   316→            if let Some(available) = state.updates.needs_update(
   317→                &result.agent_version,
   318→                &result.os_type,
   319→                &result.architecture,
   320→            ).await {
   321→                info!(
   322→                    "Agent {} needs update: {} -> {}",
   323→                    result.agent_id, result.agent_version, available.version
   324→                );
   325→
   326→                let update_id = Uuid::new_v4();
   327→
   328→                // Record update in database
   329→                if let Err(e) = db::create_agent_update(
   330→                    &state.db,
   331→                    result.agent_id,
   332→                    update_id,
   333→                    &result.agent_version,
   334→                    &available.version.to_string(),
   335→                    &available.download_url,
   336→                    &available.checksum_sha256,
   337→                ).await {
   338→                    error!("Failed to record update: {}", e);
   339→                } else {
   340→                    // Send update command
   341→                    let update_msg = ServerMessage::Update(UpdatePayload {
   342→                        update_id,
   343→                        target_version: available.version.to_string(),
   344→                        download_url: available.download_url.clone(),
   345→                        checksum_sha256: available.checksum_sha256.clone(),
   346→                        force: false,
   347→                    });
   348→
   349→                    if let Err(e) = tx.send(update_msg).await {
   350→                        error!("Failed to send update command: {}", e);
   351→                    }
   352→                }
   353→            }
   354→
   355→            result
   356→        }
   357→        Err(e) => {
   358→            error!("Authentication failed: {}", e);
   359→
   360→            let ack = ServerMessage::AuthAck(AuthAckPayload {
   361→                success: false,
   362→                agent_id: None,
   363→                error: Some(e.to_string()),
   364→            });
   365→
   366→            if let Ok(json) = serde_json::to_string(&ack) {
   367→                let _ = sender.send(Message::Text(json)).await;
   368→            }
   369→
   370→            return;
   371→        }
   372→    };
   373→
   374→    let agent_id = auth_result.agent_id;
   375→
   376→    // Spawn task to forward outgoing messages
   377→    let send_task = tokio::spawn(async move {
   378→        while let Some(msg) = rx.recv().await {
   379→            if let Ok(json) = serde_json::to_string(&msg) {
   380→                if sender.send(Message::Text(json)).await.is_err() {
   381→                    break;
   382→                }
   383→            }
   384→        }
   385→    });
   386→
   387→    // Handle incoming messages
   388→    while let Some(msg_result) = receiver.next().await {
   389→        match msg_result {
   390→            Ok(Message::Text(text)) => {
   391→                if let Err(e) = handle_agent_message(&text, agent_id, &state).await {
   392→                    error!("Error handling agent message: {}", e);
   393→                }
   394→            }
   395→            Ok(Message::Ping(data)) => {
   396→                if tx
   397→                    .send(ServerMessage::Ack { message_id: None })
   398→                    .await
   399→                    .is_err()
   400→                {
   401→                    break;
   402→                }
   403→            }
   404→            Ok(Message::Close(_)) => {
   405→                info!("Agent {} disconnected", agent_id);
   406→                break;
   407→            }
   408→            Err(e) => {
   409→                error!("WebSocket error for agent {}: {}", agent_id, e);
   410→                break;
   411→            }
   412→            _ => {}
   413→        }
   414→    }
   415→
   416→    // Cleanup
   417→    state.agents.write().await.remove(&agent_id);
   418→    let _ = db::update_agent_status(&state.db, agent_id, "offline").await;
   419→    send_task.abort();
   420→
   421→    info!("Agent {} connection closed", agent_id);
   422→}
   423→
   424→/// Authenticate an agent connection
   425→///
   426→/// Supports two modes:
   427→/// 1. Legacy: API key maps directly to an agent (api_key_hash in agents table)
   428→/// 2. Site-based: API key maps to a site, device_id identifies the specific agent
   429→async fn authenticate(
   430→    receiver: &mut futures_util::stream::SplitStream<WebSocket>,
   431→    sender: &mut futures_util::stream::SplitSink<WebSocket, Message>,
   432→    state: &AppState,
   433→) -> anyhow::Result<AuthResult> {
   434→    use tokio::time::{timeout, Duration};
   435→
   436→    // Wait for auth message with timeout
   437→    let msg = timeout(Duration::from_secs(10), receiver.next())
   438→        .await
   439→        .map_err(|_| anyhow::anyhow!("Authentication timeout"))?
   440→        .ok_or_else(|| anyhow::anyhow!("Connection closed before auth"))?
   441→        .map_err(|e| anyhow::anyhow!("WebSocket error: {}", e))?;
   442→
   443→    let text = match msg {
   444→        Message::Text(t) => t,
   445→        _ => return Err(anyhow::anyhow!("Expected text message for auth")),
   446→    };
   447→
   448→    let agent_msg: AgentMessage =
   449→        serde_json::from_str(&text).map_err(|e| anyhow::anyhow!("Invalid auth message: {}", e))?;
   450→
   451→    let auth = match agent_msg {
   452→        AgentMessage::Auth(a) => a,
   453→        _ => return Err(anyhow::anyhow!("Expected auth message")),
   454→    };
   455→
   456→    // Try site-based authentication first (if device_id is provided)
   457→    if let Some(device_id) = &auth.device_id {
   458→        // Check if api_key looks like a site code (WORD-WORD-NUMBER format)
   459→        let site = if is_site_code_format(&auth.api_key) {
   460→            info!("Attempting site code authentication: {}", auth.api_key);
   461→            db::get_site_by_code(&state.db, &auth.api_key)
   462→                .await
   463→                .map_err(|e| anyhow::anyhow!("Database error: {}", e))?
   464→        } else {
   465→            // Hash the API key and look up by hash
   466→            let api_key_hash = hash_api_key(&auth.api_key);
   467→            db::get_site_by_api_key_hash(&state.db, &api_key_hash)
   468→                .await
   469→                .map_err(|e| anyhow::anyhow!("Database error: {}", e))?
   470→        };
   471→
   472→        if let Some(site) = site {
   473→            info!("Site-based auth: site={} ({})", site.name, site.id);
   474→
   475→            // Look up or create agent by site_id + device_id
   476→            let agent = match db::get_agent_by_site_and_device(&state.db, site.id, device_id)
   477→                .await
   478→                .map_err(|e| anyhow::anyhow!("Database error: {}", e))?
   479→            {
   480→                Some(agent) => {
   481→                    // Update existing agent info
   482→                    let _ = db::update_agent_info_full(
   483→                        &state.db,
   484→                        agent.id,
   485→                        Some(&auth.hostname),
   486→                        Some(device_id),
   487→                        Some(&auth.os_version),
   488→                        Some(&auth.agent_version),
   489→                    )
   490→                    .await;
   491→                    agent
   492→                }
   493→                None => {
   494→                    // Auto-register new agent under this site
   495→                    info!(
   496→                        "Auto-registering new agent: hostname={}, device_id={}, site={}",
   497→                        auth.hostname, device_id, site.name
   498→                    );
   499→
   500→                    db::create_agent_with_site(
   501→                        &state.db,
   502→                        db::CreateAgentWithSite {
   503→                            site_id: site.id,
   504→                            device_id: device_id.clone(),
   505→                            hostname: auth.hostname.clone(),
   506→                            os_type: auth.os_type.clone(),
   507→                            os_version: Some(auth.os_version.clone()),
   508→                            agent_version: Some(auth.agent_version.clone()),
   509→                        },
   510→                    )
   511→                    .await
   512→                    .map_err(|e| anyhow::anyhow!("Failed to create agent: {}", e))?
   513→                }
   514→            };
   515→
   516→            return Ok(AuthResult {
   517→                agent_id: agent.id,
   518→                agent_version: auth.agent_version.clone(),
   519→                os_type: auth.os_type.clone(),
   520→                architecture: auth.architecture.clone(),
   521→                previous_version: auth.previous_version.clone(),
   522→                pending_update_id: auth.pending_update_id,
   523→            });
   524→        }
   525→    }
   526→
   527→    // Fall back to legacy: look up agent directly by API key hash
   528→    let api_key_hash = hash_api_key(&auth.api_key);
   529→    let agent = db::get_agent_by_api_key_hash(&state.db, &api_key_hash)
   530→        .await
   531→        .map_err(|e| anyhow::anyhow!("Database error: {}", e))?
   532→        .ok_or_else(|| anyhow::anyhow!("Invalid API key"))?;
   533→
   534→    // Update agent info (including hostname in case it changed)
   535→    let _ = db::update_agent_info(
   536→        &state.db,
   537→        agent.id,
   538→        Some(&auth.hostname),
   539→        Some(&auth.os_version),
   540→        Some(&auth.agent_version),
   541→    )
   542→    .await;
   543→
   544→    Ok(AuthResult {
   545→        agent_id: agent.id,
   546→        agent_version: auth.agent_version,
   547→        os_type: auth.os_type,
   548→        architecture: auth.architecture,
   549→        previous_version: auth.previous_version,
   550→        pending_update_id: auth.pending_update_id,
   551→    })
   552→}
   553→
   554→/// Handle a message from an authenticated agent
   555→async fn handle_agent_message(
   556→    text: &str,
   557→    agent_id: Uuid,
   558→    state: &AppState,
   559→) -> anyhow::Result<()> {
   560→    let msg: AgentMessage = serde_json::from_str(text)?;
   561→
   562→    match msg {
   563→        AgentMessage::Metrics(metrics) => {
   564→            debug!("Received metrics from agent {}: CPU={:.1}%", agent_id, metrics.cpu_percent);
   565→
   566→            // Store metrics in database
   567→            let create_metrics = db::CreateMetrics {
   568→                agent_id,
   569→                cpu_percent: Some(metrics.cpu_percent),
   570→                memory_percent: Some(metrics.memory_percent),
   571→                memory_used_bytes: Some(metrics.memory_used_bytes as i64),
   572→                disk_percent: Some(metrics.disk_percent),
   573→                disk_used_bytes: Some(metrics.disk_used_bytes as i64),
   574→                network_rx_bytes: Some(metrics.network_rx_bytes as i64),
   575→                network_tx_bytes: Some(metrics.network_tx_bytes as i64),
   576→                // Extended metrics
   577→                uptime_seconds: metrics.uptime_seconds.map(|v| v as i64),
   578→                boot_time: metrics.boot_time,
   579→                logged_in_user: metrics.logged_in_user.clone(),
   580→                user_idle_seconds: metrics.user_idle_seconds.map(|v| v as i64),
   581→                public_ip: metrics.public_ip.clone(),
   582→                memory_total_bytes: Some(metrics.memory_total_bytes as i64),
   583→                disk_total_bytes: Some(metrics.disk_total_bytes as i64),
   584→                // Temperature metrics
   585→                cpu_temp_celsius: metrics.cpu_temp_celsius,
   586→                gpu_temp_celsius: metrics.gpu_temp_celsius,
   587→                temperatures: metrics.temperatures.as_ref().map(|t| serde_json::to_value(t).ok()).flatten(),
   588→            };
   589→
   590→            db::insert_metrics(&state.db, create_metrics).await?;
   591→
   592→            // Also update agent_state for quick access to latest extended info
   593→            let _ = db::upsert_agent_state(
   594→                &state.db,
   595→                agent_id,
   596→                metrics.uptime_seconds.map(|v| v as i64),
   597→                metrics.boot_time,
   598→                metrics.logged_in_user.as_deref(),
   599→                metrics.user_idle_seconds.map(|v| v as i64),
   600→                metrics.public_ip.as_deref(),
   601→            ).await;
   602→
   603→            // Update last_seen
   604→            db::update_agent_status(&state.db, agent_id, "online").await?;
   605→        }
   606→
   607→        AgentMessage::CommandResult(result) => {
   608→            info!(
   609→                "Received command result from agent {}: command={}, exit={}",
   610→                agent_id, result.command_id, result.exit_code
   611→            );
   612→
   613→            // Update command in database
   614→            let cmd_result = db::CommandResult {
   615→                exit_code: result.exit_code,
   616→                stdout: result.stdout,
   617→                stderr: result.stderr,
   618→            };
   619→
   620→            db::update_command_result(&state.db, result.command_id, cmd_result).await?;
   621→        }
   622→
   623→        AgentMessage::WatchdogEvent(event) => {
   624→            info!(
   625→                "Received watchdog event from agent {}: {} - {}",
   626→                agent_id, event.name, event.event
   627→            );
   628→
   629→            // Store watchdog event (table exists but we'll add the insert function later)
   630→            // For now, just log it
   631→        }
   632→
   633→        AgentMessage::NetworkState(network_state) => {
   634→            debug!(
   635→                "Received network state from agent {}: {} interfaces",
   636→                agent_id,
   637→                network_state.interfaces.len()
   638→            );
   639→            // Log interface details at trace level
   640→            for iface in &network_state.interfaces {
   641→                tracing::trace!(
   642→                    "  Interface {}: IPv4={:?}, IPv6={:?}",
   643→                    iface.name,
   644→                    iface.ipv4_addresses,
   645→                    iface.ipv6_addresses
   646→                );
   647→            }
   648→            // Store network state in database
   649→            if let Ok(interfaces_json) = serde_json::to_value(&network_state.interfaces) {
   650→                let _ = db::update_agent_network_state(
   651→                    &state.db,
   652→                    agent_id,
   653→                    &interfaces_json,
   654→                    &network_state.state_hash,
   655→                ).await;
   656→            }
   657→            // Update last_seen
   658→            db::update_agent_status(&state.db, agent_id, "online").await?;
   659→        }
   660→
   661→        AgentMessage::Heartbeat => {
   662→            debug!("Received heartbeat from agent {}", agent_id);
   663→            db::update_agent_status(&state.db, agent_id, "online").await?;
   664→        }
   665→
   666→        AgentMessage::UpdateResult(result) => {
   667→            info!(
   668→                "Received update result from agent {}: update_id={}, status={}",
   669→                agent_id, result.update_id, result.status
   670→            );
   671→
   672→            // Update the agent_updates record
   673→            match result.status.as_str() {
   674→                "completed" => {
   675→                    let _ = db::complete_agent_update(&state.db, result.update_id, result.new_version.as_deref()).await;
   676→                    info!("Agent {} successfully updated to {}", agent_id, result.new_version.unwrap_or_default());
   677→                }
   678→                "failed" | "rolled_back" => {
   679→                    let _ = db::fail_agent_update(&state.db, result.update_id, result.error.as_deref()).await;
   680→                    warn!("Agent {} update failed: {}", agent_id, result.error.unwrap_or_default());
   681→                }
   682→                _ => {
   683→                    debug!("Agent {} update status: {}", agent_id, result.status);
   684→                }
   685→            }
   686→        }
   687→
   688→        AgentMessage::Auth(_) => {
   689→            warn!("Received unexpected auth message from already authenticated agent");
   690→        }
   691→    }
   692→
   693→    Ok(())
   694→}
   695→
   696→/// Hash an API key for storage/lookup
   697→pub fn hash_api_key(api_key: &str) -> String {
   698→    use sha2::{Digest, Sha256};
   699→    let mut hasher = Sha256::new();
   700→    hasher.update(api_key.as_bytes());
   701→    format!("{:x}", hasher.finalize())
   702→}
   703→
   704→/// Generate a new API key
   705→pub fn generate_api_key(prefix: &str) -> String {
   706→    use rand::Rng;
   707→    let random_bytes: [u8; 24] = rand::thread_rng().gen();
   708→    let encoded = base64::Engine::encode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, random_bytes);
   709→    format!("{}{}", prefix, encoded)
   710→}
   711→
   712→/// Check if a string looks like a site code (WORD-WORD-NUMBER format)
   713→/// Examples: SWIFT-CLOUD-6910, APPLE-GREEN-9145
   714→fn is_site_code_format(s: &str) -> bool {
   715→    let parts: Vec<&str> = s.split('-').collect();
   716→    if parts.len() != 3 {
   717→        return false;
   718→    }
   719→    // First two parts should be alphabetic (words)
   720→    // Third part should be numeric (4 digits)
   721→    parts[0].chars().all(|c| c.is_ascii_alphabetic())
   722→        && parts[1].chars().all(|c| c.is_ascii_alphabetic())
   723→        && parts[2].chars().all(|c| c.is_ascii_digit())
   724→        && parts[2].len() == 4
   725→}
   726→

<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 handler for agent connections
//!
//! Handles real-time communication with agents including:
//! - Authentication handshake
//! - Metrics ingestion
//! - Command dispatching
//! - Watchdog event handling

use std::collections::HashMap;
use std::sync::Arc;

use axum::{
    extract::{
        ws::{Message, WebSocket},
        State, WebSocketUpgrade,
    },
    response::Response,
};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tokio::sync::{mpsc, RwLock};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

use crate::db;
use crate::AppState;

/// Connected agents manager
pub struct AgentConnections {
    /// Map of agent ID to sender channel
    connections: HashMap<Uuid, mpsc::Sender<ServerMessage>>,
}

impl AgentConnections {
    pub fn new() -> Self {
        Self {
            connections: HashMap::new(),
        }
    }

    /// Add a new agent connection
    pub fn add(&mut self, agent_id: Uuid, tx: mpsc::Sender<ServerMessage>) {
        self.connections.insert(agent_id, tx);
    }

    /// Remove an agent connection
    pub fn remove(&mut self, agent_id: &Uuid) {
        self.connections.remove(agent_id);
    }

    /// Send a message to a specific agent
    pub async fn send_to(&self, agent_id: &Uuid, msg: ServerMessage) -> bool {
        if let Some(tx) = self.connections.get(agent_id) {
            tx.send(msg).await.is_ok()
        } else {
            false
        }
    }

    /// Check if an agent is connected
    pub fn is_connected(&self, agent_id: &Uuid) -> bool {
        self.connections.contains_key(agent_id)
    }

    /// Get count of connected agents
    pub fn count(&self) -> usize {
        self.connections.len()
    }
}

impl Default for AgentConnections {
    fn default() -> Self {
        Self::new()
    }
}

/// Messages from agent to server
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
#[serde(rename_all = "snake_case")]
pub enum AgentMessage {
    Auth(AuthPayload),
    Metrics(MetricsPayload),
    NetworkState(NetworkStatePayload),
    CommandResult(CommandResultPayload),
    WatchdogEvent(WatchdogEventPayload),
    UpdateResult(UpdateResultPayload),
    Heartbeat,
}

/// Messages from server to agent
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
#[serde(rename_all = "snake_case")]
pub enum ServerMessage {
    AuthAck(AuthAckPayload),
    Command(CommandPayload),
    ConfigUpdate(serde_json::Value),
    Update(UpdatePayload),
    Ack { message_id: Option<String> },
    Error { code: String, message: String },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthPayload {
    pub api_key: String,
    /// Hardware-derived device ID (for site-based registration)
    #[serde(default)]
    pub device_id: Option<String>,
    pub hostname: String,
    pub os_type: String,
    pub os_version: String,
    pub agent_version: String,
    /// Architecture (amd64, arm64, etc.)
    #[serde(default = "default_arch")]
    pub architecture: String,
    /// Previous version if reconnecting after update
    #[serde(default)]
    pub previous_version: Option<String>,
    /// Update ID if reconnecting after update
    #[serde(default)]
    pub pending_update_id: Option<Uuid>,
}

fn default_arch() -> String {
    "amd64".to_string()
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthAckPayload {
    pub success: bool,
    pub agent_id: Option<Uuid>,
    pub error: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MetricsPayload {
    pub timestamp: chrono::DateTime<chrono::Utc>,
    pub cpu_percent: f32,
    pub memory_percent: f32,
    pub memory_used_bytes: u64,
    pub memory_total_bytes: u64,
    pub disk_percent: f32,
    pub disk_used_bytes: u64,
    pub disk_total_bytes: u64,
    pub network_rx_bytes: u64,
    pub network_tx_bytes: u64,
    pub os_type: String,
    pub os_version: String,
    pub hostname: String,
    // Extended metrics (optional for backwards compatibility)
    #[serde(default)]
    pub uptime_seconds: Option<u64>,
    #[serde(default)]
    pub boot_time: Option<i64>,
    #[serde(default)]
    pub logged_in_user: Option<String>,
    #[serde(default)]
    pub user_idle_seconds: Option<u64>,
    #[serde(default)]
    pub public_ip: Option<String>,
    // Temperature metrics (optional for backwards compatibility)
    #[serde(default)]
    pub cpu_temp_celsius: Option<f32>,
    #[serde(default)]
    pub gpu_temp_celsius: Option<f32>,
    #[serde(default)]
    pub temperatures: Option<Vec<TemperatureSensor>>,
}

/// Temperature sensor reading from agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TemperatureSensor {
    pub label: String,
    pub temp_celsius: f32,
    #[serde(default)]
    pub critical_celsius: Option<f32>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommandResultPayload {
    pub command_id: Uuid,
    pub exit_code: i32,
    pub stdout: String,
    pub stderr: String,
    pub duration_ms: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WatchdogEventPayload {
    pub name: String,
    pub event: String,
    pub details: Option<String>,
}

/// Network interface information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkInterface {
    pub name: String,
    pub mac_address: Option<String>,
    pub ipv4_addresses: Vec<String>,
    pub ipv6_addresses: Vec<String>,
}

/// Network state payload from agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkStatePayload {
    pub timestamp: chrono::DateTime<chrono::Utc>,
    pub interfaces: Vec<NetworkInterface>,
    pub state_hash: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommandPayload {
    pub id: Uuid,
    pub command_type: String,
    pub command: String,
    pub timeout_seconds: Option<u64>,
    pub elevated: bool,
}

/// Update command payload from server to agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdatePayload {
    /// Unique update ID for tracking
    pub update_id: Uuid,
    /// Target version to update to
    pub target_version: String,
    /// Download URL for the new binary
    pub download_url: String,
    /// SHA256 checksum of the binary
    pub checksum_sha256: String,
    /// Whether to force update (skip version check)
    #[serde(default)]
    pub force: bool,
}

/// Update result payload from agent to server
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateResultPayload {
    /// Update ID (from the server)
    pub update_id: Uuid,
    /// Update status
    pub status: String,
    /// Old version before update
    pub old_version: String,
    /// New version after update (if successful)
    pub new_version: Option<String>,
    /// Error message if failed
    pub error: Option<String>,
}

/// Result of successful agent authentication
struct AuthResult {
    agent_id: Uuid,
    agent_version: String,
    os_type: String,
    architecture: String,
    previous_version: Option<String>,
    pending_update_id: Option<Uuid>,
}

/// WebSocket upgrade handler
pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> Response {
    ws.on_upgrade(|socket| handle_socket(socket, state))
}

/// Handle a WebSocket connection
async fn handle_socket(socket: WebSocket, state: AppState) {
    let (mut sender, mut receiver) = socket.split();

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

    // Wait for authentication message
    let auth_result = match authenticate(&mut receiver, &mut sender, &state).await {
        Ok(result) => {
            info!("Agent authenticated: {}", result.agent_id);

            // Send auth success
            let ack = ServerMessage::AuthAck(AuthAckPayload {
                success: true,
                agent_id: Some(result.agent_id),
                error: None,
            });

            if let Ok(json) = serde_json::to_string(&ack) {
                let _ = sender.send(Message::Text(json)).await;
            }

            // Register connection
            state.agents.write().await.add(result.agent_id, tx.clone());

            // Update agent status
            let _ = db::update_agent_status(&state.db, result.agent_id, "online").await;

            // Check if this is a post-update reconnection
            if let Some(prev_version) = &result.previous_version {
                if prev_version != &result.agent_version {
                    info!(
                        "Agent {} reconnected after update: {} -> {}",
                        result.agent_id, prev_version, result.agent_version
                    );
                    // Mark update as completed
                    let _ = db::complete_update_by_agent(
                        &state.db,
                        result.agent_id,
                        result.pending_update_id,
                        prev_version,
                        &result.agent_version,
                    ).await;
                }
            }

            // Check if agent needs update (auto-update enabled)
            if let Some(available) = state.updates.needs_update(
                &result.agent_version,
                &result.os_type,
                &result.architecture,
            ).await {
                info!(
                    "Agent {} needs update: {} -> {}",
                    result.agent_id, result.agent_version, available.version
                );

                let update_id = Uuid::new_v4();

                // Record update in database
                if let Err(e) = db::create_agent_update(
                    &state.db,
                    result.agent_id,
                    update_id,
                    &result.agent_version,
                    &available.version.to_string(),
                    &available.download_url,
                    &available.checksum_sha256,
                ).await {
                    error!("Failed to record update: {}", e);
                } else {
                    // Send update command
                    let update_msg = ServerMessage::Update(UpdatePayload {
                        update_id,
                        target_version: available.version.to_string(),
                        download_url: available.download_url.clone(),
                        checksum_sha256: available.checksum_sha256.clone(),
                        force: false,
                    });

                    if let Err(e) = tx.send(update_msg).await {
                        error!("Failed to send update command: {}", e);
                    }
                }
            }

            result
        }
        Err(e) => {
            error!("Authentication failed: {}", e);

            let ack = ServerMessage::AuthAck(AuthAckPayload {
                success: false,
                agent_id: None,
                error: Some(e.to_string()),
            });

            if let Ok(json) = serde_json::to_string(&ack) {
                let _ = sender.send(Message::Text(json)).await;
            }

            return;
        }
    };

    let agent_id = auth_result.agent_id;

    // Spawn task to forward outgoing messages
    let send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if let Ok(json) = serde_json::to_string(&msg) {
                if sender.send(Message::Text(json)).await.is_err() {
                    break;
                }
            }
        }
    });

    // Handle incoming messages
    while let Some(msg_result) = receiver.next().await {
        match msg_result {
            Ok(Message::Text(text)) => {
                if let Err(e) = handle_agent_message(&text, agent_id, &state).await {
                    error!("Error handling agent message: {}", e);
                }
            }
            Ok(Message::Ping(data)) => {
                if tx
                    .send(ServerMessage::Ack { message_id: None })
                    .await
                    .is_err()
                {
                    break;
                }
            }
            Ok(Message::Close(_)) => {
                info!("Agent {} disconnected", agent_id);
                break;
            }
            Err(e) => {
                error!("WebSocket error for agent {}: {}", agent_id, e);
                break;
            }
            _ => {}
        }
    }

    // Cleanup
    state.agents.write().await.remove(&agent_id);
    let _ = db::update_agent_status(&state.db, agent_id, "offline").await;
    send_task.abort();

    info!("Agent {} connection closed", agent_id);
}

/// Authenticate an agent connection
///
/// Supports two modes:
/// 1. Legacy: API key maps directly to an agent (api_key_hash in agents table)
/// 2. Site-based: API key maps to a site, device_id identifies the specific agent
async fn authenticate(
    receiver: &mut futures_util::stream::SplitStream<WebSocket>,
    sender: &mut futures_util::stream::SplitSink<WebSocket, Message>,
    state: &AppState,
) -> anyhow::Result<AuthResult> {
    use tokio::time::{timeout, Duration};

    // Wait for auth message with timeout
    let msg = timeout(Duration::from_secs(10), receiver.next())
        .await
        .map_err(|_| anyhow::anyhow!("Authentication timeout"))?
        .ok_or_else(|| anyhow::anyhow!("Connection closed before auth"))?
        .map_err(|e| anyhow::anyhow!("WebSocket error: {}", e))?;

    let text = match msg {
        Message::Text(t) => t,
        _ => return Err(anyhow::anyhow!("Expected text message for auth")),
    };

    let agent_msg: AgentMessage =
        serde_json::from_str(&text).map_err(|e| anyhow::anyhow!("Invalid auth message: {}", e))?;

    let auth = match agent_msg {
        AgentMessage::Auth(a) => a,
        _ => return Err(anyhow::anyhow!("Expected auth message")),
    };

    // Try site-based authentication first (if device_id is provided)
    if let Some(device_id) = &auth.device_id {
        // Check if api_key looks like a site code (WORD-WORD-NUMBER format)
        let site = if is_site_code_format(&auth.api_key) {
            info!("Attempting site code authentication: {}", auth.api_key);
            db::get_site_by_code(&state.db, &auth.api_key)
                .await
                .map_err(|e| anyhow::anyhow!("Database error: {}", e))?
        } else {
            // Hash the API key and look up by hash
            let api_key_hash = hash_api_key(&auth.api_key);
            db::get_site_by_api_key_hash(&state.db, &api_key_hash)
                .await
                .map_err(|e| anyhow::anyhow!("Database error: {}", e))?
        };

        if let Some(site) = site {
            info!("Site-based auth: site={} ({})", site.name, site.id);

            // Look up or create agent by site_id + device_id
            let agent = match db::get_agent_by_site_and_device(&state.db, site.id, device_id)
                .await
                .map_err(|e| anyhow::anyhow!("Database error: {}", e))?
            {
                Some(agent) => {
                    // Update existing agent info
                    let _ = db::update_agent_info_full(
                        &state.db,
                        agent.id,
                        Some(&auth.hostname),
                        Some(device_id),
                        Some(&auth.os_version),
                        Some(&auth.agent_version),
                    )
                    .await;
                    agent
                }
                None => {
                    // Auto-register new agent under this site
                    info!(
                        "Auto-registering new agent: hostname={}, device_id={}, site={}",
                        auth.hostname, device_id, site.name
                    );

                    db::create_agent_with_site(
                        &state.db,
                        db::CreateAgentWithSite {
                            site_id: site.id,
                            device_id: device_id.clone(),
                            hostname: auth.hostname.clone(),
                            os_type: auth.os_type.clone(),
                            os_version: Some(auth.os_version.clone()),
                            agent_version: Some(auth.agent_version.clone()),
                        },
                    )
                    .await
                    .map_err(|e| anyhow::anyhow!("Failed to create agent: {}", e))?
                }
            };

            return Ok(AuthResult {
                agent_id: agent.id,
                agent_version: auth.agent_version.clone(),
                os_type: auth.os_type.clone(),
                architecture: auth.architecture.clone(),
                previous_version: auth.previous_version.clone(),
                pending_update_id: auth.pending_update_id,
            });
        }
    }

    // Fall back to legacy: look up agent directly by API key hash
    let api_key_hash = hash_api_key(&auth.api_key);
    let agent = db::get_agent_by_api_key_hash(&state.db, &api_key_hash)
        .await
        .map_err(|e| anyhow::anyhow!("Database error: {}", e))?
        .ok_or_else(|| anyhow::anyhow!("Invalid API key"))?;

    // Update agent info (including hostname in case it changed)
    let _ = db::update_agent_info(
        &state.db,
        agent.id,
        Some(&auth.hostname),
        Some(&auth.os_version),
        Some(&auth.agent_version),
    )
    .await;

    Ok(AuthResult {
        agent_id: agent.id,
        agent_version: auth.agent_version,
        os_type: auth.os_type,
        architecture: auth.architecture,
        previous_version: auth.previous_version,
        pending_update_id: auth.pending_update_id,
    })
}

/// Handle a message from an authenticated agent
async fn handle_agent_message(
    text: &str,
    agent_id: Uuid,
    state: &AppState,
) -> anyhow::Result<()> {
    let msg: AgentMessage = serde_json::from_str(text)?;

    match msg {
        AgentMessage::Metrics(metrics) => {
            debug!("Received metrics from agent {}: CPU={:.1}%", agent_id, metrics.cpu_percent);

            // Store metrics in database
            let create_metrics = db::CreateMetrics {
                agent_id,
                cpu_percent: Some(metrics.cpu_percent),
                memory_percent: Some(metrics.memory_percent),
                memory_used_bytes: Some(metrics.memory_used_bytes as i64),
                disk_percent: Some(metrics.disk_percent),
                disk_used_bytes: Some(metrics.disk_used_bytes as i64),
                network_rx_bytes: Some(metrics.network_rx_bytes as i64),
                network_tx_bytes: Some(metrics.network_tx_bytes as i64),
                // Extended metrics
                uptime_seconds: metrics.uptime_seconds.map(|v| v as i64),
                boot_time: metrics.boot_time,
                logged_in_user: metrics.logged_in_user.clone(),
                user_idle_seconds: metrics.user_idle_seconds.map(|v| v as i64),
                public_ip: metrics.public_ip.clone(),
                memory_total_bytes: Some(metrics.memory_total_bytes as i64),
                disk_total_bytes: Some(metrics.disk_total_bytes as i64),
                // Temperature metrics
                cpu_temp_celsius: metrics.cpu_temp_celsius,
                gpu_temp_celsius: metrics.gpu_temp_celsius,
                temperatures: metrics.temperatures.as_ref().map(|t| serde_json::to_value(t).ok()).flatten(),
            };

            db::insert_metrics(&state.db, create_metrics).await?;

            // Also update agent_state for quick access to latest extended info
            let _ = db::upsert_agent_state(
                &state.db,
                agent_id,
                metrics.uptime_seconds.map(|v| v as i64),
                metrics.boot_time,
                metrics.logged_in_user.as_deref(),
                metrics.user_idle_seconds.map(|v| v as i64),
                metrics.public_ip.as_deref(),
            ).await;

            // Update last_seen
            db::update_agent_status(&state.db, agent_id, "online").await?;
        }

        AgentMessage::CommandResult(result) => {
            info!(
                "Received command result from agent {}: command={}, exit={}",
                agent_id, result.command_id, result.exit_code
            );

            // Update command in database
            let cmd_result = db::CommandResult {
                exit_code: result.exit_code,
                stdout: result.stdout,
                stderr: result.stderr,
            };

            db::update_command_result(&state.db, result.command_id, cmd_result).await?;
        }

        AgentMessage::WatchdogEvent(event) => {
            info!(
                "Received watchdog event from agent {}: {} - {}",
                agent_id, event.name, event.event
            );

            // Store watchdog event (table exists but we'll add the insert function later)
            // For now, just log it
        }

        AgentMessage::NetworkState(network_state) => {
            debug!(
                "Received network state from agent {}: {} interfaces",
                agent_id,
                network_state.interfaces.len()
            );
            // Log interface details at trace level
            for iface in &network_state.interfaces {
                tracing::trace!(
                    "  Interface {}: IPv4={:?}, IPv6={:?}",
                    iface.name,
                    iface.ipv4_addresses,
                    iface.ipv6_addresses
                );
            }
            // Store network state in database
            if let Ok(interfaces_json) = serde_json::to_value(&network_state.interfaces) {
                let _ = db::update_agent_network_state(
                    &state.db,
                    agent_id,
                    &interfaces_json,
                    &network_state.state_hash,
                ).await;
            }
            // Update last_seen
            db::update_agent_status(&state.db, agent_id, "online").await?;
        }

        AgentMessage::Heartbeat => {
            debug!("Received heartbeat from agent {}", agent_id);
            db::update_agent_status(&state.db, agent_id, "online").await?;
        }

        AgentMessage::UpdateResult(result) => {
            info!(
                "Received update result from agent {}: update_id={}, status={}",
                agent_id, result.update_id, result.status
            );

            // Update the agent_updates record
            match result.status.as_str() {
                "completed" => {
                    let _ = db::complete_agent_update(&state.db, result.update_id, result.new_version.as_deref()).await;
                    info!("Agent {} successfully updated to {}", agent_id, result.new_version.unwrap_or_default());
                }
                "failed" | "rolled_back" => {
                    let _ = db::fail_agent_update(&state.db, result.update_id, result.error.as_deref()).await;
                    warn!("Agent {} update failed: {}", agent_id, result.error.unwrap_or_default());
                }
                _ => {
                    debug!("Agent {} update status: {}", agent_id, result.status);
                }
            }
        }

        AgentMessage::Auth(_) => {
            warn!("Received unexpected auth message from already authenticated agent");
        }
    }

    Ok(())
}

/// Hash an API key for storage/lookup
pub fn hash_api_key(api_key: &str) -> String {
    use sha2::{Digest, Sha256};
    let mut hasher = Sha256::new();
    hasher.update(api_key.as_bytes());
    format!("{:x}", hasher.finalize())
}

/// Generate a new API key
pub fn generate_api_key(prefix: &str) -> String {
    use rand::Rng;
    let random_bytes: [u8; 24] = rand::thread_rng().gen();
    let encoded = base64::Engine::encode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, random_bytes);
    format!("{}{}", prefix, encoded)
}

/// Check if a string looks like a site code (WORD-WORD-NUMBER format)
/// Examples: SWIFT-CLOUD-6910, APPLE-GREEN-9145
fn is_site_code_format(s: &str) -> bool {
    let parts: Vec<&str> = s.split('-').collect();
    if parts.len() != 3 {
        return false;
    }
    // First two parts should be alphabetic (words)
    // Third part should be numeric (4 digits)
    parts[0].chars().all(|c| c.is_ascii_alphabetic())
        && parts[1].chars().all(|c| c.is_ascii_alphabetic())
        && parts[2].chars().all(|c| c.is_ascii_digit())
        && parts[2].len() == 4
}


     1→//! GuruRMM Server - RMM Management Server
     2→//!
     3→//! Provides the backend API and WebSocket endpoint for GuruRMM agents.
     4→//! Features:
     5→//! - Agent registration and management
     6→//! - Real-time WebSocket communication with agents
     7→//! - Metrics storage and retrieval
     8→//! - Command execution via agents
     9→//! - Dashboard authentication
    10→
    11→mod api;
    12→mod auth;
    13→mod config;
    14→mod db;
    15→mod updates;
    16→mod ws;
    17→
    18→use std::sync::Arc;
    19→
    20→use anyhow::Result;
    21→use axum::{
    22→    routing::{get, post},
    23→    Router,
    24→};
    25→use sqlx::postgres::PgPoolOptions;
    26→use tokio::sync::RwLock;
    27→use tower_http::cors::{Any, CorsLayer};
    28→use tower_http::trace::TraceLayer;
    29→use tracing::info;
    30→
    31→use crate::config::ServerConfig;
    32→use crate::updates::UpdateManager;
    33→use crate::ws::AgentConnections;
    34→
    35→/// Shared application state
    36→#[derive(Clone)]
    37→pub struct AppState {
    38→    /// Database connection pool
    39→    pub db: sqlx::PgPool,
    40→
    41→    /// Server configuration
    42→    pub config: Arc<ServerConfig>,
    43→
    44→    /// Connected agents (WebSocket connections)
    45→    pub agents: Arc<RwLock<AgentConnections>>,
    46→
    47→    /// Agent update manager
    48→    pub updates: Arc<UpdateManager>,
    49→}
    50→
    51→#[tokio::main]
    52→async fn main() -> Result<()> {
    53→    // Load environment variables from .env file
    54→    dotenvy::dotenv().ok();
    55→
    56→    // Initialize logging
    57→    tracing_subscriber::fmt()
    58→        .with_env_filter(
    59→            tracing_subscriber::EnvFilter::from_default_env()
    60→                .add_directive("gururmm_server=info".parse()?)
    61→                .add_directive("tower_http=debug".parse()?)
    62→                .add_directive("info".parse()?),
    63→        )
    64→        .init();
    65→
    66→    info!("GuruRMM Server starting...");
    67→
    68→    // Load configuration
    69→    let config = ServerConfig::from_env()?;
    70→    info!("Server configuration loaded");
    71→
    72→    // Connect to database
    73→    info!("Connecting to database...");
    74→    let db = PgPoolOptions::new()
    75→        .max_connections(config.database.max_connections)
    76→        .connect(&config.database.url)
    77→        .await?;
    78→    info!("Database connected");
    79→
    80→    // Run migrations
    81→    info!("Running database migrations...");
    82→    sqlx::migrate!("./migrations").run(&db).await?;
    83→    info!("Migrations complete");
    84→
    85→    // Initialize update manager
    86→    let update_manager = UpdateManager::new(
    87→        config.updates.downloads_dir.clone(),
    88→        config.updates.downloads_base_url.clone(),
    89→        config.updates.auto_update_enabled,
    90→        config.updates.update_timeout_secs,
    91→    );
    92→
    93→    // Initial scan for available versions
    94→    info!("Scanning for available agent versions...");
    95→    if let Err(e) = update_manager.scan_versions().await {
    96→        tracing::warn!("Failed to scan agent versions: {} (continuing without auto-update)", e);
    97→    }
    98→
    99→    let update_manager = Arc::new(update_manager);
   100→
   101→    // Spawn background scanner (handle is intentionally not awaited - runs until server shutdown)
   102→    let _scanner_handle = update_manager.spawn_scanner(config.updates.scan_interval_secs);
   103→    info!(
   104→        "Auto-update: {} (scan interval: {}s)",
   105→        if config.updates.auto_update_enabled { "enabled" } else { "disabled" },
   106→        config.updates.scan_interval_secs
   107→    );
   108→
   109→    // Create shared state
   110→    let state = AppState {
   111→        db,
   112→        config: Arc::new(config.clone()),
   113→        agents: Arc::new(RwLock::new(AgentConnections::new())),
   114→        updates: update_manager,
   115→    };
   116→
   117→    // Build router
   118→    let app = build_router(state);
   119→
   120→    // Start server
   121→    let addr = format!("{}:{}", config.server.host, config.server.port);
   122→    info!("Starting server on {}", addr);
   123→
   124→    let listener = tokio::net::TcpListener::bind(&addr).await?;
   125→    axum::serve(listener, app).await?;
   126→
   127→    Ok(())
   128→}
   129→
   130→/// Build the application router
   131→fn build_router(state: AppState) -> Router {
   132→    // CORS configuration (allow dashboard access)
   133→    let cors = CorsLayer::new()
   134→        .allow_origin(Any)
   135→        .allow_methods(Any)
   136→        .allow_headers(Any);
   137→
   138→    Router::new()
   139→        // Health check
   140→        .route("/health", get(health_check))
   141→        // WebSocket endpoint for agents
   142→        .route("/ws", get(ws::ws_handler))
   143→        // API routes
   144→        .nest("/api", api::routes())
   145→        // Middleware
   146→        .layer(TraceLayer::new_for_http())
   147→        .layer(cors)
   148→        // State
   149→        .with_state(state)
   150→}
   151→
   152→/// Health check endpoint
   153→async fn health_check() -> &'static str {
   154→    "OK"
   155→}
   156→

<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 Server - RMM Management Server
//!
//! Provides the backend API and WebSocket endpoint for GuruRMM agents.
//! Features:
//! - Agent registration and management
//! - Real-time WebSocket communication with agents
//! - Metrics storage and retrieval
//! - Command execution via agents
//! - Dashboard authentication

mod api;
mod auth;
mod config;
mod db;
mod updates;
mod ws;

use std::sync::Arc;

use anyhow::Result;
use axum::{
    routing::{get, post},
    Router,
};
use sqlx::postgres::PgPoolOptions;
use tokio::sync::RwLock;
use tower_http::cors::{Any, CorsLayer};
use tower_http::trace::TraceLayer;
use tracing::info;

use crate::config::ServerConfig;
use crate::updates::UpdateManager;
use crate::ws::AgentConnections;

/// Shared application state
#[derive(Clone)]
pub struct AppState {
    /// Database connection pool
    pub db: sqlx::PgPool,

    /// Server configuration
    pub config: Arc<ServerConfig>,

    /// Connected agents (WebSocket connections)
    pub agents: Arc<RwLock<AgentConnections>>,

    /// Agent update manager
    pub updates: Arc<UpdateManager>,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Load environment variables from .env file
    dotenvy::dotenv().ok();

    // Initialize logging
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive("gururmm_server=info".parse()?)
                .add_directive("tower_http=debug".parse()?)
                .add_directive("info".parse()?),
        )
        .init();

    info!("GuruRMM Server starting...");

    // Load configuration
    let config = ServerConfig::from_env()?;
    info!("Server configuration loaded");

    // Connect to database
    info!("Connecting to database...");
    let db = PgPoolOptions::new()
        .max_connections(config.database.max_connections)
        .connect(&config.database.url)
        .await?;
    info!("Database connected");

    // Run migrations
    info!("Running database migrations...");
    sqlx::migrate!("./migrations").run(&db).await?;
    info!("Migrations complete");

    // Initialize update manager
    let update_manager = UpdateManager::new(
        config.updates.downloads_dir.clone(),
        config.updates.downloads_base_url.clone(),
        config.updates.auto_update_enabled,
        config.updates.update_timeout_secs,
    );

    // Initial scan for available versions
    info!("Scanning for available agent versions...");
    if let Err(e) = update_manager.scan_versions().await {
        tracing::warn!("Failed to scan agent versions: {} (continuing without auto-update)", e);
    }

    let update_manager = Arc::new(update_manager);

    // Spawn background scanner (handle is intentionally not awaited - runs until server shutdown)
    let _scanner_handle = update_manager.spawn_scanner(config.updates.scan_interval_secs);
    info!(
        "Auto-update: {} (scan interval: {}s)",
        if config.updates.auto_update_enabled { "enabled" } else { "disabled" },
        config.updates.scan_interval_secs
    );

    // Create shared state
    let state = AppState {
        db,
        config: Arc::new(config.clone()),
        agents: Arc::new(RwLock::new(AgentConnections::new())),
        updates: update_manager,
    };

    // Build router
    let app = build_router(state);

    // Start server
    let addr = format!("{}:{}", config.server.host, config.server.port);
    info!("Starting server on {}", addr);

    let listener = tokio::net::TcpListener::bind(&addr).await?;
    axum::serve(listener, app).await?;

    Ok(())
}

/// Build the application router
fn build_router(state: AppState) -> Router {
    // CORS configuration (allow dashboard access)
    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods(Any)
        .allow_headers(Any);

    Router::new()
        // Health check
        .route("/health", get(health_check))
        // WebSocket endpoint for agents
        .route("/ws", get(ws::ws_handler))
        // API routes
        .nest("/api", api::routes())
        // Middleware
        .layer(TraceLayer::new_for_http())
        .layer(cors)
        // State
        .with_state(state)
}

/// Health check endpoint
async fn health_check() -> &'static str {
    "OK"
}


Now let me examine the specific API route handlers and database functions related to commands and metrics:


     1→//! Commands API endpoints
     2→//!
     3→//! All command endpoints require JWT authentication.
     4→//! Commands are logged with the user who issued them.
     5→
     6→use axum::{
     7→    extract::{Path, Query, State},
     8→    http::StatusCode,
     9→    Json,
    10→};
    11→use serde::{Deserialize, Serialize};
    12→use tracing::info;
    13→use uuid::Uuid;
    14→
    15→use crate::auth::AuthUser;
    16→use crate::db::{self, Command};
    17→use crate::ws::{CommandPayload, ServerMessage};
    18→use crate::AppState;
    19→
    20→/// Request to send a command to an agent
    21→#[derive(Debug, Deserialize)]
    22→pub struct SendCommandRequest {
    23→    /// Command type (shell, powershell, python, script)
    24→    pub command_type: String,
    25→
    26→    /// Command text to execute
    27→    pub command: String,
    28→
    29→    /// Timeout in seconds (optional, default 300)
    30→    pub timeout_seconds: Option<u64>,
    31→
    32→    /// Run as elevated/admin (optional, default false)
    33→    pub elevated: Option<bool>,
    34→}
    35→
    36→/// Response after sending a command
    37→#[derive(Debug, Serialize)]
    38→pub struct SendCommandResponse {
    39→    pub command_id: Uuid,
    40→    pub status: String,
    41→    pub message: String,
    42→}
    43→
    44→/// Query parameters for listing commands
    45→#[derive(Debug, Deserialize)]
    46→pub struct CommandsQuery {
    47→    pub limit: Option<i64>,
    48→}
    49→
    50→/// Send a command to an agent
    51→///
    52→/// Requires JWT authentication. The user ID is recorded with the command.
    53→pub async fn send_command(
    54→    State(state): State<AppState>,
    55→    auth: AuthUser,
    56→    Path(agent_id): Path<Uuid>,
    57→    Json(req): Json<SendCommandRequest>,
    58→) -> Result<Json<SendCommandResponse>, (StatusCode, String)> {
    59→    // Log the command request with user info
    60→    info!(
    61→        "Command request: user={} agent={} type={} command={}",
    62→        auth.user_id, agent_id, req.command_type,
    63→        if req.command.len() > 50 { &req.command[..50] } else { &req.command }
    64→    );
    65→
    66→    // Verify agent exists
    67→    let _agent = db::get_agent_by_id(&state.db, agent_id)
    68→        .await
    69→        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
    70→        .ok_or((StatusCode::NOT_FOUND, "Agent not found".to_string()))?;
    71→
    72→    // Create command record with authenticated user
    73→    let create = db::CreateCommand {
    74→        agent_id,
    75→        command_type: req.command_type.clone(),
    76→        command_text: req.command.clone(),
    77→        created_by: Some(auth.user_id),
    78→    };
    79→
    80→    let command = db::create_command(&state.db, create)
    81→        .await
    82→        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
    83→
    84→    // Check if agent is connected
    85→    let agents = state.agents.read().await;
    86→    if agents.is_connected(&agent_id) {
    87→        // Send command via WebSocket
    88→        let cmd_msg = ServerMessage::Command(CommandPayload {
    89→            id: command.id,
    90→            command_type: req.command_type,
    91→            command: req.command,
    92→            timeout_seconds: req.timeout_seconds,
    93→            elevated: req.elevated.unwrap_or(false),
    94→        });
    95→
    96→        if agents.send_to(&agent_id, cmd_msg).await {
    97→            // Mark as running
    98→            let _ = db::mark_command_running(&state.db, command.id).await;
    99→
   100→            return Ok(Json(SendCommandResponse {
   101→                command_id: command.id,
   102→                status: "running".to_string(),
   103→                message: "Command sent to agent".to_string(),
   104→            }));
   105→        }
   106→    }
   107→
   108→    // Agent not connected or send failed - command is queued
   109→    Ok(Json(SendCommandResponse {
   110→        command_id: command.id,
   111→        status: "pending".to_string(),
   112→        message: "Agent is offline. Command queued for execution when agent reconnects."
   113→            .to_string(),
   114→    }))
   115→}
   116→
   117→/// List recent commands
   118→///
   119→/// Requires JWT authentication.
   120→pub async fn list_commands(
   121→    State(state): State<AppState>,
   122→    _auth: AuthUser,
   123→    Query(query): Query<CommandsQuery>,
   124→) -> Result<Json<Vec<Command>>, (StatusCode, String)> {
   125→    let limit = query.limit.unwrap_or(50).min(500);
   126→
   127→    let commands = db::get_recent_commands(&state.db, limit)
   128→        .await
   129→        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
   130→
   131→    Ok(Json(commands))
   132→}
   133→
   134→/// Get a specific command by ID
   135→///
   136→/// Requires JWT authentication.
   137→pub async fn get_command(
   138→    State(state): State<AppState>,
   139→    _auth: AuthUser,
   140→    Path(id): Path<Uuid>,
   141→) -> Result<Json<Command>, (StatusCode, String)> {
   142→    let command = db::get_command_by_id(&state.db, id)
   143→        .await
   144→        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
   145→        .ok_or((StatusCode::NOT_FOUND, "Command not found".to_string()))?;
   146→
   147→    Ok(Json(command))
   148→}
   149→

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

//! Commands API endpoints
//!
//! All command endpoints require JWT authentication.
//! Commands are logged with the user who issued them.

use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    Json,
};
use serde::{Deserialize, Serialize};
use tracing::info;
use uuid::Uuid;

use crate::auth::AuthUser;
use crate::db::{self, Command};
use crate::ws::{CommandPayload, ServerMessage};
use crate::AppState;

/// Request to send a command to an agent
#[derive(Debug, Deserialize)]
pub struct SendCommandRequest {
    /// Command type (shell, powershell, python, script)
    pub command_type: String,

    /// Command text to execute
    pub command: String,

    /// Timeout in seconds (optional, default 300)
    pub timeout_seconds: Option<u64>,

    /// Run as elevated/admin (optional, default false)
    pub elevated: Option<bool>,
}

/// Response after sending a command
#[derive(Debug, Serialize)]
pub struct SendCommandResponse {
    pub command_id: Uuid,
    pub status: String,
    pub message: String,
}

/// Query parameters for listing commands
#[derive(Debug, Deserialize)]
pub struct CommandsQuery {
    pub limit: Option<i64>,
}

/// Send a command to an agent
///
/// Requires JWT authentication. The user ID is recorded with the command.
pub async fn send_command(
    State(state): State<AppState>,
    auth: AuthUser,
    Path(agent_id): Path<Uuid>,
    Json(req): Json<SendCommandRequest>,
) -> Result<Json<SendCommandResponse>, (StatusCode, String)> {
    // Log the command request with user info
    info!(
        "Command request: user={} agent={} type={} command={}",
        auth.user_id, agent_id, req.command_type,
        if req.command.len() > 50 { &req.command[..50] } else { &req.command }
    );

    // Verify agent exists
    let _agent = db::get_agent_by_id(&state.db, agent_id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
        .ok_or((StatusCode::NOT_FOUND, "Agent not found".to_string()))?;

    // Create command record with authenticated user
    let create = db::CreateCommand {
        agent_id,
        command_type: req.command_type.clone(),
        command_text: req.command.clone(),
        created_by: Some(auth.user_id),
    };

    let command = db::create_command(&state.db, create)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    // Check if agent is connected
    let agents = state.agents.read().await;
    if agents.is_connected(&agent_id) {
        // Send command via WebSocket
        let cmd_msg = ServerMessage::Command(CommandPayload {
            id: command.id,
            command_type: req.command_type,
            command: req.command,
            timeout_seconds: req.timeout_seconds,
            elevated: req.elevated.unwrap_or(false),
        });

        if agents.send_to(&agent_id, cmd_msg).await {
            // Mark as running
            let _ = db::mark_command_running(&state.db, command.id).await;

            return Ok(Json(SendCommandResponse {
                command_id: command.id,
                status: "running".to_string(),
                message: "Command sent to agent".to_string(),
            }));
        }
    }

    // Agent not connected or send failed - command is queued
    Ok(Json(SendCommandResponse {
        command_id: command.id,
        status: "pending".to_string(),
        message: "Agent is offline. Command queued for execution when agent reconnects."
            .to_string(),
    }))
}

/// List recent commands
///
/// Requires JWT authentication.
pub async fn list_commands(
    State(state): State<AppState>,
    _auth: AuthUser,
    Query(query): Query<CommandsQuery>,
) -> Result<Json<Vec<Command>>, (StatusCode, String)> {
    let limit = query.limit.unwrap_or(50).min(500);

    let commands = db::get_recent_commands(&state.db, limit)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(commands))
}

/// Get a specific command by ID
///
/// Requires JWT authentication.
pub async fn get_command(
    State(state): State<AppState>,
    _auth: AuthUser,
    Path(id): Path<Uuid>,
) -> Result<Json<Command>, (StatusCode, String)> {
    let command = db::get_command_by_id(&state.db, id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
        .ok_or((StatusCode::NOT_FOUND, "Command not found".to_string()))?;

    Ok(Json(command))
}


     1→//! Metrics API endpoints
     2→
     3→use axum::{
     4→    extract::{Path, Query, State},
     5→    http::StatusCode,
     6→    Json,
     7→};
     8→use chrono::{DateTime, Duration, Utc};
     9→use serde::Deserialize;
    10→use uuid::Uuid;
    11→
    12→use crate::db::{self, Metrics, MetricsSummary};
    13→use crate::AppState;
    14→
    15→/// Query parameters for metrics
    16→#[derive(Debug, Deserialize)]
    17→pub struct MetricsQuery {
    18→    /// Number of records to return (default: 100)
    19→    pub limit: Option<i64>,
    20→
    21→    /// Start time for range query
    22→    pub start: Option<DateTime<Utc>>,
    23→
    24→    /// End time for range query
    25→    pub end: Option<DateTime<Utc>>,
    26→}
    27→
    28→/// Get metrics for a specific agent
    29→pub async fn get_agent_metrics(
    30→    State(state): State<AppState>,
    31→    Path(id): Path<Uuid>,
    32→    Query(query): Query<MetricsQuery>,
    33→) -> Result<Json<Vec<Metrics>>, (StatusCode, String)> {
    34→    // First verify the agent exists
    35→    let agent = db::get_agent_by_id(&state.db, id)
    36→        .await
    37→        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
    38→        .ok_or((StatusCode::NOT_FOUND, "Agent not found".to_string()))?;
    39→
    40→    let metrics = if let (Some(start), Some(end)) = (query.start, query.end) {
    41→        // Range query
    42→        db::get_agent_metrics_range(&state.db, id, start, end)
    43→            .await
    44→            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
    45→    } else {
    46→        // Simple limit query
    47→        let limit = query.limit.unwrap_or(100).min(1000); // Cap at 1000
    48→        db::get_agent_metrics(&state.db, id, limit)
    49→            .await
    50→            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
    51→    };
    52→
    53→    Ok(Json(metrics))
    54→}
    55→
    56→/// Get summary metrics across all agents
    57→pub async fn get_summary(
    58→    State(state): State<AppState>,
    59→) -> Result<Json<MetricsSummary>, (StatusCode, String)> {
    60→    let summary = db::get_metrics_summary(&state.db)
    61→        .await
    62→        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
    63→
    64→    Ok(Json(summary))
    65→}
    66→

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

//! Metrics API endpoints

use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    Json,
};
use chrono::{DateTime, Duration, Utc};
use serde::Deserialize;
use uuid::Uuid;

use crate::db::{self, Metrics, MetricsSummary};
use crate::AppState;

/// Query parameters for metrics
#[derive(Debug, Deserialize)]
pub struct MetricsQuery {
    /// Number of records to return (default: 100)
    pub limit: Option<i64>,

    /// Start time for range query
    pub start: Option<DateTime<Utc>>,

    /// End time for range query
    pub end: Option<DateTime<Utc>>,
}

/// Get metrics for a specific agent
pub async fn get_agent_metrics(
    State(state): State<AppState>,
    Path(id): Path<Uuid>,
    Query(query): Query<MetricsQuery>,
) -> Result<Json<Vec<Metrics>>, (StatusCode, String)> {
    // First verify the agent exists
    let agent = db::get_agent_by_id(&state.db, id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
        .ok_or((StatusCode::NOT_FOUND, "Agent not found".to_string()))?;

    let metrics = if let (Some(start), Some(end)) = (query.start, query.end) {
        // Range query
        db::get_agent_metrics_range(&state.db, id, start, end)
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
    } else {
        // Simple limit query
        let limit = query.limit.unwrap_or(100).min(1000); // Cap at 1000
        db::get_agent_metrics(&state.db, id, limit)
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
    };

    Ok(Json(metrics))
}

/// Get summary metrics across all agents
pub async fn get_summary(
    State(state): State<AppState>,
) -> Result<Json<MetricsSummary>, (StatusCode, String)> {
    let summary = db::get_metrics_summary(&state.db)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(summary))
}