Design a stackable policy system for GuruRMM based on these requirements and codebase findings:

**Requirements:**
- Policies assignable at: Client, Site, and Agent levels
- Stack order: Client → Site → Agent (last wins for conflicts)
- Settings include: metrics collection (temps, intervals), alert thresholds (CPU > X% for Y minutes), system tray options, auto-update settings
- Alerts: Dashboard-only initially (no email/webhook)
- Delivery: Via existing ConfigUpdate WebSocket message

**Codebase Context:**
- Database: PostgreSQL with existing tables: clients, sites, agents, metrics, agent_state
- Server: Axum web framework, WebSocket in ws/mod.rs, ServerMessage::ConfigUpdate exists but unused
- Agent: Has ConfigUpdatePayload structure, TrayPolicy already shows policy pattern
- Agent config includes: MetricsConfig (interval, collect_*), WatchdogConfig (services, processes)

**Key Files:**
- server/src/db/ - Database modules (agents.rs, clients.rs, sites.rs, metrics.rs)
- server/src/api/mod.rs - API routes
- server/src/ws/mod.rs - WebSocket handling, ServerMessage enum
- agent/src/config.rs - Agent configuration
- agent/src/transport/mod.rs - Message types including ConfigUpdatePayload

**Design Requirements:**
1. Database schema for policies, policy_assignments, alerts
2. Policy structure that covers: metrics settings, thresholds, tray settings, update settings
3. Server-side policy merge logic (evaluate effective policy for an agent)
4. API endpoints for CRUD and deployment
5. Alert evaluation approach (when to check thresholds)
6. Agent-side policy application

Provide a detailed implementation plan with specific structs, tables, and code patterns.


     1→//! Database models and queries
     2→//!
     3→//! Provides database access for clients, sites, agents, metrics, commands, users, and updates.
     4→
     5→pub mod agents;
     6→pub mod clients;
     7→pub mod commands;
     8→pub mod metrics;
     9→pub mod sites;
    10→pub mod updates;
    11→pub mod users;
    12→
    13→pub use agents::*;
    14→pub use clients::*;
    15→pub use commands::*;
    16→pub use metrics::*;
    17→pub use sites::*;
    18→pub use updates::*;
    19→pub use users::*;
    20→

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

//! Database models and queries
//!
//! Provides database access for clients, sites, agents, metrics, commands, users, and updates.

pub mod agents;
pub mod clients;
pub mod commands;
pub mod metrics;
pub mod sites;
pub mod updates;
pub mod users;

pub use agents::*;
pub use clients::*;
pub use commands::*;
pub use metrics::*;
pub use sites::*;
pub use updates::*;
pub use users::*;


     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→//! Agent configuration handling
     2→//!
     3→//! Configuration is loaded from a TOML file (default: agent.toml).
     4→//! The config file defines server connection, metrics collection,
     5→//! and watchdog settings.
     6→
     7→use anyhow::{Context, Result};
     8→use serde::{Deserialize, Serialize};
     9→use std::path::Path;
    10→
    11→/// Root configuration structure
    12→#[derive(Debug, Clone, Serialize, Deserialize)]
    13→pub struct AgentConfig {
    14→    /// Server connection settings
    15→    pub server: ServerConfig,
    16→
    17→    /// Metrics collection settings
    18→    #[serde(default)]
    19→    pub metrics: MetricsConfig,
    20→
    21→    /// Watchdog settings for monitoring services/processes
    22→    #[serde(default)]
    23→    pub watchdog: WatchdogConfig,
    24→}
    25→
    26→/// Server connection configuration
    27→#[derive(Debug, Clone, Serialize, Deserialize)]
    28→pub struct ServerConfig {
    29→    /// WebSocket URL for the GuruRMM server (e.g., wss://rmm.example.com/ws)
    30→    pub url: String,
    31→
    32→    /// API key for authentication (obtained from server during registration)
    33→    pub api_key: String,
    34→
    35→    /// Optional custom hostname to report (defaults to system hostname)
    36→    pub hostname_override: Option<String>,
    37→}
    38→
    39→/// Metrics collection configuration
    40→#[derive(Debug, Clone, Serialize, Deserialize)]
    41→pub struct MetricsConfig {
    42→    /// Interval in seconds between metrics collection (default: 60)
    43→    #[serde(default = "default_metrics_interval")]
    44→    pub interval_seconds: u64,
    45→
    46→    /// Whether to collect CPU metrics
    47→    #[serde(default = "default_true")]
    48→    pub collect_cpu: bool,
    49→
    50→    /// Whether to collect memory metrics
    51→    #[serde(default = "default_true")]
    52→    pub collect_memory: bool,
    53→
    54→    /// Whether to collect disk metrics
    55→    #[serde(default = "default_true")]
    56→    pub collect_disk: bool,
    57→
    58→    /// Whether to collect network metrics
    59→    #[serde(default = "default_true")]
    60→    pub collect_network: bool,
    61→}
    62→
    63→impl Default for MetricsConfig {
    64→    fn default() -> Self {
    65→        Self {
    66→            interval_seconds: 60,
    67→            collect_cpu: true,
    68→            collect_memory: true,
    69→            collect_disk: true,
    70→            collect_network: true,
    71→        }
    72→    }
    73→}
    74→
    75→/// Watchdog configuration for service/process monitoring
    76→#[derive(Debug, Clone, Serialize, Deserialize)]
    77→pub struct WatchdogConfig {
    78→    /// Enable/disable watchdog functionality
    79→    #[serde(default)]
    80→    pub enabled: bool,
    81→
    82→    /// Interval in seconds between watchdog checks (default: 30)
    83→    #[serde(default = "default_watchdog_interval")]
    84→    pub check_interval_seconds: u64,
    85→
    86→    /// List of Windows/systemd services to monitor
    87→    #[serde(default)]
    88→    pub services: Vec<ServiceWatch>,
    89→
    90→    /// List of processes to monitor
    91→    #[serde(default)]
    92→    pub processes: Vec<ProcessWatch>,
    93→}
    94→
    95→impl Default for WatchdogConfig {
    96→    fn default() -> Self {
    97→        Self {
    98→            enabled: false,
    99→            check_interval_seconds: 30,
   100→            services: Vec::new(),
   101→            processes: Vec::new(),
   102→        }
   103→    }
   104→}
   105→
   106→/// Configuration for monitoring a service
   107→#[derive(Debug, Clone, Serialize, Deserialize)]
   108→pub struct ServiceWatch {
   109→    /// Service name (e.g., "CagService" for Datto RMM, "Syncro" for Syncro)
   110→    pub name: String,
   111→
   112→    /// Action to take when service is stopped
   113→    #[serde(default)]
   114→    pub action: WatchAction,
   115→
   116→    /// Maximum number of restart attempts before alerting (default: 3)
   117→    #[serde(default = "default_max_restarts")]
   118→    pub max_restarts: u32,
   119→
   120→    /// Cooldown period in seconds between restart attempts
   121→    #[serde(default = "default_restart_cooldown")]
   122→    pub restart_cooldown_seconds: u64,
   123→}
   124→
   125→/// Configuration for monitoring a process
   126→#[derive(Debug, Clone, Serialize, Deserialize)]
   127→pub struct ProcessWatch {
   128→    /// Process name (e.g., "AEM.exe")
   129→    pub name: String,
   130→
   131→    /// Action to take when process is not found
   132→    #[serde(default)]
   133→    pub action: WatchAction,
   134→
   135→    /// Optional path to executable to start if process is not running
   136→    pub start_command: Option<String>,
   137→}
   138→
   139→/// Action to take when a watched service/process is down
   140→#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
   141→#[serde(rename_all = "lowercase")]
   142→pub enum WatchAction {
   143→    /// Only send an alert to the server
   144→    #[default]
   145→    Alert,
   146→
   147→    /// Attempt to restart the service/process
   148→    Restart,
   149→
   150→    /// Ignore (for temporary disable without removing config)
   151→    Ignore,
   152→}
   153→
   154→// Default value functions for serde
   155→fn default_metrics_interval() -> u64 {
   156→    60
   157→}
   158→
   159→fn default_watchdog_interval() -> u64 {
   160→    30
   161→}
   162→
   163→fn default_max_restarts() -> u32 {
   164→    3
   165→}
   166→
   167→fn default_restart_cooldown() -> u64 {
   168→    60
   169→}
   170→
   171→fn default_true() -> bool {
   172→    true
   173→}
   174→
   175→impl AgentConfig {
   176→    /// Load configuration from a TOML file
   177→    pub fn load(path: &Path) -> Result<Self> {
   178→        let content = std::fs::read_to_string(path)
   179→            .with_context(|| format!("Failed to read config file: {:?}", path))?;
   180→
   181→        let config: Self = toml::from_str(&content)
   182→            .with_context(|| format!("Failed to parse config file: {:?}", path))?;
   183→
   184→        config.validate()?;
   185→        Ok(config)
   186→    }
   187→
   188→    /// Validate the configuration
   189→    fn validate(&self) -> Result<()> {
   190→        // Validate server URL
   191→        if self.server.url.is_empty() {
   192→            anyhow::bail!("Server URL cannot be empty");
   193→        }
   194→
   195→        if !self.server.url.starts_with("ws://") && !self.server.url.starts_with("wss://") {
   196→            anyhow::bail!("Server URL must start with ws:// or wss://");
   197→        }
   198→
   199→        // Validate API key
   200→        if self.server.api_key.is_empty() {
   201→            anyhow::bail!("API key cannot be empty");
   202→        }
   203→
   204→        // Validate intervals
   205→        if self.metrics.interval_seconds < 10 {
   206→            anyhow::bail!("Metrics interval must be at least 10 seconds");
   207→        }
   208→
   209→        if self.watchdog.check_interval_seconds < 5 {
   210→            anyhow::bail!("Watchdog check interval must be at least 5 seconds");
   211→        }
   212→
   213→        Ok(())
   214→    }
   215→
   216→    /// Generate a sample configuration
   217→    pub fn sample() -> Self {
   218→        Self {
   219→            server: ServerConfig {
   220→                url: "wss://rmm-api.azcomputerguru.com/ws".to_string(),
   221→                api_key: "your-api-key-here".to_string(),
   222→                hostname_override: None,
   223→            },
   224→            metrics: MetricsConfig::default(),
   225→            watchdog: WatchdogConfig {
   226→                enabled: true,
   227→                check_interval_seconds: 30,
   228→                services: vec![
   229→                    ServiceWatch {
   230→                        name: "CagService".to_string(), // Datto RMM
   231→                        action: WatchAction::Restart,
   232→                        max_restarts: 3,
   233→                        restart_cooldown_seconds: 60,
   234→                    },
   235→                    ServiceWatch {
   236→                        name: "Syncro".to_string(),
   237→                        action: WatchAction::Restart,
   238→                        max_restarts: 3,
   239→                        restart_cooldown_seconds: 60,
   240→                    },
   241→                ],
   242→                processes: vec![ProcessWatch {
   243→                    name: "AEM.exe".to_string(), // Datto AEM
   244→                    action: WatchAction::Alert,
   245→                    start_command: None,
   246→                }],
   247→            },
   248→        }
   249→    }
   250→
   251→    /// Get the hostname to report to the server
   252→    pub fn get_hostname(&self) -> String {
   253→        self.server
   254→            .hostname_override
   255→            .clone()
   256→            .unwrap_or_else(|| hostname::get().map(|h| h.to_string_lossy().to_string()).unwrap_or_else(|_| "unknown".to_string()))
   257→    }
   258→}
   259→
   260→#[cfg(test)]
   261→mod tests {
   262→    use super::*;
   263→
   264→    #[test]
   265→    fn test_sample_config_is_valid_structure() {
   266→        let sample = AgentConfig::sample();
   267→        // Sample uses placeholder values, so it won't pass full validation
   268→        // but the structure should be correct
   269→        assert!(!sample.server.url.is_empty());
   270→        assert!(!sample.server.api_key.is_empty());
   271→        assert!(sample.watchdog.enabled);
   272→        assert!(!sample.watchdog.services.is_empty());
   273→    }
   274→
   275→    #[test]
   276→    fn test_default_metrics_config() {
   277→        let config = MetricsConfig::default();
   278→        assert_eq!(config.interval_seconds, 60);
   279→        assert!(config.collect_cpu);
   280→        assert!(config.collect_memory);
   281→        assert!(config.collect_disk);
   282→        assert!(config.collect_network);
   283→    }
   284→
   285→    #[test]
   286→    fn test_watch_action_default() {
   287→        let action = WatchAction::default();
   288→        assert_eq!(action, WatchAction::Alert);
   289→    }
   290→}
   291→

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

//! Agent configuration handling
//!
//! Configuration is loaded from a TOML file (default: agent.toml).
//! The config file defines server connection, metrics collection,
//! and watchdog settings.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;

/// Root configuration structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentConfig {
    /// Server connection settings
    pub server: ServerConfig,

    /// Metrics collection settings
    #[serde(default)]
    pub metrics: MetricsConfig,

    /// Watchdog settings for monitoring services/processes
    #[serde(default)]
    pub watchdog: WatchdogConfig,
}

/// Server connection configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServerConfig {
    /// WebSocket URL for the GuruRMM server (e.g., wss://rmm.example.com/ws)
    pub url: String,

    /// API key for authentication (obtained from server during registration)
    pub api_key: String,

    /// Optional custom hostname to report (defaults to system hostname)
    pub hostname_override: Option<String>,
}

/// Metrics collection configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MetricsConfig {
    /// Interval in seconds between metrics collection (default: 60)
    #[serde(default = "default_metrics_interval")]
    pub interval_seconds: u64,

    /// Whether to collect CPU metrics
    #[serde(default = "default_true")]
    pub collect_cpu: bool,

    /// Whether to collect memory metrics
    #[serde(default = "default_true")]
    pub collect_memory: bool,

    /// Whether to collect disk metrics
    #[serde(default = "default_true")]
    pub collect_disk: bool,

    /// Whether to collect network metrics
    #[serde(default = "default_true")]
    pub collect_network: bool,
}

impl Default for MetricsConfig {
    fn default() -> Self {
        Self {
            interval_seconds: 60,
            collect_cpu: true,
            collect_memory: true,
            collect_disk: true,
            collect_network: true,
        }
    }
}

/// Watchdog configuration for service/process monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WatchdogConfig {
    /// Enable/disable watchdog functionality
    #[serde(default)]
    pub enabled: bool,

    /// Interval in seconds between watchdog checks (default: 30)
    #[serde(default = "default_watchdog_interval")]
    pub check_interval_seconds: u64,

    /// List of Windows/systemd services to monitor
    #[serde(default)]
    pub services: Vec<ServiceWatch>,

    /// List of processes to monitor
    #[serde(default)]
    pub processes: Vec<ProcessWatch>,
}

impl Default for WatchdogConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            check_interval_seconds: 30,
            services: Vec::new(),
            processes: Vec::new(),
        }
    }
}

/// Configuration for monitoring a service
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceWatch {
    /// Service name (e.g., "CagService" for Datto RMM, "Syncro" for Syncro)
    pub name: String,

    /// Action to take when service is stopped
    #[serde(default)]
    pub action: WatchAction,

    /// Maximum number of restart attempts before alerting (default: 3)
    #[serde(default = "default_max_restarts")]
    pub max_restarts: u32,

    /// Cooldown period in seconds between restart attempts
    #[serde(default = "default_restart_cooldown")]
    pub restart_cooldown_seconds: u64,
}

/// Configuration for monitoring a process
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProcessWatch {
    /// Process name (e.g., "AEM.exe")
    pub name: String,

    /// Action to take when process is not found
    #[serde(default)]
    pub action: WatchAction,

    /// Optional path to executable to start if process is not running
    pub start_command: Option<String>,
}

/// Action to take when a watched service/process is down
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum WatchAction {
    /// Only send an alert to the server
    #[default]
    Alert,

    /// Attempt to restart the service/process
    Restart,

    /// Ignore (for temporary disable without removing config)
    Ignore,
}

// Default value functions for serde
fn default_metrics_interval() -> u64 {
    60
}

fn default_watchdog_interval() -> u64 {
    30
}

fn default_max_restarts() -> u32 {
    3
}

fn default_restart_cooldown() -> u64 {
    60
}

fn default_true() -> bool {
    true
}

impl AgentConfig {
    /// Load configuration from a TOML file
    pub fn load(path: &Path) -> Result<Self> {
        let content = std::fs::read_to_string(path)
            .with_context(|| format!("Failed to read config file: {:?}", path))?;

        let config: Self = toml::from_str(&content)
            .with_context(|| format!("Failed to parse config file: {:?}", path))?;

        config.validate()?;
        Ok(config)
    }

    /// Validate the configuration
    fn validate(&self) -> Result<()> {
        // Validate server URL
        if self.server.url.is_empty() {
            anyhow::bail!("Server URL cannot be empty");
        }

        if !self.server.url.starts_with("ws://") && !self.server.url.starts_with("wss://") {
            anyhow::bail!("Server URL must start with ws:// or wss://");
        }

        // Validate API key
        if self.server.api_key.is_empty() {
            anyhow::bail!("API key cannot be empty");
        }

        // Validate intervals
        if self.metrics.interval_seconds < 10 {
            anyhow::bail!("Metrics interval must be at least 10 seconds");
        }

        if self.watchdog.check_interval_seconds < 5 {
            anyhow::bail!("Watchdog check interval must be at least 5 seconds");
        }

        Ok(())
    }

    /// Generate a sample configuration
    pub fn sample() -> Self {
        Self {
            server: ServerConfig {
                url: "wss://rmm-api.azcomputerguru.com/ws".to_string(),
                api_key: "your-api-key-here".to_string(),
                hostname_override: None,
            },
            metrics: MetricsConfig::default(),
            watchdog: WatchdogConfig {
                enabled: true,
                check_interval_seconds: 30,
                services: vec![
                    ServiceWatch {
                        name: "CagService".to_string(), // Datto RMM
                        action: WatchAction::Restart,
                        max_restarts: 3,
                        restart_cooldown_seconds: 60,
                    },
                    ServiceWatch {
                        name: "Syncro".to_string(),
                        action: WatchAction::Restart,
                        max_restarts: 3,
                        restart_cooldown_seconds: 60,
                    },
                ],
                processes: vec![ProcessWatch {
                    name: "AEM.exe".to_string(), // Datto AEM
                    action: WatchAction::Alert,
                    start_command: None,
                }],
            },
        }
    }

    /// Get the hostname to report to the server
    pub fn get_hostname(&self) -> String {
        self.server
            .hostname_override
            .clone()
            .unwrap_or_else(|| hostname::get().map(|h| h.to_string_lossy().to_string()).unwrap_or_else(|_| "unknown".to_string()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sample_config_is_valid_structure() {
        let sample = AgentConfig::sample();
        // Sample uses placeholder values, so it won't pass full validation
        // but the structure should be correct
        assert!(!sample.server.url.is_empty());
        assert!(!sample.server.api_key.is_empty());
        assert!(sample.watchdog.enabled);
        assert!(!sample.watchdog.services.is_empty());
    }

    #[test]
    fn test_default_metrics_config() {
        let config = MetricsConfig::default();
        assert_eq!(config.interval_seconds, 60);
        assert!(config.collect_cpu);
        assert!(config.collect_memory);
        assert!(config.collect_disk);
        assert!(config.collect_network);
    }

    #[test]
    fn test_watch_action_default() {
        let action = WatchAction::default();
        assert_eq!(action, WatchAction::Alert);
    }
}


     1→//! Transport layer for agent-server communication
     2→//!
     3→//! Handles WebSocket connection to the GuruRMM server with:
     4→//! - Auto-reconnection on disconnect
     5→//! - Authentication via API key
     6→//! - Sending metrics and receiving commands
     7→//! - Heartbeat to maintain connection
     8→
     9→mod websocket;
    10→
    11→pub use websocket::WebSocketClient;
    12→
    13→use serde::{Deserialize, Serialize};
    14→use uuid::Uuid;
    15→
    16→/// Messages sent from agent to server
    17→#[derive(Debug, Clone, Serialize, Deserialize)]
    18→#[serde(tag = "type", content = "payload")]
    19→#[serde(rename_all = "snake_case")]
    20→pub enum AgentMessage {
    21→    /// Authentication message (sent on connect)
    22→    Auth(AuthPayload),
    23→
    24→    /// Metrics report
    25→    Metrics(crate::metrics::SystemMetrics),
    26→
    27→    /// Network state update (sent on connect and when interfaces change)
    28→    NetworkState(crate::metrics::NetworkState),
    29→
    30→    /// Command execution result
    31→    CommandResult(CommandResultPayload),
    32→
    33→    /// Watchdog event (service stopped, restarted, etc.)
    34→    WatchdogEvent(WatchdogEventPayload),
    35→
    36→    /// Update result (success, failure, rollback)
    37→    UpdateResult(UpdateResultPayload),
    38→
    39→    /// Heartbeat to keep connection alive
    40→    Heartbeat,
    41→}
    42→
    43→/// Authentication payload
    44→#[derive(Debug, Clone, Serialize, Deserialize)]
    45→pub struct AuthPayload {
    46→    /// API key for this agent (or site)
    47→    pub api_key: String,
    48→
    49→    /// Unique device identifier (hardware-derived)
    50→    pub device_id: String,
    51→
    52→    /// Hostname of this machine
    53→    pub hostname: String,
    54→
    55→    /// Operating system type
    56→    pub os_type: String,
    57→
    58→    /// Operating system version
    59→    pub os_version: String,
    60→
    61→    /// Agent version
    62→    pub agent_version: String,
    63→
    64→    /// Architecture (amd64, arm64, etc.)
    65→    #[serde(default = "default_arch")]
    66→    pub architecture: String,
    67→
    68→    /// Previous version if reconnecting after update
    69→    #[serde(skip_serializing_if = "Option::is_none")]
    70→    pub previous_version: Option<String>,
    71→
    72→    /// Update ID if reconnecting after update
    73→    #[serde(skip_serializing_if = "Option::is_none")]
    74→    pub pending_update_id: Option<Uuid>,
    75→}
    76→
    77→fn default_arch() -> String {
    78→    #[cfg(target_arch = "x86_64")]
    79→    { "amd64".to_string() }
    80→    #[cfg(target_arch = "aarch64")]
    81→    { "arm64".to_string() }
    82→    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
    83→    { "unknown".to_string() }
    84→}
    85→
    86→/// Command execution result payload
    87→#[derive(Debug, Clone, Serialize, Deserialize)]
    88→pub struct CommandResultPayload {
    89→    /// Command ID (from the server)
    90→    pub command_id: Uuid,
    91→
    92→    /// Exit code (0 = success)
    93→    pub exit_code: i32,
    94→
    95→    /// Standard output
    96→    pub stdout: String,
    97→
    98→    /// Standard error
    99→    pub stderr: String,
   100→
   101→    /// Execution duration in milliseconds
   102→    pub duration_ms: u64,
   103→}
   104→
   105→/// Watchdog event payload
   106→#[derive(Debug, Clone, Serialize, Deserialize)]
   107→pub struct WatchdogEventPayload {
   108→    /// Service or process name
   109→    pub name: String,
   110→
   111→    /// Event type
   112→    pub event: WatchdogEvent,
   113→
   114→    /// Additional details
   115→    pub details: Option<String>,
   116→}
   117→
   118→/// Types of watchdog events
   119→#[derive(Debug, Clone, Serialize, Deserialize)]
   120→#[serde(rename_all = "snake_case")]
   121→pub enum WatchdogEvent {
   122→    /// Service/process was found stopped
   123→    Stopped,
   124→
   125→    /// Service/process was restarted by the agent
   126→    Restarted,
   127→
   128→    /// Restart attempt failed
   129→    RestartFailed,
   130→
   131→    /// Max restart attempts reached
   132→    MaxRestartsReached,
   133→
   134→    /// Service/process recovered on its own
   135→    Recovered,
   136→}
   137→
   138→/// Messages sent from server to agent
   139→#[derive(Debug, Clone, Serialize, Deserialize)]
   140→#[serde(tag = "type", content = "payload")]
   141→#[serde(rename_all = "snake_case")]
   142→pub enum ServerMessage {
   143→    /// Authentication acknowledgment
   144→    AuthAck(AuthAckPayload),
   145→
   146→    /// Command to execute
   147→    Command(CommandPayload),
   148→
   149→    /// Configuration update
   150→    ConfigUpdate(ConfigUpdatePayload),
   151→
   152→    /// Agent update command
   153→    Update(UpdatePayload),
   154→
   155→    /// Tray policy update
   156→    TrayPolicyUpdate(crate::ipc::TrayPolicy),
   157→
   158→    /// Acknowledgment of received message
   159→    Ack { message_id: Option<String> },
   160→
   161→    /// Error message
   162→    Error { code: String, message: String },
   163→}
   164→
   165→/// Authentication acknowledgment payload
   166→#[derive(Debug, Clone, Serialize, Deserialize)]
   167→pub struct AuthAckPayload {
   168→    /// Whether authentication was successful
   169→    pub success: bool,
   170→
   171→    /// Agent ID assigned by server
   172→    pub agent_id: Option<Uuid>,
   173→
   174→    /// Error message if authentication failed
   175→    pub error: Option<String>,
   176→}
   177→
   178→/// Command payload from server
   179→#[derive(Debug, Clone, Serialize, Deserialize)]
   180→pub struct CommandPayload {
   181→    /// Unique command ID
   182→    pub id: Uuid,
   183→
   184→    /// Type of command
   185→    pub command_type: CommandType,
   186→
   187→    /// Command text to execute
   188→    pub command: String,
   189→
   190→    /// Optional timeout in seconds
   191→    pub timeout_seconds: Option<u64>,
   192→
   193→    /// Whether to run as elevated/admin
   194→    pub elevated: bool,
   195→}
   196→
   197→/// Types of commands
   198→#[derive(Debug, Clone, Serialize, Deserialize)]
   199→#[serde(rename_all = "snake_case")]
   200→pub enum CommandType {
   201→    /// Shell command (cmd on Windows, bash on Unix)
   202→    Shell,
   203→
   204→    /// PowerShell command (Windows)
   205→    PowerShell,
   206→
   207→    /// Python script
   208→    Python,
   209→
   210→    /// Raw script (requires interpreter path)
   211→    Script { interpreter: String },
   212→}
   213→
   214→/// Configuration update payload
   215→#[derive(Debug, Clone, Serialize, Deserialize)]
   216→pub struct ConfigUpdatePayload {
   217→    /// New metrics interval (if changed)
   218→    pub metrics_interval_seconds: Option<u64>,
   219→
   220→    /// Updated watchdog config
   221→    pub watchdog: Option<WatchdogConfigUpdate>,
   222→}
   223→
   224→/// Watchdog configuration update
   225→#[derive(Debug, Clone, Serialize, Deserialize)]
   226→pub struct WatchdogConfigUpdate {
   227→    /// Enable/disable watchdog
   228→    pub enabled: Option<bool>,
   229→
   230→    /// Check interval
   231→    pub check_interval_seconds: Option<u64>,
   232→
   233→    // Services and processes would be included here for remote config updates
   234→}
   235→
   236→/// Update command payload from server
   237→#[derive(Debug, Clone, Serialize, Deserialize)]
   238→pub struct UpdatePayload {
   239→    /// Unique update ID for tracking
   240→    pub update_id: Uuid,
   241→
   242→    /// Target version to update to
   243→    pub target_version: String,
   244→
   245→    /// Download URL for the new binary
   246→    pub download_url: String,
   247→
   248→    /// SHA256 checksum of the binary
   249→    pub checksum_sha256: String,
   250→
   251→    /// Whether to force update (skip version check)
   252→    #[serde(default)]
   253→    pub force: bool,
   254→}
   255→
   256→/// Update result payload sent back to server
   257→#[derive(Debug, Clone, Serialize, Deserialize)]
   258→pub struct UpdateResultPayload {
   259→    /// Update ID (from the server)
   260→    pub update_id: Uuid,
   261→
   262→    /// Update status
   263→    pub status: UpdateStatus,
   264→
   265→    /// Old version before update
   266→    pub old_version: String,
   267→
   268→    /// New version after update (if successful)
   269→    pub new_version: Option<String>,
   270→
   271→    /// Error message if failed
   272→    pub error: Option<String>,
   273→}
   274→
   275→/// Update status codes
   276→#[derive(Debug, Clone, Serialize, Deserialize)]
   277→#[serde(rename_all = "snake_case")]
   278→pub enum UpdateStatus {
   279→    /// Update starting
   280→    Starting,
   281→
   282→    /// Downloading new binary
   283→    Downloading,
   284→
   285→    /// Download complete, verifying
   286→    Verifying,
   287→
   288→    /// Installing (replacing binary)
   289→    Installing,
   290→
   291→    /// Restarting service
   292→    Restarting,
   293→
   294→    /// Update completed successfully
   295→    Completed,
   296→
   297→    /// Update failed
   298→    Failed,
   299→
   300→    /// Rolled back to previous version
   301→    RolledBack,
   302→}
   303→

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

//! Transport layer for agent-server communication
//!
//! Handles WebSocket connection to the GuruRMM server with:
//! - Auto-reconnection on disconnect
//! - Authentication via API key
//! - Sending metrics and receiving commands
//! - Heartbeat to maintain connection

mod websocket;

pub use websocket::WebSocketClient;

use serde::{Deserialize, Serialize};
use uuid::Uuid;

/// Messages sent from agent to server
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
#[serde(rename_all = "snake_case")]
pub enum AgentMessage {
    /// Authentication message (sent on connect)
    Auth(AuthPayload),

    /// Metrics report
    Metrics(crate::metrics::SystemMetrics),

    /// Network state update (sent on connect and when interfaces change)
    NetworkState(crate::metrics::NetworkState),

    /// Command execution result
    CommandResult(CommandResultPayload),

    /// Watchdog event (service stopped, restarted, etc.)
    WatchdogEvent(WatchdogEventPayload),

    /// Update result (success, failure, rollback)
    UpdateResult(UpdateResultPayload),

    /// Heartbeat to keep connection alive
    Heartbeat,
}

/// Authentication payload
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthPayload {
    /// API key for this agent (or site)
    pub api_key: String,

    /// Unique device identifier (hardware-derived)
    pub device_id: String,

    /// Hostname of this machine
    pub hostname: String,

    /// Operating system type
    pub os_type: String,

    /// Operating system version
    pub os_version: String,

    /// Agent version
    pub agent_version: String,

    /// Architecture (amd64, arm64, etc.)
    #[serde(default = "default_arch")]
    pub architecture: String,

    /// Previous version if reconnecting after update
    #[serde(skip_serializing_if = "Option::is_none")]
    pub previous_version: Option<String>,

    /// Update ID if reconnecting after update
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pending_update_id: Option<Uuid>,
}

fn default_arch() -> String {
    #[cfg(target_arch = "x86_64")]
    { "amd64".to_string() }
    #[cfg(target_arch = "aarch64")]
    { "arm64".to_string() }
    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
    { "unknown".to_string() }
}

/// Command execution result payload
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommandResultPayload {
    /// Command ID (from the server)
    pub command_id: Uuid,

    /// Exit code (0 = success)
    pub exit_code: i32,

    /// Standard output
    pub stdout: String,

    /// Standard error
    pub stderr: String,

    /// Execution duration in milliseconds
    pub duration_ms: u64,
}

/// Watchdog event payload
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WatchdogEventPayload {
    /// Service or process name
    pub name: String,

    /// Event type
    pub event: WatchdogEvent,

    /// Additional details
    pub details: Option<String>,
}

/// Types of watchdog events
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WatchdogEvent {
    /// Service/process was found stopped
    Stopped,

    /// Service/process was restarted by the agent
    Restarted,

    /// Restart attempt failed
    RestartFailed,

    /// Max restart attempts reached
    MaxRestartsReached,

    /// Service/process recovered on its own
    Recovered,
}

/// Messages sent from server to agent
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
#[serde(rename_all = "snake_case")]
pub enum ServerMessage {
    /// Authentication acknowledgment
    AuthAck(AuthAckPayload),

    /// Command to execute
    Command(CommandPayload),

    /// Configuration update
    ConfigUpdate(ConfigUpdatePayload),

    /// Agent update command
    Update(UpdatePayload),

    /// Tray policy update
    TrayPolicyUpdate(crate::ipc::TrayPolicy),

    /// Acknowledgment of received message
    Ack { message_id: Option<String> },

    /// Error message
    Error { code: String, message: String },
}

/// Authentication acknowledgment payload
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthAckPayload {
    /// Whether authentication was successful
    pub success: bool,

    /// Agent ID assigned by server
    pub agent_id: Option<Uuid>,

    /// Error message if authentication failed
    pub error: Option<String>,
}

/// Command payload from server
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommandPayload {
    /// Unique command ID
    pub id: Uuid,

    /// Type of command
    pub command_type: CommandType,

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

    /// Optional timeout in seconds
    pub timeout_seconds: Option<u64>,

    /// Whether to run as elevated/admin
    pub elevated: bool,
}

/// Types of commands
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CommandType {
    /// Shell command (cmd on Windows, bash on Unix)
    Shell,

    /// PowerShell command (Windows)
    PowerShell,

    /// Python script
    Python,

    /// Raw script (requires interpreter path)
    Script { interpreter: String },
}

/// Configuration update payload
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConfigUpdatePayload {
    /// New metrics interval (if changed)
    pub metrics_interval_seconds: Option<u64>,

    /// Updated watchdog config
    pub watchdog: Option<WatchdogConfigUpdate>,
}

/// Watchdog configuration update
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WatchdogConfigUpdate {
    /// Enable/disable watchdog
    pub enabled: Option<bool>,

    /// Check interval
    pub check_interval_seconds: Option<u64>,

    // Services and processes would be included here for remote config updates
}

/// Update command payload from server
#[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 sent back to server
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateResultPayload {
    /// Update ID (from the server)
    pub update_id: Uuid,

    /// Update status
    pub status: UpdateStatus,

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

/// Update status codes
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UpdateStatus {
    /// Update starting
    Starting,

    /// Downloading new binary
    Downloading,

    /// Download complete, verifying
    Verifying,

    /// Installing (replacing binary)
    Installing,

    /// Restarting service
    Restarting,

    /// Update completed successfully
    Completed,

    /// Update failed
    Failed,

    /// Rolled back to previous version
    RolledBack,
}


Now let me explore the existing database schema and API routes to understand the current patterns better.