package queries import ( "database/sql" "fmt" "time" "github.com/Fimeg/RedFlag/aggregator-server/internal/models" "github.com/google/uuid" "github.com/jmoiron/sqlx" ) // AgentUpdateQueries handles database operations for agent update packages type AgentUpdateQueries struct { db *sqlx.DB } // NewAgentUpdateQueries creates a new AgentUpdateQueries instance func NewAgentUpdateQueries(db *sqlx.DB) *AgentUpdateQueries { return &AgentUpdateQueries{db: db} } // CreateUpdatePackage stores a new signed update package func (q *AgentUpdateQueries) CreateUpdatePackage(pkg *models.AgentUpdatePackage) error { query := ` INSERT INTO agent_update_packages ( id, version, platform, architecture, binary_path, signature, checksum, file_size, created_by, is_active ) VALUES ( :id, :version, :platform, :architecture, :binary_path, :signature, :checksum, :file_size, :created_by, :is_active ) RETURNING id, created_at ` rows, err := q.db.NamedQuery(query, pkg) if err != nil { return fmt.Errorf("failed to create update package: %w", err) } defer rows.Close() if rows.Next() { if err := rows.Scan(&pkg.ID, &pkg.CreatedAt); err != nil { return fmt.Errorf("failed to scan created package: %w", err) } } return nil } // GetUpdatePackage retrieves an update package by ID func (q *AgentUpdateQueries) GetUpdatePackage(id uuid.UUID) (*models.AgentUpdatePackage, error) { query := ` SELECT id, version, platform, architecture, binary_path, signature, checksum, file_size, created_at, created_by, is_active FROM agent_update_packages WHERE id = $1 AND is_active = true ` var pkg models.AgentUpdatePackage err := q.db.Get(&pkg, query, id) if err != nil { if err == sql.ErrNoRows { return nil, fmt.Errorf("update package not found") } return nil, fmt.Errorf("failed to get update package: %w", err) } return &pkg, nil } // GetUpdatePackageByVersion retrieves the latest update package for a version and platform func (q *AgentUpdateQueries) GetUpdatePackageByVersion(version, platform, architecture string) (*models.AgentUpdatePackage, error) { query := ` SELECT id, version, platform, architecture, binary_path, signature, checksum, file_size, created_at, created_by, is_active FROM agent_update_packages WHERE version = $1 AND platform = $2 AND architecture = $3 AND is_active = true ORDER BY created_at DESC LIMIT 1 ` var pkg models.AgentUpdatePackage err := q.db.Get(&pkg, query, version, platform, architecture) if err != nil { if err == sql.ErrNoRows { return nil, fmt.Errorf("no update package found for version %s on %s/%s", version, platform, architecture) } return nil, fmt.Errorf("failed to get update package: %w", err) } return &pkg, nil } // ListUpdatePackages retrieves all update packages with optional filtering func (q *AgentUpdateQueries) ListUpdatePackages(version, platform string, limit, offset int) ([]models.AgentUpdatePackage, error) { query := ` SELECT id, version, platform, architecture, binary_path, signature, checksum, file_size, created_at, created_by, is_active FROM agent_update_packages WHERE is_active = true ` args := []interface{}{} argIndex := 1 if version != "" { query += fmt.Sprintf(" AND version = $%d", argIndex) args = append(args, version) argIndex++ } if platform != "" { query += fmt.Sprintf(" AND platform = $%d", argIndex) args = append(args, platform) argIndex++ } query += " ORDER BY created_at DESC" if limit > 0 { query += fmt.Sprintf(" LIMIT $%d", argIndex) args = append(args, limit) argIndex++ if offset > 0 { query += fmt.Sprintf(" OFFSET $%d", argIndex) args = append(args, offset) } } var packages []models.AgentUpdatePackage err := q.db.Select(&packages, query, args...) if err != nil { return nil, fmt.Errorf("failed to list update packages: %w", err) } return packages, nil } // DeactivateUpdatePackage marks a package as inactive func (q *AgentUpdateQueries) DeactivateUpdatePackage(id uuid.UUID) error { query := `UPDATE agent_update_packages SET is_active = false WHERE id = $1` result, err := q.db.Exec(query, id) if err != nil { return fmt.Errorf("failed to deactivate update package: %w", err) } rowsAffected, err := result.RowsAffected() if err != nil { return fmt.Errorf("failed to get rows affected: %w", err) } if rowsAffected == 0 { return fmt.Errorf("no update package found to deactivate") } return nil } // UpdateAgentMachineInfo updates the machine ID and public key fingerprint for an agent func (q *AgentUpdateQueries) UpdateAgentMachineInfo(agentID uuid.UUID, machineID, publicKeyFingerprint string) error { query := ` UPDATE agents SET machine_id = $1, public_key_fingerprint = $2, updated_at = $3 WHERE id = $4 ` _, err := q.db.Exec(query, machineID, publicKeyFingerprint, time.Now().UTC(), agentID) if err != nil { return fmt.Errorf("failed to update agent machine info: %w", err) } return nil } // UpdateAgentUpdatingStatus sets the update status for an agent func (q *AgentUpdateQueries) UpdateAgentUpdatingStatus(agentID uuid.UUID, isUpdating bool, targetVersion *string) error { query := ` UPDATE agents SET is_updating = $1, updating_to_version = $2, update_initiated_at = CASE WHEN $1 = true THEN $3 ELSE update_initiated_at END, updated_at = $3 WHERE id = $4 ` now := time.Now().UTC() _, err := q.db.Exec(query, isUpdating, targetVersion, now, agentID) if err != nil { return fmt.Errorf("failed to update agent updating status: %w", err) } return nil } // GetAgentByMachineID retrieves an agent by its machine ID func (q *AgentUpdateQueries) 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, fmt.Errorf("agent not found for machine ID") } return nil, fmt.Errorf("failed to get agent by machine ID: %w", err) } return &agent, nil } // GetLatestVersion retrieves the latest available version for a platform func (q *AgentUpdateQueries) GetLatestVersion(platform string) (string, error) { query := ` SELECT version FROM agent_update_packages WHERE platform = $1 AND is_active = true ORDER BY version DESC LIMIT 1 ` var latestVersion string err := q.db.Get(&latestVersion, query, platform) if err != nil { if err == sql.ErrNoRows { return "", fmt.Errorf("no update packages available for platform %s", platform) } return "", fmt.Errorf("failed to get latest version: %w", err) } return latestVersion, nil } // GetLatestVersionByTypeAndArch retrieves the latest available version for a specific os_type and architecture func (q *AgentUpdateQueries) GetLatestVersionByTypeAndArch(osType, osArch string) (string, error) { query := ` SELECT version FROM agent_update_packages WHERE platform = $1 AND architecture = $2 AND is_active = true ORDER BY version DESC LIMIT 1 ` var latestVersion string err := q.db.Get(&latestVersion, query, osType, osArch) if err != nil { if err == sql.ErrNoRows { return "", fmt.Errorf("no update packages available for platform %s/%s", osType, osArch) } return "", fmt.Errorf("failed to get latest version: %w", err) } return latestVersion, nil } // GetPendingUpdateCommand retrieves the most recent pending update command for an agent func (q *AgentUpdateQueries) GetPendingUpdateCommand(agentID string) (*models.AgentCommand, error) { query := ` SELECT id, agent_id, command_type, params, status, source, created_at, sent_at, completed_at, result, retried_from_id FROM agent_commands WHERE agent_id = $1 AND command_type = 'install_update' AND status = 'pending' ORDER BY created_at DESC LIMIT 1 ` var command models.AgentCommand err := q.db.Get(&command, query, agentID) if err != nil { if err == sql.ErrNoRows { return nil, nil // No pending update command found } return nil, fmt.Errorf("failed to get pending update command: %w", err) } return &command, nil }