455bc75044
Critical regression fix - subsystems were hardcoded instead of reading user settings. Added CreateDefaultSubsystems to queries/subsystems.go. ConfigService now queries agent_subsystems table for actual user configuration. AgentLifecycleService creates default subsystems when creating new agents. Respects user-configured enabled/auto_run settings from UI.
311 lines
8.5 KiB
Go
311 lines
8.5 KiB
Go
package queries
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
|
|
"github.com/Fimeg/RedFlag/aggregator-server/internal/models"
|
|
"github.com/google/uuid"
|
|
"github.com/jmoiron/sqlx"
|
|
)
|
|
|
|
type SubsystemQueries struct {
|
|
db *sqlx.DB
|
|
}
|
|
|
|
func NewSubsystemQueries(db *sqlx.DB) *SubsystemQueries {
|
|
return &SubsystemQueries{db: db}
|
|
}
|
|
|
|
// GetSubsystems retrieves all subsystems for an agent
|
|
func (q *SubsystemQueries) GetSubsystems(agentID uuid.UUID) ([]models.AgentSubsystem, error) {
|
|
query := `
|
|
SELECT id, agent_id, subsystem, enabled, interval_minutes, auto_run,
|
|
last_run_at, next_run_at, created_at, updated_at
|
|
FROM agent_subsystems
|
|
WHERE agent_id = $1
|
|
ORDER BY subsystem
|
|
`
|
|
|
|
var subsystems []models.AgentSubsystem
|
|
err := q.db.Select(&subsystems, query, agentID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get subsystems: %w", err)
|
|
}
|
|
|
|
return subsystems, nil
|
|
}
|
|
|
|
// GetSubsystem retrieves a specific subsystem for an agent
|
|
func (q *SubsystemQueries) GetSubsystem(agentID uuid.UUID, subsystem string) (*models.AgentSubsystem, error) {
|
|
query := `
|
|
SELECT id, agent_id, subsystem, enabled, interval_minutes, auto_run,
|
|
last_run_at, next_run_at, created_at, updated_at
|
|
FROM agent_subsystems
|
|
WHERE agent_id = $1 AND subsystem = $2
|
|
`
|
|
|
|
var sub models.AgentSubsystem
|
|
err := q.db.Get(&sub, query, agentID, subsystem)
|
|
if err == sql.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get subsystem: %w", err)
|
|
}
|
|
|
|
return &sub, nil
|
|
}
|
|
|
|
// UpdateSubsystem updates a subsystem configuration
|
|
func (q *SubsystemQueries) UpdateSubsystem(agentID uuid.UUID, subsystem string, config models.SubsystemConfig) error {
|
|
// Build dynamic update query based on provided fields
|
|
updates := []string{}
|
|
args := []interface{}{agentID, subsystem}
|
|
argIdx := 3
|
|
|
|
if config.Enabled != nil {
|
|
updates = append(updates, fmt.Sprintf("enabled = $%d", argIdx))
|
|
args = append(args, *config.Enabled)
|
|
argIdx++
|
|
}
|
|
|
|
if config.IntervalMinutes != nil {
|
|
updates = append(updates, fmt.Sprintf("interval_minutes = $%d", argIdx))
|
|
args = append(args, *config.IntervalMinutes)
|
|
argIdx++
|
|
}
|
|
|
|
if config.AutoRun != nil {
|
|
updates = append(updates, fmt.Sprintf("auto_run = $%d", argIdx))
|
|
args = append(args, *config.AutoRun)
|
|
argIdx++
|
|
|
|
// If enabling auto_run, calculate next_run_at
|
|
if *config.AutoRun {
|
|
updates = append(updates, fmt.Sprintf("next_run_at = NOW() + INTERVAL '%d minutes'", argIdx))
|
|
}
|
|
}
|
|
|
|
if len(updates) == 0 {
|
|
return fmt.Errorf("no fields to update")
|
|
}
|
|
|
|
updates = append(updates, "updated_at = NOW()")
|
|
|
|
query := fmt.Sprintf(`
|
|
UPDATE agent_subsystems
|
|
SET %s
|
|
WHERE agent_id = $1 AND subsystem = $2
|
|
`, joinUpdates(updates))
|
|
|
|
result, err := q.db.Exec(query, args...)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to update subsystem: %w", err)
|
|
}
|
|
|
|
rows, err := result.RowsAffected()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get rows affected: %w", err)
|
|
}
|
|
|
|
if rows == 0 {
|
|
return fmt.Errorf("subsystem not found")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// UpdateLastRun updates the last_run_at timestamp for a subsystem
|
|
func (q *SubsystemQueries) UpdateLastRun(agentID uuid.UUID, subsystem string) error {
|
|
query := `
|
|
UPDATE agent_subsystems
|
|
SET last_run_at = NOW(),
|
|
next_run_at = CASE
|
|
WHEN auto_run THEN NOW() + (interval_minutes || ' minutes')::INTERVAL
|
|
ELSE next_run_at
|
|
END,
|
|
updated_at = NOW()
|
|
WHERE agent_id = $1 AND subsystem = $2
|
|
`
|
|
|
|
result, err := q.db.Exec(query, agentID, subsystem)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to update last run: %w", err)
|
|
}
|
|
|
|
rows, err := result.RowsAffected()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get rows affected: %w", err)
|
|
}
|
|
|
|
if rows == 0 {
|
|
return fmt.Errorf("subsystem not found")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetDueSubsystems retrieves all subsystems that are due to run
|
|
func (q *SubsystemQueries) GetDueSubsystems() ([]models.AgentSubsystem, error) {
|
|
query := `
|
|
SELECT id, agent_id, subsystem, enabled, interval_minutes, auto_run,
|
|
last_run_at, next_run_at, created_at, updated_at
|
|
FROM agent_subsystems
|
|
WHERE enabled = true
|
|
AND auto_run = true
|
|
AND (next_run_at IS NULL OR next_run_at <= NOW())
|
|
ORDER BY next_run_at ASC NULLS FIRST
|
|
LIMIT 1000
|
|
`
|
|
|
|
var subsystems []models.AgentSubsystem
|
|
err := q.db.Select(&subsystems, query)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get due subsystems: %w", err)
|
|
}
|
|
|
|
return subsystems, nil
|
|
}
|
|
|
|
// GetSubsystemStats retrieves statistics for a subsystem
|
|
func (q *SubsystemQueries) GetSubsystemStats(agentID uuid.UUID, subsystem string) (*models.SubsystemStats, error) {
|
|
query := `
|
|
SELECT
|
|
s.subsystem,
|
|
s.enabled,
|
|
s.last_run_at,
|
|
s.next_run_at,
|
|
s.interval_minutes,
|
|
s.auto_run,
|
|
COUNT(c.id) FILTER (WHERE c.command_type = 'scan_' || s.subsystem) as run_count,
|
|
MAX(c.status) FILTER (WHERE c.command_type = 'scan_' || s.subsystem) as last_status,
|
|
MAX(al.duration_seconds) FILTER (WHERE al.action = 'scan_' || s.subsystem) as last_duration
|
|
FROM agent_subsystems s
|
|
LEFT JOIN agent_commands c ON c.agent_id = s.agent_id
|
|
LEFT JOIN agent_logs al ON al.command_id = c.id
|
|
WHERE s.agent_id = $1 AND s.subsystem = $2
|
|
GROUP BY s.subsystem, s.enabled, s.last_run_at, s.next_run_at, s.interval_minutes, s.auto_run
|
|
`
|
|
|
|
var stats models.SubsystemStats
|
|
err := q.db.Get(&stats, query, agentID, subsystem)
|
|
if err == sql.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get subsystem stats: %w", err)
|
|
}
|
|
|
|
return &stats, nil
|
|
}
|
|
|
|
// EnableSubsystem enables a subsystem
|
|
func (q *SubsystemQueries) EnableSubsystem(agentID uuid.UUID, subsystem string) error {
|
|
enabled := true
|
|
return q.UpdateSubsystem(agentID, subsystem, models.SubsystemConfig{
|
|
Enabled: &enabled,
|
|
})
|
|
}
|
|
|
|
// DisableSubsystem disables a subsystem
|
|
func (q *SubsystemQueries) DisableSubsystem(agentID uuid.UUID, subsystem string) error {
|
|
enabled := false
|
|
return q.UpdateSubsystem(agentID, subsystem, models.SubsystemConfig{
|
|
Enabled: &enabled,
|
|
})
|
|
}
|
|
|
|
// SetAutoRun enables or disables auto-run for a subsystem
|
|
func (q *SubsystemQueries) SetAutoRun(agentID uuid.UUID, subsystem string, autoRun bool) error {
|
|
return q.UpdateSubsystem(agentID, subsystem, models.SubsystemConfig{
|
|
AutoRun: &autoRun,
|
|
})
|
|
}
|
|
|
|
// SetInterval sets the interval for a subsystem
|
|
func (q *SubsystemQueries) SetInterval(agentID uuid.UUID, subsystem string, intervalMinutes int) error {
|
|
return q.UpdateSubsystem(agentID, subsystem, models.SubsystemConfig{
|
|
IntervalMinutes: &intervalMinutes,
|
|
})
|
|
}
|
|
|
|
// CreateSubsystem creates a new subsystem configuration (used for custom subsystems)
|
|
func (q *SubsystemQueries) CreateSubsystem(sub *models.AgentSubsystem) error {
|
|
query := `
|
|
INSERT INTO agent_subsystems (agent_id, subsystem, enabled, interval_minutes, auto_run, last_run_at, next_run_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
|
RETURNING id, created_at, updated_at
|
|
`
|
|
|
|
err := q.db.QueryRow(
|
|
query,
|
|
sub.AgentID,
|
|
sub.Subsystem,
|
|
sub.Enabled,
|
|
sub.IntervalMinutes,
|
|
sub.AutoRun,
|
|
sub.LastRunAt,
|
|
sub.NextRunAt,
|
|
).Scan(&sub.ID, &sub.CreatedAt, &sub.UpdatedAt)
|
|
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create subsystem: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// DeleteSubsystem deletes a subsystem configuration
|
|
func (q *SubsystemQueries) DeleteSubsystem(agentID uuid.UUID, subsystem string) error {
|
|
query := `
|
|
DELETE FROM agent_subsystems
|
|
WHERE agent_id = $1 AND subsystem = $2
|
|
`
|
|
|
|
result, err := q.db.Exec(query, agentID, subsystem)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to delete subsystem: %w", err)
|
|
}
|
|
|
|
rows, err := result.RowsAffected()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get rows affected: %w", err)
|
|
}
|
|
|
|
if rows == 0 {
|
|
return fmt.Errorf("subsystem not found")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// CreateDefaultSubsystems creates default subsystems for a new agent
|
|
func (q *SubsystemQueries) CreateDefaultSubsystems(agentID uuid.UUID) error {
|
|
defaults := []models.AgentSubsystem{
|
|
{AgentID: agentID, Subsystem: "updates", Enabled: true, AutoRun: true, IntervalMinutes: 60},
|
|
{AgentID: agentID, Subsystem: "storage", Enabled: true, AutoRun: true, IntervalMinutes: 5},
|
|
{AgentID: agentID, Subsystem: "system", Enabled: true, AutoRun: true, IntervalMinutes: 5},
|
|
{AgentID: agentID, Subsystem: "docker", Enabled: true, AutoRun: true, IntervalMinutes: 15},
|
|
}
|
|
|
|
for _, sub := range defaults {
|
|
if err := q.CreateSubsystem(&sub); err != nil {
|
|
return fmt.Errorf("failed to create subsystem %s: %w", sub.Subsystem, err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Helper function to join update statements
|
|
func joinUpdates(updates []string) string {
|
|
result := ""
|
|
for i, update := range updates {
|
|
if i > 0 {
|
|
result += ", "
|
|
}
|
|
result += update
|
|
}
|
|
return result
|
|
}
|