Redflag/aggregator-server/cmd/server/main.go

package main

import (
	"context"
	"crypto/ed25519"
	"encoding/hex"
	"flag"
	"fmt"
	"log"
	"net/http"
	"path/filepath"
	"time"

	"github.com/Fimeg/RedFlag/aggregator-server/internal/api/handlers"
	"github.com/Fimeg/RedFlag/aggregator-server/internal/api/middleware"
	"github.com/Fimeg/RedFlag/aggregator-server/internal/config"
	"github.com/Fimeg/RedFlag/aggregator-server/internal/database"
	"github.com/Fimeg/RedFlag/aggregator-server/internal/database/queries"
	"github.com/Fimeg/RedFlag/aggregator-server/internal/logging"
	"github.com/Fimeg/RedFlag/aggregator-server/internal/scheduler"
	"github.com/Fimeg/RedFlag/aggregator-server/internal/services"
	"github.com/gin-gonic/gin"
)

// validateSigningService performs a test sign/verify to ensure the key is valid
func validateSigningService(signingService *services.SigningService) error {
	if signingService == nil {
		return fmt.Errorf("signing service is nil")
	}

	// Verify the key is accessible by getting public key and fingerprint
	publicKeyHex := signingService.GetPublicKey()
	if publicKeyHex == "" {
		return fmt.Errorf("failed to get public key from signing service")
	}

	fingerprint := signingService.GetPublicKeyFingerprint()
	if fingerprint == "" {
		return fmt.Errorf("failed to get public key fingerprint")
	}

	// Basic validation: Ed25519 public key should be 64 hex characters (32 bytes)
	if len(publicKeyHex) != 64 {
		return fmt.Errorf("invalid public key length: expected 64 hex chars, got %d", len(publicKeyHex))
	}

	return nil
}

// isSetupComplete checks if the server has been fully configured
// Returns true if all required components are ready for production
// Components checked: admin credentials, signing keys, database connectivity
func isSetupComplete(cfg *config.Config, signingService *services.SigningService, db *database.DB) bool {
	// Check if signing keys are configured
	if cfg.SigningPrivateKey == "" {
		log.Printf("Setup incomplete: Signing keys not configured")
		return false
	}

	// Check if admin password is configured (not empty)
	if cfg.Admin.Password == "" {
		log.Printf("Setup incomplete: Admin password not configured")
		return false
	}

	// Check if JWT secret is configured
	if cfg.Admin.JWTSecret == "" {
		log.Printf("Setup incomplete: JWT secret not configured")
		return false
	}

	// Check if database connection is working
	if err := db.DB.Ping(); err != nil {
		log.Printf("Setup incomplete: Database not accessible: %v", err)
		return false
	}

	// Check if database has been migrated (check for agents table)
	var agentCount int
	if err := db.DB.Get(&agentCount, "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = 'agents'"); err != nil {
		log.Printf("Setup incomplete: Database migrations not complete - agents table does not exist")
		return false
	}

	// All critical checks passed
	log.Printf("Setup validation passed: All required components configured")
	return true
}

func startWelcomeModeServer() {
	setupHandler := handlers.NewSetupHandler("/app/config")
	router := gin.Default()

	// Add CORS middleware
	router.Use(middleware.CORSMiddleware())

	// Health check (all endpoints for compatibility)
	router.GET("/health", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "waiting for configuration"})
	})
	router.GET("/api/health", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "waiting for configuration"})
	})
	router.GET("/api/v1/health", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "waiting for configuration"})
	})

	// Welcome page with setup instructions
	router.GET("/", setupHandler.ShowSetupPage)

	// Setup endpoint for web configuration
	router.POST("/api/setup/configure", setupHandler.ConfigureServer)
	router.POST("/api/setup/generate-keys", setupHandler.GenerateSigningKeys)
	router.POST("/api/setup/configure-secrets", setupHandler.ConfigureSecrets)

	// Setup endpoint for web configuration
	router.GET("/setup", setupHandler.ShowSetupPage)

	log.Printf("Welcome mode server started on :8080")
	log.Printf("Waiting for configuration...")

	if err := router.Run(":8080"); err != nil {
		log.Fatal("Failed to start welcome mode server:", err)
	}
}

func main() {
	// Parse command line flags
	var setup bool
	var migrate bool
	var version bool
	flag.BoolVar(&setup, "setup", false, "Run setup wizard")
	flag.BoolVar(&migrate, "migrate", false, "Run database migrations only")
	flag.BoolVar(&version, "version", false, "Show version information")
	flag.Parse()

	// Handle special commands
	if version {
		fmt.Printf("RedFlag Server v0.1.0-alpha\n")
		fmt.Printf("Self-hosted update management platform\n")
		return
	}

	if setup {
		if err := config.RunSetupWizard(); err != nil {
			log.Fatal("Setup failed:", err)
		}
		return
	}

	// Load configuration
	cfg, err := config.Load()
	if err != nil {
		log.Printf("Server waiting for configuration: %v", err)
		log.Printf("Run: docker-compose exec server ./redflag-server --setup")
		log.Printf("Or configure via web interface at: http://localhost:8080/setup")

		// Start welcome mode server
		startWelcomeModeServer()
		return
	}

	// Set JWT secret
	middleware.JWTSecret = cfg.Admin.JWTSecret

	// Build database URL from new config structure
	databaseURL := fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=disable",
		cfg.Database.Username, cfg.Database.Password, cfg.Database.Host, cfg.Database.Port, cfg.Database.Database)

	// Connect to database
	db, err := database.Connect(databaseURL)
	if err != nil {
		log.Fatal("Failed to connect to database:", err)
	}
	defer db.Close()

	// Handle migrate-only flag
	if migrate {
		migrationsPath := filepath.Join("internal", "database", "migrations")
		if err := db.Migrate(migrationsPath); err != nil {
			log.Fatal("Migration failed:", err)
		}
		fmt.Printf("[OK] Database migrations completed\n")
		return
	}

	// Run migrations
	migrationsPath := filepath.Join("internal", "database", "migrations")
	if err := db.Migrate(migrationsPath); err != nil {
		// For development, continue even if migrations fail
		// In production, you might want to handle this more gracefully
		fmt.Printf("Warning: Migration failed (tables may already exist): %v\n", err)
	}
	fmt.Println("[OK] Database migrations completed")

	agentQueries := queries.NewAgentQueries(db.DB)
	updateQueries := queries.NewUpdateQueries(db.DB)
	commandQueries := queries.NewCommandQueries(db.DB)
	refreshTokenQueries := queries.NewRefreshTokenQueries(db.DB)
	registrationTokenQueries := queries.NewRegistrationTokenQueries(db.DB)
	subsystemQueries := queries.NewSubsystemQueries(db.DB)
	agentUpdateQueries := queries.NewAgentUpdateQueries(db.DB)
	metricsQueries := queries.NewMetricsQueries(db.DB.DB)
	dockerQueries := queries.NewDockerQueries(db.DB.DB)
	adminQueries := queries.NewAdminQueries(db.DB)

	// Create PackageQueries for accessing signed agent update packages
	packageQueries := queries.NewPackageQueries(db.DB)

	// Initialize services
	timezoneService := services.NewTimezoneService(cfg)
	timeoutService := services.NewTimeoutService(commandQueries, updateQueries)

	// Initialize and validate signing service if private key is configured
	var signingService *services.SigningService
	if cfg.SigningPrivateKey != "" {
		var err error
		signingService, err = services.NewSigningService(cfg.SigningPrivateKey)
		if err != nil {
			log.Printf("[ERROR] Failed to initialize signing service: %v", err)
			log.Printf("[WARNING] Agent update signing is DISABLED - agents cannot be updated")
			log.Printf("[INFO] To fix: Generate signing keys at /api/setup/generate-keys and add to .env")
		} else {
			// Validate the signing key works by performing a test sign/verify
			if err := validateSigningService(signingService); err != nil {
				log.Printf("[ERROR] Signing key validation failed: %v", err)
				log.Printf("[WARNING] Agent update signing is DISABLED - key is corrupted")
				signingService = nil // Disable signing
			} else {
				log.Printf("[system] Ed25519 signing service initialized and validated")
				log.Printf("[system] Public key fingerprint: %s", signingService.GetPublicKeyFingerprint())
			}
		}
	} else {
		log.Printf("[WARNING] No signing private key configured - agent update signing disabled")
		log.Printf("[INFO] Generate keys: POST /api/setup/generate-keys")
	}
	// Initialize default security settings (critical for v0.2.x)
	fmt.Println("[OK] Initializing default security settings...")
	securitySettingsQueries := queries.NewSecuritySettingsQueries(db.DB)
	securitySettingsService, err := services.NewSecuritySettingsService(securitySettingsQueries, signingService)
	if err != nil {
		fmt.Printf("Warning: Failed to create security settings service: %v\n", err)
		fmt.Println("Security settings will need to be configured manually via the dashboard")
	} else if err := securitySettingsService.InitializeDefaultSettings(); err != nil {
		fmt.Printf("Warning: Failed to initialize default security settings: %v\n", err)
		fmt.Println("Security settings will need to be configured manually via the dashboard")
	} else {
		fmt.Println("[OK] Default security settings initialized")
	}

	// Check if setup is complete
	if !isSetupComplete(cfg, signingService, db) {
		serverAddr := cfg.Server.Host
		if serverAddr == "" {
			serverAddr = "localhost"
		}
		log.Printf("Server setup incomplete - starting welcome mode")
		log.Printf("Setup required: Admin credentials, signing keys, and database configuration")
		log.Printf("Access setup at: http://%s:%d/setup", serverAddr, cfg.Server.Port)
		startWelcomeModeServer()
		return
	}

	// Initialize admin user from .env configuration
	fmt.Println("[OK] Initializing admin user...")
	if err := adminQueries.CreateAdminIfNotExists(cfg.Admin.Username, cfg.Admin.Email, cfg.Admin.Password); err != nil {
		log.Printf("[ERROR] Failed to initialize admin user: %v", err)
	} else {
		// Update admin password from .env (runs on every startup to keep in sync)
		if err := adminQueries.UpdateAdminPassword(cfg.Admin.Username, cfg.Admin.Password); err != nil {
			log.Printf("[WARNING] Failed to update admin password: %v", err)
		} else {
			fmt.Println("[OK] Admin user initialized")
		}
	}

	// Initialize security logger
	secConfig := logging.SecurityLogConfig{
		Enabled:           true, // Could be configurable in the future
		Level:             "warning",
		LogSuccesses:      false,
		FilePath:          "/var/log/redflag/security.json",
		MaxSizeMB:         100,
		MaxFiles:          10,
		RetentionDays:     90,
		LogToDatabase:     true,
		HashIPAddresses:   true,
	}
	securityLogger, err := logging.NewSecurityLogger(secConfig, db.DB)
	if err != nil {
		log.Printf("Failed to initialize security logger: %v", err)
		securityLogger = nil
	}

	// Initialize rate limiter
	rateLimiter := middleware.NewRateLimiter()

	// Initialize handlers that don't depend on agentHandler (can be created now)
	authHandler := handlers.NewAuthHandler(cfg.Admin.JWTSecret, adminQueries)
	statsHandler := handlers.NewStatsHandler(agentQueries, updateQueries)
	settingsHandler := handlers.NewSettingsHandler(timezoneService)
	dockerHandler := handlers.NewDockerHandler(updateQueries, agentQueries, commandQueries, signingService, securityLogger)
	registrationTokenHandler := handlers.NewRegistrationTokenHandler(registrationTokenQueries, agentQueries, cfg)
	rateLimitHandler := handlers.NewRateLimitHandler(rateLimiter)
	downloadHandler := handlers.NewDownloadHandler(filepath.Join("/app"), cfg, packageQueries)
	subsystemHandler := handlers.NewSubsystemHandler(subsystemQueries, commandQueries, signingService, securityLogger)
	metricsHandler := handlers.NewMetricsHandler(metricsQueries, agentQueries, commandQueries)
	dockerReportsHandler := handlers.NewDockerReportsHandler(dockerQueries, agentQueries, commandQueries)
	agentSetupHandler := handlers.NewAgentSetupHandler(agentQueries)

	// Initialize scanner config handler (for user-configurable scanner timeouts)
	scannerConfigHandler := handlers.NewScannerConfigHandler(db.DB)

	// Initialize verification handler
	var verificationHandler *handlers.VerificationHandler
	if signingService != nil {
		verificationHandler = handlers.NewVerificationHandler(agentQueries, signingService)
	}

	// Initialize update nonce service (for version upgrade middleware)
	var updateNonceService *services.UpdateNonceService
	if signingService != nil && cfg.SigningPrivateKey != "" {
		// Decode private key for nonce service
		privateKeyBytes, err := hex.DecodeString(cfg.SigningPrivateKey)
		if err == nil && len(privateKeyBytes) == ed25519.PrivateKeySize {
			updateNonceService = services.NewUpdateNonceService(ed25519.PrivateKey(privateKeyBytes))
			log.Printf("[system] Update nonce service initialized for version upgrades")
		} else {
			log.Printf("[WARNING] Failed to initialize update nonce service: invalid private key")
		}
	}

	// Initialize system handler
	systemHandler := handlers.NewSystemHandler(signingService)

	// Initialize security handler
	securityHandler := handlers.NewSecurityHandler(signingService, agentQueries, commandQueries)

	// Initialize security settings service and handler
	securitySettingsService, err = services.NewSecuritySettingsService(securitySettingsQueries, signingService)
	if err != nil {
		log.Printf("[ERROR] Failed to initialize security settings service: %v", err)
		securitySettingsService = nil
	} else {
		log.Printf("[OK] Security settings service initialized")
	}
	// Setup router
	router := gin.Default()

	// Add CORS middleware
	router.Use(middleware.CORSMiddleware())

	// Health check
	router.GET("/health", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "healthy"})
	})
	router.GET("/api/health", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "healthy"})
	})

	// API routes
	api := router.Group("/api/v1")
	{
		// Authentication routes (with rate limiting)
		api.POST("/auth/login", rateLimiter.RateLimit("public_access", middleware.KeyByIP), authHandler.Login)
		api.POST("/auth/logout", authHandler.Logout)
		api.GET("/auth/verify", authHandler.VerifyToken)

		// Public system routes (no authentication required)
		api.GET("/public-key", rateLimiter.RateLimit("public_access", middleware.KeyByIP), systemHandler.GetPublicKey)
		api.GET("/info", rateLimiter.RateLimit("public_access", middleware.KeyByIP), systemHandler.GetSystemInfo)

		// Agent setup routes (no authentication required, with rate limiting)
		api.POST("/setup/agent", rateLimiter.RateLimit("agent_setup", middleware.KeyByIP), agentSetupHandler.SetupAgent)
		api.GET("/setup/templates", rateLimiter.RateLimit("public_access", middleware.KeyByIP), agentSetupHandler.GetTemplates)
		api.POST("/setup/validate", rateLimiter.RateLimit("agent_setup", middleware.KeyByIP), agentSetupHandler.ValidateConfiguration)

		// Build orchestrator routes (admin-only)
		buildRoutes := api.Group("/build")
		buildRoutes.Use(authHandler.WebAuthMiddleware())
		{
			buildRoutes.POST("/new", rateLimiter.RateLimit("agent_build", middleware.KeyByAgentID), handlers.NewAgentBuild)
			buildRoutes.POST("/upgrade/:agentID", rateLimiter.RateLimit("agent_build", middleware.KeyByAgentID), handlers.UpgradeAgentBuild)
			buildRoutes.POST("/detect", rateLimiter.RateLimit("agent_build", middleware.KeyByAgentID), handlers.DetectAgentInstallation)
		}

		// Public download routes (no authentication - agents need these!)
		api.GET("/downloads/:platform", rateLimiter.RateLimit("public_access", middleware.KeyByIP), downloadHandler.DownloadAgent)
		api.GET("/downloads/updates/:package_id", rateLimiter.RateLimit("public_access", middleware.KeyByIP), downloadHandler.DownloadUpdatePackage)
		api.GET("/downloads/config/:agent_id", rateLimiter.RateLimit("public_access", middleware.KeyByIP), downloadHandler.HandleConfigDownload)
		api.GET("/install/:platform", rateLimiter.RateLimit("public_access", middleware.KeyByIP), downloadHandler.InstallScript)
	}

	// Start background goroutine to mark offline agents
	// TODO: Make these values configurable via settings:
	// - Check interval (currently 2 minutes, should match agent heartbeat setting)
	// - Offline threshold (currently 10 minutes, should be based on agent check-in interval + missed checks)
	// - Missed checks before offline (default 2, so 300s agent interval * 2 = 10 minutes)
	go func() {
		ticker := time.NewTicker(2 * time.Minute) // Check every 2 minutes
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				// Mark agents as offline if they haven't checked in within 10 minutes
				if err := agentQueries.MarkOfflineAgents(10 * time.Minute); err != nil {
					log.Printf("Failed to mark offline agents: %v", err)
				}
			}
		}
	}()

	// Start timeout service
	timeoutService.Start()
	log.Println("Timeout service started")

	// Initialize and start scheduler
	schedulerConfig := scheduler.DefaultConfig()
	subsystemScheduler := scheduler.NewScheduler(schedulerConfig, agentQueries, commandQueries, subsystemQueries)

	// Initialize agentHandler now that scheduler is available
	agentHandler := handlers.NewAgentHandler(agentQueries, commandQueries, refreshTokenQueries, registrationTokenQueries, subsystemQueries, subsystemScheduler, signingService, securityLogger, cfg.CheckInInterval, cfg.LatestAgentVersion)

	// Initialize agent update handler now that agentHandler is available
	var agentUpdateHandler *handlers.AgentUpdateHandler
	if signingService != nil {
		agentUpdateHandler = handlers.NewAgentUpdateHandler(agentQueries, agentUpdateQueries, commandQueries, signingService, updateNonceService, agentHandler)
	}

	// Initialize updateHandler with the agentHandler reference
	updateHandler := handlers.NewUpdateHandler(updateQueries, agentQueries, commandQueries, agentHandler)

	// Add routes that depend on agentHandler (must be after agentHandler creation)
	api.POST("/agents/register", rateLimiter.RateLimit("agent_registration", middleware.KeyByIP), agentHandler.RegisterAgent)
	api.POST("/agents/renew", rateLimiter.RateLimit("public_access", middleware.KeyByIP), agentHandler.RenewToken)

	// Protected agent routes (with machine binding security)
	agents := api.Group("/agents")
	agents.Use(middleware.AuthMiddleware())
	agents.Use(middleware.MachineBindingMiddleware(agentQueries, cfg.MinAgentVersion)) // v0.1.22: Prevent config copying
	{
		agents.GET("/:id/commands", agentHandler.GetCommands)
		agents.GET("/:id/config", agentHandler.GetAgentConfig)
		agents.POST("/:id/updates", rateLimiter.RateLimit("agent_reports", middleware.KeyByAgentID), updateHandler.ReportUpdates)
		agents.POST("/:id/logs", rateLimiter.RateLimit("agent_reports", middleware.KeyByAgentID), updateHandler.ReportLog)
		agents.POST("/:id/dependencies", rateLimiter.RateLimit("agent_reports", middleware.KeyByAgentID), updateHandler.ReportDependencies)
		agents.POST("/:id/system-info", rateLimiter.RateLimit("agent_reports", middleware.KeyByAgentID), agentHandler.ReportSystemInfo)
		agents.POST("/:id/rapid-mode", rateLimiter.RateLimit("agent_reports", middleware.KeyByAgentID), agentHandler.SetRapidPollingMode)
		agents.POST("/:id/verify-signature", rateLimiter.RateLimit("agent_reports", middleware.KeyByAgentID), func(c *gin.Context) {
			if verificationHandler == nil {
				c.JSON(http.StatusServiceUnavailable, gin.H{"error": "signature verification service not available"})
				return
			}
			verificationHandler.VerifySignature(c)
		})
		agents.DELETE("/:id", agentHandler.UnregisterAgent)

		// New dedicated endpoints for metrics and docker images (data classification fix)
		agents.POST("/:id/metrics", rateLimiter.RateLimit("agent_reports", middleware.KeyByAgentID), metricsHandler.ReportMetrics)
		agents.POST("/:id/docker-images", rateLimiter.RateLimit("agent_reports", middleware.KeyByAgentID), dockerReportsHandler.ReportDockerImages)
	}

	// Dashboard/Web routes (protected by web auth)
	dashboard := api.Group("/")
	dashboard.Use(authHandler.WebAuthMiddleware())
	{
		dashboard.GET("/stats/summary", statsHandler.GetDashboardStats)
		dashboard.GET("/agents", agentHandler.ListAgents)
		dashboard.GET("/agents/:id", agentHandler.GetAgent)
		dashboard.POST("/agents/:id/scan", agentHandler.TriggerScan)
		dashboard.POST("/agents/:id/heartbeat", agentHandler.TriggerHeartbeat)
		dashboard.GET("/agents/:id/heartbeat", agentHandler.GetHeartbeatStatus)
		dashboard.POST("/agents/:id/reboot", agentHandler.TriggerReboot)

		// Subsystem routes for web dashboard
		dashboard.GET("/agents/:id/subsystems", subsystemHandler.GetSubsystems)
		dashboard.GET("/agents/:id/subsystems/:subsystem", subsystemHandler.GetSubsystem)
		dashboard.PATCH("/agents/:id/subsystems/:subsystem", subsystemHandler.UpdateSubsystem)
		dashboard.POST("/agents/:id/subsystems/:subsystem/enable", subsystemHandler.EnableSubsystem)
		dashboard.POST("/agents/:id/subsystems/:subsystem/disable", subsystemHandler.DisableSubsystem)
		dashboard.POST("/agents/:id/subsystems/:subsystem/trigger", subsystemHandler.TriggerSubsystem)
		dashboard.GET("/agents/:id/subsystems/:subsystem/stats", subsystemHandler.GetSubsystemStats)
		dashboard.POST("/agents/:id/subsystems/:subsystem/auto-run", subsystemHandler.SetAutoRun)
		dashboard.POST("/agents/:id/subsystems/:subsystem/interval", subsystemHandler.SetInterval)

		dashboard.GET("/updates", updateHandler.ListUpdates)
		dashboard.GET("/updates/:id", updateHandler.GetUpdate)
		dashboard.GET("/updates/:id/logs", updateHandler.GetUpdateLogs)
		dashboard.POST("/updates/:id/approve", updateHandler.ApproveUpdate)
		dashboard.POST("/updates/approve", updateHandler.ApproveUpdates)
		dashboard.POST("/updates/:id/reject", updateHandler.RejectUpdate)
		dashboard.POST("/updates/:id/install", updateHandler.InstallUpdate)
		dashboard.POST("/updates/:id/confirm-dependencies", updateHandler.ConfirmDependencies)

		// Agent update routes
		if agentUpdateHandler != nil {
			dashboard.POST("/agents/:id/update", agentUpdateHandler.UpdateAgent)
			dashboard.POST("/agents/:id/update-nonce", agentUpdateHandler.GenerateUpdateNonce)
			dashboard.POST("/agents/bulk-update", agentUpdateHandler.BulkUpdateAgents)
			dashboard.GET("/updates/packages", agentUpdateHandler.ListUpdatePackages)
			dashboard.POST("/updates/packages/sign", agentUpdateHandler.SignUpdatePackage)
			dashboard.GET("/agents/:id/updates/available", agentUpdateHandler.CheckForUpdateAvailable)
			dashboard.GET("/agents/:id/updates/status", agentUpdateHandler.GetUpdateStatus)
		}

		dashboard.GET("/logs", updateHandler.GetAllLogs)
		dashboard.GET("/logs/active", updateHandler.GetActiveOperations)

		// Command routes
		dashboard.GET("/commands/active", updateHandler.GetActiveCommands)
		dashboard.GET("/commands/recent", updateHandler.GetRecentCommands)
		dashboard.POST("/commands/:id/retry", updateHandler.RetryCommand)
		dashboard.POST("/commands/:id/cancel", updateHandler.CancelCommand)
		dashboard.DELETE("/commands/failed", updateHandler.ClearFailedCommands)

		// Settings routes
		dashboard.GET("/settings/timezone", settingsHandler.GetTimezone)
		dashboard.GET("/settings/timezones", settingsHandler.GetTimezones)
		dashboard.PUT("/settings/timezone", settingsHandler.UpdateTimezone)

		// Docker routes
		dashboard.GET("/docker/containers", dockerHandler.GetContainers)
		dashboard.GET("/docker/stats", dockerHandler.GetStats)
		dashboard.POST("/docker/containers/:container_id/images/:image_id/approve", dockerHandler.ApproveUpdate)
		dashboard.POST("/docker/containers/:container_id/images/:image_id/reject", dockerHandler.RejectUpdate)
		dashboard.POST("/docker/containers/:container_id/images/:image_id/install", dockerHandler.InstallUpdate)

		// Metrics and Docker images routes (data classification fix)
		dashboard.GET("/agents/:id/metrics", metricsHandler.GetAgentMetrics)
		dashboard.GET("/agents/:id/metrics/storage", metricsHandler.GetAgentStorageMetrics)
		dashboard.GET("/agents/:id/metrics/system", metricsHandler.GetAgentSystemMetrics)
		dashboard.GET("/agents/:id/docker-images", dockerReportsHandler.GetAgentDockerImages)
		dashboard.GET("/agents/:id/docker-info", dockerReportsHandler.GetAgentDockerInfo)

		// Admin/Registration Token routes (for agent enrollment management)
		admin := dashboard.Group("/admin")
		{
			admin.POST("/registration-tokens", rateLimiter.RateLimit("admin_token_gen", middleware.KeyByUserID), registrationTokenHandler.GenerateRegistrationToken)
			admin.GET("/registration-tokens", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), registrationTokenHandler.ListRegistrationTokens)
			admin.GET("/registration-tokens/active", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), registrationTokenHandler.GetActiveRegistrationTokens)
			admin.DELETE("/registration-tokens/:token", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), registrationTokenHandler.RevokeRegistrationToken)
			admin.DELETE("/registration-tokens/delete/:id", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), registrationTokenHandler.DeleteRegistrationToken)
			admin.POST("/registration-tokens/cleanup", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), registrationTokenHandler.CleanupExpiredTokens)
			admin.GET("/registration-tokens/stats", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), registrationTokenHandler.GetTokenStats)
			admin.GET("/registration-tokens/validate", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), registrationTokenHandler.ValidateRegistrationToken)

			// Rate Limit Management
			admin.GET("/rate-limits", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), rateLimitHandler.GetRateLimitSettings)
			admin.PUT("/rate-limits", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), rateLimitHandler.UpdateRateLimitSettings)
			admin.POST("/rate-limits/reset", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), rateLimitHandler.ResetRateLimitSettings)
			admin.GET("/rate-limits/stats", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), rateLimitHandler.GetRateLimitStats)
			admin.POST("/rate-limits/cleanup", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), rateLimitHandler.CleanupRateLimitEntries)

			// Scanner Configuration (user-configurable timeouts)
			admin.GET("/scanner-timeouts", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), scannerConfigHandler.GetScannerTimeouts)
			admin.PUT("/scanner-timeouts/:scanner_name", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), scannerConfigHandler.UpdateScannerTimeout)
			admin.POST("/scanner-timeouts/:scanner_name/reset", rateLimiter.RateLimit("admin_operations", middleware.KeyByUserID), scannerConfigHandler.ResetScannerTimeout)
		}

		// Security Health Check endpoints
		dashboard.GET("/security/overview", securityHandler.SecurityOverview)
		dashboard.GET("/security/signing", securityHandler.SigningStatus)
		dashboard.GET("/security/nonce", securityHandler.NonceValidationStatus)
		dashboard.GET("/security/commands", securityHandler.CommandValidationStatus)
		dashboard.GET("/security/machine-binding", securityHandler.MachineBindingStatus)
		dashboard.GET("/security/metrics", securityHandler.SecurityMetrics)

		// Security Settings Management endpoints (admin-only)
// 		securitySettings := dashboard.Group("/security/settings")
// 		securitySettings.Use(middleware.RequireAdmin())
// 		{
// 			securitySettings.GET("", securitySettingsHandler.GetAllSecuritySettings)
// 			securitySettings.GET("/audit", securitySettingsHandler.GetSecurityAuditTrail)
// 			securitySettings.GET("/overview", securitySettingsHandler.GetSecurityOverview)
// 			securitySettings.GET("/:category", securitySettingsHandler.GetSecuritySettingsByCategory)
// 			securitySettings.PUT("/:category/:key", securitySettingsHandler.UpdateSecuritySetting)
// 			securitySettings.POST("/validate", securitySettingsHandler.ValidateSecuritySettings)
// 			securitySettings.POST("/apply", securitySettingsHandler.ApplySecuritySettings)
// 		}
	}

	// Load subsystems into queue
	ctx := context.Background()
	if err := subsystemScheduler.LoadSubsystems(ctx); err != nil {
		log.Printf("Warning: Failed to load subsystems: %v", err)
	} else {
		log.Println("Subsystems loaded into scheduler")
	}

	// Start scheduler
	if err := subsystemScheduler.Start(); err != nil {
		log.Printf("Warning: Failed to start scheduler: %v", err)
	}

	// Add scheduler stats endpoint (after scheduler is initialized)
	router.GET("/api/v1/scheduler/stats", middleware.AuthMiddleware(), func(c *gin.Context) {
		stats := subsystemScheduler.GetStats()
		queueStats := subsystemScheduler.GetQueueStats()
		c.JSON(200, gin.H{
			"scheduler": stats,
			"queue":     queueStats,
		})
	})

	// Add graceful shutdown for services
	defer func() {
		log.Println("Shutting down services...")

		// Stop scheduler first
		if err := subsystemScheduler.Stop(); err != nil {
			log.Printf("Error stopping scheduler: %v", err)
		}

		// Stop timeout service
		timeoutService.Stop()
		log.Println("Services stopped")
	}()

	// Start server
	addr := fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port)
	fmt.Printf("\nRedFlag Aggregator Server starting on %s\n", addr)
	fmt.Printf("Admin interface: http://%s:%d/admin\n", cfg.Server.Host, cfg.Server.Port)
	fmt.Printf("Dashboard: http://%s:%d\n\n", cfg.Server.Host, cfg.Server.Port)

	if err := router.Run(addr); err != nil {
		log.Fatal("Failed to start server:", err)
	}
}