CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-07-18 18:07:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-16 19:37:01 +08:00
committed by GitHub
parent 15cc08a95b
commit d03cc47ac3
10 changed files with 1979 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/app/app.go
+14
View File
@@ -158,6 +158,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 注册漏洞记录工具
registerVulnerabilityTools(mcpServer, db, log.Logger)
registerAssetTools(mcpServer, db, log.Logger)
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
registerVisionTools(mcpServer, cfg, log.Logger)
@@ -380,6 +381,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
authHandler.SetAudit(auditSvc)
attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger)
vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger)
assetHandler := handler.NewAssetHandler(db, log.Logger)
projectHandler := handler.NewProjectHandler(db, log.Logger)
rbacHandler := handler.NewRBACHandler(db, log.Logger)
rbacHandler.SetAudit(auditSvc)
@@ -465,6 +467,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 设置漏洞工具注册器（内置工具，必须设置）
vulnerabilityRegistrar := func() error {
registerVulnerabilityTools(mcpServer, db, log.Logger)
registerAssetTools(mcpServer, db, log.Logger)
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
registerVisionTools(mcpServer, cfg, log.Logger)
return nil
@@ -554,6 +557,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
attackChainHandler,
app, // 传递 App 实例以便动态获取 knowledgeHandler
vulnerabilityHandler,
assetHandler,
projectHandler,
workflowHandler,
webshellHandler,
@@ -856,6 +860,7 @@ func setupRoutes(
attackChainHandler *handler.AttackChainHandler,
app *App, // 传递 App 实例以便动态获取 knowledgeHandler
vulnerabilityHandler *handler.VulnerabilityHandler,
assetHandler *handler.AssetHandler,
projectHandler *handler.ProjectHandler,
workflowHandler *handler.WorkflowHandler,
webshellHandler *handler.WebShellHandler,
@@ -976,6 +981,15 @@ func setupRoutes(
// 信息收集 - 自然语言解析为 FOFA 语法（需人工确认后再查询）
protected.POST("/fofa/parse", fofaHandler.ParseNaturalLanguage)
// 资产管理
protected.GET("/assets", assetHandler.List)
protected.GET("/assets/stats", assetHandler.Stats)
protected.POST("/assets/import", assetHandler.Import)
protected.POST("/assets/scan-links", assetHandler.RecordScans)
protected.PUT("/assets/project-binding", assetHandler.UpdateProjectBinding)
protected.PUT("/assets/:id", assetHandler.Update)
protected.DELETE("/assets/:id", assetHandler.Delete)
// 批量任务管理
protected.POST("/batch-tasks", agentHandler.CreateBatchQueue)
protected.GET("/batch-tasks", agentHandler.ListBatchQueues)
internal/app/asset_tools.go
+459
View File
@@ -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
}
internal/app/asset_tools_test.go
+185
View File
@@ -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)
}
}
internal/app/mcp_authorization.go
+29
View File
@@ -71,6 +71,35 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
return nil
case builtin.ToolGetVulnerability:
return resource("vulnerability:read", "vulnerability", "id")
case builtin.ToolQueryAssets:
return require("asset:read")
case builtin.ToolGetAsset:
return resource("asset:read", "asset", "id")
case builtin.ToolCreateAsset:
if err := require("asset:write"); err != nil {
return err
}
if projectID := mcpAuthorizationString(args, "project_id"); projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("asset:write"), "project", projectID)) {
return fmt.Errorf("no access to project %s", projectID)
}
return nil
case builtin.ToolUpdateAsset, builtin.ToolCompleteAssetScan:
if err := resource("asset:write", "asset", "id"); err != nil {
return err
}
if toolName == builtin.ToolCompleteAssetScan {
conversationID := mcpAuthorizationConversationID(ctx)
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("asset:write"), "conversation", conversationID) {
return fmt.Errorf("no access to conversation %s", conversationID)
}
return nil
}
if projectID := mcpAuthorizationString(args, "project_id"); projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("asset:write"), "project", projectID)) {
return fmt.Errorf("no access to project %s", projectID)
}
return nil
case builtin.ToolDeleteAsset:
return resource("asset:delete", "asset", "id")
case builtin.ToolUpsertProjectFact, builtin.ToolDeprecateProjectFact, builtin.ToolRestoreProjectFact:
return authorizeProjectTool(ctx, principal, db, "project:write")
case builtin.ToolGetProjectFact, builtin.ToolListProjectFacts, builtin.ToolSearchProjectFacts:
internal/app/mcp_authorization_test.go
+39
View File
@@ -72,6 +72,45 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
}
}
func TestMCPAssetToolAuthorizationUsesAssetPermissionsAndScope(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-asset-authz.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
user, err := db.CreateRBACUser("asset-user", "Asset User", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
owned := &database.Asset{IP: "192.0.2.10", Port: 443, Protocol: "https"}
hidden := &database.Asset{IP: "192.0.2.20", Port: 443, Protocol: "https"}
if _, err := db.UpsertAssets([]*database.Asset{owned}, user.ID); err != nil {
t.Fatal(err)
}
if _, err := db.UpsertAssets([]*database.Asset{hidden}, ""); err != nil {
t.Fatal(err)
}
permissions := map[string]bool{"asset:read": true, "asset:write": true}
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, permissions))
authorize := mcpToolAuthorizer(db)
if err := authorize(ctx, builtin.ToolQueryAssets, nil); err != nil {
t.Fatalf("asset query denied: %v", err)
}
if err := authorize(ctx, builtin.ToolGetAsset, map[string]interface{}{"id": owned.ID}); err != nil {
t.Fatalf("owned asset denied: %v", err)
}
if err := authorize(ctx, builtin.ToolGetAsset, map[string]interface{}{"id": hidden.ID}); err == nil {
t.Fatal("unassigned asset was readable")
}
if err := authorize(ctx, builtin.ToolUpdateAsset, map[string]interface{}{"id": owned.ID}); err != nil {
t.Fatalf("owned asset update denied: %v", err)
}
if err := authorize(ctx, builtin.ToolDeleteAsset, map[string]interface{}{"id": owned.ID}); err == nil {
t.Fatal("asset delete without asset:delete was allowed")
}
}
func TestExternalMCPRequiresDedicatedPermission(t *testing.T) {
authorize := externalMCPToolAuthorizer()
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
internal/database/asset.go
+843
View File
@@ -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()
}
internal/database/asset_test.go
+316
View File
@@ -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)
}
}
internal/database/database.go
+52
View File
@@ -405,6 +405,19 @@ func (db *DB) initTables() error {
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL
);`
createAssetsTable := `
CREATE TABLE IF NOT EXISTS assets (
id TEXT PRIMARY KEY,
dedup_key TEXT NOT NULL UNIQUE, project_id TEXT,
host TEXT NOT NULL DEFAULT '', ip TEXT NOT NULL DEFAULT '', port INTEGER NOT NULL DEFAULT 0,
domain TEXT NOT NULL DEFAULT '', protocol TEXT NOT NULL DEFAULT '', title TEXT NOT NULL DEFAULT '',
server TEXT NOT NULL DEFAULT '', country TEXT NOT NULL DEFAULT '', province TEXT NOT NULL DEFAULT '', city TEXT NOT NULL DEFAULT '',
source TEXT NOT NULL DEFAULT 'manual', source_query TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'active',
tags_json TEXT NOT NULL DEFAULT '[]', first_seen_at DATETIME NOT NULL, last_seen_at DATETIME NOT NULL,
created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, owner_user_id TEXT,
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL
);`
createVulnerabilityAlertSubscriptionsTable := `
CREATE TABLE IF NOT EXISTS vulnerability_alert_subscriptions (
user_id TEXT PRIMARY KEY,
@@ -715,6 +728,13 @@ func (db *DB) initTables() error {
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity);
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status);
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at);
CREATE INDEX IF NOT EXISTS idx_assets_last_seen ON assets(last_seen_at);
CREATE INDEX IF NOT EXISTS idx_assets_last_scan ON assets(last_scan_at);
CREATE INDEX IF NOT EXISTS idx_assets_ip ON assets(ip);
CREATE INDEX IF NOT EXISTS idx_assets_domain ON assets(domain);
CREATE INDEX IF NOT EXISTS idx_assets_status ON assets(status);
CREATE INDEX IF NOT EXISTS idx_assets_owner ON assets(owner_user_id);
CREATE INDEX IF NOT EXISTS idx_assets_project ON assets(project_id);
CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status);
CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at);
CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id);
@@ -823,6 +843,12 @@ func (db *DB) initTables() error {
if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
return fmt.Errorf("创建vulnerabilities表失败: %w", err)
}
if _, err := db.Exec(createAssetsTable); err != nil {
return fmt.Errorf("创建assets表失败: %w", err)
}
if err := db.migrateAssetsTable(); err != nil {
return fmt.Errorf("迁移assets表失败: %w", err)
}
if _, err := db.Exec(createBatchTaskQueuesTable); err != nil {
return fmt.Errorf("创建batch_task_queues表失败: %w", err)
@@ -950,6 +976,32 @@ func (db *DB) migrateRobotUserSessionsTable() error {
return nil
}
// migrateAssetsTable keeps databases created by the first asset-management release compatible.
func (db *DB) migrateAssetsTable() error {
columns := []struct {
name string
ddl string
}{
{"project_id", "ALTER TABLE assets ADD COLUMN project_id TEXT"},
{"last_scan_at", "ALTER TABLE assets ADD COLUMN last_scan_at DATETIME"},
{"last_scan_conversation_id", "ALTER TABLE assets ADD COLUMN last_scan_conversation_id TEXT NOT NULL DEFAULT ''"},
{"last_scan_queue_id", "ALTER TABLE assets ADD COLUMN last_scan_queue_id TEXT NOT NULL DEFAULT ''"},
{"last_scan_task_id", "ALTER TABLE assets ADD COLUMN last_scan_task_id TEXT NOT NULL DEFAULT ''"},
}
for _, column := range columns {
var count int
if err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('assets') WHERE name=?", column.name).Scan(&count); err != nil {
return err
}
if count == 0 {
if _, err := db.Exec(column.ddl); err != nil {
return err
}
}
}
return nil
}
// migrateMessagesTable 迁移 messages 表，补充 updated_at 字段。
// 语义：updated_at 表示该条消息最后一次被写入/更新的时间（例如助手占位消息在任务结束时更新正文）。
func (db *DB) migrateMessagesTable() error {
internal/database/project.go
+3
View File
@@ -268,6 +268,9 @@ func (db *DB) DeleteProject(id string) error {
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
}
if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil {
return fmt.Errorf("解除资产项目关联失败: %w", err)
}
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
if err != nil {
return fmt.Errorf("删除项目失败: %w", err)
internal/database/rbac.go
+39 -2
View File
@@ -27,6 +27,7 @@ var rbacAssignableResourceTables = map[string]string{
"project": "projects",
"conversation": "conversations",
"vulnerability": "vulnerabilities",
"asset": "assets",
"webshell": "webshell_connections",
"batch_task": "batch_task_queues",
"c2_listener": "c2_listeners",
@@ -528,6 +529,9 @@ func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID stri
if resourceType == "vulnerability" {
return db.userCanAccessVulnerabilityViaParent(userID, scope, resourceID)
}
if resourceType == "asset" {
return db.userCanAccessAssetViaParent(userID, scope, resourceID)
}
if resourceType == "conversation" {
return db.userCanAccessConversationViaParent(userID, scope, resourceID)
}
@@ -537,6 +541,15 @@ func (db *DB) UserCanAccessResource(userID, scope, resourceType, resourceID stri
return false
}
func (db *DB) userCanAccessAssetViaParent(userID, scope, assetID string) bool {
var projectID sql.NullString
if err := db.QueryRow(`SELECT project_id FROM assets WHERE id = ?`, assetID).Scan(&projectID); err != nil {
return false
}
return projectID.Valid && strings.TrimSpace(projectID.String) != "" &&
db.UserCanAccessResource(userID, scope, "project", strings.TrimSpace(projectID.String))
}
func (db *DB) userCanAccessConversationViaParent(userID, scope, conversationID string) bool {
var projectID sql.NullString
if err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, conversationID).Scan(&projectID); err != nil {
@@ -623,6 +636,8 @@ func (db *DB) userOwnsResource(userID, resourceType, resourceID string) bool {
table = "conversations"
case "vulnerability":
table = "vulnerabilities"
case "asset":
table = "assets"
case "webshell":
table = "webshell_connections"
case "batch_task":
@@ -650,6 +665,8 @@ func (db *DB) SetResourceOwner(resourceType, resourceID, userID string) error {
table = "conversations"
case "vulnerability":
table = "vulnerabilities"
case "asset":
table = "assets"
case "webshell":
table = "webshell_connections"
case "batch_task":
@@ -672,6 +689,8 @@ func (db *DB) GetResourceOwner(resourceType, resourceID string) string {
table = "conversations"
case "vulnerability":
table = "vulnerabilities"
case "asset":
table = "assets"
case "webshell":
table = "webshell_connections"
case "batch_task":
@@ -733,6 +752,10 @@ func (db *DB) ListAssignableRBACResourcesPage(resourceType, search string, limit
query = `SELECT id, title, severity FROM vulnerabilities
WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'
ORDER BY updated_at DESC LIMIT ? OFFSET ?`
case "asset":
query = `SELECT id, COALESCE(NULLIF(host,''),NULLIF(domain,''),NULLIF(ip,''),id), protocol || CASE WHEN port>0 THEN ':' || port ELSE '' END FROM assets
WHERE LOWER(host) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\' OR LOWER(ip) LIKE ? ESCAPE '\'
ORDER BY updated_at DESC LIMIT ? OFFSET ?`
case "webshell":
query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections
WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'
@@ -747,7 +770,12 @@ func (db *DB) ListAssignableRBACResourcesPage(resourceType, search string, limit
ORDER BY created_at DESC LIMIT ? OFFSET ?`
}
rows, err := db.Query(query, pattern, pattern, limit, offset)
queryArgs := []interface{}{pattern, pattern}
if resourceType == "asset" {
queryArgs = append(queryArgs, pattern)
}
queryArgs = append(queryArgs, limit, offset)
rows, err := db.Query(query, queryArgs...)
if err != nil {
return nil, err
}
@@ -781,6 +809,8 @@ func (db *DB) CountAssignableRBACResources(resourceType, search string) (int, er
query = `SELECT COUNT(*) FROM conversations WHERE LOWER(COALESCE(NULLIF(TRIM(title), ''), id)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
case "vulnerability":
query = `SELECT COUNT(*) FROM vulnerabilities WHERE LOWER(title) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
case "asset":
query = `SELECT COUNT(*) FROM assets WHERE LOWER(host) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\' OR LOWER(ip) LIKE ? ESCAPE '\'`
case "webshell":
query = `SELECT COUNT(*) FROM webshell_connections WHERE LOWER(COALESCE(NULLIF(remark, ''), url)) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
case "batch_task":
@@ -789,7 +819,11 @@ func (db *DB) CountAssignableRBACResources(resourceType, search string) (int, er
query = `SELECT COUNT(*) FROM c2_listeners WHERE LOWER(name) LIKE ? ESCAPE '\' OR LOWER(id) LIKE ? ESCAPE '\'`
}
var total int
if err := db.QueryRow(query, pattern, pattern).Scan(&total); err != nil {
queryArgs := []interface{}{pattern, pattern}
if resourceType == "asset" {
queryArgs = append(queryArgs, pattern)
}
if err := db.QueryRow(query, queryArgs...).Scan(&total); err != nil {
return 0, err
}
return total, nil
@@ -869,6 +903,8 @@ func (db *DB) lookupRBACResourceOptionsByIDs(resourceType string, ids []string)
query = `SELECT id, COALESCE(NULLIF(TRIM(title), ''), '未命名对话'), COALESCE(project_id, '') FROM conversations WHERE id IN (` + placeholders + `)`
case "vulnerability":
query = `SELECT id, title, severity FROM vulnerabilities WHERE id IN (` + placeholders + `)`
case "asset":
query = `SELECT id, COALESCE(NULLIF(host,''),NULLIF(domain,''),NULLIF(ip,''),id), protocol || CASE WHEN port>0 THEN ':' || port ELSE '' END FROM assets WHERE id IN (` + placeholders + `)`
case "webshell":
query = `SELECT id, COALESCE(NULLIF(remark, ''), url), type FROM webshell_connections WHERE id IN (` + placeholders + `)`
case "batch_task":
@@ -1039,6 +1075,7 @@ func (db *DB) AssignResourcesToUserAuto(userID string, resourceIDs []string) (in
typeTablePairs := []struct{ resourceType, table string }{
{"project", "projects"}, {"conversation", "conversations"},
{"vulnerability", "vulnerabilities"}, {"webshell", "webshell_connections"},
{"asset", "assets"},
{"batch_task", "batch_task_queues"}, {"c2_listener", "c2_listeners"},
}
detected := make(map[string]string, len(uniqueIDs))