//! Machine/Agent database operations use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use sqlx::PgPool; use uuid::Uuid; /// Machine record from database #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct Machine { pub id: Uuid, pub agent_id: String, pub hostname: String, pub os_version: Option, pub is_elevated: bool, pub is_persistent: bool, pub first_seen: DateTime, pub last_seen: DateTime, pub last_session_id: Option, pub status: String, pub created_at: DateTime, pub updated_at: DateTime, } /// Get or create a machine by agent_id (upsert) pub async fn upsert_machine( pool: &PgPool, agent_id: &str, hostname: &str, is_persistent: bool, ) -> Result { sqlx::query_as::<_, Machine>( r#" INSERT INTO connect_machines (agent_id, hostname, is_persistent, status, last_seen) VALUES ($1, $2, $3, 'online', NOW()) ON CONFLICT (agent_id) DO UPDATE SET hostname = EXCLUDED.hostname, status = 'online', last_seen = NOW() RETURNING * "#, ) .bind(agent_id) .bind(hostname) .bind(is_persistent) .fetch_one(pool) .await } /// Update machine status and info pub async fn update_machine_status( pool: &PgPool, agent_id: &str, status: &str, os_version: Option<&str>, is_elevated: bool, session_id: Option, ) -> Result<(), sqlx::Error> { sqlx::query( r#" UPDATE connect_machines SET status = $1, os_version = COALESCE($2, os_version), is_elevated = $3, last_seen = NOW(), last_session_id = COALESCE($4, last_session_id) WHERE agent_id = $5 "#, ) .bind(status) .bind(os_version) .bind(is_elevated) .bind(session_id) .bind(agent_id) .execute(pool) .await?; Ok(()) } /// Get all persistent machines (for restore on startup) pub async fn get_all_machines(pool: &PgPool) -> Result, sqlx::Error> { sqlx::query_as::<_, Machine>( "SELECT * FROM connect_machines WHERE is_persistent = true ORDER BY hostname" ) .fetch_all(pool) .await } /// Get machine by agent_id pub async fn get_machine_by_agent_id( pool: &PgPool, agent_id: &str, ) -> Result, sqlx::Error> { sqlx::query_as::<_, Machine>( "SELECT * FROM connect_machines WHERE agent_id = $1" ) .bind(agent_id) .fetch_optional(pool) .await } /// Mark machine as offline pub async fn mark_machine_offline(pool: &PgPool, agent_id: &str) -> Result<(), sqlx::Error> { sqlx::query("UPDATE connect_machines SET status = 'offline', last_seen = NOW() WHERE agent_id = $1") .bind(agent_id) .execute(pool) .await?; Ok(()) } /// Delete a machine record pub async fn delete_machine(pool: &PgPool, agent_id: &str) -> Result<(), sqlx::Error> { sqlx::query("DELETE FROM connect_machines WHERE agent_id = $1") .bind(agent_id) .execute(pool) .await?; Ok(()) } /// Update machine organization, site, and tags pub async fn update_machine_metadata( pool: &PgPool, agent_id: &str, organization: Option<&str>, site: Option<&str>, tags: &[String], ) -> Result<(), sqlx::Error> { // Only update if at least one value is provided if organization.is_none() && site.is_none() && tags.is_empty() { return Ok(()); } sqlx::query( r#" UPDATE connect_machines SET organization = COALESCE($1, organization), site = COALESCE($2, site), tags = CASE WHEN $3::text[] = '{}' THEN tags ELSE $3 END WHERE agent_id = $4 "#, ) .bind(organization) .bind(site) .bind(tags) .bind(agent_id) .execute(pool) .await?; Ok(()) }