mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-18 18:07:32 +02:00
Add files via upload
This commit is contained in:
@@ -158,6 +158,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
|||||||
|
|
||||||
// 注册漏洞记录工具
|
// 注册漏洞记录工具
|
||||||
registerVulnerabilityTools(mcpServer, db, log.Logger)
|
registerVulnerabilityTools(mcpServer, db, log.Logger)
|
||||||
|
registerAssetTools(mcpServer, db, log.Logger)
|
||||||
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
|
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
|
||||||
registerVisionTools(mcpServer, 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)
|
authHandler.SetAudit(auditSvc)
|
||||||
attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger)
|
attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger)
|
||||||
vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger)
|
vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger)
|
||||||
|
assetHandler := handler.NewAssetHandler(db, log.Logger)
|
||||||
projectHandler := handler.NewProjectHandler(db, log.Logger)
|
projectHandler := handler.NewProjectHandler(db, log.Logger)
|
||||||
rbacHandler := handler.NewRBACHandler(db, log.Logger)
|
rbacHandler := handler.NewRBACHandler(db, log.Logger)
|
||||||
rbacHandler.SetAudit(auditSvc)
|
rbacHandler.SetAudit(auditSvc)
|
||||||
@@ -465,6 +467,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
|||||||
// 设置漏洞工具注册器(内置工具,必须设置)
|
// 设置漏洞工具注册器(内置工具,必须设置)
|
||||||
vulnerabilityRegistrar := func() error {
|
vulnerabilityRegistrar := func() error {
|
||||||
registerVulnerabilityTools(mcpServer, db, log.Logger)
|
registerVulnerabilityTools(mcpServer, db, log.Logger)
|
||||||
|
registerAssetTools(mcpServer, db, log.Logger)
|
||||||
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
|
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
|
||||||
registerVisionTools(mcpServer, cfg, log.Logger)
|
registerVisionTools(mcpServer, cfg, log.Logger)
|
||||||
return nil
|
return nil
|
||||||
@@ -554,6 +557,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
|||||||
attackChainHandler,
|
attackChainHandler,
|
||||||
app, // 传递 App 实例以便动态获取 knowledgeHandler
|
app, // 传递 App 实例以便动态获取 knowledgeHandler
|
||||||
vulnerabilityHandler,
|
vulnerabilityHandler,
|
||||||
|
assetHandler,
|
||||||
projectHandler,
|
projectHandler,
|
||||||
workflowHandler,
|
workflowHandler,
|
||||||
webshellHandler,
|
webshellHandler,
|
||||||
@@ -856,6 +860,7 @@ func setupRoutes(
|
|||||||
attackChainHandler *handler.AttackChainHandler,
|
attackChainHandler *handler.AttackChainHandler,
|
||||||
app *App, // 传递 App 实例以便动态获取 knowledgeHandler
|
app *App, // 传递 App 实例以便动态获取 knowledgeHandler
|
||||||
vulnerabilityHandler *handler.VulnerabilityHandler,
|
vulnerabilityHandler *handler.VulnerabilityHandler,
|
||||||
|
assetHandler *handler.AssetHandler,
|
||||||
projectHandler *handler.ProjectHandler,
|
projectHandler *handler.ProjectHandler,
|
||||||
workflowHandler *handler.WorkflowHandler,
|
workflowHandler *handler.WorkflowHandler,
|
||||||
webshellHandler *handler.WebShellHandler,
|
webshellHandler *handler.WebShellHandler,
|
||||||
@@ -976,6 +981,15 @@ func setupRoutes(
|
|||||||
// 信息收集 - 自然语言解析为 FOFA 语法(需人工确认后再查询)
|
// 信息收集 - 自然语言解析为 FOFA 语法(需人工确认后再查询)
|
||||||
protected.POST("/fofa/parse", fofaHandler.ParseNaturalLanguage)
|
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.POST("/batch-tasks", agentHandler.CreateBatchQueue)
|
||||||
protected.GET("/batch-tasks", agentHandler.ListBatchQueues)
|
protected.GET("/batch-tasks", agentHandler.ListBatchQueues)
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -71,6 +71,35 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
|
|||||||
return nil
|
return nil
|
||||||
case builtin.ToolGetVulnerability:
|
case builtin.ToolGetVulnerability:
|
||||||
return resource("vulnerability:read", "vulnerability", "id")
|
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:
|
case builtin.ToolUpsertProjectFact, builtin.ToolDeprecateProjectFact, builtin.ToolRestoreProjectFact:
|
||||||
return authorizeProjectTool(ctx, principal, db, "project:write")
|
return authorizeProjectTool(ctx, principal, db, "project:write")
|
||||||
case builtin.ToolGetProjectFact, builtin.ToolListProjectFacts, builtin.ToolSearchProjectFacts:
|
case builtin.ToolGetProjectFact, builtin.ToolListProjectFacts, builtin.ToolSearchProjectFacts:
|
||||||
|
|||||||
@@ -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) {
|
func TestExternalMCPRequiresDedicatedPermission(t *testing.T) {
|
||||||
authorize := externalMCPToolAuthorizer()
|
authorize := externalMCPToolAuthorizer()
|
||||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
|
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
|
||||||
|
|||||||
@@ -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)<datetime('now','-30 days') THEN 1 ELSE 0 END),0),
|
||||||
|
COALESCE(SUM(CASE WHEN status='active' THEN 1 ELSE 0 END),0)
|
||||||
|
FROM assets`+where, args...)
|
||||||
|
var scanned, scanned7, scanned30, neverScanned, stale, active int
|
||||||
|
if err := coverageRow.Scan(&scanned, &scanned7, &scanned30, &neverScanned, &stale, &active); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
coverage["scanned"], coverage["scanned_7d"], coverage["scanned_30d"] = scanned, scanned7, scanned30
|
||||||
|
coverage["never_scanned"], coverage["stale"], coverage["active"] = neverScanned, stale, active
|
||||||
|
if total > 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()
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -405,6 +405,19 @@ func (db *DB) initTables() error {
|
|||||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL
|
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 := `
|
createVulnerabilityAlertSubscriptionsTable := `
|
||||||
CREATE TABLE IF NOT EXISTS vulnerability_alert_subscriptions (
|
CREATE TABLE IF NOT EXISTS vulnerability_alert_subscriptions (
|
||||||
user_id TEXT PRIMARY KEY,
|
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_severity ON vulnerabilities(severity);
|
||||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status);
|
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_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_status ON projects(status);
|
||||||
CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at);
|
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);
|
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 {
|
if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
|
||||||
return fmt.Errorf("创建vulnerabilities表失败: %w", err)
|
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 {
|
if _, err := db.Exec(createBatchTaskQueuesTable); err != nil {
|
||||||
return fmt.Errorf("创建batch_task_queues表失败: %w", err)
|
return fmt.Errorf("创建batch_task_queues表失败: %w", err)
|
||||||
@@ -950,6 +976,32 @@ func (db *DB) migrateRobotUserSessionsTable() error {
|
|||||||
return nil
|
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 字段。
|
// migrateMessagesTable 迁移 messages 表,补充 updated_at 字段。
|
||||||
// 语义:updated_at 表示该条消息最后一次被写入/更新的时间(例如助手占位消息在任务结束时更新正文)。
|
// 语义:updated_at 表示该条消息最后一次被写入/更新的时间(例如助手占位消息在任务结束时更新正文)。
|
||||||
func (db *DB) migrateMessagesTable() error {
|
func (db *DB) migrateMessagesTable() error {
|
||||||
|
|||||||
@@ -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 {
|
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
|
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)
|
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("删除项目失败: %w", err)
|
return fmt.Errorf("删除项目失败: %w", err)
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ var rbacAssignableResourceTables = map[string]string{
|
|||||||
"project": "projects",
|
"project": "projects",
|
||||||
"conversation": "conversations",
|
"conversation": "conversations",
|
||||||
"vulnerability": "vulnerabilities",
|
"vulnerability": "vulnerabilities",
|
||||||
|
"asset": "assets",
|
||||||
"webshell": "webshell_connections",
|
"webshell": "webshell_connections",
|
||||||
"batch_task": "batch_task_queues",
|
"batch_task": "batch_task_queues",
|
||||||
"c2_listener": "c2_listeners",
|
"c2_listener": "c2_listeners",
|
||||||
@@ -528,6 +529,9 @@ func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID stri
|
|||||||
if resourceType == "vulnerability" {
|
if resourceType == "vulnerability" {
|
||||||
return db.userCanAccessVulnerabilityViaParent(userID, scope, resourceID)
|
return db.userCanAccessVulnerabilityViaParent(userID, scope, resourceID)
|
||||||
}
|
}
|
||||||
|
if resourceType == "asset" {
|
||||||
|
return db.userCanAccessAssetViaParent(userID, scope, resourceID)
|
||||||
|
}
|
||||||
if resourceType == "conversation" {
|
if resourceType == "conversation" {
|
||||||
return db.userCanAccessConversationViaParent(userID, scope, resourceID)
|
return db.userCanAccessConversationViaParent(userID, scope, resourceID)
|
||||||
}
|
}
|
||||||
@@ -537,6 +541,15 @@ func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID stri
|
|||||||
return false
|
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 {
|
func (db *DB) userCanAccessConversationViaParent(userID, scope, conversationID string) bool {
|
||||||
var projectID sql.NullString
|
var projectID sql.NullString
|
||||||
if err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, conversationID).Scan(&projectID); err != nil {
|
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"
|
table = "conversations"
|
||||||
case "vulnerability":
|
case "vulnerability":
|
||||||
table = "vulnerabilities"
|
table = "vulnerabilities"
|
||||||
|
case "asset":
|
||||||
|
table = "assets"
|
||||||
case "webshell":
|
case "webshell":
|
||||||
table = "webshell_connections"
|
table = "webshell_connections"
|
||||||
case "batch_task":
|
case "batch_task":
|
||||||
@@ -650,6 +665,8 @@ func (db *DB) SetResourceOwner(resourceType, resourceID, userID string) error {
|
|||||||
table = "conversations"
|
table = "conversations"
|
||||||
case "vulnerability":
|
case "vulnerability":
|
||||||
table = "vulnerabilities"
|
table = "vulnerabilities"
|
||||||
|
case "asset":
|
||||||
|
table = "assets"
|
||||||
case "webshell":
|
case "webshell":
|
||||||
table = "webshell_connections"
|
table = "webshell_connections"
|
||||||
case "batch_task":
|
case "batch_task":
|
||||||
@@ -672,6 +689,8 @@ func (db *DB) GetResourceOwner(resourceType, resourceID string) string {
|
|||||||
table = "conversations"
|
table = "conversations"
|
||||||
case "vulnerability":
|
case "vulnerability":
|
||||||
table = "vulnerabilities"
|
table = "vulnerabilities"
|
||||||
|
case "asset":
|
||||||
|
table = "assets"
|
||||||
case "webshell":
|
case "webshell":
|
||||||
table = "webshell_connections"
|
table = "webshell_connections"
|
||||||
case "batch_task":
|
case "batch_task":
|
||||||
@@ -733,6 +752,10 @@ func (db *DB) ListAssignableRBACResourcesPage(resourceType, search string, limit
|
|||||||
query = `SELECT id, title, severity FROM vulnerabilities
|
query = `SELECT id, title, severity FROM vulnerabilities
|
||||||
WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'
|
WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'
|
||||||
ORDER BY updated_at DESC LIMIT ? OFFSET ?`
|
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":
|
case "webshell":
|
||||||
query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections
|
query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections
|
||||||
WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'
|
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 ?`
|
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 {
|
if err != nil {
|
||||||
return nil, err
|
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 '\'`
|
query = `SELECT COUNT(*) FROM conversations WHERE LOWER(COALESCE(NULLIF(TRIM(title), ''), id)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
|
||||||
case "vulnerability":
|
case "vulnerability":
|
||||||
query = `SELECT COUNT(*) FROM vulnerabilities WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
|
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":
|
case "webshell":
|
||||||
query = `SELECT COUNT(*) FROM webshell_connections WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
|
query = `SELECT COUNT(*) FROM webshell_connections WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
|
||||||
case "batch_task":
|
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 '\'`
|
query = `SELECT COUNT(*) FROM c2_listeners WHERE LOWER(name) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
|
||||||
}
|
}
|
||||||
var total int
|
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 0, err
|
||||||
}
|
}
|
||||||
return total, nil
|
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 + `)`
|
query = `SELECT id, COALESCE(NULLIF(TRIM(title), ''), '未命名对话'), COALESCE(project_id, '') FROM conversations WHERE id IN (` + placeholders + `)`
|
||||||
case "vulnerability":
|
case "vulnerability":
|
||||||
query = `SELECT id, title, severity FROM vulnerabilities WHERE id IN (` + placeholders + `)`
|
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":
|
case "webshell":
|
||||||
query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections WHERE id IN (` + placeholders + `)`
|
query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections WHERE id IN (` + placeholders + `)`
|
||||||
case "batch_task":
|
case "batch_task":
|
||||||
@@ -1039,6 +1075,7 @@ func (db *DB) AssignResourcesToUserAuto(userID string, resourceIDs []string) (in
|
|||||||
typeTablePairs := []struct{ resourceType, table string }{
|
typeTablePairs := []struct{ resourceType, table string }{
|
||||||
{"project", "projects"}, {"conversation", "conversations"},
|
{"project", "projects"}, {"conversation", "conversations"},
|
||||||
{"vulnerability", "vulnerabilities"}, {"webshell", "webshell_connections"},
|
{"vulnerability", "vulnerabilities"}, {"webshell", "webshell_connections"},
|
||||||
|
{"asset", "assets"},
|
||||||
{"batch_task", "batch_task_queues"}, {"c2_listener", "c2_listeners"},
|
{"batch_task", "batch_task_queues"}, {"c2_listener", "c2_listeners"},
|
||||||
}
|
}
|
||||||
detected := make(map[string]string, len(uniqueIDs))
|
detected := make(map[string]string, len(uniqueIDs))
|
||||||
|
|||||||
Reference in New Issue
Block a user