From e72e9fc16fdacfeb8c35c3f17720d0edf71082fd Mon Sep 17 00:00:00 2001 From: Fimeg Date: Fri, 31 Oct 2025 15:03:59 -0400 Subject: [PATCH] feat: add host restart detection and fix agent version display Potential fixes for issues #4 and #6. Agent version display: - Set CurrentVersion during registration instead of waiting for first check-in - Update UI to show "Initial Registration" instead of "Unknown" Host restart detection: - Added reboot_required, last_reboot_at, reboot_reason fields to agents table - Agent now detects pending reboots (Debian/Ubuntu via /var/run/reboot-required, RHEL/Fedora via needs-restarting) - New reboot command type with 1-minute grace period - UI shows restart alerts and adds restart button in quick actions - Restart indicator badge in agent list The reboot detection runs during system info collection and gets reported back to the server automatically. Using shutdown command for now until we make the restart mechanism user-adjustable later - need to think on that. Also need to come up with a Windows derivative outside of reading event log for detecting reboots. --- aggregator-agent/cmd/agent/main.go | 94 +++++++++++++++++++ aggregator-agent/internal/system/info.go | 69 +++++++++++++- aggregator-server/cmd/server/main.go | 1 + .../internal/api/handlers/agents.go | 61 ++++++++++++ .../migrations/013_add_reboot_tracking.up.sql | 13 +++ .../internal/database/queries/agents.go | 27 ++++++ aggregator-server/internal/models/agent.go | 16 +++- aggregator-server/internal/models/command.go | 1 + aggregator-web/src/lib/api.ts | 8 ++ aggregator-web/src/pages/Agents.tsx | 63 ++++++++++++- aggregator-web/src/types/index.ts | 4 + 11 files changed, 346 insertions(+), 11 deletions(-) create mode 100644 aggregator-server/internal/database/migrations/013_add_reboot_tracking.up.sql diff --git a/aggregator-agent/cmd/agent/main.go b/aggregator-agent/cmd/agent/main.go index 2ee0aff..67ef491 100644 --- a/aggregator-agent/cmd/agent/main.go +++ b/aggregator-agent/cmd/agent/main.go @@ -6,6 +6,7 @@ import ( "log" "math/rand" "os" + "os/exec" "runtime" "strings" "time" @@ -532,6 +533,11 @@ func runAgent(cfg *config.Config) error { log.Printf("[Heartbeat] Error disabling heartbeat: %v\n", err) } + + case "reboot": + if err := handleReboot(apiClient, cfg, cmd.ID, cmd.Params); err != nil { + log.Printf("[Reboot] Error processing reboot command: %v\n", err) + } default: log.Printf("Unknown command type: %s\n", cmd.Type) } @@ -1400,6 +1406,94 @@ func reportSystemInfo(apiClient *client.Client, cfg *config.Config) error { return nil } +// handleReboot handles reboot command +func handleReboot(apiClient *client.Client, cfg *config.Config, commandID string, params map[string]interface{}) error { + log.Println("[Reboot] Processing reboot request...") + + // Parse parameters + delayMinutes := 1 // Default to 1 minute + message := "System reboot requested by RedFlag" + + if delay, ok := params["delay_minutes"]; ok { + if delayFloat, ok := delay.(float64); ok { + delayMinutes = int(delayFloat) + } + } + if msg, ok := params["message"].(string); ok && msg != "" { + message = msg + } + + log.Printf("[Reboot] Scheduling system reboot in %d minute(s): %s", delayMinutes, message) + + var cmd *exec.Cmd + + // Execute platform-specific reboot command + if runtime.GOOS == "linux" { + // Linux: shutdown -r +MINUTES "message" + cmd = exec.Command("shutdown", "-r", fmt.Sprintf("+%d", delayMinutes), message) + } else if runtime.GOOS == "windows" { + // Windows: shutdown /r /t SECONDS /c "message" + delaySeconds := delayMinutes * 60 + cmd = exec.Command("shutdown", "/r", "/t", fmt.Sprintf("%d", delaySeconds), "/c", message) + } else { + err := fmt.Errorf("reboot not supported on platform: %s", runtime.GOOS) + log.Printf("[Reboot] Error: %v", err) + + // Report failure + logReport := client.LogReport{ + CommandID: commandID, + Action: "reboot", + Result: "failed", + Stdout: "", + Stderr: err.Error(), + ExitCode: 1, + DurationSeconds: 0, + } + apiClient.ReportLog(cfg.AgentID, logReport) + return err + } + + // Execute reboot command + output, err := cmd.CombinedOutput() + if err != nil { + log.Printf("[Reboot] Failed to schedule reboot: %v", err) + log.Printf("[Reboot] Output: %s", string(output)) + + // Report failure + logReport := client.LogReport{ + CommandID: commandID, + Action: "reboot", + Result: "failed", + Stdout: string(output), + Stderr: err.Error(), + ExitCode: 1, + DurationSeconds: 0, + } + apiClient.ReportLog(cfg.AgentID, logReport) + return err + } + + log.Printf("[Reboot] System reboot scheduled successfully") + log.Printf("[Reboot] The system will reboot in %d minute(s)", delayMinutes) + + // Report success + logReport := client.LogReport{ + CommandID: commandID, + Action: "reboot", + Result: "success", + Stdout: fmt.Sprintf("System reboot scheduled for %d minute(s) from now. Message: %s", delayMinutes, message), + Stderr: "", + ExitCode: 0, + DurationSeconds: 0, + } + + if reportErr := apiClient.ReportLog(cfg.AgentID, logReport); reportErr != nil { + log.Printf("[Reboot] Failed to report reboot command result: %v", reportErr) + } + + return nil +} + // formatTimeSince formats a duration as "X time ago" func formatTimeSince(t time.Time) string { duration := time.Since(t) diff --git a/aggregator-agent/internal/system/info.go b/aggregator-agent/internal/system/info.go index 096b031..124d81a 100644 --- a/aggregator-agent/internal/system/info.go +++ b/aggregator-agent/internal/system/info.go @@ -21,6 +21,8 @@ type SystemInfo struct { DiskInfo []DiskInfo `json:"disk_info"` RunningProcesses int `json:"running_processes"` Uptime string `json:"uptime"` + RebootRequired bool `json:"reboot_required"` + RebootReason string `json:"reboot_reason"` Metadata map[string]string `json:"metadata"` } @@ -113,6 +115,11 @@ func GetSystemInfo(agentVersion string) (*SystemInfo, error) { } } + // Check if system requires reboot + rebootRequired, rebootReason := checkRebootRequired() + info.RebootRequired = rebootRequired + info.RebootReason = rebootReason + // Add collection timestamp info.Metadata["collected_at"] = time.Now().Format(time.RFC3339) @@ -467,4 +474,64 @@ func GetLightweightMetrics() (*LightweightMetrics, error) { // In the future, we could add a background goroutine to track CPU usage return metrics, nil -} \ No newline at end of file +} +// checkRebootRequired checks if the system requires a reboot +func checkRebootRequired() (bool, string) { + if runtime.GOOS == "linux" { + return checkLinuxRebootRequired() + } else if runtime.GOOS == "windows" { + return checkWindowsRebootRequired() + } + return false, "" +} + +// checkLinuxRebootRequired checks if a Linux system requires a reboot +func checkLinuxRebootRequired() (bool, string) { + // Method 1: Check Debian/Ubuntu reboot-required file + if err := exec.Command("test", "-f", "/var/run/reboot-required").Run(); err == nil { + // File exists, reboot is required + // Try to read the packages that require reboot + if output, err := exec.Command("cat", "/var/run/reboot-required.pkgs").Output(); err == nil { + packages := strings.TrimSpace(string(output)) + if packages != "" { + // Truncate if too long + if len(packages) > 200 { + packages = packages[:200] + "..." + } + return true, "Packages: " + packages + } + } + return true, "System updates require reboot" + } + + // Method 2: Check RHEL/Fedora/Rocky using needs-restarting + cmd := exec.Command("needs-restarting", "-r") + if err := cmd.Run(); err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + // Exit code 1 means reboot is needed + if exitErr.ExitCode() == 1 { + return true, "Kernel or system libraries updated" + } + } + } + + return false, "" +} + +// checkWindowsRebootRequired checks if a Windows system requires a reboot +func checkWindowsRebootRequired() (bool, string) { + // Check Windows Update pending reboot registry keys + // HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\WindowsUpdate\Auto Update\RebootRequired + cmd := exec.Command("reg", "query", "HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\WindowsUpdate\\Auto Update\\RebootRequired") + if err := cmd.Run(); err == nil { + return true, "Windows updates require reboot" + } + + // Check Component Based Servicing pending reboot + cmd = exec.Command("reg", "query", "HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Component Based Servicing\\RebootPending") + if err := cmd.Run(); err == nil { + return true, "Component updates require reboot" + } + + return false, "" +} diff --git a/aggregator-server/cmd/server/main.go b/aggregator-server/cmd/server/main.go index c80e418..8e181a6 100644 --- a/aggregator-server/cmd/server/main.go +++ b/aggregator-server/cmd/server/main.go @@ -206,6 +206,7 @@ func main() { dashboard.POST("/agents/:id/update", agentHandler.TriggerUpdate) dashboard.POST("/agents/:id/heartbeat", agentHandler.TriggerHeartbeat) dashboard.GET("/agents/:id/heartbeat", agentHandler.GetHeartbeatStatus) + dashboard.POST("/agents/:id/reboot", agentHandler.TriggerReboot) dashboard.GET("/updates", updateHandler.ListUpdates) dashboard.GET("/updates/:id", updateHandler.GetUpdate) dashboard.GET("/updates/:id/logs", updateHandler.GetUpdateLogs) diff --git a/aggregator-server/internal/api/handlers/agents.go b/aggregator-server/internal/api/handlers/agents.go index 640fac0..c92b1e3 100644 --- a/aggregator-server/internal/api/handlers/agents.go +++ b/aggregator-server/internal/api/handlers/agents.go @@ -79,6 +79,7 @@ func (h *AgentHandler) RegisterAgent(c *gin.Context) { OSVersion: req.OSVersion, OSArchitecture: req.OSArchitecture, AgentVersion: req.AgentVersion, + CurrentVersion: req.AgentVersion, LastSeen: time.Now(), Status: "online", Metadata: models.JSONB{}, @@ -966,3 +967,63 @@ func (h *AgentHandler) SetRapidPollingMode(c *gin.Context) { "rapid_polling_until": rapidPollingUntil, }) } + +// TriggerReboot triggers a system reboot for an agent +func (h *AgentHandler) TriggerReboot(c *gin.Context) { + agentID, err := uuid.Parse(c.Param("id")) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid agent ID"}) + return + } + + // Check if agent exists + agent, err := h.agentQueries.GetAgentByID(agentID) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "agent not found"}) + return + } + + // Parse request body for optional parameters + var req struct { + DelayMinutes int `json:"delay_minutes"` + Message string `json:"message"` + } + c.ShouldBindJSON(&req) + + // Default to 1 minute delay if not specified + if req.DelayMinutes == 0 { + req.DelayMinutes = 1 + } + if req.Message == "" { + req.Message = "Reboot requested by RedFlag" + } + + // Create reboot command + cmd := &models.AgentCommand{ + ID: uuid.New(), + AgentID: agentID, + CommandType: models.CommandTypeReboot, + Params: models.JSONB{ + "delay_minutes": req.DelayMinutes, + "message": req.Message, + }, + Status: models.CommandStatusPending, + CreatedAt: time.Now(), + } + + // Save command to database + if err := h.commandQueries.CreateCommand(cmd); err != nil { + log.Printf("Failed to create reboot command: %v", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create reboot command"}) + return + } + + log.Printf("Reboot command created for agent %s (%s)", agent.Hostname, agentID) + + c.JSON(http.StatusOK, gin.H{ + "message": "reboot command sent", + "command_id": cmd.ID, + "agent_id": agentID, + "hostname": agent.Hostname, + }) +} diff --git a/aggregator-server/internal/database/migrations/013_add_reboot_tracking.up.sql b/aggregator-server/internal/database/migrations/013_add_reboot_tracking.up.sql new file mode 100644 index 0000000..c624ec1 --- /dev/null +++ b/aggregator-server/internal/database/migrations/013_add_reboot_tracking.up.sql @@ -0,0 +1,13 @@ +-- Add reboot tracking fields to agents table +ALTER TABLE agents +ADD COLUMN reboot_required BOOLEAN DEFAULT FALSE, +ADD COLUMN last_reboot_at TIMESTAMP, +ADD COLUMN reboot_reason TEXT; + +-- Add index for efficient querying of agents needing reboot +CREATE INDEX idx_agents_reboot_required ON agents(reboot_required) WHERE reboot_required = TRUE; + +-- Add comment for documentation +COMMENT ON COLUMN agents.reboot_required IS 'Whether the agent host requires a reboot to complete updates'; +COMMENT ON COLUMN agents.last_reboot_at IS 'Timestamp of the last system reboot'; +COMMENT ON COLUMN agents.reboot_reason IS 'Reason why reboot is required (e.g., kernel update, library updates)'; diff --git a/aggregator-server/internal/database/queries/agents.go b/aggregator-server/internal/database/queries/agents.go index 7a4fc2b..8fd8da6 100644 --- a/aggregator-server/internal/database/queries/agents.go +++ b/aggregator-server/internal/database/queries/agents.go @@ -204,3 +204,30 @@ func (q *AgentQueries) GetActiveAgentCount() (int, error) { 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 +} diff --git a/aggregator-server/internal/models/agent.go b/aggregator-server/internal/models/agent.go index fab1c10..5d686cf 100644 --- a/aggregator-server/internal/models/agent.go +++ b/aggregator-server/internal/models/agent.go @@ -19,11 +19,14 @@ type Agent struct { CurrentVersion string `json:"current_version" db:"current_version"` // Current running version UpdateAvailable bool `json:"update_available" db:"update_available"` // Whether update is available LastVersionCheck time.Time `json:"last_version_check" db:"last_version_check"` // Last time version was checked - LastSeen time.Time `json:"last_seen" db:"last_seen"` - Status string `json:"status" db:"status"` - Metadata JSONB `json:"metadata" db:"metadata"` - CreatedAt time.Time `json:"created_at" db:"created_at"` - UpdatedAt time.Time `json:"updated_at" db:"updated_at"` + LastSeen time.Time `json:"last_seen" db:"last_seen"` + Status string `json:"status" db:"status"` + Metadata JSONB `json:"metadata" db:"metadata"` + RebootRequired bool `json:"reboot_required" db:"reboot_required"` + LastRebootAt *time.Time `json:"last_reboot_at" db:"last_reboot_at"` + RebootReason string `json:"reboot_reason" db:"reboot_reason"` + CreatedAt time.Time `json:"created_at" db:"created_at"` + UpdatedAt time.Time `json:"updated_at" db:"updated_at"` } // AgentWithLastScan extends Agent with last scan information @@ -40,6 +43,9 @@ type AgentWithLastScan struct { LastSeen time.Time `json:"last_seen" db:"last_seen"` Status string `json:"status" db:"status"` Metadata JSONB `json:"metadata" db:"metadata"` + RebootRequired bool `json:"reboot_required" db:"reboot_required"` + LastRebootAt *time.Time `json:"last_reboot_at" db:"last_reboot_at"` + RebootReason string `json:"reboot_reason" db:"reboot_reason"` CreatedAt time.Time `json:"created_at" db:"created_at"` UpdatedAt time.Time `json:"updated_at" db:"updated_at"` LastScan *time.Time `json:"last_scan" db:"last_scan"` diff --git a/aggregator-server/internal/models/command.go b/aggregator-server/internal/models/command.go index ab64826..31584cb 100644 --- a/aggregator-server/internal/models/command.go +++ b/aggregator-server/internal/models/command.go @@ -50,6 +50,7 @@ const ( CommandTypeUpdateAgent = "update_agent" CommandTypeEnableHeartbeat = "enable_heartbeat" CommandTypeDisableHeartbeat = "disable_heartbeat" + CommandTypeReboot = "reboot" ) // Command statuses diff --git a/aggregator-web/src/lib/api.ts b/aggregator-web/src/lib/api.ts index 4cba513..702ab91 100644 --- a/aggregator-web/src/lib/api.ts +++ b/aggregator-web/src/lib/api.ts @@ -97,6 +97,14 @@ export const agentApi = { return response.data; }, + // Trigger agent reboot + rebootAgent: async (id: string, delayMinutes: number = 1, message?: string): Promise => { + await api.post(`/agents/${id}/reboot`, { + delay_minutes: delayMinutes, + message: message || 'System reboot requested by RedFlag' + }); + }, + // Unregister/remove agent unregisterAgent: async (id: string): Promise => { await api.delete(`/agents/${id}`); diff --git a/aggregator-web/src/pages/Agents.tsx b/aggregator-web/src/pages/Agents.tsx index 48a5225..bd73e18 100644 --- a/aggregator-web/src/pages/Agents.tsx +++ b/aggregator-web/src/pages/Agents.tsx @@ -21,6 +21,7 @@ import { CheckCircle, AlertCircle, XCircle, + Power, } from 'lucide-react'; import { useAgents, useAgent, useScanAgent, useScanMultipleAgents, useUnregisterAgent } from '@/hooks/useAgents'; import { useActiveCommands, useCancelCommand } from '@/hooks/useCommands'; @@ -298,6 +299,22 @@ const Agents: React.FC = () => { } }; + // Handle agent reboot + const handleRebootAgent = async (agentId: string, hostname: string) => { + if (!window.confirm( + `Schedule a system restart for agent "${hostname}"?\n\nThe system will restart in 1 minute. Any unsaved work may be lost.` + )) { + return; + } + + try { + await agentApi.rebootAgent(agentId); + toast.success(`Restart command sent to "${hostname}". System will restart in 1 minute.`); + } catch (error: any) { + toast.error(error.message || `Failed to send restart command to "${hostname}"`); + } + }; + // Handle agent removal const handleRemoveAgent = async (agentId: string, hostname: string) => { if (!window.confirm( @@ -445,7 +462,7 @@ const Agents: React.FC = () => { Version:
- {selectedAgent.current_version || 'Unknown'} + {selectedAgent.current_version || 'Initial Registration'} {selectedAgent.update_available === true && ( @@ -485,6 +502,26 @@ const Agents: React.FC = () => {
+ {/* Restart Required Alert */} + {selectedAgent.reboot_required && ( +
+
+ +
+

System Restart Required

+

+ {selectedAgent.reboot_reason || 'This system requires a restart to complete updates.'} +

+ {selectedAgent.last_reboot_at && ( +

+ Last reboot: {formatRelativeTime(selectedAgent.last_reboot_at)} +

+ )} +
+
+
+ )} + {/* Tabs */}
@@ -947,6 +984,14 @@ const Agents: React.FC = () => {
+ +