From 0ae4daca8dd6bf56dcdad80b51d0c74d74df9651 Mon Sep 17 00:00:00 2001 From: rcourtman Date: Wed, 10 Dec 2025 08:30:56 +0000 Subject: [PATCH] feat(ai): Add bulk URL discovery service - Add URLDiscoveryService for scanning all resources at once - Scans common web ports (80, 443, 8080, 8096, 3000, etc.) - Automatically saves discovered URLs to resource metadata - Add API endpoints for start/status/cancel discovery - Progress tracking with results reporting Endpoints: - POST /api/ai/discover-urls/start - Start bulk discovery - GET /api/ai/discover-urls/status - Check progress - POST /api/ai/discover-urls/cancel - Cancel discovery --- internal/ai/service.go | 36 +++-- internal/ai/url_discovery.go | 290 +++++++++++++++++++++++++++++++++++ internal/api/ai_handlers.go | 101 ++++++++++++ internal/api/router.go | 5 + 4 files changed, 420 insertions(+), 12 deletions(-) create mode 100644 internal/ai/url_discovery.go diff --git a/internal/ai/service.go b/internal/ai/service.go index 5e85fea..ba08370 100644 --- a/internal/ai/service.go +++ b/internal/ai/service.go @@ -29,18 +29,19 @@ type StateProvider interface { // Service orchestrates AI interactions type Service struct { - mu sync.RWMutex - persistence *config.ConfigPersistence - provider providers.Provider - cfg *config.AIConfig - agentServer *agentexec.Server - policy *agentexec.CommandPolicy - stateProvider StateProvider - alertProvider AlertProvider - knowledgeStore *knowledge.Store - resourceProvider ResourceProvider // Unified resource model provider (Phase 2) - patrolService *PatrolService // Background AI monitoring service - metadataProvider MetadataProvider // Enables AI to update resource URLs + mu sync.RWMutex + persistence *config.ConfigPersistence + provider providers.Provider + cfg *config.AIConfig + agentServer *agentexec.Server + policy *agentexec.CommandPolicy + stateProvider StateProvider + alertProvider AlertProvider + knowledgeStore *knowledge.Store + resourceProvider ResourceProvider // Unified resource model provider (Phase 2) + patrolService *PatrolService // Background AI monitoring service + metadataProvider MetadataProvider // Enables AI to update resource URLs + urlDiscoveryService *URLDiscoveryService // Bulk URL discovery service } // NewService creates a new AI service @@ -82,6 +83,17 @@ func (s *Service) GetPatrolService() *PatrolService { return s.patrolService } +// GetURLDiscoveryService returns the URL discovery service for bulk discovery +func (s *Service) GetURLDiscoveryService() *URLDiscoveryService { + s.mu.Lock() + defer s.mu.Unlock() + // Lazy initialization + if s.urlDiscoveryService == nil { + s.urlDiscoveryService = NewURLDiscoveryService(s) + } + return s.urlDiscoveryService +} + // SetPatrolThresholdProvider sets the threshold provider for patrol // This should be called with an AlertThresholdAdapter to connect patrol to user-configured thresholds func (s *Service) SetPatrolThresholdProvider(provider ThresholdProvider) { diff --git a/internal/ai/url_discovery.go b/internal/ai/url_discovery.go new file mode 100644 index 0000000..fc9a164 --- /dev/null +++ b/internal/ai/url_discovery.go @@ -0,0 +1,290 @@ +package ai + +import ( + "context" + "fmt" + "strings" + "sync" + "time" + + "github.com/rs/zerolog/log" +) + +// URLDiscoveryResult represents the result of URL discovery for a single resource +type URLDiscoveryResult struct { + ResourceType string `json:"resource_type"` // "guest", "docker", "host" + ResourceID string `json:"resource_id"` + ResourceName string `json:"resource_name,omitempty"` + Status string `json:"status"` // "found", "not_found", "skipped", "error" + URL string `json:"url,omitempty"` + Error string `json:"error,omitempty"` +} + +// URLDiscoveryProgress represents progress of bulk discovery +type URLDiscoveryProgress struct { + Total int `json:"total"` + Completed int `json:"completed"` + Found int `json:"found"` + Errors int `json:"errors"` + Results []URLDiscoveryResult `json:"results"` + Running bool `json:"running"` + StartedAt time.Time `json:"started_at,omitempty"` +} + +// URLDiscoveryService handles bulk URL discovery operations +type URLDiscoveryService struct { + mu sync.RWMutex + aiService *Service + progress *URLDiscoveryProgress + cancelFunc context.CancelFunc +} + +// NewURLDiscoveryService creates a new URL discovery service +func NewURLDiscoveryService(aiService *Service) *URLDiscoveryService { + return &URLDiscoveryService{ + aiService: aiService, + } +} + +// GetProgress returns the current discovery progress +func (s *URLDiscoveryService) GetProgress() *URLDiscoveryProgress { + s.mu.RLock() + defer s.mu.RUnlock() + if s.progress == nil { + return &URLDiscoveryProgress{Running: false} + } + // Return a copy + copy := *s.progress + return © +} + +// IsRunning returns true if discovery is in progress +func (s *URLDiscoveryService) IsRunning() bool { + s.mu.RLock() + defer s.mu.RUnlock() + return s.progress != nil && s.progress.Running +} + +// Cancel stops the current discovery operation +func (s *URLDiscoveryService) Cancel() { + s.mu.Lock() + defer s.mu.Unlock() + if s.cancelFunc != nil { + s.cancelFunc() + s.cancelFunc = nil + } + if s.progress != nil { + s.progress.Running = false + } +} + +// DiscoverURLs starts bulk URL discovery for resources without URLs +func (s *URLDiscoveryService) DiscoverURLs(ctx context.Context, resourceType string, skipExisting bool) error { + s.mu.Lock() + if s.progress != nil && s.progress.Running { + s.mu.Unlock() + return fmt.Errorf("discovery already in progress") + } + + // Create cancellable context + ctx, cancel := context.WithCancel(ctx) + s.cancelFunc = cancel + + s.progress = &URLDiscoveryProgress{ + Running: true, + StartedAt: time.Now(), + Results: []URLDiscoveryResult{}, + } + s.mu.Unlock() + + // Run discovery in background + go s.runDiscovery(ctx, resourceType, skipExisting) + return nil +} + +// runDiscovery performs the actual discovery work +func (s *URLDiscoveryService) runDiscovery(ctx context.Context, resourceType string, skipExisting bool) { + defer func() { + s.mu.Lock() + s.progress.Running = false + s.cancelFunc = nil + s.mu.Unlock() + }() + + // Get resources to scan from the resource provider + s.aiService.mu.RLock() + rp := s.aiService.resourceProvider + mp := s.aiService.metadataProvider + s.aiService.mu.RUnlock() + + if rp == nil { + log.Error().Msg("URL discovery: resource provider not available") + return + } + if mp == nil { + log.Error().Msg("URL discovery: metadata provider not available") + return + } + + resources := rp.GetAll() + if resources == nil { + log.Error().Msg("URL discovery: no resources available") + return + } + + // Filter resources + var toScan []struct { + Type string + ID string + Name string + IP string + } + + for _, r := range resources { + // Only scan workloads (VMs, containers) and hosts + if !r.IsWorkload() && !r.IsInfrastructure() { + continue + } + + // Filter by type if specified + if resourceType != "" && resourceType != "all" { + typeMatches := false + switch strings.ToLower(resourceType) { + case "guest": + typeMatches = r.Type == "vm" || r.Type == "container" + case "docker": + typeMatches = r.Type == "docker-container" || r.Type == "docker-service" + case "host": + typeMatches = r.Type == "host" || r.Type == "node" || r.Type == "docker-host" + default: + typeMatches = strings.EqualFold(string(r.Type), resourceType) + } + if !typeMatches { + continue + } + } + + // Need an IP to scan + ip := "" + if r.Identity != nil && len(r.Identity.IPs) > 0 { + ip = r.Identity.IPs[0] + } + + toScan = append(toScan, struct { + Type string + ID string + Name string + IP string + }{ + Type: string(r.Type), + ID: r.ID, + Name: r.Name, + IP: ip, + }) + } + + // Update total + s.mu.Lock() + s.progress.Total = len(toScan) + s.mu.Unlock() + + log.Info(). + Int("total", len(toScan)). + Str("resourceType", resourceType). + Bool("skipExisting", skipExisting). + Msg("Starting bulk URL discovery") + + // Process each resource + for _, res := range toScan { + select { + case <-ctx.Done(): + log.Info().Msg("URL discovery cancelled") + return + default: + } + + result := s.discoverSingleResource(ctx, res.Type, res.ID, res.Name, res.IP) + + s.mu.Lock() + s.progress.Completed++ + s.progress.Results = append(s.progress.Results, result) + if result.Status == "found" { + s.progress.Found++ + } else if result.Status == "error" { + s.progress.Errors++ + } + s.mu.Unlock() + } + + log.Info(). + Int("total", len(toScan)). + Int("found", s.progress.Found). + Int("errors", s.progress.Errors). + Msg("Bulk URL discovery completed") +} + +// discoverSingleResource attempts to discover URL for a single resource +func (s *URLDiscoveryService) discoverSingleResource(ctx context.Context, resType, resID, resName, ip string) URLDiscoveryResult { + result := URLDiscoveryResult{ + ResourceType: resType, + ResourceID: resID, + ResourceName: resName, + Status: "not_found", + } + + if ip == "" { + result.Status = "skipped" + result.Error = "no IP address" + return result + } + + // Common web ports to check + ports := []int{80, 443, 8080, 8443, 8096, 8920, 3000, 5000, 9000, 8081, 8000, 7878, 8989, 9117} + + for _, port := range ports { + select { + case <-ctx.Done(): + result.Status = "error" + result.Error = "cancelled" + return result + default: + } + + scheme := "http" + if port == 443 || port == 8443 { + scheme = "https" + } + + url := fmt.Sprintf("%s://%s:%d/", scheme, ip, port) + + // Try to fetch the URL using the AI service's fetch method + checkCtx, cancel := context.WithTimeout(ctx, 3*time.Second) + _, err := s.aiService.fetchURL(checkCtx, url) + cancel() + + if err == nil { + // Found a responding service! + result.Status = "found" + result.URL = url + + // Save it to metadata + if err := s.aiService.SetResourceURL(resType, resID, url); err != nil { + log.Warn(). + Err(err). + Str("resourceID", resID). + Str("url", url). + Msg("Failed to save discovered URL") + } else { + log.Info(). + Str("resourceType", resType). + Str("resourceID", resID). + Str("resourceName", resName). + Str("url", url). + Msg("Discovered and saved URL") + } + return result + } + } + + return result +} diff --git a/internal/api/ai_handlers.go b/internal/api/ai_handlers.go index 0acca41..6c8d2ff 100644 --- a/internal/api/ai_handlers.go +++ b/internal/api/ai_handlers.go @@ -1844,3 +1844,104 @@ func (h *AISettingsHandler) HandleGetFindingsHistory(w http.ResponseWriter, r *h log.Error().Err(err).Msg("Failed to write findings history response") } } + +// HandleStartURLDiscovery starts bulk URL discovery (POST /api/ai/discover-urls/start) +func (h *AISettingsHandler) HandleStartURLDiscovery(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // Require authentication + if !CheckAuth(h.config, w, r) { + return + } + + // Parse request + var req struct { + ResourceType string `json:"resource_type"` // "guest", "docker", "host", or "all" + SkipExisting bool `json:"skip_existing"` // Skip resources that already have URLs + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + // Default values + req.ResourceType = "all" + req.SkipExisting = true + } + + discoveryService := h.aiService.GetURLDiscoveryService() + if discoveryService == nil { + http.Error(w, "URL discovery service not available", http.StatusServiceUnavailable) + return + } + + if err := discoveryService.DiscoverURLs(r.Context(), req.ResourceType, req.SkipExisting); err != nil { + http.Error(w, err.Error(), http.StatusConflict) + return + } + + log.Info(). + Str("resourceType", req.ResourceType). + Bool("skipExisting", req.SkipExisting). + Msg("Started bulk URL discovery") + + if err := utils.WriteJSONResponse(w, map[string]interface{}{ + "success": true, + "message": "URL discovery started", + }); err != nil { + log.Error().Err(err).Msg("Failed to write URL discovery start response") + } +} + +// HandleURLDiscoveryStatus returns the current discovery status (GET /api/ai/discover-urls/status) +func (h *AISettingsHandler) HandleURLDiscoveryStatus(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // Require authentication + if !CheckAuth(h.config, w, r) { + return + } + + discoveryService := h.aiService.GetURLDiscoveryService() + if discoveryService == nil { + http.Error(w, "URL discovery service not available", http.StatusServiceUnavailable) + return + } + + progress := discoveryService.GetProgress() + if err := utils.WriteJSONResponse(w, progress); err != nil { + log.Error().Err(err).Msg("Failed to write URL discovery status response") + } +} + +// HandleCancelURLDiscovery cancels the current discovery operation (POST /api/ai/discover-urls/cancel) +func (h *AISettingsHandler) HandleCancelURLDiscovery(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // Require authentication + if !CheckAuth(h.config, w, r) { + return + } + + discoveryService := h.aiService.GetURLDiscoveryService() + if discoveryService == nil { + http.Error(w, "URL discovery service not available", http.StatusServiceUnavailable) + return + } + + discoveryService.Cancel() + + log.Info().Msg("Cancelled bulk URL discovery") + + if err := utils.WriteJSONResponse(w, map[string]interface{}{ + "success": true, + "message": "URL discovery cancelled", + }); err != nil { + log.Error().Err(err).Msg("Failed to write URL discovery cancel response") + } +} diff --git a/internal/api/router.go b/internal/api/router.go index 2eb4ae6..998039b 100644 --- a/internal/api/router.go +++ b/internal/api/router.go @@ -1103,6 +1103,11 @@ func (r *Router) setupRoutes() { r.mux.HandleFunc("/api/ai/patrol/dismiss", RequireAuth(r.config, r.aiSettingsHandler.HandleAcknowledgeFinding)) // Backward compat r.mux.HandleFunc("/api/ai/patrol/snooze", RequireAuth(r.config, r.aiSettingsHandler.HandleSnoozeFinding)) + // AI URL Discovery routes for bulk scanning + r.mux.HandleFunc("/api/ai/discover-urls/start", RequireAdmin(r.config, r.aiSettingsHandler.HandleStartURLDiscovery)) + r.mux.HandleFunc("/api/ai/discover-urls/status", RequireAuth(r.config, r.aiSettingsHandler.HandleURLDiscoveryStatus)) + r.mux.HandleFunc("/api/ai/discover-urls/cancel", RequireAdmin(r.config, r.aiSettingsHandler.HandleCancelURLDiscovery)) + // Agent WebSocket for AI command execution r.mux.HandleFunc("/api/agent/ws", r.handleAgentWebSocket)