package queries import ( "database/sql" "fmt" "time" "github.com/Fimeg/RedFlag/aggregator-server/internal/models" "github.com/google/uuid" "github.com/jmoiron/sqlx" ) type AgentQueries struct { db *sqlx.DB } func NewAgentQueries(db *sqlx.DB) *AgentQueries { return &AgentQueries{db: db} } // CreateAgent inserts a new agent into the database func (q *AgentQueries) CreateAgent(agent *models.Agent) error { query := ` INSERT INTO agents ( id, hostname, os_type, os_version, os_architecture, agent_version, last_seen, status, metadata ) VALUES ( :id, :hostname, :os_type, :os_version, :os_architecture, :agent_version, :last_seen, :status, :metadata ) ` _, err := q.db.NamedExec(query, agent) return err } // GetAgentByID retrieves an agent by ID func (q *AgentQueries) GetAgentByID(id uuid.UUID) (*models.Agent, error) { var agent models.Agent query := `SELECT * FROM agents WHERE id = $1` err := q.db.Get(&agent, query, id) if err != nil { return nil, err } return &agent, nil } // UpdateAgentLastSeen updates the agent's last_seen timestamp func (q *AgentQueries) UpdateAgentLastSeen(id uuid.UUID) error { query := `UPDATE agents SET last_seen = $1, status = 'online' WHERE id = $2` _, err := q.db.Exec(query, time.Now().UTC(), id) return err } // UpdateAgent updates an agent's full record including metadata func (q *AgentQueries) UpdateAgent(agent *models.Agent) error { query := ` UPDATE agents SET hostname = :hostname, os_type = :os_type, os_version = :os_version, os_architecture = :os_architecture, agent_version = :agent_version, last_seen = :last_seen, status = :status, metadata = :metadata WHERE id = :id ` _, err := q.db.NamedExec(query, agent) return err } // UpdateAgentMetadata updates only the metadata, last_seen, and status fields // Used for metrics updates to avoid overwriting version tracking func (q *AgentQueries) UpdateAgentMetadata(id uuid.UUID, metadata models.JSONB, status string, lastSeen time.Time) error { query := ` UPDATE agents SET last_seen = $1, status = $2, metadata = $3 WHERE id = $4 ` _, err := q.db.Exec(query, lastSeen, status, metadata, id) return err } // ListAgents returns all agents with optional filtering func (q *AgentQueries) ListAgents(status, osType string) ([]models.Agent, error) { var agents []models.Agent query := `SELECT * FROM agents WHERE 1=1` args := []interface{}{} argIdx := 1 if status != "" { query += ` AND status = $` + string(rune(argIdx+'0')) args = append(args, status) argIdx++ } if osType != "" { query += ` AND os_type = $` + string(rune(argIdx+'0')) args = append(args, osType) } query += ` ORDER BY last_seen DESC` err := q.db.Select(&agents, query, args...) return agents, err } // MarkOfflineAgents marks agents as offline if they haven't checked in recently func (q *AgentQueries) MarkOfflineAgents(threshold time.Duration) error { query := ` UPDATE agents SET status = 'offline' WHERE last_seen < $1 AND status = 'online' ` _, err := q.db.Exec(query, time.Now().Add(-threshold)) return err } // GetAgentLastScan gets the last scan time from update events func (q *AgentQueries) GetAgentLastScan(id uuid.UUID) (*time.Time, error) { var lastScan time.Time query := `SELECT MAX(created_at) FROM update_events WHERE agent_id = $1` err := q.db.Get(&lastScan, query, id) if err != nil { return nil, err } return &lastScan, nil } // GetAgentWithLastScan gets agent information including last scan time func (q *AgentQueries) GetAgentWithLastScan(id uuid.UUID) (*models.AgentWithLastScan, error) { var agent models.AgentWithLastScan query := ` SELECT a.*, (SELECT MAX(created_at) FROM update_events WHERE agent_id = a.id) as last_scan FROM agents a WHERE a.id = $1` err := q.db.Get(&agent, query, id) if err != nil { return nil, err } return &agent, nil } // ListAgentsWithLastScan returns all agents with their last scan times func (q *AgentQueries) ListAgentsWithLastScan(status, osType string) ([]models.AgentWithLastScan, error) { var agents []models.AgentWithLastScan query := ` SELECT a.*, (SELECT MAX(created_at) FROM update_events WHERE agent_id = a.id) as last_scan FROM agents a WHERE 1=1` args := []interface{}{} argIdx := 1 if status != "" { query += ` AND a.status = $` + string(rune(argIdx+'0')) args = append(args, status) argIdx++ } if osType != "" { query += ` AND a.os_type = $` + string(rune(argIdx+'0')) args = append(args, osType) argIdx++ } query += ` ORDER BY a.last_seen DESC` err := q.db.Select(&agents, query, args...) return agents, err } // UpdateAgentVersion updates the agent's version information and checks for updates func (q *AgentQueries) UpdateAgentVersion(id uuid.UUID, currentVersion string) error { query := ` UPDATE agents SET current_version = $1, last_version_check = $2 WHERE id = $3 ` _, err := q.db.Exec(query, currentVersion, time.Now().UTC(), id) return err } // UpdateAgentUpdateAvailable sets whether an update is available for an agent func (q *AgentQueries) UpdateAgentUpdateAvailable(id uuid.UUID, updateAvailable bool) error { query := ` UPDATE agents SET update_available = $1 WHERE id = $2 ` _, err := q.db.Exec(query, updateAvailable, id) return err } // DeleteAgent removes an agent and all associated data func (q *AgentQueries) DeleteAgent(id uuid.UUID) error { // Start a transaction for atomic deletion tx, err := q.db.Beginx() if err != nil { return err } defer tx.Rollback() // Delete the agent (CASCADE will handle related records) _, err = tx.Exec("DELETE FROM agents WHERE id = $1", id) if err != nil { return err } // Commit the transaction return tx.Commit() } // GetActiveAgentCount returns the count of active (online) agents func (q *AgentQueries) GetActiveAgentCount() (int, error) { var count int query := `SELECT COUNT(*) FROM agents WHERE status = 'online'` err := q.db.Get(&count, query) return count, err } // UpdateAgentRebootStatus updates the reboot status for an agent func (q *AgentQueries) UpdateAgentRebootStatus(id uuid.UUID, required bool, reason string) error { query := ` UPDATE agents SET reboot_required = $1, reboot_reason = $2, updated_at = $3 WHERE id = $4 ` _, err := q.db.Exec(query, required, reason, time.Now(), id) return err } // UpdateAgentLastReboot updates the last reboot timestamp for an agent func (q *AgentQueries) UpdateAgentLastReboot(id uuid.UUID, rebootTime time.Time) error { query := ` UPDATE agents SET last_reboot_at = $1, reboot_required = FALSE, reboot_reason = '', updated_at = $2 WHERE id = $3 ` _, err := q.db.Exec(query, rebootTime, time.Now(), id) return err } // GetAgentByMachineID retrieves an agent by its machine ID func (q *AgentQueries) GetAgentByMachineID(machineID string) (*models.Agent, error) { query := ` SELECT id, hostname, os_type, os_version, os_architecture, agent_version, current_version, update_available, last_version_check, machine_id, public_key_fingerprint, is_updating, updating_to_version, update_initiated_at, last_seen, status, metadata, reboot_required, last_reboot_at, reboot_reason, created_at, updated_at FROM agents WHERE machine_id = $1 ` var agent models.Agent err := q.db.Get(&agent, query, machineID) if err != nil { if err == sql.ErrNoRows { return nil, nil // Return nil if not found (not an error) } return nil, fmt.Errorf("failed to get agent by machine ID: %w", err) } return &agent, nil } // UpdateAgentUpdatingStatus updates the agent's update status func (q *AgentQueries) UpdateAgentUpdatingStatus(id uuid.UUID, isUpdating bool, updatingToVersion *string) error { query := ` UPDATE agents SET is_updating = $1, updating_to_version = $2, update_initiated_at = CASE WHEN $1 = true THEN $3 ELSE NULL END, updated_at = $3 WHERE id = $4 ` var versionPtr *string if updatingToVersion != nil { versionPtr = updatingToVersion } _, err := q.db.Exec(query, isUpdating, versionPtr, time.Now(), id) return err }