claudetools/projects/msp-tools/guru-connect/server/src/db/machines.rs

//! 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<String>,
    pub is_elevated: bool,
    pub is_persistent: bool,
    pub first_seen: DateTime<Utc>,
    pub last_seen: DateTime<Utc>,
    pub last_session_id: Option<Uuid>,
    pub status: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

/// 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<Machine, sqlx::Error> {
    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<Uuid>,
) -> 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<Vec<Machine>, 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<Option<Machine>, 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(())
}