From 1f3472955c02eddd318887d3c1e1546348750302 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:11:16 +0800 Subject: [PATCH] Add files via upload --- internal/handler/asset.go | 247 +++++++++++++++++++++++++- internal/handler/asset_test.go | 83 +++++++++ internal/handler/config.go | 16 +- internal/handler/fofa.go | 57 +++--- internal/handler/fofa_test.go | 86 +++++++++ internal/handler/openapi.go | 82 +++++++++ internal/handler/openapi_i18n.go | 8 +- internal/handler/openapi_i18n_test.go | 46 +++++ 8 files changed, 594 insertions(+), 31 deletions(-) create mode 100644 internal/handler/asset_test.go create mode 100644 internal/handler/fofa_test.go create mode 100644 internal/handler/openapi_i18n_test.go diff --git a/internal/handler/asset.go b/internal/handler/asset.go index c7c58047..0eb878db 100644 --- a/internal/handler/asset.go +++ b/internal/handler/asset.go @@ -18,6 +18,11 @@ type AssetHandler struct { logger *zap.Logger } +const ( + maxAssetImportBatch = 100000 + maxAssetOperationBatch = 10000 +) + func NewAssetHandler(db *database.DB, logger *zap.Logger) *AssetHandler { return &AssetHandler{db: db, logger: logger} } @@ -29,6 +34,13 @@ func assetAccess(c *gin.Context) database.RBACListAccess { return database.RBACListAccess{} } +func assetAccessForPermission(c *gin.Context, permission string) database.RBACListAccess { + if session, ok := security.CurrentSession(c); ok { + return database.RBACListAccess{UserID: session.UserID, Scope: session.ScopeFor(permission)} + } + return database.RBACListAccess{} +} + type importAssetsRequest struct { Assets []*database.Asset `json:"assets" binding:"required"` Source string `json:"source"` @@ -51,14 +63,35 @@ type updateAssetsProjectRequest struct { ProjectID string `json:"project_id"` } +type bulkUpdateAssetsRequest struct { + AssetIDs []string `json:"asset_ids" binding:"required"` + Status *string `json:"status"` + ResponsiblePerson *string `json:"responsible_person"` + Department *string `json:"department"` + BusinessSystem *string `json:"business_system"` + Environment *string `json:"environment"` + Criticality *string `json:"criticality"` + AddTags []string `json:"add_tags"` + RemoveTags []string `json:"remove_tags"` +} + +type assetIDsRequest struct { + AssetIDs []string `json:"asset_ids" binding:"required"` +} + +type mergeAssetsRequest struct { + AssetIDs []string `json:"asset_ids" binding:"required"` + PrimaryID string `json:"primary_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 之间"}) + if len(req.Assets) == 0 || len(req.Assets) > maxAssetImportBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "assets 数量必须在 1-100000 之间"}) return } owner := "" @@ -132,6 +165,11 @@ func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error) 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"))), + RiskLevel: strings.ToLower(strings.TrimSpace(c.Query("risk_level"))), + Country: strings.TrimSpace(c.Query("country")), Province: strings.TrimSpace(c.Query("province")), City: strings.TrimSpace(c.Query("city")), + ResponsiblePerson: strings.TrimSpace(c.Query("responsible_person")), Department: strings.TrimSpace(c.Query("department")), + BusinessSystem: strings.TrimSpace(c.Query("business_system")), Environment: strings.ToLower(strings.TrimSpace(c.Query("environment"))), + Criticality: strings.ToLower(strings.TrimSpace(c.Query("criticality"))), } if raw := strings.TrimSpace(c.Query("port")); raw != "" { port, err := strconv.Atoi(raw) @@ -140,6 +178,21 @@ func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error) } filter.Port = &port } + for field, target := range map[string]**int{ + "min_vulnerabilities": &filter.MinVulnerabilities, + "max_vulnerabilities": &filter.MaxVulnerabilities, + "scan_overdue_days": &filter.ScanOverdueDays, + } { + raw := strings.TrimSpace(c.Query(field)) + if raw == "" { + continue + } + value, err := strconv.Atoi(raw) + if err != nil || value < 0 || (field == "scan_overdue_days" && value == 0) { + return filter, &assetQueryError{field: field, value: raw} + } + *target = &value + } var err error if filter.LastScanBefore, err = parseAssetQueryTime("last_scan_before", c.Query("last_scan_before")); err != nil { return filter, err @@ -147,6 +200,12 @@ func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error) if filter.LastScanAfter, err = parseAssetQueryTime("last_scan_after", c.Query("last_scan_after")); err != nil { return filter, err } + if filter.FirstSeenBefore, err = parseAssetQueryTime("first_seen_before", c.Query("first_seen_before")); err != nil { + return filter, err + } + if filter.FirstSeenAfter, err = parseAssetQueryTime("first_seen_after", c.Query("first_seen_after")); err != nil { + return filter, err + } if filter.LastSeenBefore, err = parseAssetQueryTime("last_seen_before", c.Query("last_seen_before")); err != nil { return filter, err } @@ -156,6 +215,21 @@ func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error) return filter, nil } +// Selection resolves all assets matching the current filter for cross-page actions. +func (h *AssetHandler) Selection(c *gin.Context) { + filter, err := assetListFilterFromQuery(c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + assets, total, err := h.db.ListAssetsForOperation(maxAssetOperationBatch, filter, assetAccess(c)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error(), "total": total}) + return + } + c.JSON(http.StatusOK, gin.H{"assets": assets, "total": total}) +} + type assetQueryError struct{ field, value string } func (e *assetQueryError) Error() string { @@ -200,8 +274,8 @@ func (h *AssetHandler) RecordScans(c *gin.Context) { 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 之间"}) + if len(req.Scans) == 0 || len(req.Scans) > maxAssetOperationBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "scans 数量必须在 1-10000 之间"}) return } access := assetAccess(c) @@ -273,8 +347,8 @@ func (h *AssetHandler) UpdateProjectBinding(c *gin.Context) { 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 之间"}) + if len(req.AssetIDs) == 0 || len(req.AssetIDs) > maxAssetOperationBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"}) return } req.ProjectID = strings.TrimSpace(req.ProjectID) @@ -296,6 +370,167 @@ func (h *AssetHandler) UpdateProjectBinding(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"updated": updated, "project_id": req.ProjectID}) } +func (h *AssetHandler) BulkUpdate(c *gin.Context) { + var req bulkUpdateAssetsRequest + 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) > maxAssetOperationBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"}) + return + } + updated, err := h.db.UpdateAssetsBulk(req.AssetIDs, database.AssetBulkPatch{ + Status: req.Status, ResponsiblePerson: req.ResponsiblePerson, Department: req.Department, + BusinessSystem: req.BusinessSystem, Environment: req.Environment, Criticality: req.Criticality, + AddTags: req.AddTags, RemoveTags: req.RemoveTags, + }, assetAccess(c)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"updated": updated}) +} + +func (h *AssetHandler) BatchDelete(c *gin.Context) { + var req assetIDsRequest + 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) > maxAssetOperationBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"}) + return + } + deleted, err := h.db.DeleteAssets(req.AssetIDs, assetAccess(c)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"deleted": deleted}) +} + +func assetIdentityKeys(asset *database.Asset) map[string]struct{} { + keys := map[string]struct{}{} + if value := strings.ToLower(strings.TrimSpace(asset.Domain)); value != "" { + keys["domain:"+value] = struct{}{} + } + if value := strings.ToLower(strings.Trim(strings.TrimSpace(asset.IP), "[]")); value != "" { + keys["ip:"+value] = struct{}{} + } + if value := strings.ToLower(strings.TrimSpace(asset.Host)); value != "" { + keys["host:"+value] = struct{}{} + } + return keys +} + +func shareAssetIdentity(left, right *database.Asset) bool { + for key := range assetIdentityKeys(left) { + if _, ok := assetIdentityKeys(right)[key]; ok { + return true + } + } + return false +} + +// Merge keeps the selected primary asset and safely combines compatible duplicate metadata. +func (h *AssetHandler) Merge(c *gin.Context) { + var req mergeAssetsRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.AssetIDs) < 2 || len(req.AssetIDs) > 100 { + c.JSON(http.StatusBadRequest, gin.H{"error": "合并资产数量必须在 2-100 之间"}) + return + } + writeAccess := assetAccessForPermission(c, "asset:write") + deleteAccess := assetAccessForPermission(c, "asset:delete") + primaryID := strings.TrimSpace(req.PrimaryID) + if primaryID == "" { + primaryID = strings.TrimSpace(req.AssetIDs[0]) + } + primary, err := h.db.GetAsset(primaryID, writeAccess) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "主资产不存在或无权访问"}) + return + } + others := make([]*database.Asset, 0, len(req.AssetIDs)-1) + seen := map[string]struct{}{primaryID: {}} + for _, id := range req.AssetIDs { + id = strings.TrimSpace(id) + if id == "" || id == primaryID { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + item, err := h.db.GetAsset(id, writeAccess) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "部分资产不存在或无权访问"}) + return + } + if !shareAssetIdentity(primary, item) { + c.JSON(http.StatusBadRequest, gin.H{"error": "所选资产没有共同域名、IP 或 Host，不能判定为重复资产"}) + return + } + others = append(others, item) + } + if len(others) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "至少需要两个不同资产"}) + return + } + mergeText := func(dst *string, src string) { + if strings.TrimSpace(*dst) == "" && strings.TrimSpace(src) != "" { + *dst = src + } + } + tagSet := map[string]struct{}{} + for _, tag := range primary.Tags { + tagSet[tag] = struct{}{} + } + for _, item := range others { + mergeText(&primary.ProjectID, item.ProjectID) + mergeText(&primary.Host, item.Host) + mergeText(&primary.IP, item.IP) + mergeText(&primary.Domain, item.Domain) + mergeText(&primary.Protocol, item.Protocol) + mergeText(&primary.Title, item.Title) + mergeText(&primary.Server, item.Server) + mergeText(&primary.Country, item.Country) + mergeText(&primary.Province, item.Province) + mergeText(&primary.City, item.City) + mergeText(&primary.ResponsiblePerson, item.ResponsiblePerson) + mergeText(&primary.Department, item.Department) + mergeText(&primary.BusinessSystem, item.BusinessSystem) + mergeText(&primary.Environment, item.Environment) + mergeText(&primary.Criticality, item.Criticality) + for _, tag := range item.Tags { + tagSet[tag] = struct{}{} + } + } + primary.Tags = primary.Tags[:0] + for tag := range tagSet { + primary.Tags = append(primary.Tags, tag) + } + if len(primary.Tags) > 30 { + c.JSON(http.StatusBadRequest, gin.H{"error": "合并后标签超过 30 个"}) + return + } + ids := make([]string, 0, len(others)) + for _, item := range others { + ids = append(ids, item.ID) + } + merged, err := h.db.MergeAssets(primary, ids, writeAccess, deleteAccess) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + updated, _ := h.db.GetAsset(primary.ID, writeAccess) + c.JSON(http.StatusOK, gin.H{"merged": merged, "asset": updated}) +} + 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": "资产不存在或无权删除"}) diff --git a/internal/handler/asset_test.go b/internal/handler/asset_test.go new file mode 100644 index 00000000..b45cebb3 --- /dev/null +++ b/internal/handler/asset_test.go @@ -0,0 +1,83 @@ +package handler + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestAssetListPaginatesWithinProject(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "asset-list-pagination.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + project, err := db.CreateProject(&database.Project{Name: "Paged Project", Status: "active"}) + if err != nil { + t.Fatal(err) + } + otherProject, err := db.CreateProject(&database.Project{Name: "Other Project", Status: "active"}) + if err != nil { + t.Fatal(err) + } + + assets := make([]*database.Asset, 0, 8) + for i := 1; i <= 7; i++ { + assets = append(assets, &database.Asset{ + ProjectID: project.ID, + IP: fmt.Sprintf("192.0.2.%d", i), + Port: 80, + Protocol: "http", + }) + } + assets = append(assets, &database.Asset{ + ProjectID: otherProject.ID, + IP: "198.51.100.1", + Port: 443, + Protocol: "https", + }) + if _, err := db.UpsertAssets(assets, "", true); err != nil { + t.Fatal(err) + } + + router := gin.New() + router.GET("/api/assets", NewAssetHandler(db, zap.NewNop()).List) + request := httptest.NewRequest(http.MethodGet, "/api/assets?project_id="+project.ID+"&page=2&page_size=3", nil) + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + if response.Code != http.StatusOK { + t.Fatalf("unexpected status %d: %s", response.Code, response.Body.String()) + } + var payload struct { + Assets []*database.Asset `json:"assets"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"page_size"` + TotalPages int `json:"total_pages"` + } + if err := json.Unmarshal(response.Body.Bytes(), &payload); err != nil { + t.Fatal(err) + } + if payload.Total != 7 || payload.Page != 2 || payload.PageSize != 3 || payload.TotalPages != 3 { + t.Fatalf("unexpected pagination: total=%d page=%d page_size=%d total_pages=%d", + payload.Total, payload.Page, payload.PageSize, payload.TotalPages) + } + if len(payload.Assets) != 3 { + t.Fatalf("expected 3 assets on page 2, got %d", len(payload.Assets)) + } + for _, asset := range payload.Assets { + if asset.ProjectID != project.ID { + t.Fatalf("asset from another project leaked into page: %#v", asset) + } + } +} diff --git a/internal/handler/config.go b/internal/handler/config.go index 9087939f..6a1fc776 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -775,7 +775,7 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) { // 更新FOFA配置 if req.FOFA != nil { h.config.FOFA = *req.FOFA - h.logger.Info("更新FOFA配置", zap.String("email", h.config.FOFA.Email)) + h.logger.Info("更新FOFA配置", zap.String("base_url", h.config.FOFA.BaseURL)) } // 更新MCP配置 @@ -1788,7 +1788,7 @@ func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) { root := doc.Content[0] fofaNode := ensureMap(root, "fofa") setStringInMap(fofaNode, "base_url", cfg.BaseURL) - setStringInMap(fofaNode, "email", cfg.Email) + removeKeyFromMap(fofaNode, "email") setStringInMap(fofaNode, "api_key", cfg.APIKey) } @@ -2102,6 +2102,18 @@ func setStringInMap(mapNode *yaml.Node, key, value string) { valueNode.Value = value } +func removeKeyFromMap(mapNode *yaml.Node, key string) { + if mapNode == nil || mapNode.Kind != yaml.MappingNode { + return + } + for i := 0; i+1 < len(mapNode.Content); i += 2 { + if mapNode.Content[i].Value == key { + mapNode.Content = append(mapNode.Content[:i], mapNode.Content[i+2:]...) + return + } + } +} + func setStringSliceInMap(mapNode *yaml.Node, key string, values []string) { _, valueNode := ensureKeyValue(mapNode, key) valueNode.Kind = yaml.SequenceNode diff --git a/internal/handler/fofa.go b/internal/handler/fofa.go index 84ec8131..e3eae38b 100644 --- a/internal/handler/fofa.go +++ b/internal/handler/fofa.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "net" "net/http" "net/url" "os" @@ -36,7 +37,7 @@ func NewFofaHandler(cfg *config.Config, logger *zap.Logger) *FofaHandler { return &FofaHandler{ cfg: cfg, logger: logger, - client: &http.Client{Timeout: 30 * time.Second}, + client: &http.Client{Timeout: 60 * time.Second}, openAIClient: openaiClient.NewClient(llmCfg, llmHTTPClient, logger), } } @@ -80,22 +81,15 @@ type fofaSearchResponse struct { Results []map[string]interface{} `json:"results"` } -func (h *FofaHandler) resolveCredentials() (email, apiKey string) { - // 优先环境变量（便于容器部署），其次配置文件 - email = strings.TrimSpace(os.Getenv("FOFA_EMAIL")) - apiKey = strings.TrimSpace(os.Getenv("FOFA_API_KEY")) - if email != "" && apiKey != "" { - return email, apiKey +func (h *FofaHandler) resolveAPIKey() string { + // 优先环境变量（便于容器部署），其次配置文件。 + if apiKey := strings.TrimSpace(os.Getenv("FOFA_API_KEY")); apiKey != "" { + return apiKey } if h.cfg != nil { - if email == "" { - email = strings.TrimSpace(h.cfg.FOFA.Email) - } - if apiKey == "" { - apiKey = strings.TrimSpace(h.cfg.FOFA.APIKey) - } + return strings.TrimSpace(h.cfg.FOFA.APIKey) } - return email, apiKey + return "" } func (h *FofaHandler) resolveBaseURL() string { @@ -357,12 +351,12 @@ func (h *FofaHandler) Search(c *gin.Context) { req.Fields = "host,ip,port,domain,title,protocol,country,province,city,server" } - email, apiKey := h.resolveCredentials() - if email == "" || apiKey == "" { + apiKey := h.resolveAPIKey() + if apiKey == "" { c.JSON(http.StatusBadRequest, gin.H{ - "error": "FOFA 未配置：请在系统设置中填写 FOFA Email/API Key，或设置环境变量 FOFA_EMAIL/FOFA_API_KEY", - "need": []string{"fofa.email", "fofa.api_key"}, - "env_key": []string{"FOFA_EMAIL", "FOFA_API_KEY"}, + "error": "FOFA 未配置：请在系统设置的资产管理中填写 FOFA API Key，或设置环境变量 FOFA_API_KEY", + "need": []string{"fofa.api_key"}, + "env_key": []string{"FOFA_API_KEY"}, }) return } @@ -377,7 +371,6 @@ func (h *FofaHandler) Search(c *gin.Context) { } params := u.Query() - params.Set("email", email) params.Set("key", apiKey) params.Set("qbase64", qb64) params.Set("size", fmt.Sprintf("%d", req.Size)) @@ -396,10 +389,18 @@ func (h *FofaHandler) Search(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": "创建请求失败: " + err.Error()}) return } + httpReq.Header.Set("User-Agent", "CyberStrikeAI/1.7.4") + httpReq.Header.Set("Accept", "application/json") resp, err := h.client.Do(httpReq) if err != nil { - c.JSON(http.StatusBadGateway, gin.H{"error": "请求 FOFA 失败: " + err.Error()}) + status, message, timeout := safeFofaRequestError(err) + h.logger.Warn("请求 FOFA 失败", + zap.String("host", u.Host), + zap.Bool("timeout", timeout), + zap.String("error_type", fmt.Sprintf("%T", err)), + ) + c.JSON(status, gin.H{"error": message}) return } defer resp.Body.Close() @@ -448,6 +449,20 @@ func (h *FofaHandler) Search(c *gin.Context) { }) } +func safeFofaRequestError(err error) (status int, message string, timeout bool) { + var netErr net.Error + timeout = errors.Is(err, context.DeadlineExceeded) || + (errors.As(err, &netErr) && netErr.Timeout()) + if timeout { + return http.StatusGatewayTimeout, + "FOFA 请求超时（60 秒）：请稍后重试，或减少返回数量和返回字段", + true + } + return http.StatusBadGateway, + "无法连接 FOFA 服务，请检查服务器网络或代理配置", + false +} + func splitAndCleanCSV(s string) []string { parts := strings.Split(s, ",") out := make([]string, 0, len(parts)) diff --git a/internal/handler/fofa_test.go b/internal/handler/fofa_test.go new file mode 100644 index 00000000..a35fd0d8 --- /dev/null +++ b/internal/handler/fofa_test.go @@ -0,0 +1,86 @@ +package handler + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestFofaSearchUsesAPIKeyWithoutEmail(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("FOFA_API_KEY", "") + t.Setenv("FOFA_EMAIL", "legacy@example.com") + + var receivedEmail string + var receivedKey string + fofaServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedEmail = r.URL.Query().Get("email") + receivedKey = r.URL.Query().Get("key") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"error":false,"size":1,"page":1,"results":[["https://example.com"]]}`)) + })) + defer fofaServer.Close() + + h := NewFofaHandler(&config.Config{ + FOFA: config.FofaConfig{ + BaseURL: fofaServer.URL, + APIKey: "test-api-key", + }, + }, zap.NewNop()) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + body := `{"query":"domain=\"example.com\"","fields":"host"}` + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + h.Search(ctx) + + if recorder.Code != http.StatusOK { + t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String()) + } + if receivedEmail != "" { + t.Fatalf("FOFA request unexpectedly included email = %q", receivedEmail) + } + if receivedKey != "test-api-key" { + t.Fatalf("FOFA request key = %q, want %q", receivedKey, "test-api-key") + } + + var response fofaSearchResponse + if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if response.ResultsCount != 1 { + t.Fatalf("results_count = %d, want 1", response.ResultsCount) + } +} + +func TestSafeFofaRequestErrorDoesNotExposeURLOrAPIKey(t *testing.T) { + const secretURL = "https://fofa.info/api/v1/search/all?key=secret-api-key" + err := &url.Error{ + Op: http.MethodGet, + URL: secretURL, + Err: context.DeadlineExceeded, + } + + status, message, timeout := safeFofaRequestError(err) + + if status != http.StatusGatewayTimeout { + t.Fatalf("status = %d, want %d", status, http.StatusGatewayTimeout) + } + if !timeout { + t.Fatal("timeout = false, want true") + } + if strings.Contains(message, "secret-api-key") || strings.Contains(message, secretURL) { + t.Fatalf("safe error exposed request URL or API key: %q", message) + } +} diff --git a/internal/handler/openapi.go b/internal/handler/openapi.go index 4ce3fe5f..9832601c 100644 --- a/internal/handler/openapi.go +++ b/internal/handler/openapi.go @@ -241,6 +241,58 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) { }, }, }, + "AssetImportItem": map[string]interface{}{ + "type": "object", + "description": "待导入资产；host、ip、domain 至少一项非空", + "properties": map[string]interface{}{ + "project_id": map[string]interface{}{"type": "string", "description": "所属项目 ID；调用者必须有权访问"}, + "host": map[string]interface{}{"type": "string", "maxLength": 500, "example": "https://app.example.com:443"}, + "ip": map[string]interface{}{"type": "string", "example": "192.0.2.10"}, + "port": map[string]interface{}{"type": "integer", "minimum": 0, "maximum": 65535, "example": 443}, + "domain": map[string]interface{}{"type": "string", "example": "app.example.com"}, + "protocol": map[string]interface{}{"type": "string", "example": "https"}, + "title": map[string]interface{}{"type": "string", "maxLength": 500}, + "server": map[string]interface{}{"type": "string", "maxLength": 255, "example": "nginx"}, + "country": map[string]interface{}{"type": "string"}, + "province": map[string]interface{}{"type": "string"}, + "city": map[string]interface{}{"type": "string"}, + "responsible_person": map[string]interface{}{"type": "string", "maxLength": 255, "description": "资产负责人"}, + "department": map[string]interface{}{"type": "string", "maxLength": 255, "description": "所属部门"}, + "business_system": map[string]interface{}{"type": "string", "maxLength": 255, "description": "所属业务系统"}, + "environment": map[string]interface{}{"type": "string", "enum": []string{"production", "staging", "testing", "development", "other"}}, + "criticality": map[string]interface{}{"type": "string", "enum": []string{"critical", "high", "medium", "low"}}, + "source": map[string]interface{}{"type": "string"}, + "source_query": map[string]interface{}{"type": "string"}, + "status": map[string]interface{}{"type": "string", "enum": []string{"active", "inactive"}, "default": "active"}, + "tags": map[string]interface{}{ + "type": "array", + "maxItems": 30, + "items": map[string]interface{}{"type": "string", "maxLength": 64}, + }, + }, + }, + "AssetImportRequest": map[string]interface{}{ + "type": "object", + "required": []string{"assets"}, + "properties": map[string]interface{}{ + "assets": map[string]interface{}{ + "type": "array", + "minItems": 1, + "maxItems": 100000, + "items": map[string]interface{}{"$ref": "#/components/schemas/AssetImportItem"}, + }, + "source": map[string]interface{}{"type": "string", "description": "未在资产中填写来源时使用的默认来源"}, + "source_query": map[string]interface{}{"type": "string", "description": "默认来源查询或导入文件名"}, + }, + }, + "AssetImportResult": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "created": map[string]interface{}{"type": "integer", "description": "新建数量", "example": 120}, + "updated": map[string]interface{}{"type": "integer", "description": "去重合并数量", "example": 8}, + "skipped": map[string]interface{}{"type": "integer", "description": "跳过数量", "example": 2}, + }, + }, "ExecutionResult": map[string]interface{}{ "type": "object", "properties": map[string]interface{}{ @@ -2434,6 +2486,36 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) { }, }, }, + "/api/assets/import": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"资产管理"}, + "summary": "批量导入资产", + "description": "新增或按“目标 + 端口 + 协议”去重更新资产。接收 JSON，不直接接收 XLSX/CSV 文件；单次最多 100000 条，需要 asset:write 权限。", + "operationId": "importAssets", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{"$ref": "#/components/schemas/AssetImportRequest"}, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "导入完成", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{"$ref": "#/components/schemas/AssetImportResult"}, + }, + }, + }, + "400": map[string]interface{}{"description": "数量或资产字段校验失败"}, + "401": map[string]interface{}{"description": "未授权"}, + "403": map[string]interface{}{"description": "缺少 asset:write 权限或无权访问指定项目"}, + "500": map[string]interface{}{"description": "导入事务失败"}, + }, + }, + }, "/api/projects": map[string]interface{}{ "get": map[string]interface{}{ "tags": []string{"项目管理"}, diff --git a/internal/handler/openapi_i18n.go b/internal/handler/openapi_i18n.go index 11ce5733..21286e49 100644 --- a/internal/handler/openapi_i18n.go +++ b/internal/handler/openapi_i18n.go @@ -11,7 +11,7 @@ var apiDocI18nTagToKey = map[string]string{ "知识库": "knowledgeBase", "MCP": "mcp", "FOFA信息收集": "fofaRecon", "终端": "terminal", "WebShell管理": "webshellManagement", "对话附件": "chatUploads", "机器人集成": "robotIntegration", "多代理Markdown": "markdownAgents", - "项目管理": "projectManagement", + "项目管理": "projectManagement", "资产管理": "assetManagement", } var apiDocI18nSummaryToKey = map[string]string{ @@ -71,8 +71,9 @@ var apiDocI18nSummaryToKey = map[string]string{ "列出技能包文件": "listSkillPackageFiles", "获取技能包文件内容": "getSkillPackageFile", "写入技能包文件": "putSkillPackageFile", "批量获取工具名称": "batchGetToolNames", "获取知识库统计": "getKnowledgeStats", - "列出项目": "listProjects", "创建项目": "createProject", "获取项目": "getProject", + "列出项目": "listProjects", "创建项目": "createProject", "获取项目": "getProject", "更新项目": "updateProject", "删除项目": "deleteProject", + "批量导入资产": "importAssets", "列出或按 key 获取事实": "listProjectFacts", "创建/更新事实": "upsertProjectFact", "获取项目事实攻击路径图": "getProjectFactGraph", "列出项目全部事实边": "listProjectFactEdges", "添加事实边": "createProjectFactEdge", "删除事实边": "deleteProjectFactEdge", @@ -109,6 +110,9 @@ var apiDocI18nResponseDescToKey = map[string]string{ "成功": "success", "nodes + edges": "factGraphNodesEdges", "边列表": "edgeList", "边已创建": "edgeCreated", "沉淀结果（facts/edges/graph）": "promoteAttackChainResult", + "导入完成": "assetImportCompleted", "数量或资产字段校验失败": "assetImportValidationFailed", + "缺少 asset:write 权限或无权访问指定项目": "assetImportForbidden", + "导入事务失败": "assetImportTransactionFailed", } // enrichSpecWithI18nKeys 在 spec 的每个 operation 上写入 x-i18n-tags、x-i18n-summary， diff --git a/internal/handler/openapi_i18n_test.go b/internal/handler/openapi_i18n_test.go new file mode 100644 index 00000000..12f08a0d --- /dev/null +++ b/internal/handler/openapi_i18n_test.go @@ -0,0 +1,46 @@ +package handler + +import "testing" + +func TestEnrichSpecWithI18nKeysForAssetImport(t *testing.T) { + responses := map[string]interface{}{ + "200": map[string]interface{}{"description": "导入完成"}, + "400": map[string]interface{}{"description": "数量或资产字段校验失败"}, + "403": map[string]interface{}{"description": "缺少 asset:write 权限或无权访问指定项目"}, + "500": map[string]interface{}{"description": "导入事务失败"}, + } + operation := map[string]interface{}{ + "tags": []string{"资产管理"}, + "summary": "批量导入资产", + "responses": responses, + } + spec := map[string]interface{}{ + "paths": map[string]interface{}{ + "/api/assets/import": map[string]interface{}{ + "post": operation, + }, + }, + } + + enrichSpecWithI18nKeys(spec) + + tagKeys, ok := operation["x-i18n-tags"].([]string) + if !ok || len(tagKeys) != 1 || tagKeys[0] != "assetManagement" { + t.Fatalf("unexpected asset tag i18n keys: %#v", operation["x-i18n-tags"]) + } + if got := operation["x-i18n-summary"]; got != "importAssets" { + t.Fatalf("unexpected asset summary i18n key: %#v", got) + } + expectedResponseKeys := map[string]string{ + "200": "assetImportCompleted", + "400": "assetImportValidationFailed", + "403": "assetImportForbidden", + "500": "assetImportTransactionFailed", + } + for status, want := range expectedResponseKeys { + response := responses[status].(map[string]interface{}) + if got := response["x-i18n-description"]; got != want { + t.Errorf("unexpected asset response i18n key for %s: got %#v, want %q", status, got, want) + } + } +}