From d03cc47ac3ceb9d993cf558671a2255e32cc09eb 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:37:01 +0800 Subject: [PATCH] Add files via upload --- internal/app/app.go | 14 + internal/app/asset_tools.go | 459 ++++++++++++++ internal/app/asset_tools_test.go | 185 ++++++ internal/app/mcp_authorization.go | 29 + internal/app/mcp_authorization_test.go | 39 ++ internal/database/asset.go | 843 +++++++++++++++++++++++++ internal/database/asset_test.go | 316 +++++++++ internal/database/database.go | 52 ++ internal/database/project.go | 3 + internal/database/rbac.go | 41 +- 10 files changed, 1979 insertions(+), 2 deletions(-) create mode 100644 internal/app/asset_tools.go create mode 100644 internal/app/asset_tools_test.go create mode 100644 internal/database/asset.go create mode 100644 internal/database/asset_test.go diff --git a/internal/app/app.go b/internal/app/app.go index 1c0c4b49..d0ef977e 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -158,6 +158,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error // 注册漏洞记录工具 registerVulnerabilityTools(mcpServer, db, log.Logger) + registerAssetTools(mcpServer, db, log.Logger) registerProjectFactTools(mcpServer, db, cfg, log.Logger) registerVisionTools(mcpServer, cfg, log.Logger) @@ -380,6 +381,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error authHandler.SetAudit(auditSvc) attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger) vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger) + assetHandler := handler.NewAssetHandler(db, log.Logger) projectHandler := handler.NewProjectHandler(db, log.Logger) rbacHandler := handler.NewRBACHandler(db, log.Logger) rbacHandler.SetAudit(auditSvc) @@ -465,6 +467,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error // 设置漏洞工具注册器（内置工具，必须设置） vulnerabilityRegistrar := func() error { registerVulnerabilityTools(mcpServer, db, log.Logger) + registerAssetTools(mcpServer, db, log.Logger) registerProjectFactTools(mcpServer, db, cfg, log.Logger) registerVisionTools(mcpServer, cfg, log.Logger) return nil @@ -554,6 +557,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error attackChainHandler, app, // 传递 App 实例以便动态获取 knowledgeHandler vulnerabilityHandler, + assetHandler, projectHandler, workflowHandler, webshellHandler, @@ -856,6 +860,7 @@ func setupRoutes( attackChainHandler *handler.AttackChainHandler, app *App, // 传递 App 实例以便动态获取 knowledgeHandler vulnerabilityHandler *handler.VulnerabilityHandler, + assetHandler *handler.AssetHandler, projectHandler *handler.ProjectHandler, workflowHandler *handler.WorkflowHandler, webshellHandler *handler.WebShellHandler, @@ -976,6 +981,15 @@ func setupRoutes( // 信息收集 - 自然语言解析为 FOFA 语法（需人工确认后再查询） protected.POST("/fofa/parse", fofaHandler.ParseNaturalLanguage) + // 资产管理 + protected.GET("/assets", assetHandler.List) + protected.GET("/assets/stats", assetHandler.Stats) + protected.POST("/assets/import", assetHandler.Import) + protected.POST("/assets/scan-links", assetHandler.RecordScans) + protected.PUT("/assets/project-binding", assetHandler.UpdateProjectBinding) + protected.PUT("/assets/:id", assetHandler.Update) + protected.DELETE("/assets/:id", assetHandler.Delete) + // 批量任务管理 protected.POST("/batch-tasks", agentHandler.CreateBatchQueue) protected.GET("/batch-tasks", agentHandler.ListBatchQueues) diff --git a/internal/app/asset_tools.go b/internal/app/asset_tools.go new file mode 100644 index 00000000..6f7f3820 --- /dev/null +++ b/internal/app/asset_tools.go @@ -0,0 +1,459 @@ +package app + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strings" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" + + "go.uber.org/zap" +) + +const agentAssetPageSizeMax = 50 + +func registerAssetTools(server *mcp.Server, db *database.DB, logger *zap.Logger) { + if server == nil || db == nil { + return + } + properties := assetMutationProperties() + + server.RegisterTool(mcp.Tool{ + Name: builtin.ToolCreateAsset, ShortDescription: "新增或去重更新资产", + Description: "向资产库新增资产。按目标+端口+协议去重；若资产已存在则更新非空字段。至少提供 host、ip、domain 之一。", + InputSchema: map[string]interface{}{"type": "object", "properties": properties, "anyOf": []interface{}{ + map[string]interface{}{"required": []string{"host"}}, map[string]interface{}{"required": []string{"ip"}}, map[string]interface{}{"required": []string{"domain"}}, + }}, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + asset, err := assetFromCreateArgs(args) + if err != nil { + return textResult("错误: "+err.Error(), true), nil + } + access, owner, global := assetAccessFromToolContext(ctx, "asset:write") + result, err := db.UpsertAssets([]*database.Asset{asset}, owner, global) + if err != nil { + logger.Error("Agent 保存资产失败", zap.Error(err)) + return textResult("错误: "+err.Error(), true), nil + } + if result.Skipped > 0 || asset.ID == "" { + return textResult("资产未保存：同一资产已存在但当前用户无权更新，或目标字段为空", true), nil + } + saved, err := db.GetAsset(asset.ID, access) + if err != nil { + return textResult("资产已保存，但无法读取结果: "+err.Error(), true), nil + } + action := "created" + if result.Updated > 0 { + action = "updated" + } + return assetJSONResult(map[string]interface{}{"action": action, "asset": assetToolDetail(saved)}) + }) + + server.RegisterTool(mcp.Tool{ + Name: builtin.ToolGetAsset, ShortDescription: "按 ID 查看资产详情", Description: "按资产 ID 返回完整资产详情。查询列表时先用 query_assets，避免一次拉取过多详情。", + InputSchema: map[string]interface{}{"type": "object", "properties": map[string]interface{}{"id": map[string]interface{}{"type": "string", "description": "资产 ID"}}, "required": []string{"id"}}, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + projectID, projectScoped, err := agentAssetProjectScope(db, ctx) + if err != nil { + return textResult("错误: "+err.Error(), true), nil + } + asset, err := db.GetAsset(strings.TrimSpace(strArg(args, "id")), assetAccessOnly(ctx, "asset:read")) + if err != nil { + if err == sql.ErrNoRows { + return textResult("错误: 资产不存在或无权查看", true), nil + } + return textResult("错误: "+err.Error(), true), nil + } + if projectScoped && strings.TrimSpace(asset.ProjectID) != projectID { + return textResult("错误: 资产不存在或不属于当前对话绑定的项目", true), nil + } + return assetJSONResult(assetToolDetail(asset)) + }) + + server.RegisterTool(mcp.Tool{ + Name: builtin.ToolQueryAssets, ShortDescription: "灵活分页查询资产", + Description: "分页查询资产。支持精确字段、时间范围、扫描状态和白名单排序。查最久未扫描资产请使用 sort_by=last_scan_at、sort_order=asc；从未扫描资产会排在最前。默认每页 20 条，最大 50 条，返回精简摘要；使用 get_asset 获取单条详情。", + InputSchema: assetQuerySchema(), + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + filter, page, pageSize, err := assetFilterFromToolArgs(args) + if err != nil { + return textResult("错误: "+err.Error(), true), nil + } + projectID, projectScoped, err := agentAssetProjectScope(db, ctx) + if err != nil { + return textResult("错误: "+err.Error(), true), nil + } + if projectScoped { + // 对话绑定项目后，项目范围是服务端强制边界；不能通过工具参数扩大或切换范围。 + filter.ProjectID = projectID + } + items, total, err := db.ListAssets(pageSize, (page-1)*pageSize, filter, assetAccessOnly(ctx, "asset:read")) + if err != nil { + return textResult("错误: "+err.Error(), true), nil + } + totalPages := (total + pageSize - 1) / pageSize + if totalPages < 1 { + totalPages = 1 + } + var b strings.Builder + b.WriteString(fmt.Sprintf("资产查询：第 %d/%d 页，本页 %d 条，共 %d 条，page_size=%d\n", page, totalPages, len(items), total, pageSize)) + for _, asset := range items { + b.WriteString(formatAssetListItem(asset)) + b.WriteByte('\n') + } + if page < totalPages { + b.WriteString(fmt.Sprintf("下一页：保持筛选条件并设置 page=%d。", page+1)) + } + return textResult(b.String(), false), nil + }) + + updateProperties := assetMutationProperties() + updateProperties["id"] = map[string]interface{}{"type": "string", "description": "资产 ID"} + server.RegisterTool(mcp.Tool{ + Name: builtin.ToolUpdateAsset, ShortDescription: "局部更新资产", + Description: "按 ID 局部更新资产，只修改显式传入的字段；可传空 project_id 清除项目绑定，可传空 tags 清空标签。", + InputSchema: map[string]interface{}{"type": "object", "properties": updateProperties, "required": []string{"id"}}, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + id := strings.TrimSpace(strArg(args, "id")) + access := assetAccessOnly(ctx, "asset:write") + asset, err := db.GetAsset(id, access) + if err != nil { + return textResult("错误: 资产不存在或无权更新", true), nil + } + if err := applyAssetPatch(asset, args); err != nil { + return textResult("错误: "+err.Error(), true), nil + } + if err := db.UpdateAsset(id, asset, access); err != nil { + return textResult("错误: "+err.Error(), true), nil + } + updated, err := db.GetAsset(id, access) + if err != nil { + return textResult("资产已更新，但无法读取结果: "+err.Error(), true), nil + } + return assetJSONResult(map[string]interface{}{"action": "updated", "asset": assetToolDetail(updated)}) + }) + + server.RegisterTool(mcp.Tool{ + Name: builtin.ToolDeleteAsset, ShortDescription: "删除资产", Description: "按 ID 永久删除资产记录。仅在用户明确要求删除时调用。", + InputSchema: map[string]interface{}{"type": "object", "properties": map[string]interface{}{"id": map[string]interface{}{"type": "string", "description": "资产 ID"}}, "required": []string{"id"}}, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + id := strings.TrimSpace(strArg(args, "id")) + if err := db.DeleteAsset(id, assetAccessOnly(ctx, "asset:delete")); err != nil { + return textResult("错误: 资产不存在或无权删除", true), nil + } + return textResult("资产已删除: "+id, false), nil + }) + + server.RegisterTool(mcp.Tool{ + Name: builtin.ToolCompleteAssetScan, + ShortDescription: "完成资产扫描并回写结果", + Description: "目标扫描完成后调用：把资产的上次扫描时间更新为当前时间，并关联当前对话。相关漏洞数量不手填，而是自动统计当前扫描对话中通过 record_vulnerability 保存的漏洞。应在漏洞均已落库后调用；一个扫描对话建议只对应一个资产。", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{"type": "string", "description": "已完成扫描的资产 ID"}, + }, + "required": []string{"id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + id := strings.TrimSpace(strArg(args, "id")) + conversationID := conversationIDFromToolCtx(ctx) + if conversationID == "" { + return textResult("错误: 无法确定当前扫描对话", true), nil + } + access := assetAccessOnly(ctx, "asset:write") + if err := db.CompleteAssetScan(id, conversationID, access); err != nil { + if err == sql.ErrNoRows { + return textResult("错误: 资产不存在或无权回写扫描结果", true), nil + } + return textResult("错误: "+err.Error(), true), nil + } + updated, err := db.GetAsset(id, access) + if err != nil { + return textResult("扫描结果已回写，但无法读取资产: "+err.Error(), true), nil + } + return assetJSONResult(map[string]interface{}{ + "action": "scan_completed", + "message": "上次扫描时间已更新；相关漏洞数由当前扫描对话中已保存的漏洞自动计算", + "asset": assetToolDetail(updated), + }) + }) +} + +func assetMutationProperties() map[string]interface{} { + return map[string]interface{}{ + "project_id": map[string]interface{}{"type": "string"}, "host": map[string]interface{}{"type": "string"}, + "ip": map[string]interface{}{"type": "string"}, "port": map[string]interface{}{"type": "integer", "minimum": 0, "maximum": 65535}, + "domain": map[string]interface{}{"type": "string"}, "protocol": map[string]interface{}{"type": "string"}, + "title": map[string]interface{}{"type": "string"}, "server": map[string]interface{}{"type": "string"}, + "country": map[string]interface{}{"type": "string"}, "province": map[string]interface{}{"type": "string"}, "city": map[string]interface{}{"type": "string"}, + "source": map[string]interface{}{"type": "string"}, "source_query": map[string]interface{}{"type": "string"}, + "status": map[string]interface{}{"type": "string", "enum": []string{"active", "inactive"}}, + "tags": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "maxItems": 50}, + } +} + +func assetQuerySchema() map[string]interface{} { + properties := map[string]interface{}{ + "q": map[string]interface{}{"type": "string", "description": "模糊搜索 host、IP、域名、标题、服务和标签"}, + "project_id": map[string]interface{}{"type": "string"}, "status": map[string]interface{}{"type": "string", "enum": []string{"active", "inactive"}}, + "protocol": map[string]interface{}{"type": "string"}, "source": map[string]interface{}{"type": "string"}, "tag": map[string]interface{}{"type": "string"}, + "host": map[string]interface{}{"type": "string"}, "ip": map[string]interface{}{"type": "string"}, "domain": map[string]interface{}{"type": "string"}, + "port": map[string]interface{}{"type": "integer", "minimum": 0, "maximum": 65535}, + "scan_state": map[string]interface{}{"type": "string", "enum": []string{"never", "scanned"}, "description": "never=从未扫描，scanned=扫描过"}, + "last_scan_before": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"}, + "last_scan_after": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"}, + "last_seen_before": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"}, + "last_seen_after": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"}, + "sort_by": map[string]interface{}{"type": "string", "enum": []string{"last_seen_at", "last_scan_at", "first_seen_at", "created_at", "updated_at", "host", "port"}}, + "sort_order": map[string]interface{}{"type": "string", "enum": []string{"asc", "desc"}}, + "page": map[string]interface{}{"type": "integer", "minimum": 1}, + "page_size": map[string]interface{}{"type": "integer", "minimum": 1, "maximum": agentAssetPageSizeMax}, + } + return map[string]interface{}{"type": "object", "properties": properties} +} + +func assetFromCreateArgs(args map[string]interface{}) (*database.Asset, error) { + asset := &database.Asset{} + if err := applyAssetPatch(asset, args); err != nil { + return nil, err + } + if strings.TrimSpace(asset.Host) == "" && strings.TrimSpace(asset.IP) == "" && strings.TrimSpace(asset.Domain) == "" { + return nil, fmt.Errorf("host、ip、domain 至少需要一个") + } + return asset, nil +} + +func applyAssetPatch(asset *database.Asset, args map[string]interface{}) error { + setString := func(key string, dst *string) { + if _, ok := args[key]; ok { + *dst = strings.TrimSpace(strArg(args, key)) + } + } + setString("project_id", &asset.ProjectID) + setString("host", &asset.Host) + setString("ip", &asset.IP) + setString("domain", &asset.Domain) + setString("protocol", &asset.Protocol) + setString("title", &asset.Title) + setString("server", &asset.Server) + setString("country", &asset.Country) + setString("province", &asset.Province) + setString("city", &asset.City) + setString("source", &asset.Source) + setString("source_query", &asset.SourceQuery) + setString("status", &asset.Status) + if _, ok := args["port"]; ok { + port := intArg(args, "port", -1) + if port < 0 || port > 65535 { + return fmt.Errorf("port 必须在 0-65535 之间") + } + asset.Port = port + } + if raw, ok := args["tags"]; ok { + tags, err := stringSliceArg(raw) + if err != nil { + return fmt.Errorf("tags: %w", err) + } + asset.Tags = tags + } + return nil +} + +func assetFilterFromToolArgs(args map[string]interface{}) (database.AssetListFilter, int, int, error) { + filter := database.AssetListFilter{ + Search: strings.TrimSpace(strArg(args, "q")), ProjectID: strings.TrimSpace(strArg(args, "project_id")), Status: strings.ToLower(strings.TrimSpace(strArg(args, "status"))), + Protocol: strings.ToLower(strings.TrimSpace(strArg(args, "protocol"))), Source: strings.TrimSpace(strArg(args, "source")), Tag: strings.TrimSpace(strArg(args, "tag")), + Host: strings.TrimSpace(strArg(args, "host")), IP: strings.TrimSpace(strArg(args, "ip")), Domain: strings.TrimSpace(strArg(args, "domain")), + ScanState: strings.ToLower(strings.TrimSpace(strArg(args, "scan_state"))), SortBy: strings.ToLower(strings.TrimSpace(strArg(args, "sort_by"))), + SortOrder: strings.ToLower(strings.TrimSpace(strArg(args, "sort_order"))), + } + if !oneOfOrEmpty(filter.Status, "active", "inactive") { + return filter, 0, 0, fmt.Errorf("status 仅支持 active 或 inactive") + } + if !oneOfOrEmpty(filter.ScanState, "never", "scanned") { + return filter, 0, 0, fmt.Errorf("scan_state 仅支持 never 或 scanned") + } + if !oneOfOrEmpty(filter.SortBy, "last_seen_at", "last_scan_at", "first_seen_at", "created_at", "updated_at", "host", "port") { + return filter, 0, 0, fmt.Errorf("sort_by 不受支持") + } + if !oneOfOrEmpty(filter.SortOrder, "asc", "desc") { + return filter, 0, 0, fmt.Errorf("sort_order 仅支持 asc 或 desc") + } + if _, ok := args["port"]; ok { + port := intArg(args, "port", -1) + if port < 0 || port > 65535 { + return filter, 0, 0, fmt.Errorf("port 必须在 0-65535 之间") + } + filter.Port = &port + } + var err error + if filter.LastScanBefore, err = parseAssetToolTime("last_scan_before", strArg(args, "last_scan_before")); err != nil { + return filter, 0, 0, err + } + if filter.LastScanAfter, err = parseAssetToolTime("last_scan_after", strArg(args, "last_scan_after")); err != nil { + return filter, 0, 0, err + } + if filter.LastSeenBefore, err = parseAssetToolTime("last_seen_before", strArg(args, "last_seen_before")); err != nil { + return filter, 0, 0, err + } + if filter.LastSeenAfter, err = parseAssetToolTime("last_seen_after", strArg(args, "last_seen_after")); err != nil { + return filter, 0, 0, err + } + page := intArg(args, "page", 1) + pageSize := intArg(args, "page_size", 20) + if page < 1 || page > 1_000_000 { + return filter, 0, 0, fmt.Errorf("page 必须在 1-1000000 之间") + } + if pageSize < 1 || pageSize > agentAssetPageSizeMax { + return filter, 0, 0, fmt.Errorf("page_size 必须在 1-%d 之间", agentAssetPageSizeMax) + } + return filter, page, pageSize, nil +} + +func oneOfOrEmpty(value string, allowed ...string) bool { + if value == "" { + return true + } + for _, candidate := range allowed { + if value == candidate { + return true + } + } + return false +} + +func parseAssetToolTime(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, fmt.Errorf("%s 必须是 RFC3339 时间或 YYYY-MM-DD", field) +} + +func stringSliceArg(raw interface{}) ([]string, error) { + values := []string{} + switch typed := raw.(type) { + case []string: + values = append(values, typed...) + case []interface{}: + for _, item := range typed { + value, ok := item.(string) + if !ok { + return nil, fmt.Errorf("必须是字符串数组") + } + values = append(values, value) + } + default: + return nil, fmt.Errorf("必须是字符串数组") + } + if len(values) > 50 { + return nil, fmt.Errorf("最多 50 个标签") + } + return values, nil +} + +func assetAccessOnly(ctx context.Context, permission string) database.RBACListAccess { + principal, ok := authctx.PrincipalFromContext(ctx) + if !ok { + return database.RBACListAccess{} + } + return database.RBACListAccess{UserID: principal.UserID, Scope: principal.ScopeFor(permission)} +} + +func assetAccessFromToolContext(ctx context.Context, permission string) (database.RBACListAccess, string, bool) { + principal, ok := authctx.PrincipalFromContext(ctx) + if !ok { + return database.RBACListAccess{}, "", false + } + access := database.RBACListAccess{UserID: principal.UserID, Scope: principal.ScopeFor(permission)} + return access, principal.UserID, access.Scope == database.RBACScopeAll +} + +// agentAssetProjectScope returns the hard asset-read boundary implied by the +// current conversation. An unbound conversation (or a tool call outside a +// conversation) keeps the existing all-accessible-assets behavior. A bound +// conversation can only read assets assigned to that exact project. +func agentAssetProjectScope(db *database.DB, ctx context.Context) (projectID string, scoped bool, err error) { + conversationID := conversationIDFromToolCtx(ctx) + if conversationID == "" { + return "", false, nil + } + projectID, err = db.GetConversationProjectID(conversationID) + if err != nil { + return "", false, fmt.Errorf("无法确定当前对话的项目范围") + } + projectID = strings.TrimSpace(projectID) + return projectID, projectID != "", nil +} + +func formatAssetListItem(asset *database.Asset) string { + target := asset.Domain + if target == "" { + target = asset.IP + } + if target == "" { + target = asset.Host + } + if asset.Port > 0 { + target = fmt.Sprintf("%s:%d", target, asset.Port) + } + lastScan := "never" + if asset.LastScanAt != nil { + lastScan = asset.LastScanAt.Format(time.RFC3339) + } + return fmt.Sprintf("- id=%s | target=%s | protocol=%s | status=%s | last_scan_at=%s | risk=%s | vulnerabilities=%d", asset.ID, truncateRunes(target, 120), truncateRunes(asset.Protocol, 30), truncateRunes(asset.Status, 30), lastScan, asset.RiskLevel, asset.VulnerabilityCount) +} + +// assetToolDetail keeps even a single unusually large imported record from +// consuming the model context. The database and HTTP API retain full values. +func assetToolDetail(asset *database.Asset) map[string]interface{} { + if asset == nil { + return nil + } + tags := make([]string, 0, len(asset.Tags)) + for i, tag := range asset.Tags { + if i >= 50 { + break + } + tags = append(tags, truncateRunes(tag, 100)) + } + detail := map[string]interface{}{ + "id": asset.ID, "project_id": asset.ProjectID, "project_name": truncateRunes(asset.ProjectName, 200), + "host": truncateRunes(asset.Host, 500), "ip": truncateRunes(asset.IP, 100), "port": asset.Port, + "domain": truncateRunes(asset.Domain, 255), "protocol": truncateRunes(asset.Protocol, 50), + "title": truncateRunes(asset.Title, 500), "server": truncateRunes(asset.Server, 500), + "country": truncateRunes(asset.Country, 100), "province": truncateRunes(asset.Province, 100), "city": truncateRunes(asset.City, 100), + "source": truncateRunes(asset.Source, 100), "source_query": truncateRunes(asset.SourceQuery, 2000), + "status": truncateRunes(asset.Status, 50), "tags": tags, + "first_seen_at": asset.FirstSeenAt, "last_seen_at": asset.LastSeenAt, "created_at": asset.CreatedAt, "updated_at": asset.UpdatedAt, + "last_scan_conversation_id": asset.LastScanConversationID, "last_scan_queue_id": asset.LastScanQueueID, "last_scan_task_id": asset.LastScanTaskID, + "vulnerability_count": asset.VulnerabilityCount, "risk_level": asset.RiskLevel, + } + if asset.LastScanAt != nil { + detail["last_scan_at"] = asset.LastScanAt + } + if len(asset.Tags) > len(tags) { + detail["tags_truncated"] = true + } + return detail +} + +func assetJSONResult(value interface{}) (*mcp.ToolResult, error) { + encoded, err := json.MarshalIndent(value, "", " ") + if err != nil { + return textResult("错误: "+err.Error(), true), nil + } + return textResult(string(encoded), false), nil +} diff --git a/internal/app/asset_tools_test.go b/internal/app/asset_tools_test.go new file mode 100644 index 00000000..50cc7222 --- /dev/null +++ b/internal/app/asset_tools_test.go @@ -0,0 +1,185 @@ +package app + +import ( + "context" + "path/filepath" + "strings" + "testing" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" + + "go.uber.org/zap" +) + +func TestAssetToolsCRUDQueryAndPageLimit(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "asset-tools.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + user, err := db.CreateRBACUser("asset-agent", "Asset Agent", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{ + "asset:read": true, "asset:write": true, "asset:delete": true, + }) + ctx := authctx.WithPrincipal(context.Background(), principal) + server := mcp.NewServer(zap.NewNop()) + server.SetToolAuthorizer(mcpToolAuthorizer(db)) + registerAssetTools(server, db, zap.NewNop()) + + wantTools := map[string]bool{ + builtin.ToolCreateAsset: false, builtin.ToolGetAsset: false, builtin.ToolQueryAssets: false, + builtin.ToolUpdateAsset: false, builtin.ToolDeleteAsset: false, builtin.ToolCompleteAssetScan: false, + } + for _, tool := range server.GetAllTools() { + if _, ok := wantTools[tool.Name]; ok { + wantTools[tool.Name] = true + } + } + for name, found := range wantTools { + if !found { + t.Fatalf("asset tool not registered: %s", name) + } + } + + result, _, err := server.CallTool(ctx, builtin.ToolCreateAsset, map[string]interface{}{ + "ip": "192.0.2.42", "port": 443, "protocol": "https", "title": "Before", "tags": []interface{}{"prod"}, + }) + if err != nil || result == nil || result.IsError { + t.Fatalf("create asset result=%#v err=%v", result, err) + } + assets, total, err := db.ListAssets(20, 0, database.AssetListFilter{}, database.RBACListAccess{UserID: user.ID, Scope: database.RBACScopeAssigned}) + if err != nil || total != 1 || len(assets) != 1 { + t.Fatalf("saved assets total=%d len=%d err=%v", total, len(assets), err) + } + id := assets[0].ID + + result, _, err = server.CallTool(ctx, builtin.ToolUpdateAsset, map[string]interface{}{"id": id, "title": "After"}) + if err != nil || result == nil || result.IsError { + t.Fatalf("update asset result=%#v err=%v", result, err) + } + updated, err := db.GetAsset(id, database.RBACListAccess{UserID: user.ID, Scope: database.RBACScopeAssigned}) + if err != nil || updated.Title != "After" || updated.IP != "192.0.2.42" { + t.Fatalf("partial update lost fields: %#v err=%v", updated, err) + } + + result, _, err = server.CallTool(ctx, builtin.ToolQueryAssets, map[string]interface{}{ + "sort_by": "last_scan_at", "sort_order": "asc", "page": 1, "page_size": 1, + }) + if err != nil || result == nil || result.IsError || !strings.Contains(toolResultText(result), "第 1/1 页") || !strings.Contains(toolResultText(result), "last_scan_at=never") { + t.Fatalf("query asset result=%#v err=%v", result, err) + } + result, _, err = server.CallTool(ctx, builtin.ToolQueryAssets, map[string]interface{}{"page_size": agentAssetPageSizeMax + 1}) + if err != nil || result == nil || !result.IsError { + t.Fatalf("oversized page was accepted: result=%#v err=%v", result, err) + } + + conversation, err := db.CreateConversation("asset scan", database.ConversationCreateMeta{}) + if err != nil { + t.Fatal(err) + } + if err := db.AssignResourceToUser(user.ID, "conversation", conversation.ID); err != nil { + t.Fatal(err) + } + if _, err := db.CreateVulnerability(&database.Vulnerability{ConversationID: conversation.ID, Title: "finding", Severity: "high", Target: "192.0.2.42"}); err != nil { + t.Fatal(err) + } + scanCtx := mcp.WithMCPConversationID(ctx, conversation.ID) + result, _, err = server.CallTool(scanCtx, builtin.ToolCompleteAssetScan, map[string]interface{}{"id": id}) + if err != nil || result == nil || result.IsError { + t.Fatalf("complete scan result=%#v err=%v", result, err) + } + scanned, err := db.GetAsset(id, database.RBACListAccess{UserID: user.ID, Scope: database.RBACScopeAssigned}) + if err != nil || scanned.LastScanAt == nil || scanned.LastScanConversationID != conversation.ID || scanned.VulnerabilityCount != 1 { + t.Fatalf("scan fields not updated: %#v err=%v", scanned, err) + } + + result, _, err = server.CallTool(ctx, builtin.ToolDeleteAsset, map[string]interface{}{"id": id}) + if err != nil || result == nil || result.IsError { + t.Fatalf("delete asset result=%#v err=%v", result, err) + } + if _, err := db.GetAsset(id, database.RBACListAccess{Scope: database.RBACScopeAll}); err == nil { + t.Fatal("asset still exists after delete") + } +} + +func toolResultText(result *mcp.ToolResult) string { + var b strings.Builder + if result == nil { + return "" + } + for _, content := range result.Content { + b.WriteString(content.Text) + } + return b.String() +} + +func TestAssetReadToolsRespectConversationProjectScope(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "asset-project-scope.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + projectA, err := db.CreateProject(&database.Project{Name: "Project A"}) + if err != nil { + t.Fatal(err) + } + projectB, err := db.CreateProject(&database.Project{Name: "Project B"}) + if err != nil { + t.Fatal(err) + } + assets := []*database.Asset{ + {ProjectID: projectA.ID, IP: "192.0.2.10", Protocol: "https"}, + {ProjectID: projectB.ID, IP: "192.0.2.20", Protocol: "https"}, + {IP: "192.0.2.30", Protocol: "https"}, + } + if result, err := db.UpsertAssets(assets, "", true); err != nil || result.Created != len(assets) { + t.Fatalf("seed assets result=%#v err=%v", result, err) + } + + bound, err := db.CreateConversation("bound", database.ConversationCreateMeta{ProjectID: projectA.ID}) + if err != nil { + t.Fatal(err) + } + unbound, err := db.CreateConversation("unbound", database.ConversationCreateMeta{}) + if err != nil { + t.Fatal(err) + } + principal := authctx.NewPrincipal("admin", "admin", database.RBACScopeAll, map[string]bool{"asset:read": true}) + ctx := authctx.WithPrincipal(context.Background(), principal) + server := mcp.NewServer(zap.NewNop()) + server.SetToolAuthorizer(mcpToolAuthorizer(db)) + registerAssetTools(server, db, zap.NewNop()) + + boundCtx := mcp.WithMCPConversationID(ctx, bound.ID) + result, _, err := server.CallTool(boundCtx, builtin.ToolQueryAssets, map[string]interface{}{}) + text := toolResultText(result) + if err != nil || result == nil || result.IsError || !strings.Contains(text, assets[0].ID) || strings.Contains(text, assets[1].ID) || strings.Contains(text, assets[2].ID) { + t.Fatalf("bound query escaped project scope: result=%#v text=%q err=%v", result, text, err) + } + + // Even an explicit foreign project_id cannot override the conversation boundary. + result, _, err = server.CallTool(boundCtx, builtin.ToolQueryAssets, map[string]interface{}{"project_id": projectB.ID}) + text = toolResultText(result) + if err != nil || result == nil || result.IsError || !strings.Contains(text, assets[0].ID) || strings.Contains(text, assets[1].ID) { + t.Fatalf("project_id overrode conversation scope: result=%#v text=%q err=%v", result, text, err) + } + + result, _, err = server.CallTool(boundCtx, builtin.ToolGetAsset, map[string]interface{}{"id": assets[1].ID}) + if err != nil || result == nil || !result.IsError { + t.Fatalf("bound get read a foreign-project asset: result=%#v err=%v", result, err) + } + + unboundCtx := mcp.WithMCPConversationID(ctx, unbound.ID) + result, _, err = server.CallTool(unboundCtx, builtin.ToolQueryAssets, map[string]interface{}{"page_size": 10}) + text = toolResultText(result) + if err != nil || result == nil || result.IsError || !strings.Contains(text, assets[0].ID) || !strings.Contains(text, assets[1].ID) || !strings.Contains(text, assets[2].ID) { + t.Fatalf("unbound query did not retain all-assets behavior: result=%#v text=%q err=%v", result, text, err) + } +} diff --git a/internal/app/mcp_authorization.go b/internal/app/mcp_authorization.go index 3d3c815a..a92b6728 100644 --- a/internal/app/mcp_authorization.go +++ b/internal/app/mcp_authorization.go @@ -71,6 +71,35 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string return nil case builtin.ToolGetVulnerability: return resource("vulnerability:read", "vulnerability", "id") + case builtin.ToolQueryAssets: + return require("asset:read") + case builtin.ToolGetAsset: + return resource("asset:read", "asset", "id") + case builtin.ToolCreateAsset: + if err := require("asset:write"); err != nil { + return err + } + if projectID := mcpAuthorizationString(args, "project_id"); projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("asset:write"), "project", projectID)) { + return fmt.Errorf("no access to project %s", projectID) + } + return nil + case builtin.ToolUpdateAsset, builtin.ToolCompleteAssetScan: + if err := resource("asset:write", "asset", "id"); err != nil { + return err + } + if toolName == builtin.ToolCompleteAssetScan { + conversationID := mcpAuthorizationConversationID(ctx) + if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("asset:write"), "conversation", conversationID) { + return fmt.Errorf("no access to conversation %s", conversationID) + } + return nil + } + if projectID := mcpAuthorizationString(args, "project_id"); projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("asset:write"), "project", projectID)) { + return fmt.Errorf("no access to project %s", projectID) + } + return nil + case builtin.ToolDeleteAsset: + return resource("asset:delete", "asset", "id") case builtin.ToolUpsertProjectFact, builtin.ToolDeprecateProjectFact, builtin.ToolRestoreProjectFact: return authorizeProjectTool(ctx, principal, db, "project:write") case builtin.ToolGetProjectFact, builtin.ToolListProjectFacts, builtin.ToolSearchProjectFacts: diff --git a/internal/app/mcp_authorization_test.go b/internal/app/mcp_authorization_test.go index 0324ca26..6b185b1a 100644 --- a/internal/app/mcp_authorization_test.go +++ b/internal/app/mcp_authorization_test.go @@ -72,6 +72,45 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) { } } +func TestMCPAssetToolAuthorizationUsesAssetPermissionsAndScope(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-asset-authz.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + user, err := db.CreateRBACUser("asset-user", "Asset User", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + owned := &database.Asset{IP: "192.0.2.10", Port: 443, Protocol: "https"} + hidden := &database.Asset{IP: "192.0.2.20", Port: 443, Protocol: "https"} + if _, err := db.UpsertAssets([]*database.Asset{owned}, user.ID); err != nil { + t.Fatal(err) + } + if _, err := db.UpsertAssets([]*database.Asset{hidden}, ""); err != nil { + t.Fatal(err) + } + + permissions := map[string]bool{"asset:read": true, "asset:write": true} + ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, permissions)) + authorize := mcpToolAuthorizer(db) + if err := authorize(ctx, builtin.ToolQueryAssets, nil); err != nil { + t.Fatalf("asset query denied: %v", err) + } + if err := authorize(ctx, builtin.ToolGetAsset, map[string]interface{}{"id": owned.ID}); err != nil { + t.Fatalf("owned asset denied: %v", err) + } + if err := authorize(ctx, builtin.ToolGetAsset, map[string]interface{}{"id": hidden.ID}); err == nil { + t.Fatal("unassigned asset was readable") + } + if err := authorize(ctx, builtin.ToolUpdateAsset, map[string]interface{}{"id": owned.ID}); err != nil { + t.Fatalf("owned asset update denied: %v", err) + } + if err := authorize(ctx, builtin.ToolDeleteAsset, map[string]interface{}{"id": owned.ID}); err == nil { + t.Fatal("asset delete without asset:delete was allowed") + } +} + func TestExternalMCPRequiresDedicatedPermission(t *testing.T) { authorize := externalMCPToolAuthorizer() ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true})) diff --git a/internal/database/asset.go b/internal/database/asset.go new file mode 100644 index 00000000..c90e78d1 --- /dev/null +++ b/internal/database/asset.go @@ -0,0 +1,843 @@ +package database + +import ( + "database/sql" + "encoding/json" + "fmt" + "net" + "net/url" + "regexp" + "strconv" + "strings" + "time" + "unicode/utf8" + + "github.com/google/uuid" + "golang.org/x/net/idna" +) + +// Asset is a persistent, deduplicated target discovered manually or by recon providers. +type Asset struct { + ID string `json:"id"` + ProjectID string `json:"project_id,omitempty"` + ProjectName string `json:"project_name,omitempty"` + Host string `json:"host"` + IP string `json:"ip"` + Port int `json:"port"` + Domain string `json:"domain"` + Protocol string `json:"protocol"` + Title string `json:"title"` + Server string `json:"server"` + Country string `json:"country"` + Province string `json:"province"` + City string `json:"city"` + Source string `json:"source"` + SourceQuery string `json:"source_query"` + Status string `json:"status"` + Tags []string `json:"tags"` + FirstSeenAt time.Time `json:"first_seen_at"` + LastSeenAt time.Time `json:"last_seen_at"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + LastScanAt *time.Time `json:"last_scan_at,omitempty"` + LastScanConversationID string `json:"last_scan_conversation_id,omitempty"` + LastScanQueueID string `json:"last_scan_queue_id,omitempty"` + LastScanTaskID string `json:"last_scan_task_id,omitempty"` + VulnerabilityCount int `json:"vulnerability_count"` + RiskLevel string `json:"risk_level"` + OwnerUserID string `json:"-"` +} + +type AssetListFilter struct { + Search string + Status string + Protocol string + ProjectID string + Source string + Tag string + Host string + IP string + Domain string + Port *int + ScanState string + LastScanBefore *time.Time + LastScanAfter *time.Time + LastSeenBefore *time.Time + LastSeenAfter *time.Time + SortBy string + SortOrder string +} + +type AssetImportResult struct { + Created int `json:"created"` + Updated int `json:"updated"` + Skipped int `json:"skipped"` +} + +func normalizeAsset(a *Asset) { + a.Host = strings.TrimSpace(a.Host) + a.IP = strings.ToLower(strings.TrimSpace(a.IP)) + a.Domain = strings.ToLower(strings.TrimSpace(a.Domain)) + a.Protocol = strings.ToLower(strings.TrimSpace(a.Protocol)) + a.Title = strings.TrimSpace(a.Title) + a.Server = strings.TrimSpace(a.Server) + a.Country = strings.TrimSpace(a.Country) + a.Province = strings.TrimSpace(a.Province) + a.City = strings.TrimSpace(a.City) + a.Source = strings.TrimSpace(a.Source) + a.SourceQuery = strings.TrimSpace(a.SourceQuery) + a.ProjectID = strings.TrimSpace(a.ProjectID) + a.Status = strings.ToLower(strings.TrimSpace(a.Status)) + if a.Status == "" { + a.Status = "active" + } + if a.Source == "" { + a.Source = "manual" + } + seen := map[string]bool{} + tags := make([]string, 0, len(a.Tags)) + for _, tag := range a.Tags { + tag = strings.TrimSpace(tag) + if tag != "" && !seen[tag] { + seen[tag] = true + tags = append(tags, tag) + } + } + a.Tags = tags + // URL 型 Host 是常见输入。缺失的结构化字段在服务端同样补齐，确保 + // API、MCP 与 Web 端产生一致的去重键，而不依赖某个客户端正确解析。 + if strings.Contains(a.Host, "://") { + if parsed, err := url.Parse(a.Host); err == nil && parsed.Hostname() != "" && parsed.User == nil { + hostname := strings.Trim(strings.ToLower(parsed.Hostname()), "[]") + if net.ParseIP(hostname) != nil && a.IP == "" { + a.IP = hostname + } else if a.Domain == "" { + if ascii, err := idna.Lookup.ToASCII(hostname); err == nil { + a.Domain = strings.ToLower(ascii) + } + } + if a.Protocol == "" { + a.Protocol = strings.ToLower(parsed.Scheme) + } + if a.Port == 0 { + if parsed.Port() != "" { + a.Port, _ = strconv.Atoi(parsed.Port()) + } else if a.Protocol == "https" { + a.Port = 443 + } else if a.Protocol == "http" { + a.Port = 80 + } + } + } + } + // Recon providers occasionally return placeholders, multiple values, or + // provider-specific identifiers in structured fields. They are optional + // enrichment; a valid Host must not make the entire batch fail because one + // of those fields is dirty. + if strings.EqualFold(a.Source, "fofa") { + if a.IP != "" && net.ParseIP(strings.Trim(a.IP, "[]")) == nil { + a.IP = "" + } + if a.Domain != "" { + ascii, err := idna.Lookup.ToASCII(strings.TrimSuffix(a.Domain, ".")) + if err != nil || !validAssetDomain(ascii) { + a.Domain = "" + } else { + a.Domain = strings.ToLower(ascii) + } + } + if a.Protocol != "" && !assetProtocolPattern.MatchString(a.Protocol) { + a.Protocol = "" + } + } +} + +var assetProtocolPattern = regexp.MustCompile(`^[a-z][a-z0-9+.-]{0,31}$`) + +// AssetValidationError distinguishes user-correctable asset data from storage failures. +type AssetValidationError struct{ Message string } + +func (e *AssetValidationError) Error() string { return e.Message } + +func assetValidationErrorf(format string, args ...interface{}) error { + return &AssetValidationError{Message: fmt.Sprintf(format, args...)} +} + +func validateAsset(a *Asset) error { + if a == nil { + return assetValidationErrorf("资产不能为空") + } + if a.Host == "" && a.IP == "" && a.Domain == "" { + return assetValidationErrorf("资产目标不能为空") + } + if a.Port < 0 || a.Port > 65535 { + return assetValidationErrorf("端口必须在 0-65535 之间") + } + if a.IP != "" && net.ParseIP(strings.Trim(a.IP, "[]")) == nil { + return assetValidationErrorf("IP 地址格式无效") + } + if a.Domain != "" { + ascii, err := idna.Lookup.ToASCII(strings.TrimSuffix(a.Domain, ".")) + if err != nil || !validAssetDomain(ascii) { + return assetValidationErrorf("域名格式无效") + } + a.Domain = strings.ToLower(ascii) + } + if a.Protocol != "" && !assetProtocolPattern.MatchString(a.Protocol) { + return assetValidationErrorf("协议格式无效") + } + if a.Status != "active" && a.Status != "inactive" { + return assetValidationErrorf("资产状态必须为 active 或 inactive") + } + for name, value := range map[string]string{ + "Host": a.Host, "域名": a.Domain, "协议": a.Protocol, "页面标题": a.Title, + "服务指纹": a.Server, "国家/地区": a.Country, "省份/州": a.Province, "城市": a.City, + } { + limit := 255 + if name == "Host" || name == "页面标题" { + limit = 500 + } + if utf8.RuneCountInString(value) > limit { + return assetValidationErrorf("%s不能超过 %d 个字符", name, limit) + } + } + if len(a.Tags) > 30 { + return assetValidationErrorf("标签不能超过 30 个") + } + for _, tag := range a.Tags { + if utf8.RuneCountInString(tag) > 64 { + return assetValidationErrorf("单个标签不能超过 64 个字符") + } + } + return nil +} + +func validAssetDomain(domain string) bool { + domain = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(domain)), ".") + if domain == "" || len(domain) > 253 || net.ParseIP(domain) != nil { + return false + } + for _, label := range strings.Split(domain, ".") { + if len(label) == 0 || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' { + return false + } + for _, r := range label { + if (r < 'a' || r > 'z') && (r < '0' || r > '9') && r != '-' { + return false + } + } + } + return true +} + +func assetDedupKey(a *Asset) string { + target := a.Domain + if target == "" { + target = a.IP + } + if target == "" { + target = strings.ToLower(a.Host) + } + return strings.Join([]string{target, strconv.Itoa(a.Port), a.Protocol}, "|") +} + +func appendAssetAccess(query string, args []interface{}, access RBACListAccess, alias string) (string, []interface{}) { + if strings.TrimSpace(access.UserID) == "" || access.Scope == RBACScopeAll { + return query, args + } + prefix := "" + if alias != "" { + prefix = alias + "." + } + query += ` AND (` + prefix + `owner_user_id = ? OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'asset' AND ra.resource_id = ` + prefix + `id + ) OR (` + prefix + `project_id IS NOT NULL AND ` + prefix + `project_id <> '' AND ( + EXISTS (SELECT 1 FROM projects ap WHERE ap.id=` + prefix + `project_id AND ap.owner_user_id=?) + OR EXISTS (SELECT 1 FROM rbac_resource_assignments pra WHERE pra.user_id=? AND pra.resource_type='project' AND pra.resource_id=` + prefix + `project_id) + )))` + return query, append(args, access.UserID, access.UserID, access.UserID, access.UserID) +} + +func (db *DB) UpsertAssets(assets []*Asset, ownerUserID string, allowGlobal ...bool) (AssetImportResult, error) { + result := AssetImportResult{} + tx, err := db.Begin() + if err != nil { + return result, err + } + defer tx.Rollback() + now := time.Now() + for _, asset := range assets { + if asset == nil { + result.Skipped++ + continue + } + normalizeAsset(asset) + if err := validateAsset(asset); err != nil { + return result, fmt.Errorf("第 %d 个资产无效: %w", result.Created+result.Updated+result.Skipped+1, err) + } + key := assetDedupKey(asset) + if key == "|0|" { + result.Skipped++ + continue + } + var existingID string + var existingOwner sql.NullString + err := tx.QueryRow(`SELECT id,owner_user_id FROM assets WHERE dedup_key = ?`, key).Scan(&existingID, &existingOwner) + tagsJSON, _ := json.Marshal(asset.Tags) + if err == sql.ErrNoRows { + asset.ID = uuid.NewString() + asset.FirstSeenAt, asset.LastSeenAt, asset.CreatedAt, asset.UpdatedAt = now, now, now, now + _, err = tx.Exec(`INSERT INTO assets ( + id,dedup_key,project_id,host,ip,port,domain,protocol,title,server,country,province,city,source,source_query,status,tags_json, + first_seen_at,last_seen_at,created_at,updated_at,owner_user_id + ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`, + asset.ID, key, nullIfEmpty(asset.ProjectID), asset.Host, asset.IP, asset.Port, asset.Domain, asset.Protocol, asset.Title, asset.Server, + asset.Country, asset.Province, asset.City, asset.Source, asset.SourceQuery, asset.Status, string(tagsJSON), + now, now, now, now, nullIfEmpty(ownerUserID)) + if err != nil { + return result, fmt.Errorf("创建资产失败: %w", err) + } + if ownerUserID != "" { + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_resource_assignments (id,user_id,resource_type,resource_id,created_at) SELECT ?,id,?,?,? FROM rbac_users WHERE id=?`, uuid.NewString(), "asset", asset.ID, now, ownerUserID); err != nil { + return result, fmt.Errorf("授权新资产失败: %w", err) + } + } + result.Created++ + continue + } + if err != nil { + return result, fmt.Errorf("检查资产去重键失败: %w", err) + } + asset.ID = existingID + global := len(allowGlobal) > 0 && allowGlobal[0] + if !global && existingOwner.Valid && strings.TrimSpace(existingOwner.String) != "" && strings.TrimSpace(existingOwner.String) != strings.TrimSpace(ownerUserID) { + result.Skipped++ + continue + } + _, err = tx.Exec(`UPDATE assets SET + host=CASE WHEN ?<>'' THEN ? ELSE host END, ip=CASE WHEN ?<>'' THEN ? ELSE ip END, + domain=CASE WHEN ?<>'' THEN ? ELSE domain END, protocol=CASE WHEN ?<>'' THEN ? ELSE protocol END, + title=CASE WHEN ?<>'' THEN ? ELSE title END, server=CASE WHEN ?<>'' THEN ? ELSE server END, + country=CASE WHEN ?<>'' THEN ? ELSE country END, province=CASE WHEN ?<>'' THEN ? ELSE province END, + city=CASE WHEN ?<>'' THEN ? ELSE city END, source=CASE WHEN ?<>'' THEN ? ELSE source END, + source_query=CASE WHEN ?<>'' THEN ? ELSE source_query END, project_id=CASE WHEN ?<>'' THEN ? ELSE project_id END, + tags_json=CASE WHEN ?<>'[]' THEN ? ELSE tags_json END, + last_seen_at=?, updated_at=? WHERE id=?`, + asset.Host, asset.Host, asset.IP, asset.IP, asset.Domain, asset.Domain, asset.Protocol, asset.Protocol, + asset.Title, asset.Title, asset.Server, asset.Server, asset.Country, asset.Country, asset.Province, asset.Province, + asset.City, asset.City, asset.Source, asset.Source, asset.SourceQuery, asset.SourceQuery, asset.ProjectID, nullIfEmpty(asset.ProjectID), string(tagsJSON), string(tagsJSON), + now, now, existingID) + if err != nil { + return result, fmt.Errorf("更新资产失败: %w", err) + } + if ownerUserID != "" && (!existingOwner.Valid || strings.TrimSpace(existingOwner.String) == ownerUserID) { + if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_resource_assignments (id,user_id,resource_type,resource_id,created_at) SELECT ?,id,?,?,? FROM rbac_users WHERE id=?`, uuid.NewString(), "asset", existingID, now, ownerUserID); err != nil { + return result, fmt.Errorf("授权资产失败: %w", err) + } + } + result.Updated++ + } + if err := tx.Commit(); err != nil { + return result, err + } + return result, nil +} + +func assetWhere(filter AssetListFilter, access RBACListAccess) (string, []interface{}) { + query := " WHERE 1=1" + args := []interface{}{} + if q := strings.TrimSpace(filter.Search); q != "" { + pattern := "%" + escapeAssetLike(strings.ToLower(q)) + "%" + query += ` AND (LOWER(host) LIKE ? ESCAPE '\' OR LOWER(ip) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\' + OR LOWER(title) LIKE ? ESCAPE '\' OR LOWER(server) LIKE ? ESCAPE '\' OR LOWER(tags_json) LIKE ? ESCAPE '\')` + for i := 0; i < 6; i++ { + args = append(args, pattern) + } + } + if filter.Status != "" { + query += " AND status = ?" + args = append(args, filter.Status) + } + if filter.Protocol != "" { + query += " AND protocol = ?" + args = append(args, filter.Protocol) + } + if filter.ProjectID != "" { + query += " AND assets.project_id = ?" + args = append(args, filter.ProjectID) + } + if filter.Source != "" { + query += " AND LOWER(assets.source) = LOWER(?)" + args = append(args, strings.TrimSpace(filter.Source)) + } + if tag := strings.TrimSpace(filter.Tag); tag != "" { + pattern := "%\"" + escapeAssetLike(strings.ToLower(tag)) + "\"%" + query += ` AND LOWER(assets.tags_json) LIKE ? ESCAPE '\'` + args = append(args, pattern) + } + if filter.Host != "" { + query += " AND LOWER(assets.host) = LOWER(?)" + args = append(args, strings.TrimSpace(filter.Host)) + } + if filter.IP != "" { + query += " AND LOWER(assets.ip) = LOWER(?)" + args = append(args, strings.TrimSpace(filter.IP)) + } + if filter.Domain != "" { + query += " AND LOWER(assets.domain) = LOWER(?)" + args = append(args, strings.TrimSpace(filter.Domain)) + } + if filter.Port != nil { + query += " AND assets.port = ?" + args = append(args, *filter.Port) + } + switch strings.ToLower(strings.TrimSpace(filter.ScanState)) { + case "never": + query += " AND " + assetEffectiveLastScanExpr + " IS NULL" + case "scanned": + query += " AND " + assetEffectiveLastScanExpr + " IS NOT NULL" + } + if filter.LastScanBefore != nil { + query += " AND " + assetEffectiveLastScanExpr + " < ?" + args = append(args, *filter.LastScanBefore) + } + if filter.LastScanAfter != nil { + query += " AND " + assetEffectiveLastScanExpr + " > ?" + args = append(args, *filter.LastScanAfter) + } + if filter.LastSeenBefore != nil { + query += " AND assets.last_seen_at < ?" + args = append(args, *filter.LastSeenBefore) + } + if filter.LastSeenAfter != nil { + query += " AND assets.last_seen_at > ?" + args = append(args, *filter.LastSeenAfter) + } + return appendAssetAccess(query, args, access, "assets") +} + +func escapeAssetLike(value string) string { + value = strings.ReplaceAll(value, `\`, `\\`) + value = strings.ReplaceAll(value, `%`, `\%`) + return strings.ReplaceAll(value, `_`, `\_`) +} + +func scanAsset(scanner interface{ Scan(...interface{}) error }) (*Asset, error) { + var a Asset + var tags string + var lastScanAt interface{} + err := scanner.Scan(&a.ID, &a.ProjectID, &a.ProjectName, &a.Host, &a.IP, &a.Port, &a.Domain, &a.Protocol, &a.Title, &a.Server, &a.Country, + &a.Province, &a.City, &a.Source, &a.SourceQuery, &a.Status, &tags, &a.FirstSeenAt, &a.LastSeenAt, &a.CreatedAt, &a.UpdatedAt, + &lastScanAt, &a.LastScanConversationID, &a.LastScanQueueID, &a.LastScanTaskID, &a.VulnerabilityCount, &a.RiskLevel) + if err != nil { + return nil, err + } + if parsed, ok := parseAssetScanTime(lastScanAt); ok { + a.LastScanAt = &parsed + } + _ = json.Unmarshal([]byte(tags), &a.Tags) + return &a, nil +} + +func parseAssetScanTime(value interface{}) (time.Time, bool) { + if value == nil { + return time.Time{}, false + } + if parsed, ok := value.(time.Time); ok { + return parsed, true + } + var raw string + switch typed := value.(type) { + case string: + raw = typed + case []byte: + raw = string(typed) + default: + raw = fmt.Sprint(typed) + } + for _, layout := range []string{ + time.RFC3339Nano, + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999Z07:00", + "2006-01-02 15:04:05-07:00", + "2006-01-02 15:04:05", + } { + if parsed, err := time.Parse(layout, strings.TrimSpace(raw)); err == nil { + return parsed, true + } + } + return time.Time{}, false +} + +const assetEffectiveLastScanExpr = `COALESCE( + (SELECT bt.completed_at FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.completed_at IS NOT NULL LIMIT 1), + (SELECT MAX(m.updated_at) FROM messages m WHERE m.conversation_id=assets.last_scan_conversation_id AND m.role='assistant'), + assets.last_scan_at + )` + +const assetSelectColumns = `assets.id,COALESCE(assets.project_id,''),COALESCE(p.name,''),assets.host,assets.ip,assets.port,assets.domain,assets.protocol,assets.title,assets.server,assets.country, + assets.province,assets.city,assets.source,assets.source_query,assets.status,assets.tags_json,assets.first_seen_at,assets.last_seen_at,assets.created_at,assets.updated_at, + ` + assetEffectiveLastScanExpr + `,COALESCE(assets.last_scan_conversation_id,''),COALESCE(assets.last_scan_queue_id,''),COALESCE(assets.last_scan_task_id,''), + (SELECT COUNT(DISTINCT v.id) FROM vulnerabilities v WHERE + (COALESCE(assets.last_scan_conversation_id,'')<>'' AND v.conversation_id=assets.last_scan_conversation_id) + OR (COALESCE(assets.last_scan_task_id,'')<>'' AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=v.conversation_id))), + CASE WHEN assets.last_scan_at IS NULL THEN 'unassessed' ELSE CASE COALESCE(( + SELECT MAX(CASE LOWER(COALESCE(v.severity,'')) WHEN 'critical' THEN 5 WHEN 'high' THEN 4 WHEN 'medium' THEN 3 WHEN 'low' THEN 2 WHEN 'info' THEN 1 ELSE 0 END) + FROM vulnerabilities v WHERE + LOWER(COALESCE(v.status,'open')) NOT IN ('fixed','false_positive','ignored') AND ( + (COALESCE(assets.last_scan_conversation_id,'')<>'' AND v.conversation_id=assets.last_scan_conversation_id) + OR (COALESCE(assets.last_scan_task_id,'')<>'' AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=v.conversation_id)) + ) + ),0) WHEN 5 THEN 'critical' WHEN 4 THEN 'high' WHEN 3 THEN 'medium' WHEN 2 THEN 'low' WHEN 1 THEN 'info' ELSE 'normal' END END` + +// MarkAssetScanned links an asset to the conversation or batch subtask created from it. +// The link lets the asset list show the latest scan time and vulnerabilities produced by that scan. +func (db *DB) MarkAssetScanned(id, conversationID, queueID, taskID string, access RBACListAccess) error { + where, args := appendAssetAccess(" WHERE id = ?", []interface{}{strings.TrimSpace(id)}, access, "assets") + res, err := db.Exec(`UPDATE assets SET last_scan_at=?,last_scan_conversation_id=?,last_scan_queue_id=?,last_scan_task_id=?,updated_at=?`+where, + append([]interface{}{time.Now(), strings.TrimSpace(conversationID), strings.TrimSpace(queueID), strings.TrimSpace(taskID), time.Now()}, args...)...) + if err != nil { + return err + } + n, _ := res.RowsAffected() + if n == 0 { + return sql.ErrNoRows + } + return nil +} + +// CompleteAssetScan records completion from inside an Agent conversation. If +// the asset was launched as a batch task, keep its task/queue link only when +// that task belongs to the current conversation; a later ad-hoc chat scan must +// not retain stale task associations. +func (db *DB) CompleteAssetScan(id, conversationID string, access RBACListAccess) error { + id = strings.TrimSpace(id) + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return fmt.Errorf("扫描对话不能为空") + } + where, args := appendAssetAccess(" WHERE id = ?", []interface{}{id}, access, "assets") + now := time.Now() + res, err := db.Exec(`UPDATE assets SET + last_scan_at=?,last_scan_conversation_id=?, + last_scan_queue_id=CASE WHEN EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=?) THEN last_scan_queue_id ELSE '' END, + last_scan_task_id=CASE WHEN EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=?) THEN last_scan_task_id ELSE '' END, + updated_at=?`+where, + append([]interface{}{now, conversationID, conversationID, conversationID, now}, args...)...) + if err != nil { + return err + } + n, _ := res.RowsAffected() + if n == 0 { + return sql.ErrNoRows + } + return nil +} + +func (db *DB) BatchTaskBelongsToQueue(taskID, queueID string) bool { + var count int + err := db.QueryRow(`SELECT COUNT(*) FROM batch_tasks WHERE id=? AND queue_id=?`, strings.TrimSpace(taskID), strings.TrimSpace(queueID)).Scan(&count) + return err == nil && count > 0 +} + +func (db *DB) ListAssets(limit, offset int, filter AssetListFilter, access RBACListAccess) ([]*Asset, int, error) { + if limit < 1 { + limit = 20 + } + if limit > 100 { + limit = 100 + } + if offset < 0 { + offset = 0 + } + where, args := assetWhere(filter, access) + var total int + if err := db.QueryRow("SELECT COUNT(*) FROM assets"+where, args...).Scan(&total); err != nil { + return nil, 0, err + } + orderBy := assetOrderBy(filter.SortBy, filter.SortOrder) + rows, err := db.Query("SELECT "+assetSelectColumns+" FROM assets LEFT JOIN projects p ON p.id=assets.project_id"+where+" ORDER BY "+orderBy+" LIMIT ? OFFSET ?", append(args, limit, offset)...) + if err != nil { + return nil, 0, err + } + defer rows.Close() + items := []*Asset{} + for rows.Next() { + a, err := scanAsset(rows) + if err != nil { + return nil, 0, err + } + items = append(items, a) + } + return items, total, rows.Err() +} + +func assetOrderBy(sortBy, sortOrder string) string { + direction := "DESC" + if strings.EqualFold(strings.TrimSpace(sortOrder), "asc") { + direction = "ASC" + } + var expression string + switch strings.ToLower(strings.TrimSpace(sortBy)) { + case "last_scan_at": + expression = assetEffectiveLastScanExpr + // For oldest-first queries, assets that have never been scanned are the + // most overdue and intentionally appear first. NULLs stay last for DESC. + if direction == "ASC" { + return "CASE WHEN " + expression + " IS NULL THEN 0 ELSE 1 END ASC, " + expression + " ASC, assets.id ASC" + } + return "CASE WHEN " + expression + " IS NULL THEN 1 ELSE 0 END ASC, " + expression + " DESC, assets.id ASC" + case "first_seen_at": + expression = "assets.first_seen_at" + case "created_at": + expression = "assets.created_at" + case "updated_at": + expression = "assets.updated_at" + case "host": + expression = "LOWER(assets.host)" + case "port": + expression = "assets.port" + default: + expression = "assets.last_seen_at" + } + return expression + " " + direction + ", assets.id ASC" +} + +func (db *DB) GetAsset(id string, access RBACListAccess) (*Asset, error) { + query, args := appendAssetAccess("SELECT "+assetSelectColumns+" FROM assets LEFT JOIN projects p ON p.id=assets.project_id WHERE assets.id = ?", []interface{}{id}, access, "assets") + return scanAsset(db.QueryRow(query, args...)) +} + +func (db *DB) UpdateAsset(id string, a *Asset, access RBACListAccess) error { + normalizeAsset(a) + if err := validateAsset(a); err != nil { + return err + } + key := assetDedupKey(a) + if key == "|0|" { + return fmt.Errorf("资产目标不能为空") + } + tags, _ := json.Marshal(a.Tags) + where, args := appendAssetAccess(" WHERE id = ?", []interface{}{id}, access, "assets") + res, err := db.Exec(`UPDATE assets SET dedup_key=?,project_id=?,host=?,ip=?,port=?,domain=?,protocol=?,title=?,server=?,country=?,province=?,city=?,source=?,source_query=?,status=?,tags_json=?,updated_at=?`+where, + append([]interface{}{key, nullIfEmpty(a.ProjectID), a.Host, a.IP, a.Port, a.Domain, a.Protocol, a.Title, a.Server, a.Country, a.Province, a.City, a.Source, a.SourceQuery, a.Status, string(tags), time.Now()}, args...)...) + if err != nil { + return err + } + n, _ := res.RowsAffected() + if n == 0 { + return sql.ErrNoRows + } + return nil +} + +// UpdateAssetsProject atomically replaces the project binding for every asset. +// It refuses the whole update when any requested asset is missing or outside +// the caller's access scope, so a bulk action can never partially succeed. +func (db *DB) UpdateAssetsProject(ids []string, projectID string, access RBACListAccess) (int, error) { + unique := make([]string, 0, len(ids)) + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + id = strings.TrimSpace(id) + if id == "" { + continue + } + if _, exists := seen[id]; exists { + continue + } + seen[id] = struct{}{} + unique = append(unique, id) + } + if len(unique) == 0 { + return 0, fmt.Errorf("资产列表不能为空") + } + + tx, err := db.Begin() + if err != nil { + return 0, err + } + defer tx.Rollback() + + placeholders := strings.TrimSuffix(strings.Repeat("?,", len(unique)), ",") + idArgs := make([]interface{}, len(unique)) + for i, id := range unique { + idArgs[i] = id + } + countQuery, countArgs := appendAssetAccess("SELECT COUNT(*) FROM assets WHERE id IN ("+placeholders+")", idArgs, access, "assets") + var accessible int + if err := tx.QueryRow(countQuery, countArgs...).Scan(&accessible); err != nil { + return 0, err + } + if accessible != len(unique) { + return 0, fmt.Errorf("部分资产不存在或无权更新") + } + + updateArgs := []interface{}{nullIfEmpty(strings.TrimSpace(projectID)), time.Now()} + updateArgs = append(updateArgs, idArgs...) + updateQuery, updateArgs := appendAssetAccess("UPDATE assets SET project_id=?,updated_at=? WHERE id IN ("+placeholders+")", updateArgs, access, "assets") + result, err := tx.Exec(updateQuery, updateArgs...) + if err != nil { + return 0, err + } + updated, err := result.RowsAffected() + if err != nil { + return 0, err + } + if int(updated) != len(unique) { + return 0, fmt.Errorf("批量更新资产失败") + } + if err := tx.Commit(); err != nil { + return 0, err + } + return int(updated), nil +} + +func (db *DB) DeleteAsset(id string, access RBACListAccess) error { + where, args := appendAssetAccess(" WHERE id = ?", []interface{}{id}, access, "assets") + res, err := db.Exec("DELETE FROM assets"+where, args...) + if err != nil { + return err + } + n, _ := res.RowsAffected() + if n == 0 { + return sql.ErrNoRows + } + return nil +} + +func (db *DB) GetAssetStats(access RBACListAccess, requestedDays ...int) (map[string]interface{}, error) { + days := 30 + if len(requestedDays) > 0 && (requestedDays[0] == 7 || requestedDays[0] == 30 || requestedDays[0] == 90) { + days = requestedDays[0] + } + where, args := appendAssetAccess(" WHERE 1=1", nil, access, "assets") + stats := map[string]interface{}{} + row := db.QueryRow(`SELECT COUNT(*),COUNT(DISTINCT NULLIF(ip,'')),COUNT(DISTINCT NULLIF(domain,'')), + COUNT(DISTINCT CASE WHEN port>0 THEN CAST(port AS TEXT) END), + COALESCE(SUM(CASE WHEN datetime(last_seen_at)>=datetime('now','-7 days') THEN 1 ELSE 0 END),0) FROM assets`+where, args...) + var total, ips, domains, ports, recent int + if err := row.Scan(&total, &ips, &domains, &ports, &recent); err != nil { + return nil, err + } + stats["total"], stats["ips"], stats["domains"], stats["ports"], stats["recent"] = total, ips, domains, ports, recent + rows, err := db.Query(`SELECT CASE WHEN protocol='' THEN 'unknown' ELSE protocol END,COUNT(*) FROM assets`+where+` GROUP BY protocol ORDER BY COUNT(*) DESC LIMIT 8`, args...) + if err != nil { + return nil, err + } + defer rows.Close() + dist := []map[string]interface{}{} + for rows.Next() { + var name string + var count int + if err := rows.Scan(&name, &count); err != nil { + return nil, err + } + dist = append(dist, map[string]interface{}{"name": name, "count": count}) + } + stats["protocols"] = dist + stats["period_days"] = days + + coverage := map[string]interface{}{} + coverageRow := db.QueryRow(`SELECT + COALESCE(SUM(CASE WHEN last_scan_at IS NOT NULL THEN 1 ELSE 0 END),0), + COALESCE(SUM(CASE WHEN datetime(last_scan_at)>=datetime('now','-7 days') THEN 1 ELSE 0 END),0), + COALESCE(SUM(CASE WHEN datetime(last_scan_at)>=datetime('now','-30 days') THEN 1 ELSE 0 END),0), + COALESCE(SUM(CASE WHEN last_scan_at IS NULL THEN 1 ELSE 0 END),0), + COALESCE(SUM(CASE WHEN last_scan_at IS NOT NULL AND datetime(last_scan_at) 0 { + coverage["rate"] = int(float64(scanned) / float64(total) * 100) + coverage["recent_rate"] = int(float64(scanned30) / float64(total) * 100) + } else { + coverage["rate"], coverage["recent_rate"] = 0, 0 + } + stats["coverage"] = coverage + + assetDaily := map[string]map[string]int{} + trendWhere, trendArgs := appendAssetAccess(" WHERE datetime(first_seen_at)>=datetime('now',?)", []interface{}{fmt.Sprintf("-%d days", days-1)}, access, "assets") + trendRows, err := db.Query(`SELECT date(first_seen_at), COUNT(*) + FROM assets`+trendWhere+` GROUP BY date(first_seen_at) ORDER BY date(first_seen_at)`, trendArgs...) + if err != nil { + return nil, err + } + for trendRows.Next() { + var day string + var added int + if err := trendRows.Scan(&day, &added); err != nil { + trendRows.Close() + return nil, err + } + assetDaily[day] = map[string]int{"added": added, "inactive": 0} + } + if err := trendRows.Close(); err != nil { + return nil, err + } + inactiveWhere, inactiveArgs := appendAssetAccess(" WHERE status='inactive' AND datetime(updated_at)>=datetime('now',?)", []interface{}{fmt.Sprintf("-%d days", days-1)}, access, "assets") + inactiveRows, err := db.Query(`SELECT date(updated_at), COUNT(*) FROM assets`+inactiveWhere+` GROUP BY date(updated_at) ORDER BY date(updated_at)`, inactiveArgs...) + if err != nil { + return nil, err + } + for inactiveRows.Next() { + var day string + var inactive int + if err := inactiveRows.Scan(&day, &inactive); err != nil { + inactiveRows.Close() + return nil, err + } + if _, ok := assetDaily[day]; !ok { + assetDaily[day] = map[string]int{"added": 0, "inactive": 0} + } + assetDaily[day]["inactive"] = inactive + } + if err := inactiveRows.Close(); err != nil { + return nil, err + } + + riskDaily := map[string]map[string]int{} + riskWhere, riskArgs := appendVulnerabilityAccessFilter(" WHERE datetime(created_at)>=datetime('now',?)", []interface{}{fmt.Sprintf("-%d days", days-1)}, access) + riskRows, err := db.Query(`SELECT date(created_at), COUNT(*), + COALESCE(SUM(CASE WHEN LOWER(severity) IN ('critical','high') THEN 1 ELSE 0 END),0) + FROM vulnerabilities`+riskWhere+` GROUP BY date(created_at) ORDER BY date(created_at)`, riskArgs...) + if err != nil { + return nil, err + } + for riskRows.Next() { + var day string + var discovered, highRisk int + if err := riskRows.Scan(&day, &discovered, &highRisk); err != nil { + riskRows.Close() + return nil, err + } + riskDaily[day] = map[string]int{"discovered": discovered, "high_risk": highRisk} + } + if err := riskRows.Close(); err != nil { + return nil, err + } + + assetTrend := make([]map[string]interface{}, 0, days) + riskTrend := make([]map[string]interface{}, 0, days) + start := time.Now().UTC().Truncate(24*time.Hour).AddDate(0, 0, -(days - 1)) + for i := 0; i < days; i++ { + day := start.AddDate(0, 0, i).Format("2006-01-02") + assetPoint := map[string]interface{}{"date": day, "added": 0, "inactive": 0} + if values, ok := assetDaily[day]; ok { + assetPoint["added"], assetPoint["inactive"] = values["added"], values["inactive"] + } + assetTrend = append(assetTrend, assetPoint) + riskPoint := map[string]interface{}{"date": day, "discovered": 0, "high_risk": 0} + if values, ok := riskDaily[day]; ok { + riskPoint["discovered"], riskPoint["high_risk"] = values["discovered"], values["high_risk"] + } + riskTrend = append(riskTrend, riskPoint) + } + stats["asset_trend"], stats["risk_trend"] = assetTrend, riskTrend + return stats, rows.Err() +} diff --git a/internal/database/asset_test.go b/internal/database/asset_test.go new file mode 100644 index 00000000..8390c24e --- /dev/null +++ b/internal/database/asset_test.go @@ -0,0 +1,316 @@ +package database + +import ( + "path/filepath" + "strings" + "testing" + "time" + + "go.uber.org/zap" +) + +func TestAssetURLNormalizationAndValidation(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "asset-validation.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + asset := &Asset{Host: "https://例子.测试/path", Tags: []string{" prod ", "prod"}} + result, err := db.UpsertAssets([]*Asset{asset}, "") + if err != nil || result.Created != 1 { + t.Fatalf("URL asset was not created: result=%#v err=%v", result, err) + } + if asset.Domain != "xn--fsqu00a.xn--0zwm56d" || asset.Protocol != "https" || asset.Port != 443 { + t.Fatalf("URL fields were not normalized: %#v", asset) + } + if len(asset.Tags) != 1 || asset.Tags[0] != "prod" { + t.Fatalf("tags were not normalized: %#v", asset.Tags) + } + + invalid := []*Asset{ + {IP: "999.1.1.1", Status: "active"}, + {Domain: "bad_domain.example", Status: "active"}, + {Domain: "example.com", Port: 70000, Status: "active"}, + {Domain: "example.com", Protocol: "HTTP 1.1", Status: "active"}, + {Domain: "example.com", Status: "deleted"}, + } + for _, candidate := range invalid { + if _, err := db.UpsertAssets([]*Asset{candidate}, ""); err == nil { + t.Fatalf("invalid asset unexpectedly accepted: %#v", candidate) + } + } + + for _, host := range []string{"123", "not a formal target", "https://", "https://user:password@example.com"} { + result, err := db.UpsertAssets([]*Asset{{Host: host}}, "") + if err != nil || result.Created != 1 { + t.Fatalf("opaque asset address %q was not accepted: result=%#v err=%v", host, result, err) + } + } +} + +func TestAssetValidationRejectsOversizedTags(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "asset-tag-validation.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + _, err = db.UpsertAssets([]*Asset{{Domain: "example.com", Tags: []string{strings.Repeat("x", 65)}}}, "") + if err == nil || !strings.Contains(err.Error(), "标签") { + t.Fatalf("expected tag validation error, got %v", err) + } +} + +func TestFofaAssetIgnoresInvalidOptionalStructuredFields(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "fofa-asset-validation.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + asset := &Asset{ + Host: "https://203.0.113.59:8443", + IP: "203.0.113.59", + Domain: "provider_specific_invalid_domain_59", + Port: 8443, + Protocol: "https", + Source: "fofa", + } + result, err := db.UpsertAssets([]*Asset{asset}, "") + if err != nil || result.Created != 1 { + t.Fatalf("FOFA asset with dirty optional domain was not created: result=%#v err=%v", result, err) + } + if asset.Domain != "" || asset.IP != "203.0.113.59" { + t.Fatalf("FOFA structured fields were not sanitized: %#v", asset) + } +} + +func TestAssetUpsertDeduplicatesAndUpdates(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "assets.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + first := &Asset{Host: "https://example.com", Domain: "Example.COM", Port: 443, Protocol: "HTTPS", Title: "Old", Source: "fofa"} + result, err := db.UpsertAssets([]*Asset{first}, "user-a") + if err != nil || result.Created != 1 || result.Updated != 0 { + t.Fatalf("first upsert = %#v, %v", result, err) + } + second := &Asset{Domain: "example.com", Port: 443, Protocol: "https", Title: "New", Server: "nginx", Source: "fofa"} + result, err = db.UpsertAssets([]*Asset{second}, "user-a") + if err != nil || result.Created != 0 || result.Updated != 1 { + t.Fatalf("second upsert = %#v, %v", result, err) + } + assets, total, err := db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll}) + if err != nil || total != 1 || len(assets) != 1 { + t.Fatalf("list assets total=%d len=%d err=%v", total, len(assets), err) + } + if assets[0].Title != "New" || assets[0].Server != "nginx" || assets[0].Protocol != "https" { + t.Fatalf("asset not refreshed: %#v", assets[0]) + } + stats, err := db.GetAssetStats(RBACListAccess{Scope: RBACScopeAll}) + if err != nil || stats["total"] != 1 { + t.Fatalf("stats=%#v err=%v", stats, err) + } + coverage, ok := stats["coverage"].(map[string]interface{}) + if !ok || coverage["never_scanned"] != 1 || coverage["rate"] != 0 { + t.Fatalf("coverage=%#v", stats["coverage"]) + } + assetTrend, ok := stats["asset_trend"].([]map[string]interface{}) + if !ok || len(assetTrend) != 30 { + t.Fatalf("asset trend=%#v", stats["asset_trend"]) + } + riskTrend, ok := stats["risk_trend"].([]map[string]interface{}) + if !ok || len(riskTrend) != 30 { + t.Fatalf("risk trend=%#v", stats["risk_trend"]) + } +} + +func TestAssetAccessFiltersOwners(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "assets-access.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + now := time.Now() + if _, err := db.Exec(`INSERT INTO rbac_users (id,username,display_name,password_hash,enabled,is_builtin,created_at,updated_at) VALUES ('user-a','user-a','User A','hash',1,0,?,?)`, now, now); err != nil { + t.Fatal(err) + } + if _, err := db.UpsertAssets([]*Asset{{IP: "10.0.0.1", Port: 80, Protocol: "http"}}, "user-a"); err != nil { + t.Fatal(err) + } + _, total, err := db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{UserID: "user-b", Scope: RBACScopeAssigned}) + if err != nil || total != 0 { + t.Fatalf("unexpected cross-user assets: total=%d err=%v", total, err) + } + _, total, err = db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{UserID: "user-a", Scope: RBACScopeOwn}) + if err != nil || total != 1 { + t.Fatalf("owner cannot list asset: total=%d err=%v", total, err) + } + assets, _, err := db.ListAssets(1, 0, AssetListFilter{}, RBACListAccess{UserID: "user-a", Scope: RBACScopeAssigned}) + if err != nil || len(assets) != 1 || !db.UserCanAccessResource("user-a", RBACScopeAssigned, "asset", assets[0].ID) { + t.Fatalf("creator assignment missing: assets=%d err=%v", len(assets), err) + } + options, err := db.ListAssignableRBACResources("asset", "10.0.0.1", 10) + if err != nil || len(options) != 1 { + t.Fatalf("asset resource picker: options=%#v err=%v", options, err) + } + project, err := db.CreateProject(&Project{Name: "Alpha", Status: "active"}) + if err != nil { + t.Fatal(err) + } + if err := db.SetResourceOwner("project", project.ID, "user-b"); err != nil { + t.Fatal(err) + } + asset := assets[0] + asset.ProjectID = project.ID + if err := db.UpdateAsset(asset.ID, asset, RBACListAccess{Scope: RBACScopeAll}); err != nil { + t.Fatal(err) + } + projectAssets, total, err := db.ListAssets(20, 0, AssetListFilter{ProjectID: project.ID}, RBACListAccess{UserID: "user-b", Scope: RBACScopeOwn}) + if err != nil || total != 1 || len(projectAssets) != 1 || projectAssets[0].ProjectName != "Alpha" { + t.Fatalf("project-bound asset access failed: total=%d assets=%#v err=%v", total, projectAssets, err) + } + if !db.UserCanAccessResource("user-b", RBACScopeOwn, "asset", asset.ID) { + t.Fatal("project owner cannot access bound asset") + } +} + +func TestUpdateAssetsProjectIsAtomicAndScoped(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "asset-batch-project.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + project, err := db.CreateProject(&Project{Name: "Batch Project", Status: "active"}) + if err != nil { + t.Fatal(err) + } + if _, err := db.UpsertAssets([]*Asset{ + {IP: "192.0.2.1", Port: 80, Protocol: "http"}, + {IP: "192.0.2.2", Port: 443, Protocol: "https"}, + }, "owner-a"); err != nil { + t.Fatal(err) + } + assets, _, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll}) + if err != nil || len(assets) != 2 { + t.Fatalf("list assets: len=%d err=%v", len(assets), err) + } + ids := []string{assets[0].ID, assets[1].ID} + updated, err := db.UpdateAssetsProject(ids, project.ID, RBACListAccess{UserID: "owner-a", Scope: RBACScopeOwn}) + if err != nil || updated != 2 { + t.Fatalf("batch bind: updated=%d err=%v", updated, err) + } + for _, id := range ids { + asset, err := db.GetAsset(id, RBACListAccess{Scope: RBACScopeAll}) + if err != nil || asset.ProjectID != project.ID { + t.Fatalf("asset %s was not bound: asset=%#v err=%v", id, asset, err) + } + } + + if _, err := db.UpdateAssetsProject([]string{ids[0], "missing"}, "", RBACListAccess{Scope: RBACScopeAll}); err == nil { + t.Fatal("partial batch update unexpectedly succeeded") + } + asset, err := db.GetAsset(ids[0], RBACListAccess{Scope: RBACScopeAll}) + if err != nil || asset.ProjectID != project.ID { + t.Fatalf("failed batch changed an asset: asset=%#v err=%v", asset, err) + } + + updated, err = db.UpdateAssetsProject(ids, "", RBACListAccess{Scope: RBACScopeAll}) + if err != nil || updated != 2 { + t.Fatalf("batch unbind: updated=%d err=%v", updated, err) + } +} + +func TestAssetScanLinkReturnsTimeAndRelatedVulnerabilities(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "asset-scan.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + if _, err := db.UpsertAssets([]*Asset{{IP: "192.0.2.10", Port: 443, Protocol: "https"}}, ""); err != nil { + t.Fatal(err) + } + assets, _, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll}) + if err != nil || len(assets) != 1 { + t.Fatalf("list assets: len=%d err=%v", len(assets), err) + } + conv, err := db.CreateConversation("asset scan", ConversationCreateMeta{}) + if err != nil { + t.Fatal(err) + } + if err := db.MarkAssetScanned(assets[0].ID, conv.ID, "", "", RBACListAccess{Scope: RBACScopeAll}); err != nil { + t.Fatal(err) + } + if _, err := db.CreateVulnerability(&Vulnerability{ConversationID: conv.ID, Title: "finding", Severity: "high", Target: "192.0.2.10"}); err != nil { + t.Fatal(err) + } + linked, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll}) + if err != nil { + t.Fatal(err) + } + if linked.LastScanAt == nil || linked.LastScanConversationID != conv.ID || linked.VulnerabilityCount != 1 || linked.RiskLevel != "high" { + t.Fatalf("unexpected scan metadata: %#v", linked) + } + if _, err := db.Exec(`UPDATE vulnerabilities SET status='fixed' WHERE conversation_id=?`, conv.ID); err != nil { + t.Fatal(err) + } + resolved, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll}) + if err != nil { + t.Fatal(err) + } + if resolved.VulnerabilityCount != 1 || resolved.RiskLevel != "normal" { + t.Fatalf("resolved finding should remain in history without raising current risk: %#v", resolved) + } +} + +func TestAssetListFlexibleFiltersAndOldestScanPagination(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "asset-query.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + assets := []*Asset{ + {IP: "192.0.2.1", Port: 443, Protocol: "https", Source: "fofa", Tags: []string{"prod"}}, + {IP: "192.0.2.2", Port: 80, Protocol: "http", Source: "manual", Tags: []string{"prod", "legacy"}}, + {Domain: "never.example.com", Port: 443, Protocol: "https", Source: "manual", Tags: []string{"prod"}}, + } + if _, err := db.UpsertAssets(assets, ""); err != nil { + t.Fatal(err) + } + old := time.Now().Add(-90 * 24 * time.Hour).UTC() + recent := time.Now().Add(-24 * time.Hour).UTC() + if _, err := db.Exec(`UPDATE assets SET last_scan_at=? WHERE id=?`, old, assets[0].ID); err != nil { + t.Fatal(err) + } + if _, err := db.Exec(`UPDATE assets SET last_scan_at=? WHERE id=?`, recent, assets[1].ID); err != nil { + t.Fatal(err) + } + + access := RBACListAccess{Scope: RBACScopeAll} + firstPage, total, err := db.ListAssets(2, 0, AssetListFilter{Tag: "prod", SortBy: "last_scan_at", SortOrder: "asc"}, access) + if err != nil || total != 3 || len(firstPage) != 2 { + t.Fatalf("oldest scan page: total=%d len=%d err=%v", total, len(firstPage), err) + } + if firstPage[0].ID != assets[2].ID || firstPage[0].LastScanAt != nil || firstPage[1].ID != assets[0].ID { + t.Fatalf("expected never-scanned then oldest scanned asset, got %#v", firstPage) + } + secondPage, _, err := db.ListAssets(2, 2, AssetListFilter{Tag: "prod", SortBy: "last_scan_at", SortOrder: "asc"}, access) + if err != nil || len(secondPage) != 1 || secondPage[0].ID != assets[1].ID { + t.Fatalf("unexpected second page: %#v err=%v", secondPage, err) + } + + never, total, err := db.ListAssets(20, 0, AssetListFilter{ScanState: "never"}, access) + if err != nil || total != 1 || len(never) != 1 || never[0].ID != assets[2].ID { + t.Fatalf("never-scanned filter: total=%d assets=%#v err=%v", total, never, err) + } + port := 443 + filtered, total, err := db.ListAssets(20, 0, AssetListFilter{Source: "fofa", Port: &port, LastScanBefore: &recent}, access) + if err != nil || total != 1 || len(filtered) != 1 || filtered[0].ID != assets[0].ID { + t.Fatalf("structured filters: total=%d assets=%#v err=%v", total, filtered, err) + } +} diff --git a/internal/database/database.go b/internal/database/database.go index d80c58fc..713a1df2 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -405,6 +405,19 @@ func (db *DB) initTables() error { FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL );` + createAssetsTable := ` + CREATE TABLE IF NOT EXISTS assets ( + id TEXT PRIMARY KEY, + dedup_key TEXT NOT NULL UNIQUE, project_id TEXT, + host TEXT NOT NULL DEFAULT '', ip TEXT NOT NULL DEFAULT '', port INTEGER NOT NULL DEFAULT 0, + domain TEXT NOT NULL DEFAULT '', protocol TEXT NOT NULL DEFAULT '', title TEXT NOT NULL DEFAULT '', + server TEXT NOT NULL DEFAULT '', country TEXT NOT NULL DEFAULT '', province TEXT NOT NULL DEFAULT '', city TEXT NOT NULL DEFAULT '', + source TEXT NOT NULL DEFAULT 'manual', source_query TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'active', + tags_json TEXT NOT NULL DEFAULT '[]', first_seen_at DATETIME NOT NULL, last_seen_at DATETIME NOT NULL, + created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, owner_user_id TEXT, + FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL + );` + createVulnerabilityAlertSubscriptionsTable := ` CREATE TABLE IF NOT EXISTS vulnerability_alert_subscriptions ( user_id TEXT PRIMARY KEY, @@ -715,6 +728,13 @@ func (db *DB) initTables() error { CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity); CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status); CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at); + CREATE INDEX IF NOT EXISTS idx_assets_last_seen ON assets(last_seen_at); + CREATE INDEX IF NOT EXISTS idx_assets_last_scan ON assets(last_scan_at); + CREATE INDEX IF NOT EXISTS idx_assets_ip ON assets(ip); + CREATE INDEX IF NOT EXISTS idx_assets_domain ON assets(domain); + CREATE INDEX IF NOT EXISTS idx_assets_status ON assets(status); + CREATE INDEX IF NOT EXISTS idx_assets_owner ON assets(owner_user_id); + CREATE INDEX IF NOT EXISTS idx_assets_project ON assets(project_id); CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status); CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at); CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id); @@ -823,6 +843,12 @@ func (db *DB) initTables() error { if _, err := db.Exec(createVulnerabilitiesTable); err != nil { return fmt.Errorf("创建vulnerabilities表失败: %w", err) } + if _, err := db.Exec(createAssetsTable); err != nil { + return fmt.Errorf("创建assets表失败: %w", err) + } + if err := db.migrateAssetsTable(); err != nil { + return fmt.Errorf("迁移assets表失败: %w", err) + } if _, err := db.Exec(createBatchTaskQueuesTable); err != nil { return fmt.Errorf("创建batch_task_queues表失败: %w", err) @@ -950,6 +976,32 @@ func (db *DB) migrateRobotUserSessionsTable() error { return nil } +// migrateAssetsTable keeps databases created by the first asset-management release compatible. +func (db *DB) migrateAssetsTable() error { + columns := []struct { + name string + ddl string + }{ + {"project_id", "ALTER TABLE assets ADD COLUMN project_id TEXT"}, + {"last_scan_at", "ALTER TABLE assets ADD COLUMN last_scan_at DATETIME"}, + {"last_scan_conversation_id", "ALTER TABLE assets ADD COLUMN last_scan_conversation_id TEXT NOT NULL DEFAULT ''"}, + {"last_scan_queue_id", "ALTER TABLE assets ADD COLUMN last_scan_queue_id TEXT NOT NULL DEFAULT ''"}, + {"last_scan_task_id", "ALTER TABLE assets ADD COLUMN last_scan_task_id TEXT NOT NULL DEFAULT ''"}, + } + for _, column := range columns { + var count int + if err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('assets') WHERE name=?", column.name).Scan(&count); err != nil { + return err + } + if count == 0 { + if _, err := db.Exec(column.ddl); err != nil { + return err + } + } + } + return nil +} + // migrateMessagesTable 迁移 messages 表，补充 updated_at 字段。 // 语义：updated_at 表示该条消息最后一次被写入/更新的时间（例如助手占位消息在任务结束时更新正文）。 func (db *DB) migrateMessagesTable() error { diff --git a/internal/database/project.go b/internal/database/project.go index d163aee8..75be4d69 100644 --- a/internal/database/project.go +++ b/internal/database/project.go @@ -268,6 +268,9 @@ func (db *DB) DeleteProject(id string) error { if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil { return fmt.Errorf("解除漏洞项目关联失败: %w", err) } + if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil { + return fmt.Errorf("解除资产项目关联失败: %w", err) + } _, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id) if err != nil { return fmt.Errorf("删除项目失败: %w", err) diff --git a/internal/database/rbac.go b/internal/database/rbac.go index aa15b5e6..28e86502 100644 --- a/internal/database/rbac.go +++ b/internal/database/rbac.go @@ -27,6 +27,7 @@ var rbacAssignableResourceTables = map[string]string{ "project": "projects", "conversation": "conversations", "vulnerability": "vulnerabilities", + "asset": "assets", "webshell": "webshell_connections", "batch_task": "batch_task_queues", "c2_listener": "c2_listeners", @@ -528,6 +529,9 @@ func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID stri if resourceType == "vulnerability" { return db.userCanAccessVulnerabilityViaParent(userID, scope, resourceID) } + if resourceType == "asset" { + return db.userCanAccessAssetViaParent(userID, scope, resourceID) + } if resourceType == "conversation" { return db.userCanAccessConversationViaParent(userID, scope, resourceID) } @@ -537,6 +541,15 @@ func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID stri return false } +func (db *DB) userCanAccessAssetViaParent(userID, scope, assetID string) bool { + var projectID sql.NullString + if err := db.QueryRow(`SELECT project_id FROM assets WHERE id = ?`, assetID).Scan(&projectID); err != nil { + return false + } + return projectID.Valid && strings.TrimSpace(projectID.String) != "" && + db.UserCanAccessResource(userID, scope, "project", strings.TrimSpace(projectID.String)) +} + func (db *DB) userCanAccessConversationViaParent(userID, scope, conversationID string) bool { var projectID sql.NullString if err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, conversationID).Scan(&projectID); err != nil { @@ -623,6 +636,8 @@ func (db *DB) userOwnsResource(userID, resourceType, resourceID string) bool { table = "conversations" case "vulnerability": table = "vulnerabilities" + case "asset": + table = "assets" case "webshell": table = "webshell_connections" case "batch_task": @@ -650,6 +665,8 @@ func (db *DB) SetResourceOwner(resourceType, resourceID, userID string) error { table = "conversations" case "vulnerability": table = "vulnerabilities" + case "asset": + table = "assets" case "webshell": table = "webshell_connections" case "batch_task": @@ -672,6 +689,8 @@ func (db *DB) GetResourceOwner(resourceType, resourceID string) string { table = "conversations" case "vulnerability": table = "vulnerabilities" + case "asset": + table = "assets" case "webshell": table = "webshell_connections" case "batch_task": @@ -733,6 +752,10 @@ func (db *DB) ListAssignableRBACResourcesPage(resourceType, search string, limit query = `SELECT id, title, severity FROM vulnerabilities WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\' ORDER BY updated_at DESC LIMIT ? OFFSET ?` + case "asset": + query = `SELECT id, COALESCE(NULLIF(host,''),NULLIF(domain,''),NULLIF(ip,''),id), protocol || CASE WHEN port>0 THEN ':' || port ELSE '' END FROM assets + WHERE LOWER(host) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\' OR LOWER(ip) LIKE ? ESCAPE '\' + ORDER BY updated_at DESC LIMIT ? OFFSET ?` case "webshell": query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\' @@ -747,7 +770,12 @@ func (db *DB) ListAssignableRBACResourcesPage(resourceType, search string, limit ORDER BY created_at DESC LIMIT ? OFFSET ?` } - rows, err := db.Query(query, pattern, pattern, limit, offset) + queryArgs := []interface{}{pattern, pattern} + if resourceType == "asset" { + queryArgs = append(queryArgs, pattern) + } + queryArgs = append(queryArgs, limit, offset) + rows, err := db.Query(query, queryArgs...) if err != nil { return nil, err } @@ -781,6 +809,8 @@ func (db *DB) CountAssignableRBACResources(resourceType, search string) (int, er query = `SELECT COUNT(*) FROM conversations WHERE LOWER(COALESCE(NULLIF(TRIM(title), ''), id)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'` case "vulnerability": query = `SELECT COUNT(*) FROM vulnerabilities WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'` + case "asset": + query = `SELECT COUNT(*) FROM assets WHERE LOWER(host) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\' OR LOWER(ip) LIKE ? ESCAPE '\'` case "webshell": query = `SELECT COUNT(*) FROM webshell_connections WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'` case "batch_task": @@ -789,7 +819,11 @@ func (db *DB) CountAssignableRBACResources(resourceType, search string) (int, er query = `SELECT COUNT(*) FROM c2_listeners WHERE LOWER(name) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'` } var total int - if err := db.QueryRow(query, pattern, pattern).Scan(&total); err != nil { + queryArgs := []interface{}{pattern, pattern} + if resourceType == "asset" { + queryArgs = append(queryArgs, pattern) + } + if err := db.QueryRow(query, queryArgs...).Scan(&total); err != nil { return 0, err } return total, nil @@ -869,6 +903,8 @@ func (db *DB) lookupRBACResourceOptionsByIDs(resourceType string, ids []string) query = `SELECT id, COALESCE(NULLIF(TRIM(title), ''), '未命名对话'), COALESCE(project_id, '') FROM conversations WHERE id IN (` + placeholders + `)` case "vulnerability": query = `SELECT id, title, severity FROM vulnerabilities WHERE id IN (` + placeholders + `)` + case "asset": + query = `SELECT id, COALESCE(NULLIF(host,''),NULLIF(domain,''),NULLIF(ip,''),id), protocol || CASE WHEN port>0 THEN ':' || port ELSE '' END FROM assets WHERE id IN (` + placeholders + `)` case "webshell": query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections WHERE id IN (` + placeholders + `)` case "batch_task": @@ -1039,6 +1075,7 @@ func (db *DB) AssignResourcesToUserAuto(userID string, resourceIDs []string) (in typeTablePairs := []struct{ resourceType, table string }{ {"project", "projects"}, {"conversation", "conversations"}, {"vulnerability", "vulnerabilities"}, {"webshell", "webshell_connections"}, + {"asset", "assets"}, {"batch_task", "batch_task_queues"}, {"c2_listener", "c2_listeners"}, } detected := make(map[string]string, len(uniqueIDs))