From b7e45a319fa78d4c1019e62c6ccbacb8c4fbc157 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:42:13 +0800 Subject: [PATCH] Add files via upload --- internal/handler/asset.go | 305 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 305 insertions(+) create mode 100644 internal/handler/asset.go diff --git a/internal/handler/asset.go b/internal/handler/asset.go new file mode 100644 index 00000000..c7c58047 --- /dev/null +++ b/internal/handler/asset.go @@ -0,0 +1,305 @@ +package handler + +import ( + "errors" + "net/http" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type AssetHandler struct { + db *database.DB + logger *zap.Logger +} + +func NewAssetHandler(db *database.DB, logger *zap.Logger) *AssetHandler { + return &AssetHandler{db: db, logger: logger} +} + +func assetAccess(c *gin.Context) database.RBACListAccess { + if session, ok := security.CurrentSession(c); ok { + return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope} + } + return database.RBACListAccess{} +} + +type importAssetsRequest struct { + Assets []*database.Asset `json:"assets" binding:"required"` + Source string `json:"source"` + SourceQuery string `json:"source_query"` +} + +type assetScanLink struct { + AssetID string `json:"asset_id" binding:"required"` + ConversationID string `json:"conversation_id"` + QueueID string `json:"queue_id"` + TaskID string `json:"task_id"` +} + +type recordAssetScansRequest struct { + Scans []assetScanLink `json:"scans" binding:"required"` +} + +type updateAssetsProjectRequest struct { + AssetIDs []string `json:"asset_ids" binding:"required"` + ProjectID string `json:"project_id"` +} + +func (h *AssetHandler) Import(c *gin.Context) { + var req importAssetsRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.Assets) == 0 || len(req.Assets) > 1000 { + c.JSON(http.StatusBadRequest, gin.H{"error": "assets 数量必须在 1-1000 之间"}) + return + } + owner := "" + allowGlobal := false + if session, ok := security.CurrentSession(c); ok { + owner = session.UserID + allowGlobal = session.Scope == database.RBACScopeAll + } + for _, asset := range req.Assets { + if asset == nil { + continue + } + if strings.TrimSpace(asset.ProjectID) != "" { + if session, ok := security.CurrentSession(c); ok && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", strings.TrimSpace(asset.ProjectID)) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权绑定该项目"}) + return + } + } + if strings.TrimSpace(asset.Source) == "" { + asset.Source = strings.TrimSpace(req.Source) + } + if strings.TrimSpace(asset.SourceQuery) == "" { + asset.SourceQuery = strings.TrimSpace(req.SourceQuery) + } + } + result, err := h.db.UpsertAssets(req.Assets, owner, allowGlobal) + if err != nil { + var validationErr *database.AssetValidationError + if errors.As(err, &validationErr) { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + h.logger.Error("导入资产失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, result) +} + +func (h *AssetHandler) List(c *gin.Context) { + page, _ := strconv.Atoi(c.DefaultQuery("page", "1")) + pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "20")) + if page < 1 { + page = 1 + } + if pageSize < 1 || pageSize > 100 { + pageSize = 20 + } + filter, err := assetListFilterFromQuery(c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + assets, total, err := h.db.ListAssets(pageSize, (page-1)*pageSize, filter, assetAccess(c)) + if err != nil { + h.logger.Error("加载资产失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + totalPages := (total + pageSize - 1) / pageSize + if totalPages < 1 { + totalPages = 1 + } + c.JSON(http.StatusOK, gin.H{"assets": assets, "total": total, "page": page, "page_size": pageSize, "total_pages": totalPages}) +} + +func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error) { + filter := database.AssetListFilter{ + Search: strings.TrimSpace(c.Query("q")), Status: strings.ToLower(strings.TrimSpace(c.Query("status"))), + Protocol: strings.ToLower(strings.TrimSpace(c.Query("protocol"))), ProjectID: strings.TrimSpace(c.Query("project_id")), + Source: strings.TrimSpace(c.Query("source")), Tag: strings.TrimSpace(c.Query("tag")), Host: strings.TrimSpace(c.Query("host")), + IP: strings.TrimSpace(c.Query("ip")), Domain: strings.TrimSpace(c.Query("domain")), ScanState: strings.ToLower(strings.TrimSpace(c.Query("scan_state"))), + SortBy: strings.ToLower(strings.TrimSpace(c.Query("sort_by"))), SortOrder: strings.ToLower(strings.TrimSpace(c.Query("sort_order"))), + } + if raw := strings.TrimSpace(c.Query("port")); raw != "" { + port, err := strconv.Atoi(raw) + if err != nil || port < 0 || port > 65535 { + return filter, &assetQueryError{field: "port", value: raw} + } + filter.Port = &port + } + var err error + if filter.LastScanBefore, err = parseAssetQueryTime("last_scan_before", c.Query("last_scan_before")); err != nil { + return filter, err + } + if filter.LastScanAfter, err = parseAssetQueryTime("last_scan_after", c.Query("last_scan_after")); err != nil { + return filter, err + } + if filter.LastSeenBefore, err = parseAssetQueryTime("last_seen_before", c.Query("last_seen_before")); err != nil { + return filter, err + } + if filter.LastSeenAfter, err = parseAssetQueryTime("last_seen_after", c.Query("last_seen_after")); err != nil { + return filter, err + } + return filter, nil +} + +type assetQueryError struct{ field, value string } + +func (e *assetQueryError) Error() string { + return e.field + " 参数无效: " + e.value +} + +func parseAssetQueryTime(field, value string) (*time.Time, error) { + value = strings.TrimSpace(value) + if value == "" { + return nil, nil + } + for _, layout := range []string{time.RFC3339, "2006-01-02"} { + if parsed, err := time.Parse(layout, value); err == nil { + return &parsed, nil + } + } + return nil, &assetQueryError{field: field, value: value} +} + +func (h *AssetHandler) Stats(c *gin.Context) { + days := 30 + if raw := strings.TrimSpace(c.Query("days")); raw != "" { + parsed, err := strconv.Atoi(raw) + if err != nil || (parsed != 7 && parsed != 30 && parsed != 90) { + c.JSON(http.StatusBadRequest, gin.H{"error": "days 仅支持 7、30 或 90"}) + return + } + days = parsed + } + stats, err := h.db.GetAssetStats(assetAccess(c), days) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, stats) +} + +// RecordScans stores the execution link created by the asset-library scan action. +func (h *AssetHandler) RecordScans(c *gin.Context) { + var req recordAssetScansRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.Scans) == 0 || len(req.Scans) > 1000 { + c.JSON(http.StatusBadRequest, gin.H{"error": "scans 数量必须在 1-1000 之间"}) + return + } + access := assetAccess(c) + for _, scan := range req.Scans { + conversationID := strings.TrimSpace(scan.ConversationID) + queueID := strings.TrimSpace(scan.QueueID) + taskID := strings.TrimSpace(scan.TaskID) + if conversationID == "" && taskID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "conversation_id 或 task_id 至少需要一个"}) + return + } + if taskID != "" && (queueID == "" || !h.db.BatchTaskBelongsToQueue(taskID, queueID)) { + c.JSON(http.StatusBadRequest, gin.H{"error": "任务不属于指定队列"}) + return + } + if _, err := h.db.GetAsset(strings.TrimSpace(scan.AssetID), access); err != nil { + c.JSON(http.StatusForbidden, gin.H{"error": "资产不存在或无权扫描"}) + return + } + if session, ok := security.CurrentSession(c); ok { + if id := conversationID; id != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权关联该对话"}) + return + } + if id := queueID; id != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "batch_task", id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权关联该任务队列"}) + return + } + } + } + for _, scan := range req.Scans { + if err := h.db.MarkAssetScanned(scan.AssetID, scan.ConversationID, scan.QueueID, scan.TaskID, access); err != nil { + h.logger.Error("记录资产扫描失败", zap.String("asset_id", scan.AssetID), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + c.JSON(http.StatusOK, gin.H{"updated": len(req.Scans)}) +} + +func (h *AssetHandler) Update(c *gin.Context) { + var asset database.Asset + if err := c.ShouldBindJSON(&asset); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if asset.ProjectID != "" { + if session, ok := security.CurrentSession(c); ok && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", asset.ProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权绑定该项目"}) + return + } + } + if err := h.db.UpdateAsset(c.Param("id"), &asset, assetAccess(c)); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + updated, err := h.db.GetAsset(c.Param("id"), assetAccess(c)) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "资产不存在"}) + return + } + c.JSON(http.StatusOK, updated) +} + +// UpdateProjectBinding replaces the project binding for a selected asset set. +func (h *AssetHandler) UpdateProjectBinding(c *gin.Context) { + var req updateAssetsProjectRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.AssetIDs) == 0 || len(req.AssetIDs) > 1000 { + c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-1000 之间"}) + return + } + req.ProjectID = strings.TrimSpace(req.ProjectID) + if req.ProjectID != "" { + if _, err := h.db.GetProject(req.ProjectID); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "项目不存在"}) + return + } + if session, ok := security.CurrentSession(c); ok && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", req.ProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权绑定该项目"}) + return + } + } + updated, err := h.db.UpdateAssetsProject(req.AssetIDs, req.ProjectID, assetAccess(c)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"updated": updated, "project_id": req.ProjectID}) +} + +func (h *AssetHandler) Delete(c *gin.Context) { + if err := h.db.DeleteAsset(c.Param("id"), assetAccess(c)); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "资产不存在或无权删除"}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true}) +}