Fimeg/Redflag
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add host restart detection and fix agent version display

Browse Source 
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.
This commit is contained in:
Fimeg
2025-10-31 15:03:59 -04:00
parent 3f9164c7ca
commit e72e9fc16f
11 changed files with 346 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
aggregator-agent/cmd/agent/main.go
View File

@@ -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)

69
aggregator-agent/internal/system/info.go
View File

@@ -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
}
}
// 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, ""
}

1
aggregator-server/cmd/server/main.go
View File

@@ -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)

61
aggregator-server/internal/api/handlers/agents.go
View File

@@ -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,
})
}

13
aggregator-server/internal/database/migrations/013_add_reboot_tracking.up.sql Normal file
View File

@@ -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)';

27
aggregator-server/internal/database/queries/agents.go
View File

@@ -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
}

16
aggregator-server/internal/models/agent.go
View File

@@ -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"`

1
aggregator-server/internal/models/command.go
View File

@@ -50,6 +50,7 @@ const (
CommandTypeUpdateAgent = "update_agent"
CommandTypeEnableHeartbeat = "enable_heartbeat"
CommandTypeDisableHeartbeat = "disable_heartbeat"
CommandTypeReboot = "reboot"
)
// Command statuses

8
aggregator-web/src/lib/api.ts
View File

@@ -97,6 +97,14 @@ export const agentApi = {
return response.data;
},
// Trigger agent reboot
rebootAgent: async (id: string, delayMinutes: number = 1, message?: string): Promise<void> => {
await api.post(`/agents/${id}/reboot`, {
delay_minutes: delayMinutes,
message: message || 'System reboot requested by RedFlag'
});
},
// Unregister/remove agent
unregisterAgent: async (id: string): Promise<void> => {
await api.delete(`/agents/${id}`);

63
aggregator-web/src/pages/Agents.tsx
View File

@@ -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 = () => {
<span className="text-gray-500">Version:</span>
<div className="flex items-center space-x-1">
<span className="font-medium text-gray-900">
{selectedAgent.current_version || 'Unknown'}
{selectedAgent.current_version || 'Initial Registration'}
</span>
{selectedAgent.update_available === true && (
<span className="flex items-center text-xs text-amber-600 bg-amber-50 px-2 py-0.5 rounded-full">
@@ -485,6 +502,26 @@ const Agents: React.FC = () => {
</div>
</div>
{/* Restart Required Alert */}
{selectedAgent.reboot_required && (
<div className="mb-6 bg-amber-50 border border-amber-200 rounded-lg p-4">
<div className="flex items-start">
<AlertCircle className="h-5 w-5 text-amber-600 mt-0.5 mr-3 flex-shrink-0" />
<div className="flex-1">
<h3 className="text-sm font-medium text-amber-900">System Restart Required</h3>
<p className="text-sm text-amber-700 mt-1">
{selectedAgent.reboot_reason || 'This system requires a restart to complete updates.'}
</p>
{selectedAgent.last_reboot_at && (
<p className="text-xs text-amber-600 mt-1">
Last reboot: {formatRelativeTime(selectedAgent.last_reboot_at)}
</p>
)}
</div>
</div>
</div>
)}
{/* Tabs */}
<div className="mb-6">
<div className="border-b border-gray-200">
@@ -947,6 +984,14 @@ const Agents: React.FC = () => {
</div>
</div>
<button
onClick={() => handleRebootAgent(selectedAgent.id, selectedAgent.hostname)}
className="w-full btn btn-warning"
>
<Power className="h-4 w-4 mr-2" />
Restart Host
</button>
<button
onClick={() => handleRemoveAgent(selectedAgent.id, selectedAgent.hostname)}
disabled={unregisterAgentMutation.isPending}
@@ -1157,14 +1202,22 @@ const Agents: React.FC = () => {
</div>
</td>
<td className="table-cell">
<span className={cn('badge', getStatusColor(isOnline(agent.last_seen) ? 'online' : 'offline'))}>
{isOnline(agent.last_seen) ? 'Online' : 'Offline'}
</span>
<div className="flex flex-col space-y-1">
<span className={cn('badge', getStatusColor(isOnline(agent.last_seen) ? 'online' : 'offline'))}>
{isOnline(agent.last_seen) ? 'Online' : 'Offline'}
</span>
{agent.reboot_required && (
<span className="flex items-center text-xs text-amber-700 bg-amber-50 px-1.5 py-0.5 rounded-full w-fit">
<Power className="h-3 w-3 mr-1" />
Restart
</span>
)}
</div>
</td>
<td className="table-cell">
<div className="flex items-center space-x-2">
<span className="text-sm text-gray-900">
{agent.current_version || 'Unknown'}
{agent.current_version || 'Initial Registration'}
</span>
{agent.update_available === true && (
<span className="flex items-center text-xs text-amber-600 bg-amber-50 px-1.5 py-0.5 rounded-full">

4
aggregator-web/src/types/index.ts
View File

@@ -21,6 +21,10 @@ export interface Agent {
status: 'online' | 'offline';
created_at: string;
updated_at: string;
current_version?: string;
reboot_required?: boolean;
last_reboot_at?: string | null;
reboot_reason?: string;
metadata?: Record<string, any>;
// Note: ip_address not available from API yet
}