xingrin/go-backend/internal/handler/vulnerability.go

package handler

import (
	"errors"
	"strconv"

	"github.com/gin-gonic/gin"
	"github.com/xingrin/go-backend/internal/dto"
	"github.com/xingrin/go-backend/internal/model"
	"github.com/xingrin/go-backend/internal/service"
)

// VulnerabilityHandler handles vulnerability endpoints
type VulnerabilityHandler struct {
	svc *service.VulnerabilityService
}

// NewVulnerabilityHandler creates a new vulnerability handler
func NewVulnerabilityHandler(svc *service.VulnerabilityService) *VulnerabilityHandler {
	return &VulnerabilityHandler{svc: svc}
}

// toVulnerabilityResponse converts model to response DTO
func toVulnerabilityResponse(v *model.Vulnerability) dto.VulnerabilityResponse {
	return dto.VulnerabilityResponse{
		ID:          v.ID,
		TargetID:    v.TargetID,
		URL:         v.URL,
		VulnType:    v.VulnType,
		Severity:    v.Severity,
		Source:      v.Source,
		CVSSScore:   v.CVSSScore,
		Description: v.Description,
		RawOutput:   v.RawOutput,
		IsReviewed:  v.IsReviewed,
		ReviewedAt:  v.ReviewedAt,
		CreatedAt:   v.CreatedAt,
	}
}

// ListAll returns paginated vulnerabilities for all targets
// GET /api/assets/vulnerabilities/
func (h *VulnerabilityHandler) ListAll(c *gin.Context) {
	var query dto.VulnerabilityListQuery
	if !dto.BindQuery(c, &query) {
		return
	}

	vulnerabilities, total, err := h.svc.ListAll(&query)
	if err != nil {
		dto.InternalError(c, "Failed to list vulnerabilities")
		return
	}

	// Convert to response
	var resp []dto.VulnerabilityResponse
	for _, v := range vulnerabilities {
		resp = append(resp, toVulnerabilityResponse(&v))
	}

	dto.Paginated(c, resp, total, query.GetPage(), query.GetPageSize())
}

// GetByID returns a vulnerability by ID
// GET /api/assets/vulnerabilities/:id/
func (h *VulnerabilityHandler) GetByID(c *gin.Context) {
	id, err := strconv.Atoi(c.Param("id"))
	if err != nil {
		dto.BadRequest(c, "Invalid vulnerability ID")
		return
	}

	vulnerability, err := h.svc.GetByID(id)
	if err != nil {
		if errors.Is(err, service.ErrVulnerabilityNotFound) {
			dto.NotFound(c, "Vulnerability not found")
			return
		}
		dto.InternalError(c, "Failed to get vulnerability")
		return
	}

	dto.Success(c, toVulnerabilityResponse(vulnerability))
}

// ListByTarget returns paginated vulnerabilities for a target
// GET /api/targets/:id/vulnerabilities/
func (h *VulnerabilityHandler) ListByTarget(c *gin.Context) {
	targetID, err := strconv.Atoi(c.Param("id"))
	if err != nil {
		dto.BadRequest(c, "Invalid target ID")
		return
	}

	var query dto.VulnerabilityListQuery
	if !dto.BindQuery(c, &query) {
		return
	}

	vulnerabilities, total, err := h.svc.ListByTarget(targetID, &query)
	if err != nil {
		if errors.Is(err, service.ErrTargetNotFound) {
			dto.NotFound(c, "Target not found")
			return
		}
		dto.InternalError(c, "Failed to list vulnerabilities")
		return
	}

	// Convert to response
	var resp []dto.VulnerabilityResponse
	for _, v := range vulnerabilities {
		resp = append(resp, toVulnerabilityResponse(&v))
	}

	dto.Paginated(c, resp, total, query.GetPage(), query.GetPageSize())
}

// BulkCreate creates multiple vulnerabilities for a target
// POST /api/targets/:id/vulnerabilities/bulk-create/
func (h *VulnerabilityHandler) BulkCreate(c *gin.Context) {
	targetID, err := strconv.Atoi(c.Param("id"))
	if err != nil {
		dto.BadRequest(c, "Invalid target ID")
		return
	}

	var req dto.BulkCreateVulnerabilitiesRequest
	if !dto.BindJSON(c, &req) {
		return
	}

	createdCount, err := h.svc.BulkCreate(targetID, req.Vulnerabilities)
	if err != nil {
		if errors.Is(err, service.ErrTargetNotFound) {
			dto.NotFound(c, "Target not found")
			return
		}
		dto.InternalError(c, "Failed to create vulnerabilities")
		return
	}

	dto.Success(c, dto.BulkCreateVulnerabilitiesResponse{
		CreatedCount: int(createdCount),
	})
}

// BulkDelete deletes multiple vulnerabilities by IDs
// POST /api/vulnerabilities/bulk-delete/
func (h *VulnerabilityHandler) BulkDelete(c *gin.Context) {
	var req dto.BulkDeleteRequest
	if !dto.BindJSON(c, &req) {
		return
	}

	deletedCount, err := h.svc.BulkDelete(req.IDs)
	if err != nil {
		dto.InternalError(c, "Failed to delete vulnerabilities")
		return
	}

	dto.Success(c, dto.BulkDeleteResponse{DeletedCount: deletedCount})
}

// MarkAsReviewed marks a vulnerability as reviewed
// PATCH /api/vulnerabilities/:id/review/
func (h *VulnerabilityHandler) MarkAsReviewed(c *gin.Context) {
	id, err := strconv.Atoi(c.Param("id"))
	if err != nil {
		dto.BadRequest(c, "Invalid vulnerability ID")
		return
	}

	vulnerability, err := h.svc.MarkAsReviewed(id)
	if err != nil {
		if errors.Is(err, service.ErrVulnerabilityNotFound) {
			dto.NotFound(c, "Vulnerability not found")
			return
		}
		dto.InternalError(c, "Failed to mark vulnerability as reviewed")
		return
	}

	dto.Success(c, toVulnerabilityResponse(vulnerability))
}

// MarkAsUnreviewed marks a vulnerability as pending (unreview)
// PATCH /api/vulnerabilities/:id/unreview/
func (h *VulnerabilityHandler) MarkAsUnreviewed(c *gin.Context) {
	id, err := strconv.Atoi(c.Param("id"))
	if err != nil {
		dto.BadRequest(c, "Invalid vulnerability ID")
		return
	}

	vulnerability, err := h.svc.MarkAsUnreviewed(id)
	if err != nil {
		if errors.Is(err, service.ErrVulnerabilityNotFound) {
			dto.NotFound(c, "Vulnerability not found")
			return
		}
		dto.InternalError(c, "Failed to mark vulnerability as pending")
		return
	}

	dto.Success(c, toVulnerabilityResponse(vulnerability))
}

// BulkMarkAsReviewed marks multiple vulnerabilities as reviewed
// POST /api/vulnerabilities/bulk-review/
func (h *VulnerabilityHandler) BulkMarkAsReviewed(c *gin.Context) {
	var req dto.BulkReviewRequest
	if !dto.BindJSON(c, &req) {
		return
	}

	updatedCount, err := h.svc.BulkMarkAsReviewed(req.IDs)
	if err != nil {
		dto.InternalError(c, "Failed to mark vulnerabilities as reviewed")
		return
	}

	dto.Success(c, dto.BulkReviewResponse{UpdatedCount: updatedCount})
}

// BulkMarkAsUnreviewed marks multiple vulnerabilities as pending
// POST /api/vulnerabilities/bulk-unreview/
func (h *VulnerabilityHandler) BulkMarkAsUnreviewed(c *gin.Context) {
	var req dto.BulkReviewRequest
	if !dto.BindJSON(c, &req) {
		return
	}

	updatedCount, err := h.svc.BulkMarkAsUnreviewed(req.IDs)
	if err != nil {
		dto.InternalError(c, "Failed to mark vulnerabilities as pending")
		return
	}

	dto.Success(c, dto.BulkReviewResponse{UpdatedCount: updatedCount})
}