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 = () => {
+ +