Redflag/aggregator-server/internal/api/handlers/docker_reports.go

package handlers

import (
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/Fimeg/RedFlag/aggregator-server/internal/database/queries"
	"github.com/Fimeg/RedFlag/aggregator-server/internal/models"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
)

// DockerReportsHandler handles Docker image reports from agents
type DockerReportsHandler struct {
	dockerQueries  *queries.DockerQueries
	agentQueries   *queries.AgentQueries
	commandQueries *queries.CommandQueries
}

func NewDockerReportsHandler(dq *queries.DockerQueries, aq *queries.AgentQueries, cq *queries.CommandQueries) *DockerReportsHandler {
	return &DockerReportsHandler{
		dockerQueries:  dq,
		agentQueries:   aq,
		commandQueries: cq,
	}
}

// ReportDockerImages handles Docker image reports from agents using event sourcing
func (h *DockerReportsHandler) ReportDockerImages(c *gin.Context) {
	agentID := c.MustGet("agent_id").(uuid.UUID)

	// Update last_seen timestamp
	if err := h.agentQueries.UpdateAgentLastSeen(agentID); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update last seen"})
		return
	}

	var req models.DockerReportRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	// Validate command exists and belongs to agent
	commandID, err := uuid.Parse(req.CommandID)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid command ID format"})
		return
	}

	command, err := h.commandQueries.GetCommandByID(commandID)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "command not found"})
		return
	}

	if command.AgentID != agentID {
		c.JSON(http.StatusForbidden, gin.H{"error": "unauthorized command"})
		return
	}

	// Convert Docker images to events
	events := make([]models.StoredDockerImage, 0, len(req.Images))
	for _, item := range req.Images {
		event := models.StoredDockerImage{
			ID:               uuid.New(),
			AgentID:          agentID,
			PackageType:      item.PackageType,
			PackageName:      item.PackageName,
			CurrentVersion:   item.CurrentVersion,
			AvailableVersion: item.AvailableVersion,
			Severity:         item.Severity,
			RepositorySource: item.RepositorySource,
			Metadata:         models.JSONB(item.Metadata),
			EventType:        "discovered",
			CreatedAt:        req.Timestamp,
		}
		events = append(events, event)
	}

	// Store events in batch with error isolation
	if err := h.dockerQueries.CreateDockerEventsBatch(events); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to record docker image events"})
		return
	}

	// Update command status to completed
	result := models.JSONB{
		"docker_images_count": len(req.Images),
		"logged_at":          time.Now(),
	}

	if err := h.commandQueries.MarkCommandCompleted(commandID, result); err != nil {
		fmt.Printf("Warning: Failed to mark docker command %s as completed: %v\n", commandID, err)
	}

	c.JSON(http.StatusOK, gin.H{
		"message":    "docker image events recorded",
		"count":      len(events),
		"command_id": req.CommandID,
	})
}

// GetAgentDockerImages retrieves Docker image updates for a specific agent
func (h *DockerReportsHandler) GetAgentDockerImages(c *gin.Context) {
	agentIDStr := c.Param("agentId")
	agentID, err := uuid.Parse(agentIDStr)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid agent ID"})
		return
	}

	// Parse query parameters
	page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
	pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "50"))
	if page < 1 {
		page = 1
	}
	if pageSize < 1 || pageSize > 100 {
		pageSize = 50
	}

	imageName := c.Query("image_name")
	registry := c.Query("registry")
	severity := c.Query("severity")
	hasUpdatesStr := c.Query("has_updates")

	// Build filter
	filter := &models.DockerFilter{
		AgentID:    &agentID,
		ImageName:  nil,
		Registry:   nil,
		Severity:   nil,
		HasUpdates: nil,
		Limit:      &pageSize,
		Offset:     &((page - 1) * pageSize),
	}

	if imageName != "" {
		filter.ImageName = &imageName
	}
	if registry != "" {
		filter.Registry = &registry
	}
	if severity != "" {
		filter.Severity = &severity
	}
	if hasUpdatesStr != "" {
		hasUpdates := hasUpdatesStr == "true"
		filter.HasUpdates = &hasUpdates
	}

	// Fetch Docker images
	result, err := h.dockerQueries.GetDockerImages(filter)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to fetch docker images"})
		return
	}

	c.JSON(http.StatusOK, result)
}

// GetAgentDockerInfo retrieves detailed Docker information for an agent
func (h *DockerReportsHandler) GetAgentDockerInfo(c *gin.Context) {
	agentIDStr := c.Param("agentId")
	agentID, err := uuid.Parse(agentIDStr)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid agent ID"})
		return
	}

	// Get all Docker images for this agent
	pageSize := 100
	offset := 0

	filter := &models.DockerFilter{
		AgentID: &agentID,
		Limit:   &pageSize,
		Offset:  &offset,
	}

	result, err := h.dockerQueries.GetDockerImages(filter)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to fetch docker images"})
		return
	}

	// Convert to detailed format
	dockerInfo := make([]models.DockerImageInfo, 0, len(result.Images))
	for _, image := range result.Images {
		info := models.DockerImageInfo{
			Name:          image.PackageName,
			Tag:           extractTag(image.PackageName),
			ImageID:       image.CurrentVersion,
			Size:          parseImageSize(image.Metadata),
			CreatedAt:     image.CreatedAt,
			Registry:      image.RepositorySource,
			HasUpdate:     image.AvailableVersion != image.CurrentVersion,
			LatestImageID: image.AvailableVersion,
			Severity:      image.Severity,
			Labels:        extractLabels(image.Metadata),
			LastScanned:   image.CreatedAt,
		}
		dockerInfo = append(dockerInfo, info)
	}

	c.JSON(http.StatusOK, gin.H{
		"docker_images": dockerInfo,
		"is_live":       isDockerRecentlyUpdated(result.Images),
		"total":         len(dockerInfo),
		"updates_available": countUpdates(dockerInfo),
	})
}

// Helper function to extract tag from image name
func extractTag(imageName string) string {
	// Simple implementation - split by ":" and return last part
	parts := strings.Split(imageName, ":")
	if len(parts) > 1 {
		return parts[len(parts)-1]
	}
	return "latest"
}

// Helper function to parse image size from metadata
func parseImageSize(metadata models.JSONB) int64 {
	// Check if size is stored in metadata
	if sizeStr, ok := metadata["size"].(string); ok {
		if size, err := strconv.ParseInt(sizeStr, 10, 64); err == nil {
			return size
		}
	}
	return 0
}

// Helper function to extract labels from metadata
func extractLabels(metadata models.JSONB) map[string]string {
	labels := make(map[string]string)
	if labelsData, ok := metadata["labels"].(map[string]interface{}); ok {
		for k, v := range labelsData {
			if str, ok := v.(string); ok {
				labels[k] = str
			}
		}
	}
	return labels
}

// Helper function to check if Docker images are recently updated
func isDockerRecentlyUpdated(images []models.StoredDockerImage) bool {
	if len(images) == 0 {
		return false
	}

	// Check if any image was updated in the last 5 minutes
	now := time.Now()
	for _, image := range images {
		if now.Sub(image.CreatedAt) < 5*time.Minute {
			return true
		}
	}
	return false
}

// Helper function to count available updates
func countUpdates(images []models.DockerImageInfo) int {
	count := 0
	for _, image := range images {
		if image.HasUpdate {
			count++
		}
	}
	return count
}