     1→//! Agent database operations
     2→
     3→use chrono::{DateTime, Utc};
     4→use serde::{Deserialize, Serialize};
     5→use sqlx::PgPool;
     6→use uuid::Uuid;
     7→
     8→/// Agent record from database
     9→#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
    10→pub struct Agent {
    11→    pub id: Uuid,
    12→    pub hostname: String,
    13→    #[serde(skip_serializing)]
    14→    pub api_key_hash: Option<String>,  // Nullable: new agents use site's api_key
    15→    pub os_type: String,
    16→    pub os_version: Option<String>,
    17→    pub agent_version: Option<String>,
    18→    pub last_seen: Option<DateTime<Utc>>,
    19→    pub status: String,
    20→    pub created_at: DateTime<Utc>,
    21→    pub updated_at: DateTime<Utc>,
    22→    // New fields for site-based registration
    23→    pub device_id: Option<String>,     // Hardware-derived unique ID
    24→    pub site_id: Option<Uuid>,         // Which site this agent belongs to
    25→}
    26→
    27→/// Agent without sensitive fields (for API responses)
    28→#[derive(Debug, Clone, Serialize, Deserialize)]
    29→pub struct AgentResponse {
    30→    pub id: Uuid,
    31→    pub hostname: String,
    32→    pub os_type: String,
    33→    pub os_version: Option<String>,
    34→    pub agent_version: Option<String>,
    35→    pub last_seen: Option<DateTime<Utc>>,
    36→    pub status: String,
    37→    pub created_at: DateTime<Utc>,
    38→    pub device_id: Option<String>,
    39→    pub site_id: Option<Uuid>,
    40→    pub site_name: Option<String>,
    41→    pub client_name: Option<String>,
    42→}
    43→
    44→impl From<Agent> for AgentResponse {
    45→    fn from(agent: Agent) -> Self {
    46→        Self {
    47→            id: agent.id,
    48→            hostname: agent.hostname,
    49→            os_type: agent.os_type,
    50→            os_version: agent.os_version,
    51→            agent_version: agent.agent_version,
    52→            last_seen: agent.last_seen,
    53→            status: agent.status,
    54→            created_at: agent.created_at,
    55→            device_id: agent.device_id,
    56→            site_id: agent.site_id,
    57→            site_name: None,
    58→            client_name: None,
    59→        }
    60→    }
    61→}
    62→
    63→/// Create a new agent registration
    64→#[derive(Debug, Clone, Deserialize)]
    65→pub struct CreateAgent {
    66→    pub hostname: String,
    67→    pub api_key_hash: String,
    68→    pub os_type: String,
    69→    pub os_version: Option<String>,
    70→    pub agent_version: Option<String>,
    71→}
    72→
    73→/// Insert a new agent into the database
    74→pub async fn create_agent(pool: &PgPool, agent: CreateAgent) -> Result<Agent, sqlx::Error> {
    75→    sqlx::query_as::<_, Agent>(
    76→        r#"
    77→        INSERT INTO agents (hostname, api_key_hash, os_type, os_version, agent_version, status)
    78→        VALUES ($1, $2, $3, $4, $5, 'offline')
    79→        RETURNING *
    80→        "#,
    81→    )
    82→    .bind(&agent.hostname)
    83→    .bind(&agent.api_key_hash)
    84→    .bind(&agent.os_type)
    85→    .bind(&agent.os_version)
    86→    .bind(&agent.agent_version)
    87→    .fetch_one(pool)
    88→    .await
    89→}
    90→
    91→/// Get an agent by ID
    92→pub async fn get_agent_by_id(pool: &PgPool, id: Uuid) -> Result<Option<Agent>, sqlx::Error> {
    93→    sqlx::query_as::<_, Agent>("SELECT * FROM agents WHERE id = $1")
    94→        .bind(id)
    95→        .fetch_optional(pool)
    96→        .await
    97→}
    98→
    99→/// Get an agent by API key hash
   100→pub async fn get_agent_by_api_key_hash(
   101→    pool: &PgPool,
   102→    api_key_hash: &str,
   103→) -> Result<Option<Agent>, sqlx::Error> {
   104→    sqlx::query_as::<_, Agent>("SELECT * FROM agents WHERE api_key_hash = $1")
   105→        .bind(api_key_hash)
   106→        .fetch_optional(pool)
   107→        .await
   108→}
   109→
   110→/// Get all agents
   111→pub async fn get_all_agents(pool: &PgPool) -> Result<Vec<Agent>, sqlx::Error> {
   112→    sqlx::query_as::<_, Agent>("SELECT * FROM agents ORDER BY hostname")
   113→        .fetch_all(pool)
   114→        .await
   115→}
   116→
   117→/// Get agents by status
   118→pub async fn get_agents_by_status(
   119→    pool: &PgPool,
   120→    status: &str,
   121→) -> Result<Vec<Agent>, sqlx::Error> {
   122→    sqlx::query_as::<_, Agent>("SELECT * FROM agents WHERE status = $1 ORDER BY hostname")
   123→        .bind(status)
   124→        .fetch_all(pool)
   125→        .await
   126→}
   127→
   128→/// Update agent status and last_seen
   129→pub async fn update_agent_status(
   130→    pool: &PgPool,
   131→    id: Uuid,
   132→    status: &str,
   133→) -> Result<(), sqlx::Error> {
   134→    sqlx::query("UPDATE agents SET status = $1, last_seen = NOW() WHERE id = $2")
   135→        .bind(status)
   136→        .bind(id)
   137→        .execute(pool)
   138→        .await?;
   139→    Ok(())
   140→}
   141→
   142→/// Update agent info (on connection)
   143→pub async fn update_agent_info(
   144→    pool: &PgPool,
   145→    id: Uuid,
   146→    hostname: Option<&str>,
   147→    os_version: Option<&str>,
   148→    agent_version: Option<&str>,
   149→) -> Result<(), sqlx::Error> {
   150→    sqlx::query(
   151→        r#"
   152→        UPDATE agents
   153→        SET hostname = COALESCE($1, hostname),
   154→            os_version = COALESCE($2, os_version),
   155→            agent_version = COALESCE($3, agent_version),
   156→            last_seen = NOW(),
   157→            status = 'online'
   158→        WHERE id = $4
   159→        "#,
   160→    )
   161→    .bind(hostname)
   162→    .bind(os_version)
   163→    .bind(agent_version)
   164→    .bind(id)
   165→    .execute(pool)
   166→    .await?;
   167→    Ok(())
   168→}
   169→
   170→/// Delete an agent
   171→pub async fn delete_agent(pool: &PgPool, id: Uuid) -> Result<bool, sqlx::Error> {
   172→    let result = sqlx::query("DELETE FROM agents WHERE id = $1")
   173→        .bind(id)
   174→        .execute(pool)
   175→        .await?;
   176→    Ok(result.rows_affected() > 0)
   177→}
   178→
   179→/// Get agent count statistics
   180→#[derive(Debug, Clone, Serialize, Deserialize)]
   181→pub struct AgentStats {
   182→    pub total: i64,
   183→    pub online: i64,
   184→    pub offline: i64,
   185→}
   186→
   187→pub async fn get_agent_stats(pool: &PgPool) -> Result<AgentStats, sqlx::Error> {
   188→    let total: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM agents")
   189→        .fetch_one(pool)
   190→        .await?;
   191→
   192→    let online: (i64,) =
   193→        sqlx::query_as("SELECT COUNT(*) FROM agents WHERE status = 'online'")
   194→            .fetch_one(pool)
   195→            .await?;
   196→
   197→    Ok(AgentStats {
   198→        total: total.0,
   199→        online: online.0,
   200→        offline: total.0 - online.0,
   201→    })
   202→}
   203→
   204→// ============================================================================
   205→// Site-based agent operations
   206→// ============================================================================
   207→
   208→/// Data for creating an agent via site registration
   209→#[derive(Debug, Clone)]
   210→pub struct CreateAgentWithSite {
   211→    pub site_id: Uuid,
   212→    pub device_id: String,
   213→    pub hostname: String,
   214→    pub os_type: String,
   215→    pub os_version: Option<String>,
   216→    pub agent_version: Option<String>,
   217→}
   218→
   219→/// Create a new agent under a site (site-based registration)
   220→pub async fn create_agent_with_site(
   221→    pool: &PgPool,
   222→    agent: CreateAgentWithSite,
   223→) -> Result<Agent, sqlx::Error> {
   224→    sqlx::query_as::<_, Agent>(
   225→        r#"
   226→        INSERT INTO agents (site_id, device_id, hostname, os_type, os_version, agent_version, status)
   227→        VALUES ($1, $2, $3, $4, $5, $6, 'offline')
   228→        RETURNING *
   229→        "#,
   230→    )
   231→    .bind(&agent.site_id)
   232→    .bind(&agent.device_id)
   233→    .bind(&agent.hostname)
   234→    .bind(&agent.os_type)
   235→    .bind(&agent.os_version)
   236→    .bind(&agent.agent_version)
   237→    .fetch_one(pool)
   238→    .await
   239→}
   240→
   241→/// Get an agent by site_id and device_id (for site-based auth)
   242→pub async fn get_agent_by_site_and_device(
   243→    pool: &PgPool,
   244→    site_id: Uuid,
   245→    device_id: &str,
   246→) -> Result<Option<Agent>, sqlx::Error> {
   247→    sqlx::query_as::<_, Agent>(
   248→        "SELECT * FROM agents WHERE site_id = $1 AND device_id = $2"
   249→    )
   250→    .bind(site_id)
   251→    .bind(device_id)
   252→    .fetch_optional(pool)
   253→    .await
   254→}
   255→
   256→/// Get all agents for a site
   257→pub async fn get_agents_by_site(pool: &PgPool, site_id: Uuid) -> Result<Vec<Agent>, sqlx::Error> {
   258→    sqlx::query_as::<_, Agent>(
   259→        "SELECT * FROM agents WHERE site_id = $1 ORDER BY hostname"
   260→    )
   261→    .bind(site_id)
   262→    .fetch_all(pool)
   263→    .await
   264→}
   265→
   266→/// Agent with site and client details for API responses
   267→#[derive(Debug, Clone, sqlx::FromRow, Serialize)]
   268→pub struct AgentWithDetails {
   269→    pub id: Uuid,
   270→    pub hostname: String,
   271→    pub os_type: String,
   272→    pub os_version: Option<String>,
   273→    pub agent_version: Option<String>,
   274→    pub last_seen: Option<DateTime<Utc>>,
   275→    pub status: String,
   276→    pub created_at: DateTime<Utc>,
   277→    pub device_id: Option<String>,
   278→    pub site_id: Option<Uuid>,
   279→    pub site_name: Option<String>,
   280→    pub client_id: Option<Uuid>,
   281→    pub client_name: Option<String>,
   282→}
   283→
   284→/// Get all agents with site/client details
   285→pub async fn get_all_agents_with_details(pool: &PgPool) -> Result<Vec<AgentWithDetails>, sqlx::Error> {
   286→    sqlx::query_as::<_, AgentWithDetails>(
   287→        r#"
   288→        SELECT
   289→            a.id, a.hostname, a.os_type, a.os_version, a.agent_version,
   290→            a.last_seen, a.status, a.created_at, a.device_id, a.site_id,
   291→            s.name as site_name,
   292→            c.id as client_id,
   293→            c.name as client_name
   294→        FROM agents a
   295→        LEFT JOIN sites s ON a.site_id = s.id
   296→        LEFT JOIN clients c ON s.client_id = c.id
   297→        ORDER BY c.name, s.name, a.hostname
   298→        "#,
   299→    )
   300→    .fetch_all(pool)
   301→    .await
   302→}
   303→
   304→/// Update agent info including device_id (on connection)
   305→pub async fn update_agent_info_full(
   306→    pool: &PgPool,
   307→    id: Uuid,
   308→    hostname: Option<&str>,
   309→    device_id: Option<&str>,
   310→    os_version: Option<&str>,
   311→    agent_version: Option<&str>,
   312→) -> Result<(), sqlx::Error> {
   313→    sqlx::query(
   314→        r#"
   315→        UPDATE agents
   316→        SET hostname = COALESCE($1, hostname),
   317→            device_id = COALESCE($2, device_id),
   318→            os_version = COALESCE($3, os_version),
   319→            agent_version = COALESCE($4, agent_version),
   320→            last_seen = NOW(),
   321→            status = 'online'
   322→        WHERE id = $5
   323→        "#,
   324→    )
   325→    .bind(hostname)
   326→    .bind(device_id)
   327→    .bind(os_version)
   328→    .bind(agent_version)
   329→    .bind(id)
   330→    .execute(pool)
   331→    .await?;
   332→    Ok(())
   333→}
   334→
   335→/// Move an agent to a different site (or remove from site if site_id is None)
   336→pub async fn move_agent_to_site(
   337→    pool: &PgPool,
   338→    agent_id: Uuid,
   339→    site_id: Option<Uuid>,
   340→) -> Result<Option<Agent>, sqlx::Error> {
   341→    sqlx::query_as::<_, Agent>(
   342→        r#"
   343→        UPDATE agents
   344→        SET site_id = $1,
   345→            updated_at = NOW()
   346→        WHERE id = $2
   347→        RETURNING *
   348→        "#,
   349→    )
   350→    .bind(site_id)
   351→    .bind(agent_id)
   352→    .fetch_optional(pool)
   353→    .await
   354→}
   355→
   356→/// Get agents that are not assigned to any site
   357→pub async fn get_unassigned_agents(pool: &PgPool) -> Result<Vec<Agent>, sqlx::Error> {
   358→    sqlx::query_as::<_, Agent>(
   359→        "SELECT * FROM agents WHERE site_id IS NULL ORDER BY hostname"
   360→    )
   361→    .fetch_all(pool)
   362→    .await
   363→}
   364→
   365→/// Update agent last_seen timestamp (for heartbeat)
   366→pub async fn update_agent_last_seen(
   367→    pool: &PgPool,
   368→    id: Uuid,
   369→) -> Result<(), sqlx::Error> {
   370→    sqlx::query("UPDATE agents SET last_seen = NOW(), status = 'online' WHERE id = $1")
   371→        .bind(id)
   372→        .execute(pool)
   373→        .await?;
   374→    Ok(())
   375→}
   376→

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

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use uuid::Uuid;

/// Agent record from database
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct Agent {
    pub id: Uuid,
    pub hostname: String,
    #[serde(skip_serializing)]
    pub api_key_hash: Option<String>,  // Nullable: new agents use site's api_key
    pub os_type: String,
    pub os_version: Option<String>,
    pub agent_version: Option<String>,
    pub last_seen: Option<DateTime<Utc>>,
    pub status: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    // New fields for site-based registration
    pub device_id: Option<String>,     // Hardware-derived unique ID
    pub site_id: Option<Uuid>,         // Which site this agent belongs to
}

/// Agent without sensitive fields (for API responses)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentResponse {
    pub id: Uuid,
    pub hostname: String,
    pub os_type: String,
    pub os_version: Option<String>,
    pub agent_version: Option<String>,
    pub last_seen: Option<DateTime<Utc>>,
    pub status: String,
    pub created_at: DateTime<Utc>,
    pub device_id: Option<String>,
    pub site_id: Option<Uuid>,
    pub site_name: Option<String>,
    pub client_name: Option<String>,
}

impl From<Agent> for AgentResponse {
    fn from(agent: Agent) -> Self {
        Self {
            id: agent.id,
            hostname: agent.hostname,
            os_type: agent.os_type,
            os_version: agent.os_version,
            agent_version: agent.agent_version,
            last_seen: agent.last_seen,
            status: agent.status,
            created_at: agent.created_at,
            device_id: agent.device_id,
            site_id: agent.site_id,
            site_name: None,
            client_name: None,
        }
    }
}

/// Create a new agent registration
#[derive(Debug, Clone, Deserialize)]
pub struct CreateAgent {
    pub hostname: String,
    pub api_key_hash: String,
    pub os_type: String,
    pub os_version: Option<String>,
    pub agent_version: Option<String>,
}

/// Insert a new agent into the database
pub async fn create_agent(pool: &PgPool, agent: CreateAgent) -> Result<Agent, sqlx::Error> {
    sqlx::query_as::<_, Agent>(
        r#"
        INSERT INTO agents (hostname, api_key_hash, os_type, os_version, agent_version, status)
        VALUES ($1, $2, $3, $4, $5, 'offline')
        RETURNING *
        "#,
    )
    .bind(&agent.hostname)
    .bind(&agent.api_key_hash)
    .bind(&agent.os_type)
    .bind(&agent.os_version)
    .bind(&agent.agent_version)
    .fetch_one(pool)
    .await
}

/// Get an agent by ID
pub async fn get_agent_by_id(pool: &PgPool, id: Uuid) -> Result<Option<Agent>, sqlx::Error> {
    sqlx::query_as::<_, Agent>("SELECT * FROM agents WHERE id = $1")
        .bind(id)
        .fetch_optional(pool)
        .await
}

/// Get an agent by API key hash
pub async fn get_agent_by_api_key_hash(
    pool: &PgPool,
    api_key_hash: &str,
) -> Result<Option<Agent>, sqlx::Error> {
    sqlx::query_as::<_, Agent>("SELECT * FROM agents WHERE api_key_hash = $1")
        .bind(api_key_hash)
        .fetch_optional(pool)
        .await
}

/// Get all agents
pub async fn get_all_agents(pool: &PgPool) -> Result<Vec<Agent>, sqlx::Error> {
    sqlx::query_as::<_, Agent>("SELECT * FROM agents ORDER BY hostname")
        .fetch_all(pool)
        .await
}

/// Get agents by status
pub async fn get_agents_by_status(
    pool: &PgPool,
    status: &str,
) -> Result<Vec<Agent>, sqlx::Error> {
    sqlx::query_as::<_, Agent>("SELECT * FROM agents WHERE status = $1 ORDER BY hostname")
        .bind(status)
        .fetch_all(pool)
        .await
}

/// Update agent status and last_seen
pub async fn update_agent_status(
    pool: &PgPool,
    id: Uuid,
    status: &str,
) -> Result<(), sqlx::Error> {
    sqlx::query("UPDATE agents SET status = $1, last_seen = NOW() WHERE id = $2")
        .bind(status)
        .bind(id)
        .execute(pool)
        .await?;
    Ok(())
}

/// Update agent info (on connection)
pub async fn update_agent_info(
    pool: &PgPool,
    id: Uuid,
    hostname: Option<&str>,
    os_version: Option<&str>,
    agent_version: Option<&str>,
) -> Result<(), sqlx::Error> {
    sqlx::query(
        r#"
        UPDATE agents
        SET hostname = COALESCE($1, hostname),
            os_version = COALESCE($2, os_version),
            agent_version = COALESCE($3, agent_version),
            last_seen = NOW(),
            status = 'online'
        WHERE id = $4
        "#,
    )
    .bind(hostname)
    .bind(os_version)
    .bind(agent_version)
    .bind(id)
    .execute(pool)
    .await?;
    Ok(())
}

/// Delete an agent
pub async fn delete_agent(pool: &PgPool, id: Uuid) -> Result<bool, sqlx::Error> {
    let result = sqlx::query("DELETE FROM agents WHERE id = $1")
        .bind(id)
        .execute(pool)
        .await?;
    Ok(result.rows_affected() > 0)
}

/// Get agent count statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentStats {
    pub total: i64,
    pub online: i64,
    pub offline: i64,
}

pub async fn get_agent_stats(pool: &PgPool) -> Result<AgentStats, sqlx::Error> {
    let total: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM agents")
        .fetch_one(pool)
        .await?;

    let online: (i64,) =
        sqlx::query_as("SELECT COUNT(*) FROM agents WHERE status = 'online'")
            .fetch_one(pool)
            .await?;

    Ok(AgentStats {
        total: total.0,
        online: online.0,
        offline: total.0 - online.0,
    })
}

// ============================================================================
// Site-based agent operations
// ============================================================================

/// Data for creating an agent via site registration
#[derive(Debug, Clone)]
pub struct CreateAgentWithSite {
    pub site_id: Uuid,
    pub device_id: String,
    pub hostname: String,
    pub os_type: String,
    pub os_version: Option<String>,
    pub agent_version: Option<String>,
}

/// Create a new agent under a site (site-based registration)
pub async fn create_agent_with_site(
    pool: &PgPool,
    agent: CreateAgentWithSite,
) -> Result<Agent, sqlx::Error> {
    sqlx::query_as::<_, Agent>(
        r#"
        INSERT INTO agents (site_id, device_id, hostname, os_type, os_version, agent_version, status)
        VALUES ($1, $2, $3, $4, $5, $6, 'offline')
        RETURNING *
        "#,
    )
    .bind(&agent.site_id)
    .bind(&agent.device_id)
    .bind(&agent.hostname)
    .bind(&agent.os_type)
    .bind(&agent.os_version)
    .bind(&agent.agent_version)
    .fetch_one(pool)
    .await
}

/// Get an agent by site_id and device_id (for site-based auth)
pub async fn get_agent_by_site_and_device(
    pool: &PgPool,
    site_id: Uuid,
    device_id: &str,
) -> Result<Option<Agent>, sqlx::Error> {
    sqlx::query_as::<_, Agent>(
        "SELECT * FROM agents WHERE site_id = $1 AND device_id = $2"
    )
    .bind(site_id)
    .bind(device_id)
    .fetch_optional(pool)
    .await
}

/// Get all agents for a site
pub async fn get_agents_by_site(pool: &PgPool, site_id: Uuid) -> Result<Vec<Agent>, sqlx::Error> {
    sqlx::query_as::<_, Agent>(
        "SELECT * FROM agents WHERE site_id = $1 ORDER BY hostname"
    )
    .bind(site_id)
    .fetch_all(pool)
    .await
}

/// Agent with site and client details for API responses
#[derive(Debug, Clone, sqlx::FromRow, Serialize)]
pub struct AgentWithDetails {
    pub id: Uuid,
    pub hostname: String,
    pub os_type: String,
    pub os_version: Option<String>,
    pub agent_version: Option<String>,
    pub last_seen: Option<DateTime<Utc>>,
    pub status: String,
    pub created_at: DateTime<Utc>,
    pub device_id: Option<String>,
    pub site_id: Option<Uuid>,
    pub site_name: Option<String>,
    pub client_id: Option<Uuid>,
    pub client_name: Option<String>,
}

/// Get all agents with site/client details
pub async fn get_all_agents_with_details(pool: &PgPool) -> Result<Vec<AgentWithDetails>, sqlx::Error> {
    sqlx::query_as::<_, AgentWithDetails>(
        r#"
        SELECT
            a.id, a.hostname, a.os_type, a.os_version, a.agent_version,
            a.last_seen, a.status, a.created_at, a.device_id, a.site_id,
            s.name as site_name,
            c.id as client_id,
            c.name as client_name
        FROM agents a
        LEFT JOIN sites s ON a.site_id = s.id
        LEFT JOIN clients c ON s.client_id = c.id
        ORDER BY c.name, s.name, a.hostname
        "#,
    )
    .fetch_all(pool)
    .await
}

/// Update agent info including device_id (on connection)
pub async fn update_agent_info_full(
    pool: &PgPool,
    id: Uuid,
    hostname: Option<&str>,
    device_id: Option<&str>,
    os_version: Option<&str>,
    agent_version: Option<&str>,
) -> Result<(), sqlx::Error> {
    sqlx::query(
        r#"
        UPDATE agents
        SET hostname = COALESCE($1, hostname),
            device_id = COALESCE($2, device_id),
            os_version = COALESCE($3, os_version),
            agent_version = COALESCE($4, agent_version),
            last_seen = NOW(),
            status = 'online'
        WHERE id = $5
        "#,
    )
    .bind(hostname)
    .bind(device_id)
    .bind(os_version)
    .bind(agent_version)
    .bind(id)
    .execute(pool)
    .await?;
    Ok(())
}

/// Move an agent to a different site (or remove from site if site_id is None)
pub async fn move_agent_to_site(
    pool: &PgPool,
    agent_id: Uuid,
    site_id: Option<Uuid>,
) -> Result<Option<Agent>, sqlx::Error> {
    sqlx::query_as::<_, Agent>(
        r#"
        UPDATE agents
        SET site_id = $1,
            updated_at = NOW()
        WHERE id = $2
        RETURNING *
        "#,
    )
    .bind(site_id)
    .bind(agent_id)
    .fetch_optional(pool)
    .await
}

/// Get agents that are not assigned to any site
pub async fn get_unassigned_agents(pool: &PgPool) -> Result<Vec<Agent>, sqlx::Error> {
    sqlx::query_as::<_, Agent>(
        "SELECT * FROM agents WHERE site_id IS NULL ORDER BY hostname"
    )
    .fetch_all(pool)
    .await
}

/// Update agent last_seen timestamp (for heartbeat)
pub async fn update_agent_last_seen(
    pool: &PgPool,
    id: Uuid,
) -> Result<(), sqlx::Error> {
    sqlx::query("UPDATE agents SET last_seen = NOW(), status = 'online' WHERE id = $1")
        .bind(id)
        .execute(pool)
        .await?;
    Ok(())
}


     1→//! Client (organization) database operations
     2→
     3→use serde::{Deserialize, Serialize};
     4→use sqlx::PgPool;
     5→use uuid::Uuid;
     6→
     7→/// Client record from database
     8→#[derive(Debug, Clone, sqlx::FromRow)]
     9→pub struct Client {
    10→    pub id: Uuid,
    11→    pub name: String,
    12→    pub code: Option<String>,
    13→    pub notes: Option<String>,
    14→    pub is_active: bool,
    15→    pub created_at: chrono::DateTime<chrono::Utc>,
    16→    pub updated_at: chrono::DateTime<chrono::Utc>,
    17→}
    18→
    19→/// Client response for API
    20→#[derive(Debug, Clone, Serialize)]
    21→pub struct ClientResponse {
    22→    pub id: Uuid,
    23→    pub name: String,
    24→    pub code: Option<String>,
    25→    pub notes: Option<String>,
    26→    pub is_active: bool,
    27→    pub created_at: chrono::DateTime<chrono::Utc>,
    28→    pub site_count: Option<i64>,
    29→    pub agent_count: Option<i64>,
    30→}
    31→
    32→impl From<Client> for ClientResponse {
    33→    fn from(c: Client) -> Self {
    34→        ClientResponse {
    35→            id: c.id,
    36→            name: c.name,
    37→            code: c.code,
    38→            notes: c.notes,
    39→            is_active: c.is_active,
    40→            created_at: c.created_at,
    41→            site_count: None,
    42→            agent_count: None,
    43→        }
    44→    }
    45→}
    46→
    47→/// Data for creating a new client
    48→#[derive(Debug, Deserialize)]
    49→pub struct CreateClient {
    50→    pub name: String,
    51→    pub code: Option<String>,
    52→    pub notes: Option<String>,
    53→}
    54→
    55→/// Data for updating a client
    56→#[derive(Debug, Deserialize)]
    57→pub struct UpdateClient {
    58→    pub name: Option<String>,
    59→    pub code: Option<String>,
    60→    pub notes: Option<String>,
    61→    pub is_active: Option<bool>,
    62→}
    63→
    64→/// Create a new client
    65→pub async fn create_client(pool: &PgPool, client: CreateClient) -> Result<Client, sqlx::Error> {
    66→    sqlx::query_as::<_, Client>(
    67→        r#"
    68→        INSERT INTO clients (name, code, notes)
    69→        VALUES ($1, $2, $3)
    70→        RETURNING *
    71→        "#,
    72→    )
    73→    .bind(&client.name)
    74→    .bind(&client.code)
    75→    .bind(&client.notes)
    76→    .fetch_one(pool)
    77→    .await
    78→}
    79→
    80→/// Get a client by ID
    81→pub async fn get_client_by_id(pool: &PgPool, id: Uuid) -> Result<Option<Client>, sqlx::Error> {
    82→    sqlx::query_as::<_, Client>("SELECT * FROM clients WHERE id = $1")
    83→        .bind(id)
    84→        .fetch_optional(pool)
    85→        .await
    86→}
    87→
    88→/// Get all clients
    89→pub async fn get_all_clients(pool: &PgPool) -> Result<Vec<Client>, sqlx::Error> {
    90→    sqlx::query_as::<_, Client>("SELECT * FROM clients ORDER BY name")
    91→        .fetch_all(pool)
    92→        .await
    93→}
    94→
    95→/// Get all clients with counts
    96→#[derive(Debug, Clone, sqlx::FromRow)]
    97→pub struct ClientWithCounts {
    98→    pub id: Uuid,
    99→    pub name: String,
   100→    pub code: Option<String>,
   101→    pub notes: Option<String>,
   102→    pub is_active: bool,
   103→    pub created_at: chrono::DateTime<chrono::Utc>,
   104→    pub updated_at: chrono::DateTime<chrono::Utc>,
   105→    pub site_count: i64,
   106→    pub agent_count: i64,
   107→}
   108→
   109→pub async fn get_all_clients_with_counts(pool: &PgPool) -> Result<Vec<ClientWithCounts>, sqlx::Error> {
   110→    sqlx::query_as::<_, ClientWithCounts>(
   111→        r#"
   112→        SELECT
   113→            c.*,
   114→            COALESCE((SELECT COUNT(*) FROM sites WHERE client_id = c.id), 0) as site_count,
   115→            COALESCE((SELECT COUNT(*) FROM agents a JOIN sites s ON a.site_id = s.id WHERE s.client_id = c.id), 0) as agent_count
   116→        FROM clients c
   117→        ORDER BY c.name
   118→        "#,
   119→    )
   120→    .fetch_all(pool)
   121→    .await
   122→}
   123→
   124→/// Update a client
   125→pub async fn update_client(
   126→    pool: &PgPool,
   127→    id: Uuid,
   128→    update: UpdateClient,
   129→) -> Result<Option<Client>, sqlx::Error> {
   130→    sqlx::query_as::<_, Client>(
   131→        r#"
   132→        UPDATE clients
   133→        SET name = COALESCE($1, name),
   134→            code = COALESCE($2, code),
   135→            notes = COALESCE($3, notes),
   136→            is_active = COALESCE($4, is_active)
   137→        WHERE id = $5
   138→        RETURNING *
   139→        "#,
   140→    )
   141→    .bind(&update.name)
   142→    .bind(&update.code)
   143→    .bind(&update.notes)
   144→    .bind(&update.is_active)
   145→    .bind(id)
   146→    .fetch_optional(pool)
   147→    .await
   148→}
   149→
   150→/// Delete a client
   151→pub async fn delete_client(pool: &PgPool, id: Uuid) -> Result<bool, sqlx::Error> {
   152→    let result = sqlx::query("DELETE FROM clients WHERE id = $1")
   153→        .bind(id)
   154→        .execute(pool)
   155→        .await?;
   156→    Ok(result.rows_affected() > 0)
   157→}
   158→

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

//! Client (organization) database operations

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

/// Client record from database
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct Client {
    pub id: Uuid,
    pub name: String,
    pub code: Option<String>,
    pub notes: Option<String>,
    pub is_active: bool,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub updated_at: chrono::DateTime<chrono::Utc>,
}

/// Client response for API
#[derive(Debug, Clone, Serialize)]
pub struct ClientResponse {
    pub id: Uuid,
    pub name: String,
    pub code: Option<String>,
    pub notes: Option<String>,
    pub is_active: bool,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub site_count: Option<i64>,
    pub agent_count: Option<i64>,
}

impl From<Client> for ClientResponse {
    fn from(c: Client) -> Self {
        ClientResponse {
            id: c.id,
            name: c.name,
            code: c.code,
            notes: c.notes,
            is_active: c.is_active,
            created_at: c.created_at,
            site_count: None,
            agent_count: None,
        }
    }
}

/// Data for creating a new client
#[derive(Debug, Deserialize)]
pub struct CreateClient {
    pub name: String,
    pub code: Option<String>,
    pub notes: Option<String>,
}

/// Data for updating a client
#[derive(Debug, Deserialize)]
pub struct UpdateClient {
    pub name: Option<String>,
    pub code: Option<String>,
    pub notes: Option<String>,
    pub is_active: Option<bool>,
}

/// Create a new client
pub async fn create_client(pool: &PgPool, client: CreateClient) -> Result<Client, sqlx::Error> {
    sqlx::query_as::<_, Client>(
        r#"
        INSERT INTO clients (name, code, notes)
        VALUES ($1, $2, $3)
        RETURNING *
        "#,
    )
    .bind(&client.name)
    .bind(&client.code)
    .bind(&client.notes)
    .fetch_one(pool)
    .await
}

/// Get a client by ID
pub async fn get_client_by_id(pool: &PgPool, id: Uuid) -> Result<Option<Client>, sqlx::Error> {
    sqlx::query_as::<_, Client>("SELECT * FROM clients WHERE id = $1")
        .bind(id)
        .fetch_optional(pool)
        .await
}

/// Get all clients
pub async fn get_all_clients(pool: &PgPool) -> Result<Vec<Client>, sqlx::Error> {
    sqlx::query_as::<_, Client>("SELECT * FROM clients ORDER BY name")
        .fetch_all(pool)
        .await
}

/// Get all clients with counts
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct ClientWithCounts {
    pub id: Uuid,
    pub name: String,
    pub code: Option<String>,
    pub notes: Option<String>,
    pub is_active: bool,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub updated_at: chrono::DateTime<chrono::Utc>,
    pub site_count: i64,
    pub agent_count: i64,
}

pub async fn get_all_clients_with_counts(pool: &PgPool) -> Result<Vec<ClientWithCounts>, sqlx::Error> {
    sqlx::query_as::<_, ClientWithCounts>(
        r#"
        SELECT
            c.*,
            COALESCE((SELECT COUNT(*) FROM sites WHERE client_id = c.id), 0) as site_count,
            COALESCE((SELECT COUNT(*) FROM agents a JOIN sites s ON a.site_id = s.id WHERE s.client_id = c.id), 0) as agent_count
        FROM clients c
        ORDER BY c.name
        "#,
    )
    .fetch_all(pool)
    .await
}

/// Update a client
pub async fn update_client(
    pool: &PgPool,
    id: Uuid,
    update: UpdateClient,
) -> Result<Option<Client>, sqlx::Error> {
    sqlx::query_as::<_, Client>(
        r#"
        UPDATE clients
        SET name = COALESCE($1, name),
            code = COALESCE($2, code),
            notes = COALESCE($3, notes),
            is_active = COALESCE($4, is_active)
        WHERE id = $5
        RETURNING *
        "#,
    )
    .bind(&update.name)
    .bind(&update.code)
    .bind(&update.notes)
    .bind(&update.is_active)
    .bind(id)
    .fetch_optional(pool)
    .await
}

/// Delete a client
pub async fn delete_client(pool: &PgPool, id: Uuid) -> Result<bool, sqlx::Error> {
    let result = sqlx::query("DELETE FROM clients WHERE id = $1")
        .bind(id)
        .execute(pool)
        .await?;
    Ok(result.rows_affected() > 0)
}


     1→//! Site database operations
     2→
     3→use rand::Rng;
     4→use serde::{Deserialize, Serialize};
     5→use sqlx::PgPool;
     6→use uuid::Uuid;
     7→
     8→/// Word lists for generating site codes
     9→const ADJECTIVES: &[&str] = &[
    10→    "BLUE", "GREEN", "RED", "GOLD", "SILVER", "IRON", "COPPER", "BRONZE",
    11→    "SWIFT", "BRIGHT", "DARK", "LIGHT", "BOLD", "CALM", "WILD", "WARM",
    12→    "NORTH", "SOUTH", "EAST", "WEST", "UPPER", "LOWER", "INNER", "OUTER",
    13→];
    14→
    15→const NOUNS: &[&str] = &[
    16→    "HAWK", "EAGLE", "TIGER", "LION", "WOLF", "BEAR", "FALCON", "PHOENIX",
    17→    "PEAK", "VALLEY", "RIVER", "OCEAN", "STORM", "CLOUD", "STAR", "MOON",
    18→    "TOWER", "BRIDGE", "GATE", "FORGE", "CASTLE", "HARBOR", "MEADOW", "GROVE",
    19→];
    20→
    21→/// Generate a human-friendly site code (e.g., "BLUE-TIGER-4829")
    22→pub fn generate_site_code() -> String {
    23→    let mut rng = rand::thread_rng();
    24→    let adj = ADJECTIVES[rng.gen_range(0..ADJECTIVES.len())];
    25→    let noun = NOUNS[rng.gen_range(0..NOUNS.len())];
    26→    let num: u16 = rng.gen_range(1000..9999);
    27→    format!("{}-{}-{}", adj, noun, num)
    28→}
    29→
    30→/// Site record from database
    31→#[derive(Debug, Clone, sqlx::FromRow)]
    32→pub struct Site {
    33→    pub id: Uuid,
    34→    pub client_id: Uuid,
    35→    pub name: String,
    36→    pub site_code: String,
    37→    pub api_key_hash: String,
    38→    pub address: Option<String>,
    39→    pub notes: Option<String>,
    40→    pub is_active: bool,
    41→    pub created_at: chrono::DateTime<chrono::Utc>,
    42→    pub updated_at: chrono::DateTime<chrono::Utc>,
    43→}
    44→
    45→/// Site response for API
    46→#[derive(Debug, Clone, Serialize)]
    47→pub struct SiteResponse {
    48→    pub id: Uuid,
    49→    pub client_id: Uuid,
    50→    pub client_name: Option<String>,
    51→    pub name: String,
    52→    pub site_code: String,
    53→    pub address: Option<String>,
    54→    pub notes: Option<String>,
    55→    pub is_active: bool,
    56→    pub created_at: chrono::DateTime<chrono::Utc>,
    57→    pub agent_count: Option<i64>,
    58→}
    59→
    60→impl From<Site> for SiteResponse {
    61→    fn from(s: Site) -> Self {
    62→        SiteResponse {
    63→            id: s.id,
    64→            client_id: s.client_id,
    65→            client_name: None,
    66→            name: s.name,
    67→            site_code: s.site_code,
    68→            address: s.address,
    69→            notes: s.notes,
    70→            is_active: s.is_active,
    71→            created_at: s.created_at,
    72→            agent_count: None,
    73→        }
    74→    }
    75→}
    76→
    77→/// Data for creating a new site
    78→#[derive(Debug, Deserialize)]
    79→pub struct CreateSite {
    80→    pub client_id: Uuid,
    81→    pub name: String,
    82→    pub address: Option<String>,
    83→    pub notes: Option<String>,
    84→}
    85→
    86→/// Internal create with all fields
    87→pub struct CreateSiteInternal {
    88→    pub client_id: Uuid,
    89→    pub name: String,
    90→    pub site_code: String,
    91→    pub api_key_hash: String,
    92→    pub address: Option<String>,
    93→    pub notes: Option<String>,
    94→}
    95→
    96→/// Data for updating a site
    97→#[derive(Debug, Deserialize)]
    98→pub struct UpdateSite {
    99→    pub name: Option<String>,
   100→    pub address: Option<String>,
   101→    pub notes: Option<String>,
   102→    pub is_active: Option<bool>,
   103→}
   104→
   105→/// Create a new site
   106→pub async fn create_site(pool: &PgPool, site: CreateSiteInternal) -> Result<Site, sqlx::Error> {
   107→    sqlx::query_as::<_, Site>(
   108→        r#"
   109→        INSERT INTO sites (client_id, name, site_code, api_key_hash, address, notes)
   110→        VALUES ($1, $2, $3, $4, $5, $6)
   111→        RETURNING *
   112→        "#,
   113→    )
   114→    .bind(&site.client_id)
   115→    .bind(&site.name)
   116→    .bind(&site.site_code)
   117→    .bind(&site.api_key_hash)
   118→    .bind(&site.address)
   119→    .bind(&site.notes)
   120→    .fetch_one(pool)
   121→    .await
   122→}
   123→
   124→/// Get a site by ID
   125→pub async fn get_site_by_id(pool: &PgPool, id: Uuid) -> Result<Option<Site>, sqlx::Error> {
   126→    sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = $1")
   127→        .bind(id)
   128→        .fetch_optional(pool)
   129→        .await
   130→}
   131→
   132→/// Get a site by site code
   133→pub async fn get_site_by_code(pool: &PgPool, site_code: &str) -> Result<Option<Site>, sqlx::Error> {
   134→    sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE site_code = $1 AND is_active = true")
   135→        .bind(site_code.to_uppercase())
   136→        .fetch_optional(pool)
   137→        .await
   138→}
   139→
   140→/// Get a site by API key hash
   141→pub async fn get_site_by_api_key_hash(pool: &PgPool, api_key_hash: &str) -> Result<Option<Site>, sqlx::Error> {
   142→    sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE api_key_hash = $1 AND is_active = true")
   143→        .bind(api_key_hash)
   144→        .fetch_optional(pool)
   145→        .await
   146→}
   147→
   148→/// Get all sites for a client
   149→pub async fn get_sites_by_client(pool: &PgPool, client_id: Uuid) -> Result<Vec<Site>, sqlx::Error> {
   150→    sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE client_id = $1 ORDER BY name")
   151→        .bind(client_id)
   152→        .fetch_all(pool)
   153→        .await
   154→}
   155→
   156→/// Get all sites
   157→pub async fn get_all_sites(pool: &PgPool) -> Result<Vec<Site>, sqlx::Error> {
   158→    sqlx::query_as::<_, Site>("SELECT * FROM sites ORDER BY name")
   159→        .fetch_all(pool)
   160→        .await
   161→}
   162→
   163→/// Site with client name and agent count
   164→#[derive(Debug, Clone, sqlx::FromRow)]
   165→pub struct SiteWithDetails {
   166→    pub id: Uuid,
   167→    pub client_id: Uuid,
   168→    pub name: String,
   169→    pub site_code: String,
   170→    pub api_key_hash: String,
   171→    pub address: Option<String>,
   172→    pub notes: Option<String>,
   173→    pub is_active: bool,
   174→    pub created_at: chrono::DateTime<chrono::Utc>,
   175→    pub updated_at: chrono::DateTime<chrono::Utc>,
   176→    pub client_name: String,
   177→    pub agent_count: i64,
   178→}
   179→
   180→pub async fn get_all_sites_with_details(pool: &PgPool) -> Result<Vec<SiteWithDetails>, sqlx::Error> {
   181→    sqlx::query_as::<_, SiteWithDetails>(
   182→        r#"
   183→        SELECT
   184→            s.*,
   185→            c.name as client_name,
   186→            COALESCE((SELECT COUNT(*) FROM agents WHERE site_id = s.id), 0) as agent_count
   187→        FROM sites s
   188→        JOIN clients c ON s.client_id = c.id
   189→        ORDER BY c.name, s.name
   190→        "#,
   191→    )
   192→    .fetch_all(pool)
   193→    .await
   194→}
   195→
   196→/// Update a site
   197→pub async fn update_site(
   198→    pool: &PgPool,
   199→    id: Uuid,
   200→    update: UpdateSite,
   201→) -> Result<Option<Site>, sqlx::Error> {
   202→    sqlx::query_as::<_, Site>(
   203→        r#"
   204→        UPDATE sites
   205→        SET name = COALESCE($1, name),
   206→            address = COALESCE($2, address),
   207→            notes = COALESCE($3, notes),
   208→            is_active = COALESCE($4, is_active)
   209→        WHERE id = $5
   210→        RETURNING *
   211→        "#,
   212→    )
   213→    .bind(&update.name)
   214→    .bind(&update.address)
   215→    .bind(&update.notes)
   216→    .bind(&update.is_active)
   217→    .bind(id)
   218→    .fetch_optional(pool)
   219→    .await
   220→}
   221→
   222→/// Regenerate API key for a site
   223→pub async fn regenerate_site_api_key(
   224→    pool: &PgPool,
   225→    id: Uuid,
   226→    new_api_key_hash: &str,
   227→) -> Result<bool, sqlx::Error> {
   228→    let result = sqlx::query("UPDATE sites SET api_key_hash = $1 WHERE id = $2")
   229→        .bind(new_api_key_hash)
   230→        .bind(id)
   231→        .execute(pool)
   232→        .await?;
   233→    Ok(result.rows_affected() > 0)
   234→}
   235→
   236→/// Delete a site
   237→pub async fn delete_site(pool: &PgPool, id: Uuid) -> Result<bool, sqlx::Error> {
   238→    let result = sqlx::query("DELETE FROM sites WHERE id = $1")
   239→        .bind(id)
   240→        .execute(pool)
   241→        .await?;
   242→    Ok(result.rows_affected() > 0)
   243→}
   244→
   245→/// Check if a site code is unique
   246→pub async fn is_site_code_unique(pool: &PgPool, site_code: &str) -> Result<bool, sqlx::Error> {
   247→    let result: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM sites WHERE site_code = $1")
   248→        .bind(site_code.to_uppercase())
   249→        .fetch_one(pool)
   250→        .await?;
   251→    Ok(result.0 == 0)
   252→}
   253→
   254→/// Generate a unique site code (tries up to 10 times)
   255→pub async fn generate_unique_site_code(pool: &PgPool) -> Result<String, sqlx::Error> {
   256→    for _ in 0..10 {
   257→        let code = generate_site_code();
   258→        if is_site_code_unique(pool, &code).await? {
   259→            return Ok(code);
   260→        }
   261→    }
   262→    // Fallback: add random suffix
   263→    Ok(format!("{}-{}", generate_site_code(), rand::thread_rng().gen_range(100..999)))
   264→}
   265→

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

//! Site database operations

use rand::Rng;
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use uuid::Uuid;

/// Word lists for generating site codes
const ADJECTIVES: &[&str] = &[
    "BLUE", "GREEN", "RED", "GOLD", "SILVER", "IRON", "COPPER", "BRONZE",
    "SWIFT", "BRIGHT", "DARK", "LIGHT", "BOLD", "CALM", "WILD", "WARM",
    "NORTH", "SOUTH", "EAST", "WEST", "UPPER", "LOWER", "INNER", "OUTER",
];

const NOUNS: &[&str] = &[
    "HAWK", "EAGLE", "TIGER", "LION", "WOLF", "BEAR", "FALCON", "PHOENIX",
    "PEAK", "VALLEY", "RIVER", "OCEAN", "STORM", "CLOUD", "STAR", "MOON",
    "TOWER", "BRIDGE", "GATE", "FORGE", "CASTLE", "HARBOR", "MEADOW", "GROVE",
];

/// Generate a human-friendly site code (e.g., "BLUE-TIGER-4829")
pub fn generate_site_code() -> String {
    let mut rng = rand::thread_rng();
    let adj = ADJECTIVES[rng.gen_range(0..ADJECTIVES.len())];
    let noun = NOUNS[rng.gen_range(0..NOUNS.len())];
    let num: u16 = rng.gen_range(1000..9999);
    format!("{}-{}-{}", adj, noun, num)
}

/// Site record from database
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct Site {
    pub id: Uuid,
    pub client_id: Uuid,
    pub name: String,
    pub site_code: String,
    pub api_key_hash: String,
    pub address: Option<String>,
    pub notes: Option<String>,
    pub is_active: bool,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub updated_at: chrono::DateTime<chrono::Utc>,
}

/// Site response for API
#[derive(Debug, Clone, Serialize)]
pub struct SiteResponse {
    pub id: Uuid,
    pub client_id: Uuid,
    pub client_name: Option<String>,
    pub name: String,
    pub site_code: String,
    pub address: Option<String>,
    pub notes: Option<String>,
    pub is_active: bool,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub agent_count: Option<i64>,
}

impl From<Site> for SiteResponse {
    fn from(s: Site) -> Self {
        SiteResponse {
            id: s.id,
            client_id: s.client_id,
            client_name: None,
            name: s.name,
            site_code: s.site_code,
            address: s.address,
            notes: s.notes,
            is_active: s.is_active,
            created_at: s.created_at,
            agent_count: None,
        }
    }
}

/// Data for creating a new site
#[derive(Debug, Deserialize)]
pub struct CreateSite {
    pub client_id: Uuid,
    pub name: String,
    pub address: Option<String>,
    pub notes: Option<String>,
}

/// Internal create with all fields
pub struct CreateSiteInternal {
    pub client_id: Uuid,
    pub name: String,
    pub site_code: String,
    pub api_key_hash: String,
    pub address: Option<String>,
    pub notes: Option<String>,
}

/// Data for updating a site
#[derive(Debug, Deserialize)]
pub struct UpdateSite {
    pub name: Option<String>,
    pub address: Option<String>,
    pub notes: Option<String>,
    pub is_active: Option<bool>,
}

/// Create a new site
pub async fn create_site(pool: &PgPool, site: CreateSiteInternal) -> Result<Site, sqlx::Error> {
    sqlx::query_as::<_, Site>(
        r#"
        INSERT INTO sites (client_id, name, site_code, api_key_hash, address, notes)
        VALUES ($1, $2, $3, $4, $5, $6)
        RETURNING *
        "#,
    )
    .bind(&site.client_id)
    .bind(&site.name)
    .bind(&site.site_code)
    .bind(&site.api_key_hash)
    .bind(&site.address)
    .bind(&site.notes)
    .fetch_one(pool)
    .await
}

/// Get a site by ID
pub async fn get_site_by_id(pool: &PgPool, id: Uuid) -> Result<Option<Site>, sqlx::Error> {
    sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = $1")
        .bind(id)
        .fetch_optional(pool)
        .await
}

/// Get a site by site code
pub async fn get_site_by_code(pool: &PgPool, site_code: &str) -> Result<Option<Site>, sqlx::Error> {
    sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE site_code = $1 AND is_active = true")
        .bind(site_code.to_uppercase())
        .fetch_optional(pool)
        .await
}

/// Get a site by API key hash
pub async fn get_site_by_api_key_hash(pool: &PgPool, api_key_hash: &str) -> Result<Option<Site>, sqlx::Error> {
    sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE api_key_hash = $1 AND is_active = true")
        .bind(api_key_hash)
        .fetch_optional(pool)
        .await
}

/// Get all sites for a client
pub async fn get_sites_by_client(pool: &PgPool, client_id: Uuid) -> Result<Vec<Site>, sqlx::Error> {
    sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE client_id = $1 ORDER BY name")
        .bind(client_id)
        .fetch_all(pool)
        .await
}

/// Get all sites
pub async fn get_all_sites(pool: &PgPool) -> Result<Vec<Site>, sqlx::Error> {
    sqlx::query_as::<_, Site>("SELECT * FROM sites ORDER BY name")
        .fetch_all(pool)
        .await
}

/// Site with client name and agent count
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct SiteWithDetails {
    pub id: Uuid,
    pub client_id: Uuid,
    pub name: String,
    pub site_code: String,
    pub api_key_hash: String,
    pub address: Option<String>,
    pub notes: Option<String>,
    pub is_active: bool,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub updated_at: chrono::DateTime<chrono::Utc>,
    pub client_name: String,
    pub agent_count: i64,
}

pub async fn get_all_sites_with_details(pool: &PgPool) -> Result<Vec<SiteWithDetails>, sqlx::Error> {
    sqlx::query_as::<_, SiteWithDetails>(
        r#"
        SELECT
            s.*,
            c.name as client_name,
            COALESCE((SELECT COUNT(*) FROM agents WHERE site_id = s.id), 0) as agent_count
        FROM sites s
        JOIN clients c ON s.client_id = c.id
        ORDER BY c.name, s.name
        "#,
    )
    .fetch_all(pool)
    .await
}

/// Update a site
pub async fn update_site(
    pool: &PgPool,
    id: Uuid,
    update: UpdateSite,
) -> Result<Option<Site>, sqlx::Error> {
    sqlx::query_as::<_, Site>(
        r#"
        UPDATE sites
        SET name = COALESCE($1, name),
            address = COALESCE($2, address),
            notes = COALESCE($3, notes),
            is_active = COALESCE($4, is_active)
        WHERE id = $5
        RETURNING *
        "#,
    )
    .bind(&update.name)
    .bind(&update.address)
    .bind(&update.notes)
    .bind(&update.is_active)
    .bind(id)
    .fetch_optional(pool)
    .await
}

/// Regenerate API key for a site
pub async fn regenerate_site_api_key(
    pool: &PgPool,
    id: Uuid,
    new_api_key_hash: &str,
) -> Result<bool, sqlx::Error> {
    let result = sqlx::query("UPDATE sites SET api_key_hash = $1 WHERE id = $2")
        .bind(new_api_key_hash)
        .bind(id)
        .execute(pool)
        .await?;
    Ok(result.rows_affected() > 0)
}

/// Delete a site
pub async fn delete_site(pool: &PgPool, id: Uuid) -> Result<bool, sqlx::Error> {
    let result = sqlx::query("DELETE FROM sites WHERE id = $1")
        .bind(id)
        .execute(pool)
        .await?;
    Ok(result.rows_affected() > 0)
}

/// Check if a site code is unique
pub async fn is_site_code_unique(pool: &PgPool, site_code: &str) -> Result<bool, sqlx::Error> {
    let result: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM sites WHERE site_code = $1")
        .bind(site_code.to_uppercase())
        .fetch_one(pool)
        .await?;
    Ok(result.0 == 0)
}

/// Generate a unique site code (tries up to 10 times)
pub async fn generate_unique_site_code(pool: &PgPool) -> Result<String, sqlx::Error> {
    for _ in 0..10 {
        let code = generate_site_code();
        if is_site_code_unique(pool, &code).await? {
            return Ok(code);
        }
    }
    // Fallback: add random suffix
    Ok(format!("{}-{}", generate_site_code(), rand::thread_rng().gen_range(100..999)))
}


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

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

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

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

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

use crate::AppState;

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


Let me also examine the existing metrics structure and the agent's IPC/tray module to understand the TrayPolicy pattern mentioned.


     1→//! Metrics database operations
     2→
     3→use chrono::{DateTime, Utc};
     4→use serde::{Deserialize, Serialize};
     5→use sqlx::PgPool;
     6→use uuid::Uuid;
     7→
     8→/// Metrics record from database
     9→#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
    10→pub struct Metrics {
    11→    pub id: i64,
    12→    pub agent_id: Uuid,
    13→    pub timestamp: DateTime<Utc>,
    14→    pub cpu_percent: Option<f32>,
    15→    pub memory_percent: Option<f32>,
    16→    pub memory_used_bytes: Option<i64>,
    17→    pub disk_percent: Option<f32>,
    18→    pub disk_used_bytes: Option<i64>,
    19→    pub network_rx_bytes: Option<i64>,
    20→    pub network_tx_bytes: Option<i64>,
    21→    // Extended metrics
    22→    pub uptime_seconds: Option<i64>,
    23→    pub boot_time: Option<i64>,
    24→    pub logged_in_user: Option<String>,
    25→    pub user_idle_seconds: Option<i64>,
    26→    pub public_ip: Option<String>,
    27→    pub memory_total_bytes: Option<i64>,
    28→    pub disk_total_bytes: Option<i64>,
    29→    // Temperature metrics
    30→    pub cpu_temp_celsius: Option<f32>,
    31→    pub gpu_temp_celsius: Option<f32>,
    32→    pub temperatures: Option<serde_json::Value>,
    33→}
    34→
    35→/// Create metrics data from agent
    36→#[derive(Debug, Clone, Deserialize)]
    37→pub struct CreateMetrics {
    38→    pub agent_id: Uuid,
    39→    pub cpu_percent: Option<f32>,
    40→    pub memory_percent: Option<f32>,
    41→    pub memory_used_bytes: Option<i64>,
    42→    pub disk_percent: Option<f32>,
    43→    pub disk_used_bytes: Option<i64>,
    44→    pub network_rx_bytes: Option<i64>,
    45→    pub network_tx_bytes: Option<i64>,
    46→    // Extended metrics
    47→    pub uptime_seconds: Option<i64>,
    48→    pub boot_time: Option<i64>,
    49→    pub logged_in_user: Option<String>,
    50→    pub user_idle_seconds: Option<i64>,
    51→    pub public_ip: Option<String>,
    52→    pub memory_total_bytes: Option<i64>,
    53→    pub disk_total_bytes: Option<i64>,
    54→    // Temperature metrics
    55→    pub cpu_temp_celsius: Option<f32>,
    56→    pub gpu_temp_celsius: Option<f32>,
    57→    pub temperatures: Option<serde_json::Value>,
    58→}
    59→
    60→/// Insert metrics into the database
    61→pub async fn insert_metrics(pool: &PgPool, metrics: CreateMetrics) -> Result<i64, sqlx::Error> {
    62→    let result: (i64,) = sqlx::query_as(
    63→        r#"
    64→        INSERT INTO metrics (
    65→            agent_id, cpu_percent, memory_percent, memory_used_bytes,
    66→            disk_percent, disk_used_bytes, network_rx_bytes, network_tx_bytes,
    67→            uptime_seconds, boot_time, logged_in_user, user_idle_seconds,
    68→            public_ip, memory_total_bytes, disk_total_bytes,
    69→            cpu_temp_celsius, gpu_temp_celsius, temperatures
    70→        )
    71→        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18)
    72→        RETURNING id
    73→        "#,
    74→    )
    75→    .bind(metrics.agent_id)
    76→    .bind(metrics.cpu_percent)
    77→    .bind(metrics.memory_percent)
    78→    .bind(metrics.memory_used_bytes)
    79→    .bind(metrics.disk_percent)
    80→    .bind(metrics.disk_used_bytes)
    81→    .bind(metrics.network_rx_bytes)
    82→    .bind(metrics.network_tx_bytes)
    83→    .bind(metrics.uptime_seconds)
    84→    .bind(metrics.boot_time)
    85→    .bind(&metrics.logged_in_user)
    86→    .bind(metrics.user_idle_seconds)
    87→    .bind(&metrics.public_ip)
    88→    .bind(metrics.memory_total_bytes)
    89→    .bind(metrics.disk_total_bytes)
    90→    .bind(metrics.cpu_temp_celsius)
    91→    .bind(metrics.gpu_temp_celsius)
    92→    .bind(&metrics.temperatures)
    93→    .fetch_one(pool)
    94→    .await?;
    95→
    96→    Ok(result.0)
    97→}
    98→
    99→/// Get recent metrics for an agent
   100→pub async fn get_agent_metrics(
   101→    pool: &PgPool,
   102→    agent_id: Uuid,
   103→    limit: i64,
   104→) -> Result<Vec<Metrics>, sqlx::Error> {
   105→    sqlx::query_as::<_, Metrics>(
   106→        r#"
   107→        SELECT * FROM metrics
   108→        WHERE agent_id = $1
   109→        ORDER BY timestamp DESC
   110→        LIMIT $2
   111→        "#,
   112→    )
   113→    .bind(agent_id)
   114→    .bind(limit)
   115→    .fetch_all(pool)
   116→    .await
   117→}
   118→
   119→/// Get metrics for an agent within a time range
   120→pub async fn get_agent_metrics_range(
   121→    pool: &PgPool,
   122→    agent_id: Uuid,
   123→    start: DateTime<Utc>,
   124→    end: DateTime<Utc>,
   125→) -> Result<Vec<Metrics>, sqlx::Error> {
   126→    sqlx::query_as::<_, Metrics>(
   127→        r#"
   128→        SELECT * FROM metrics
   129→        WHERE agent_id = $1 AND timestamp >= $2 AND timestamp <= $3
   130→        ORDER BY timestamp ASC
   131→        "#,
   132→    )
   133→    .bind(agent_id)
   134→    .bind(start)
   135→    .bind(end)
   136→    .fetch_all(pool)
   137→    .await
   138→}
   139→
   140→/// Get latest metrics for an agent
   141→pub async fn get_latest_metrics(
   142→    pool: &PgPool,
   143→    agent_id: Uuid,
   144→) -> Result<Option<Metrics>, sqlx::Error> {
   145→    sqlx::query_as::<_, Metrics>(
   146→        r#"
   147→        SELECT * FROM metrics
   148→        WHERE agent_id = $1
   149→        ORDER BY timestamp DESC
   150→        LIMIT 1
   151→        "#,
   152→    )
   153→    .bind(agent_id)
   154→    .fetch_optional(pool)
   155→    .await
   156→}
   157→
   158→/// Delete old metrics (for cleanup jobs)
   159→pub async fn delete_old_metrics(
   160→    pool: &PgPool,
   161→    older_than: DateTime<Utc>,
   162→) -> Result<u64, sqlx::Error> {
   163→    let result = sqlx::query("DELETE FROM metrics WHERE timestamp < $1")
   164→        .bind(older_than)
   165→        .execute(pool)
   166→        .await?;
   167→
   168→    Ok(result.rows_affected())
   169→}
   170→
   171→/// Summary statistics for dashboard
   172→#[derive(Debug, Clone, Serialize, Deserialize)]
   173→pub struct MetricsSummary {
   174→    pub avg_cpu: f32,
   175→    pub avg_memory: f32,
   176→    pub avg_disk: f32,
   177→    pub total_network_rx: i64,
   178→    pub total_network_tx: i64,
   179→}
   180→
   181→/// Get summary metrics across all agents (last hour)
   182→pub async fn get_metrics_summary(pool: &PgPool) -> Result<MetricsSummary, sqlx::Error> {
   183→    let result: (Option<f64>, Option<f64>, Option<f64>, Option<i64>, Option<i64>) = sqlx::query_as(
   184→        r#"
   185→        SELECT
   186→            AVG(cpu_percent)::float8,
   187→            AVG(memory_percent)::float8,
   188→            AVG(disk_percent)::float8,
   189→            SUM(network_rx_bytes),
   190→            SUM(network_tx_bytes)
   191→        FROM metrics
   192→        WHERE timestamp > NOW() - INTERVAL '1 hour'
   193→        "#,
   194→    )
   195→    .fetch_one(pool)
   196→    .await?;
   197→
   198→    Ok(MetricsSummary {
   199→        avg_cpu: result.0.unwrap_or(0.0) as f32,
   200→        avg_memory: result.1.unwrap_or(0.0) as f32,
   201→        avg_disk: result.2.unwrap_or(0.0) as f32,
   202→        total_network_rx: result.3.unwrap_or(0),
   203→        total_network_tx: result.4.unwrap_or(0),
   204→    })
   205→}
   206→
   207→// ============================================================================
   208→// Agent State (network interfaces, extended info snapshot)
   209→// ============================================================================
   210→
   211→/// Agent state record from database
   212→#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
   213→pub struct AgentState {
   214→    pub agent_id: Uuid,
   215→    pub network_interfaces: Option<serde_json::Value>,
   216→    pub network_state_hash: Option<String>,
   217→    pub uptime_seconds: Option<i64>,
   218→    pub boot_time: Option<i64>,
   219→    pub logged_in_user: Option<String>,
   220→    pub user_idle_seconds: Option<i64>,
   221→    pub public_ip: Option<String>,
   222→    pub network_updated_at: Option<DateTime<Utc>>,
   223→    pub metrics_updated_at: Option<DateTime<Utc>>,
   224→}
   225→
   226→/// Update or insert agent state (upsert)
   227→pub async fn upsert_agent_state(
   228→    pool: &PgPool,
   229→    agent_id: Uuid,
   230→    uptime_seconds: Option<i64>,
   231→    boot_time: Option<i64>,
   232→    logged_in_user: Option<&str>,
   233→    user_idle_seconds: Option<i64>,
   234→    public_ip: Option<&str>,
   235→) -> Result<(), sqlx::Error> {
   236→    sqlx::query(
   237→        r#"
   238→        INSERT INTO agent_state (
   239→            agent_id, uptime_seconds, boot_time, logged_in_user,
   240→            user_idle_seconds, public_ip, metrics_updated_at
   241→        )
   242→        VALUES ($1, $2, $3, $4, $5, $6, NOW())
   243→        ON CONFLICT (agent_id) DO UPDATE SET
   244→            uptime_seconds = EXCLUDED.uptime_seconds,
   245→            boot_time = EXCLUDED.boot_time,
   246→            logged_in_user = EXCLUDED.logged_in_user,
   247→            user_idle_seconds = EXCLUDED.user_idle_seconds,
   248→            public_ip = EXCLUDED.public_ip,
   249→            metrics_updated_at = NOW()
   250→        "#,
   251→    )
   252→    .bind(agent_id)
   253→    .bind(uptime_seconds)
   254→    .bind(boot_time)
   255→    .bind(logged_in_user)
   256→    .bind(user_idle_seconds)
   257→    .bind(public_ip)
   258→    .execute(pool)
   259→    .await?;
   260→
   261→    Ok(())
   262→}
   263→
   264→/// Update network state for an agent
   265→pub async fn update_agent_network_state(
   266→    pool: &PgPool,
   267→    agent_id: Uuid,
   268→    interfaces: &serde_json::Value,
   269→    state_hash: &str,
   270→) -> Result<(), sqlx::Error> {
   271→    sqlx::query(
   272→        r#"
   273→        INSERT INTO agent_state (agent_id, network_interfaces, network_state_hash, network_updated_at)
   274→        VALUES ($1, $2, $3, NOW())
   275→        ON CONFLICT (agent_id) DO UPDATE SET
   276→            network_interfaces = EXCLUDED.network_interfaces,
   277→            network_state_hash = EXCLUDED.network_state_hash,
   278→            network_updated_at = NOW()
   279→        "#,
   280→    )
   281→    .bind(agent_id)
   282→    .bind(interfaces)
   283→    .bind(state_hash)
   284→    .execute(pool)
   285→    .await?;
   286→
   287→    Ok(())
   288→}
   289→
   290→/// Get agent state by agent ID
   291→pub async fn get_agent_state(pool: &PgPool, agent_id: Uuid) -> Result<Option<AgentState>, sqlx::Error> {
   292→    sqlx::query_as::<_, AgentState>("SELECT * FROM agent_state WHERE agent_id = $1")
   293→        .bind(agent_id)
   294→        .fetch_optional(pool)
   295→        .await
   296→}
   297→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! Metrics database operations

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use uuid::Uuid;

/// Metrics record from database
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct Metrics {
    pub id: i64,
    pub agent_id: Uuid,
    pub timestamp: DateTime<Utc>,
    pub cpu_percent: Option<f32>,
    pub memory_percent: Option<f32>,
    pub memory_used_bytes: Option<i64>,
    pub disk_percent: Option<f32>,
    pub disk_used_bytes: Option<i64>,
    pub network_rx_bytes: Option<i64>,
    pub network_tx_bytes: Option<i64>,
    // Extended metrics
    pub uptime_seconds: Option<i64>,
    pub boot_time: Option<i64>,
    pub logged_in_user: Option<String>,
    pub user_idle_seconds: Option<i64>,
    pub public_ip: Option<String>,
    pub memory_total_bytes: Option<i64>,
    pub disk_total_bytes: Option<i64>,
    // Temperature metrics
    pub cpu_temp_celsius: Option<f32>,
    pub gpu_temp_celsius: Option<f32>,
    pub temperatures: Option<serde_json::Value>,
}

/// Create metrics data from agent
#[derive(Debug, Clone, Deserialize)]
pub struct CreateMetrics {
    pub agent_id: Uuid,
    pub cpu_percent: Option<f32>,
    pub memory_percent: Option<f32>,
    pub memory_used_bytes: Option<i64>,
    pub disk_percent: Option<f32>,
    pub disk_used_bytes: Option<i64>,
    pub network_rx_bytes: Option<i64>,
    pub network_tx_bytes: Option<i64>,
    // Extended metrics
    pub uptime_seconds: Option<i64>,
    pub boot_time: Option<i64>,
    pub logged_in_user: Option<String>,
    pub user_idle_seconds: Option<i64>,
    pub public_ip: Option<String>,
    pub memory_total_bytes: Option<i64>,
    pub disk_total_bytes: Option<i64>,
    // Temperature metrics
    pub cpu_temp_celsius: Option<f32>,
    pub gpu_temp_celsius: Option<f32>,
    pub temperatures: Option<serde_json::Value>,
}

/// Insert metrics into the database
pub async fn insert_metrics(pool: &PgPool, metrics: CreateMetrics) -> Result<i64, sqlx::Error> {
    let result: (i64,) = sqlx::query_as(
        r#"
        INSERT INTO metrics (
            agent_id, cpu_percent, memory_percent, memory_used_bytes,
            disk_percent, disk_used_bytes, network_rx_bytes, network_tx_bytes,
            uptime_seconds, boot_time, logged_in_user, user_idle_seconds,
            public_ip, memory_total_bytes, disk_total_bytes,
            cpu_temp_celsius, gpu_temp_celsius, temperatures
        )
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18)
        RETURNING id
        "#,
    )
    .bind(metrics.agent_id)
    .bind(metrics.cpu_percent)
    .bind(metrics.memory_percent)
    .bind(metrics.memory_used_bytes)
    .bind(metrics.disk_percent)
    .bind(metrics.disk_used_bytes)
    .bind(metrics.network_rx_bytes)
    .bind(metrics.network_tx_bytes)
    .bind(metrics.uptime_seconds)
    .bind(metrics.boot_time)
    .bind(&metrics.logged_in_user)
    .bind(metrics.user_idle_seconds)
    .bind(&metrics.public_ip)
    .bind(metrics.memory_total_bytes)
    .bind(metrics.disk_total_bytes)
    .bind(metrics.cpu_temp_celsius)
    .bind(metrics.gpu_temp_celsius)
    .bind(&metrics.temperatures)
    .fetch_one(pool)
    .await?;

    Ok(result.0)
}

/// Get recent metrics for an agent
pub async fn get_agent_metrics(
    pool: &PgPool,
    agent_id: Uuid,
    limit: i64,
) -> Result<Vec<Metrics>, sqlx::Error> {
    sqlx::query_as::<_, Metrics>(
        r#"
        SELECT * FROM metrics
        WHERE agent_id = $1
        ORDER BY timestamp DESC
        LIMIT $2
        "#,
    )
    .bind(agent_id)
    .bind(limit)
    .fetch_all(pool)
    .await
}

/// Get metrics for an agent within a time range
pub async fn get_agent_metrics_range(
    pool: &PgPool,
    agent_id: Uuid,
    start: DateTime<Utc>,
    end: DateTime<Utc>,
) -> Result<Vec<Metrics>, sqlx::Error> {
    sqlx::query_as::<_, Metrics>(
        r#"
        SELECT * FROM metrics
        WHERE agent_id = $1 AND timestamp >= $2 AND timestamp <= $3
        ORDER BY timestamp ASC
        "#,
    )
    .bind(agent_id)
    .bind(start)
    .bind(end)
    .fetch_all(pool)
    .await
}

/// Get latest metrics for an agent
pub async fn get_latest_metrics(
    pool: &PgPool,
    agent_id: Uuid,
) -> Result<Option<Metrics>, sqlx::Error> {
    sqlx::query_as::<_, Metrics>(
        r#"
        SELECT * FROM metrics
        WHERE agent_id = $1
        ORDER BY timestamp DESC
        LIMIT 1
        "#,
    )
    .bind(agent_id)
    .fetch_optional(pool)
    .await
}

/// Delete old metrics (for cleanup jobs)
pub async fn delete_old_metrics(
    pool: &PgPool,
    older_than: DateTime<Utc>,
) -> Result<u64, sqlx::Error> {
    let result = sqlx::query("DELETE FROM metrics WHERE timestamp < $1")
        .bind(older_than)
        .execute(pool)
        .await?;

    Ok(result.rows_affected())
}

/// Summary statistics for dashboard
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MetricsSummary {
    pub avg_cpu: f32,
    pub avg_memory: f32,
    pub avg_disk: f32,
    pub total_network_rx: i64,
    pub total_network_tx: i64,
}

/// Get summary metrics across all agents (last hour)
pub async fn get_metrics_summary(pool: &PgPool) -> Result<MetricsSummary, sqlx::Error> {
    let result: (Option<f64>, Option<f64>, Option<f64>, Option<i64>, Option<i64>) = sqlx::query_as(
        r#"
        SELECT
            AVG(cpu_percent)::float8,
            AVG(memory_percent)::float8,
            AVG(disk_percent)::float8,
            SUM(network_rx_bytes),
            SUM(network_tx_bytes)
        FROM metrics
        WHERE timestamp > NOW() - INTERVAL '1 hour'
        "#,
    )
    .fetch_one(pool)
    .await?;

    Ok(MetricsSummary {
        avg_cpu: result.0.unwrap_or(0.0) as f32,
        avg_memory: result.1.unwrap_or(0.0) as f32,
        avg_disk: result.2.unwrap_or(0.0) as f32,
        total_network_rx: result.3.unwrap_or(0),
        total_network_tx: result.4.unwrap_or(0),
    })
}

// ============================================================================
// Agent State (network interfaces, extended info snapshot)
// ============================================================================

/// Agent state record from database
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct AgentState {
    pub agent_id: Uuid,
    pub network_interfaces: Option<serde_json::Value>,
    pub network_state_hash: Option<String>,
    pub uptime_seconds: Option<i64>,
    pub boot_time: Option<i64>,
    pub logged_in_user: Option<String>,
    pub user_idle_seconds: Option<i64>,
    pub public_ip: Option<String>,
    pub network_updated_at: Option<DateTime<Utc>>,
    pub metrics_updated_at: Option<DateTime<Utc>>,
}

/// Update or insert agent state (upsert)
pub async fn upsert_agent_state(
    pool: &PgPool,
    agent_id: Uuid,
    uptime_seconds: Option<i64>,
    boot_time: Option<i64>,
    logged_in_user: Option<&str>,
    user_idle_seconds: Option<i64>,
    public_ip: Option<&str>,
) -> Result<(), sqlx::Error> {
    sqlx::query(
        r#"
        INSERT INTO agent_state (
            agent_id, uptime_seconds, boot_time, logged_in_user,
            user_idle_seconds, public_ip, metrics_updated_at
        )
        VALUES ($1, $2, $3, $4, $5, $6, NOW())
        ON CONFLICT (agent_id) DO UPDATE SET
            uptime_seconds = EXCLUDED.uptime_seconds,
            boot_time = EXCLUDED.boot_time,
            logged_in_user = EXCLUDED.logged_in_user,
            user_idle_seconds = EXCLUDED.user_idle_seconds,
            public_ip = EXCLUDED.public_ip,
            metrics_updated_at = NOW()
        "#,
    )
    .bind(agent_id)
    .bind(uptime_seconds)
    .bind(boot_time)
    .bind(logged_in_user)
    .bind(user_idle_seconds)
    .bind(public_ip)
    .execute(pool)
    .await?;

    Ok(())
}

/// Update network state for an agent
pub async fn update_agent_network_state(
    pool: &PgPool,
    agent_id: Uuid,
    interfaces: &serde_json::Value,
    state_hash: &str,
) -> Result<(), sqlx::Error> {
    sqlx::query(
        r#"
        INSERT INTO agent_state (agent_id, network_interfaces, network_state_hash, network_updated_at)
        VALUES ($1, $2, $3, NOW())
        ON CONFLICT (agent_id) DO UPDATE SET
            network_interfaces = EXCLUDED.network_interfaces,
            network_state_hash = EXCLUDED.network_state_hash,
            network_updated_at = NOW()
        "#,
    )
    .bind(agent_id)
    .bind(interfaces)
    .bind(state_hash)
    .execute(pool)
    .await?;

    Ok(())
}

/// Get agent state by agent ID
pub async fn get_agent_state(pool: &PgPool, agent_id: Uuid) -> Result<Option<AgentState>, sqlx::Error> {
    sqlx::query_as::<_, AgentState>("SELECT * FROM agent_state WHERE agent_id = $1")
        .bind(agent_id)
        .fetch_optional(pool)
        .await
}


C:\Users\MikeSwanson\claude-projects\gururmm\agent\src\ipc.rs
C:\Users\MikeSwanson\claude-projects\gururmm\tray\src\ipc.rs