CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-07-20 10:51:09 +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}))