Add PostgreSQL database persistence
- Add connect_machines, connect_sessions, connect_session_events, connect_support_codes tables - Implement db module with connection pooling (sqlx) - Add machine persistence across server restarts - Add audit logging for session/viewer events - Support codes now persisted to database 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
118
server/src/db/machines.rs
Normal file
118
server/src/db/machines.rs
Normal file
@@ -0,0 +1,118 @@
|
||||
//! 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(())
|
||||
}
|
||||
Reference in New Issue
Block a user