mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-17 16:37:18 +02:00
Delete internal directory
This commit is contained in:
公明
-2256
File diff suppressed because it is too large
Load Diff
@@ -1,511 +0,0 @@
|
||||
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 之一。",
|
||||
// Bedrock rejects tool schemas with top-level oneOf/allOf/anyOf. The
|
||||
// host/ip/domain requirement is enforced by assetFromCreateArgs below.
|
||||
InputSchema: map[string]interface{}{"type": "object", "properties": properties},
|
||||
}, 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"},
|
||||
"responsible_person": map[string]interface{}{"type": "string"}, "department": map[string]interface{}{"type": "string"},
|
||||
"business_system": map[string]interface{}{"type": "string"},
|
||||
"environment": map[string]interface{}{"type": "string", "enum": []string{"production", "staging", "testing", "development", "other"}},
|
||||
"criticality": map[string]interface{}{"type": "string", "enum": []string{"critical", "high", "medium", "low"}},
|
||||
"source": map[string]interface{}{"type": "string"}, "source_query": map[string]interface{}{"type": "string"},
|
||||
"status": map[string]interface{}{"type": "string", "enum": []string{"active", "inactive"}},
|
||||
"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},
|
||||
"risk_level": map[string]interface{}{"type": "string", "enum": []string{"unassessed", "critical", "high", "medium", "low", "info", "normal"}},
|
||||
"min_vulnerabilities": map[string]interface{}{"type": "integer", "minimum": 0},
|
||||
"max_vulnerabilities": map[string]interface{}{"type": "integer", "minimum": 0},
|
||||
"country": map[string]interface{}{"type": "string"}, "province": map[string]interface{}{"type": "string"}, "city": map[string]interface{}{"type": "string"},
|
||||
"responsible_person": map[string]interface{}{"type": "string"}, "department": map[string]interface{}{"type": "string"},
|
||||
"business_system": map[string]interface{}{"type": "string"}, "environment": map[string]interface{}{"type": "string"}, "criticality": map[string]interface{}{"type": "string"},
|
||||
"scan_state": map[string]interface{}{"type": "string", "enum": []string{"never", "scanned"}, "description": "never=从未扫描，scanned=扫描过"},
|
||||
"scan_overdue_days": map[string]interface{}{"type": "integer", "minimum": 1},
|
||||
"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"},
|
||||
"first_seen_before": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
|
||||
"first_seen_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", "risk_level", "vulnerability_count"}},
|
||||
"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("responsible_person", &asset.ResponsiblePerson)
|
||||
setString("department", &asset.Department)
|
||||
setString("business_system", &asset.BusinessSystem)
|
||||
setString("environment", &asset.Environment)
|
||||
setString("criticality", &asset.Criticality)
|
||||
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"))),
|
||||
RiskLevel: strings.ToLower(strings.TrimSpace(strArg(args, "risk_level"))),
|
||||
Country: strings.TrimSpace(strArg(args, "country")), Province: strings.TrimSpace(strArg(args, "province")), City: strings.TrimSpace(strArg(args, "city")),
|
||||
ResponsiblePerson: strings.TrimSpace(strArg(args, "responsible_person")), Department: strings.TrimSpace(strArg(args, "department")),
|
||||
BusinessSystem: strings.TrimSpace(strArg(args, "business_system")), Environment: strings.ToLower(strings.TrimSpace(strArg(args, "environment"))),
|
||||
Criticality: strings.ToLower(strings.TrimSpace(strArg(args, "criticality"))),
|
||||
}
|
||||
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", "risk_level", "vulnerability_count") {
|
||||
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
|
||||
}
|
||||
if _, ok := args["min_vulnerabilities"]; ok {
|
||||
value := intArg(args, "min_vulnerabilities", -1)
|
||||
if value < 0 {
|
||||
return filter, 0, 0, fmt.Errorf("min_vulnerabilities 不能小于 0")
|
||||
}
|
||||
filter.MinVulnerabilities = &value
|
||||
}
|
||||
if _, ok := args["max_vulnerabilities"]; ok {
|
||||
value := intArg(args, "max_vulnerabilities", -1)
|
||||
if value < 0 {
|
||||
return filter, 0, 0, fmt.Errorf("max_vulnerabilities 不能小于 0")
|
||||
}
|
||||
filter.MaxVulnerabilities = &value
|
||||
}
|
||||
if _, ok := args["scan_overdue_days"]; ok {
|
||||
value := intArg(args, "scan_overdue_days", 0)
|
||||
if value < 1 {
|
||||
return filter, 0, 0, fmt.Errorf("scan_overdue_days 必须大于 0")
|
||||
}
|
||||
filter.ScanOverdueDays = &value
|
||||
}
|
||||
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.FirstSeenBefore, err = parseAssetToolTime("first_seen_before", strArg(args, "first_seen_before")); err != nil {
|
||||
return filter, 0, 0, err
|
||||
}
|
||||
if filter.FirstSeenAfter, err = parseAssetToolTime("first_seen_after", strArg(args, "first_seen_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),
|
||||
"responsible_person": truncateRunes(asset.ResponsiblePerson, 255), "department": truncateRunes(asset.Department, 255),
|
||||
"business_system": truncateRunes(asset.BusinessSystem, 255), "environment": asset.Environment, "criticality": asset.Criticality,
|
||||
"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
|
||||
}
|
||||
@@ -1,201 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
for _, tool := range server.GetAllTools() {
|
||||
if tool.Name != builtin.ToolCreateAsset {
|
||||
continue
|
||||
}
|
||||
for _, keyword := range []string{"oneOf", "allOf", "anyOf"} {
|
||||
if _, exists := tool.InputSchema[keyword]; exists {
|
||||
t.Fatalf("create asset schema contains Bedrock-incompatible top-level %s", keyword)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result, _, err := server.CallTool(ctx, builtin.ToolCreateAsset, map[string]interface{}{"title": "Missing target"})
|
||||
if err != nil || result == nil || !result.IsError {
|
||||
t.Fatalf("create asset accepted missing host/ip/domain: result=%#v err=%v", result, err)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -1,228 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// C2HITLBridge 实现 C2 Manager 的 HITLBridge 接口，将危险任务桥接到现有 HITL 审批流。
|
||||
// 审批记录写入 hitl_interrupts 表，与现有 HITL 系统共享前端审批 UI。
|
||||
type C2HITLBridge struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
timeout time.Duration
|
||||
getConvID func() string
|
||||
}
|
||||
|
||||
// NewC2HITLBridge 创建 C2 HITL 桥
|
||||
func NewC2HITLBridge(db *database.DB, logger *zap.Logger) *C2HITLBridge {
|
||||
return &C2HITLBridge{
|
||||
db: db,
|
||||
logger: logger,
|
||||
timeout: 5 * time.Minute,
|
||||
getConvID: func() string { return "" },
|
||||
}
|
||||
}
|
||||
|
||||
// SetConversationIDGetter 设置获取当前对话 ID 的函数
|
||||
func (b *C2HITLBridge) SetConversationIDGetter(fn func() string) {
|
||||
b.getConvID = fn
|
||||
}
|
||||
|
||||
// SetTimeout 设置审批超时（0 表示不超时）
|
||||
func (b *C2HITLBridge) SetTimeout(d time.Duration) {
|
||||
b.timeout = d
|
||||
}
|
||||
|
||||
// RequestApproval 实现 HITLBridge 接口：写入 hitl_interrupts 表并轮询等待审批结果
|
||||
func (b *C2HITLBridge) RequestApproval(ctx context.Context, req c2.HITLApprovalRequest) error {
|
||||
interruptID := "hitl_c2_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
now := time.Now()
|
||||
|
||||
convID := req.ConversationID
|
||||
if convID == "" {
|
||||
convID = b.getConvID()
|
||||
}
|
||||
if convID == "" {
|
||||
convID = "c2_system"
|
||||
}
|
||||
|
||||
payload, _ := json.Marshal(map[string]interface{}{
|
||||
"task_id": req.TaskID,
|
||||
"session_id": req.SessionID,
|
||||
"task_type": req.TaskType,
|
||||
"payload": req.PayloadJSON,
|
||||
"source": req.Source,
|
||||
"reason": req.Reason,
|
||||
"c2_operation": true,
|
||||
})
|
||||
|
||||
_, err := b.db.Exec(`INSERT INTO hitl_interrupts
|
||||
(id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', ?)`,
|
||||
interruptID, convID, "", "approval",
|
||||
c2.MCPToolC2Task, req.TaskID,
|
||||
string(payload), now,
|
||||
)
|
||||
if err != nil {
|
||||
b.logger.Error("C2 HITL: 创建审批记录失败，拒绝执行", zap.Error(err))
|
||||
return fmt.Errorf("C2 HITL 审批记录创建失败，安全起见拒绝执行: %w", err)
|
||||
}
|
||||
|
||||
b.logger.Info("C2 HITL: 等待人工审批",
|
||||
zap.String("interrupt_id", interruptID),
|
||||
zap.String("task_id", req.TaskID),
|
||||
zap.String("task_type", req.TaskType),
|
||||
)
|
||||
|
||||
// Poll DB waiting for decision
|
||||
ticker := time.NewTicker(500 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
var deadline <-chan time.Time
|
||||
if b.timeout > 0 {
|
||||
timer := time.NewTimer(b.timeout)
|
||||
defer timer.Stop()
|
||||
deadline = timer.C
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
_, _ = b.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject',
|
||||
decision_comment='context cancelled', decided_at=? WHERE id=? AND status='pending'`,
|
||||
time.Now(), interruptID)
|
||||
return ctx.Err()
|
||||
|
||||
case <-deadline:
|
||||
_, _ = b.db.Exec(`UPDATE hitl_interrupts SET status='timeout', decision='reject',
|
||||
decision_comment='C2 HITL timeout auto-reject for safety', decided_at=? WHERE id=? AND status='pending'`,
|
||||
time.Now(), interruptID)
|
||||
b.logger.Warn("C2 HITL: 审批超时，安全起见拒绝执行", zap.String("interrupt_id", interruptID))
|
||||
return fmt.Errorf("C2 HITL 审批超时，危险任务已被自动拒绝")
|
||||
|
||||
case <-ticker.C:
|
||||
var status, decision string
|
||||
err := b.db.QueryRow(`SELECT status, COALESCE(decision, '') FROM hitl_interrupts WHERE id = ?`,
|
||||
interruptID).Scan(&status, &decision)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch status {
|
||||
case "decided", "timeout":
|
||||
if decision == "reject" {
|
||||
return fmt.Errorf("C2 危险任务被人工拒绝")
|
||||
}
|
||||
return nil
|
||||
case "cancelled":
|
||||
return fmt.Errorf("C2 审批已取消")
|
||||
case "pending":
|
||||
continue
|
||||
default:
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// C2HooksConfig 配置 C2 Manager 的 Hooks
|
||||
type C2HooksConfig struct {
|
||||
DB *database.DB
|
||||
Logger *zap.Logger
|
||||
AttackChainRecord func(session *database.C2Session, phase string, description string)
|
||||
VulnRecord func(session *database.C2Session, title string, severity string)
|
||||
}
|
||||
|
||||
// SetupC2Hooks 设置 C2 Manager 的业务钩子
|
||||
func SetupC2Hooks(cfg *C2HooksConfig) c2.Hooks {
|
||||
return c2.Hooks{
|
||||
OnSessionFirstSeen: func(session *database.C2Session) {
|
||||
// 新会话上线
|
||||
cfg.Logger.Info("C2 Session first seen",
|
||||
zap.String("session_id", session.ID),
|
||||
zap.String("hostname", session.Hostname),
|
||||
zap.String("os", session.OS),
|
||||
zap.String("arch", session.Arch),
|
||||
)
|
||||
|
||||
// 记录漏洞（初始访问点）
|
||||
if cfg.VulnRecord != nil {
|
||||
cfg.VulnRecord(session, fmt.Sprintf("C2 Session Established: %s@%s", session.Username, session.Hostname), "high")
|
||||
}
|
||||
|
||||
// 记录攻击链（Initial Access）
|
||||
if cfg.AttackChainRecord != nil {
|
||||
cfg.AttackChainRecord(session, "initial-access", fmt.Sprintf("Implant beacon from %s/%s", session.Hostname, session.InternalIP))
|
||||
}
|
||||
},
|
||||
OnTaskCompleted: func(task *database.C2Task, sessionID string) {
|
||||
// 任务完成
|
||||
cfg.Logger.Debug("C2 Task completed",
|
||||
zap.String("task_id", task.ID),
|
||||
zap.String("task_type", task.TaskType),
|
||||
zap.String("status", task.Status),
|
||||
)
|
||||
|
||||
// 根据任务类型记录攻击链
|
||||
if cfg.AttackChainRecord != nil {
|
||||
session, _ := cfg.DB.GetC2Session(sessionID)
|
||||
if session != nil {
|
||||
phase := taskToAttackPhase(task.TaskType)
|
||||
if phase != "" {
|
||||
cfg.AttackChainRecord(session, phase, fmt.Sprintf("Task %s: %s", task.TaskType, task.Status))
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// taskToAttackPhase 将任务类型映射到 ATT&CK 阶段
|
||||
func taskToAttackPhase(taskType string) string {
|
||||
switch taskType {
|
||||
case "exec", "shell":
|
||||
return "execution"
|
||||
case "upload":
|
||||
return "persistence"
|
||||
case "download":
|
||||
return "exfiltration"
|
||||
case "screenshot":
|
||||
return "collection"
|
||||
case "kill_proc":
|
||||
return "impact"
|
||||
case "port_fwd", "socks_start":
|
||||
return "lateral-movement"
|
||||
case "load_assembly":
|
||||
return "defense-evasion"
|
||||
case "persist":
|
||||
return "persistence"
|
||||
case "self_delete":
|
||||
return "defense-evasion"
|
||||
default:
|
||||
return "execution"
|
||||
}
|
||||
}
|
||||
|
||||
// SetupC2HITLBridgeWithAgent 设置 HITL 桥接器
|
||||
// 这个函数将由 App 调用，注入必要的依赖
|
||||
func SetupC2HITLBridgeWithAgent(db *database.DB, logger *zap.Logger) c2.HITLBridge {
|
||||
return &C2HITLBridge{
|
||||
db: db,
|
||||
logger: logger,
|
||||
timeout: 5 * time.Minute,
|
||||
getConvID: func() string { return "" },
|
||||
}
|
||||
}
|
||||
@@ -1,104 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/handler"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// setupC2Runtime 创建 C2 Manager、看门狗与取消函数；不注册 MCP 工具（由 Apply 统一 ClearTools 后注册）。
|
||||
func setupC2Runtime(
|
||||
cfg *config.Config,
|
||||
db *database.DB,
|
||||
agentHandler *handler.AgentHandler,
|
||||
logger *zap.Logger,
|
||||
) (*c2.Manager, *c2.SessionWatchdog, context.CancelFunc) {
|
||||
if !cfg.C2.EnabledEffective() {
|
||||
return nil, nil, nil
|
||||
}
|
||||
c2Manager := c2.NewManager(db, logger, "tmp/c2")
|
||||
c2Manager.Registry().Register(string(c2.ListenerTypeTCPReverse), c2.NewTCPReverseListener)
|
||||
c2Manager.Registry().Register(string(c2.ListenerTypeHTTPBeacon), c2.NewHTTPBeaconListener)
|
||||
c2Manager.Registry().Register(string(c2.ListenerTypeHTTPSBeacon), c2.NewHTTPSBeaconListener)
|
||||
c2Manager.Registry().Register(string(c2.ListenerTypeWebSocket), c2.NewWebSocketListener)
|
||||
c2HITLBridge := NewC2HITLBridge(db, logger)
|
||||
c2Manager.SetHITLBridge(c2HITLBridge)
|
||||
c2Manager.SetHITLDangerousGate(func(conversationID, toolName string) bool {
|
||||
return agentHandler.HITLNeedsToolApproval(conversationID, toolName)
|
||||
})
|
||||
c2Hooks := SetupC2Hooks(&C2HooksConfig{
|
||||
DB: db,
|
||||
Logger: logger,
|
||||
AttackChainRecord: func(session *database.C2Session, phase string, description string) {
|
||||
logger.Info("C2 Attack Chain",
|
||||
zap.String("session_id", session.ID),
|
||||
zap.String("phase", phase),
|
||||
zap.String("desc", description),
|
||||
)
|
||||
},
|
||||
VulnRecord: func(session *database.C2Session, title string, severity string) {
|
||||
logger.Info("C2 Vulnerability",
|
||||
zap.String("session_id", session.ID),
|
||||
zap.String("title", title),
|
||||
zap.String("severity", severity),
|
||||
)
|
||||
},
|
||||
})
|
||||
c2Manager.SetHooks(c2Hooks)
|
||||
c2Manager.RestoreRunningListeners()
|
||||
c2Watchdog := c2.NewSessionWatchdog(c2Manager)
|
||||
watchdogCtx, watchdogCancel := context.WithCancel(context.Background())
|
||||
go c2Watchdog.Run(watchdogCtx)
|
||||
return c2Manager, c2Watchdog, watchdogCancel
|
||||
}
|
||||
|
||||
// ReconcileC2AfterConfigApply 根据当前内存配置启停 C2（不写盘；在 Apply 中 ClearTools 之前调用）。
|
||||
func (a *App) ReconcileC2AfterConfigApply() error {
|
||||
if !a.config.C2.EnabledEffective() {
|
||||
a.shutdownC2()
|
||||
return nil
|
||||
}
|
||||
if a.c2Manager != nil {
|
||||
return nil
|
||||
}
|
||||
if a.db == nil || a.agentHandler == nil {
|
||||
return nil
|
||||
}
|
||||
m, wd, cancel := setupC2Runtime(a.config, a.db, a.agentHandler, a.logger.Logger)
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
a.c2Manager = m
|
||||
a.c2Watchdog = wd
|
||||
a.c2WatchdogCancel = cancel
|
||||
if a.c2Handler != nil {
|
||||
a.c2Handler.SetManager(m)
|
||||
}
|
||||
a.logger.Info("C2 子系统已按配置启动")
|
||||
return nil
|
||||
}
|
||||
|
||||
// shutdownC2 停止看门狗与所有监听器，并断开 Handler 引用。
|
||||
func (a *App) shutdownC2() {
|
||||
had := a.c2WatchdogCancel != nil || a.c2Manager != nil
|
||||
if a.c2WatchdogCancel != nil {
|
||||
a.c2WatchdogCancel()
|
||||
a.c2WatchdogCancel = nil
|
||||
}
|
||||
a.c2Watchdog = nil
|
||||
if a.c2Manager != nil {
|
||||
a.c2Manager.Close()
|
||||
a.c2Manager = nil
|
||||
}
|
||||
if a.c2Handler != nil {
|
||||
a.c2Handler.SetManager(nil)
|
||||
}
|
||||
if had {
|
||||
a.logger.Info("C2 子系统已关闭")
|
||||
}
|
||||
}
|
||||
@@ -1,919 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// registerC2Tools 注册所有 C2 MCP 工具（合并同类项，减少工具数量以节省上下文 token）。
|
||||
// webListenPort 为本进程 Web/API 监听端口（配置 server.port，启动时已加载），用于 MCP 描述中提示勿与 C2 bind_port 冲突。
|
||||
func registerC2Tools(mcpServer *mcp.Server, c2Manager *c2.Manager, logger *zap.Logger, webListenPort int) {
|
||||
registerC2ListenerTool(mcpServer, c2Manager, logger, webListenPort)
|
||||
registerC2SessionTool(mcpServer, c2Manager, logger)
|
||||
registerC2TaskTool(mcpServer, c2Manager, logger)
|
||||
registerC2TaskManageTool(mcpServer, c2Manager, logger)
|
||||
registerC2PayloadTool(mcpServer, c2Manager, logger, webListenPort)
|
||||
registerC2EventTool(mcpServer, c2Manager, logger)
|
||||
registerC2ProfileTool(mcpServer, c2Manager, logger)
|
||||
registerC2FileTool(mcpServer, c2Manager, logger)
|
||||
logger.Debug("C2 MCP tools registered (8 unified tools)")
|
||||
}
|
||||
|
||||
func makeC2Result(data interface{}, err error) (*mcp.ToolResult, error) {
|
||||
if err != nil {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: err.Error()}},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
text, _ := json.Marshal(data)
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: string(text)}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_listener — 监听器统一工具
|
||||
// ============================================================================
|
||||
|
||||
func registerC2ListenerTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListenPort int) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Listener,
|
||||
Description: fmt.Sprintf(`C2 监听器管理。通过 action 参数选择操作：
|
||||
- list: 列出所有监听器
|
||||
- get: 获取监听器详情（需 listener_id）
|
||||
- create: 创建监听器（需 name, type, bind_port）。成功时除 listener 外会返回 implant_token（仅此一次，用于 X-Implant-Token / oneliner；list/get/start 不再返回）
|
||||
- update: 更新监听器配置（需 listener_id，可改 name/bind_host/bind_port/remark/config/callback_host）
|
||||
- start: 启动监听器（需 listener_id）
|
||||
- stop: 停止监听器（需 listener_id）
|
||||
- delete: 删除监听器（需 listener_id）
|
||||
监听器类型: tcp_reverse, http_beacon, https_beacon, websocket
|
||||
tcp_reverse 默认仅接受 CSB1 加密 Beacon（AES-GCM + ImplantToken）才登记会话；经典 bash/nc 反弹需在 config.allow_legacy_shell=true（公网不推荐）。
|
||||
端口约束：create/update 的 bind_port 禁止与本平台 Web/API 所用端口相同。当前本服务该端口为 %d（配置项 server.port，随进程启动从配置文件加载）。若 bind_port 与此相同会导致本服务或监听器 bind 失败、Beacon/oneliner 误连到 Web 而非 C2。请为监听器另选空闲端口。`, webListenPort),
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/start/stop/delete", "enum": []string{"list", "get", "create", "update", "start", "stop", "delete"}},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID（get/update/start/stop/delete 需要）"},
|
||||
"name": map[string]interface{}{"type": "string", "description": "监听器名称（create/update）"},
|
||||
"type": map[string]interface{}{"type": "string", "description": "监听器类型（create）", "enum": []string{"tcp_reverse", "http_beacon", "https_beacon", "websocket"}},
|
||||
"bind_host": map[string]interface{}{"type": "string", "description": "绑定地址，默认 127.0.0.1；外网监听常用 0.0.0.0"},
|
||||
"callback_host": map[string]interface{}{"type": "string", "description": "可选：植入端/Payload 回连主机名（公网 IP 或域名）。写入 config_json；生成 oneliner/beacon 时优先于 bind_host。update 时传入空字符串可清除"},
|
||||
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口（create 必填）。须 ≠ %d（当前本服务 Web/API 端口，配置 server.port）", webListenPort), "minimum": 1, "maximum": 65535},
|
||||
"project_id": map[string]interface{}{"type": "string", "description": "所属项目 ID。create 省略时默认使用当前对话绑定项目；未绑定项目的对话则创建未绑定监听器"},
|
||||
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
|
||||
"remark": map[string]interface{}{"type": "string", "description": "备注"},
|
||||
"config": map[string]interface{}{"type": "object", "description": "高级配置（beacon 路径/TLS/OPSEC 等），create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell（默认 false）"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
id := getString(params, "listener_id")
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx), mcpEffectiveProjectFilter(ctx, m.DB()))
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
for _, li := range listeners {
|
||||
li.EncryptionKey = ""
|
||||
li.ImplantToken = ""
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"listeners": listeners, "count": len(listeners)}, nil)
|
||||
|
||||
case "get":
|
||||
listener, err := m.DB().GetC2Listener(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if listener == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("listener not found"))
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
return makeC2Result(map[string]interface{}{"listener": listener}, nil)
|
||||
|
||||
case "create":
|
||||
var cfg *c2.ListenerConfig
|
||||
if cfgRaw, ok := params["config"]; ok && cfgRaw != nil {
|
||||
cfgBytes, _ := json.Marshal(cfgRaw)
|
||||
cfg = &c2.ListenerConfig{}
|
||||
_ = json.Unmarshal(cfgBytes, cfg)
|
||||
}
|
||||
projectID := strings.TrimSpace(getString(params, "project_id"))
|
||||
if projectID == "" {
|
||||
projectID = mcpEffectiveProjectFilter(ctx, m.DB())
|
||||
if projectID == database.ProjectFilterUnbound {
|
||||
projectID = ""
|
||||
}
|
||||
}
|
||||
input := c2.CreateListenerInput{
|
||||
Name: getString(params, "name"),
|
||||
Type: getString(params, "type"),
|
||||
BindHost: getString(params, "bind_host"),
|
||||
BindPort: int(getFloat64(params, "bind_port")),
|
||||
ProfileID: getString(params, "profile_id"),
|
||||
Remark: getString(params, "remark"),
|
||||
ProjectID: projectID,
|
||||
Config: cfg,
|
||||
CallbackHost: getString(params, "callback_host"),
|
||||
}
|
||||
listener, err := m.CreateListener(input)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = m.DB().SetResourceOwner("c2_listener", listener.ID, principal.UserID)
|
||||
_ = m.DB().AssignResourceToUser(principal.UserID, "c2_listener", listener.ID)
|
||||
}
|
||||
implantToken := listener.ImplantToken
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
return makeC2Result(map[string]interface{}{
|
||||
"listener": listener,
|
||||
"implant_token": implantToken,
|
||||
}, nil)
|
||||
|
||||
case "update":
|
||||
listener, err := m.DB().GetC2Listener(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if listener == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("listener not found"))
|
||||
}
|
||||
if m.IsListenerRunning(id) {
|
||||
newHost := getString(params, "bind_host")
|
||||
newPort := int(getFloat64(params, "bind_port"))
|
||||
if (newHost != "" && newHost != listener.BindHost) || (newPort > 0 && newPort != listener.BindPort) {
|
||||
return makeC2Result(nil, fmt.Errorf("cannot modify bind address while listener is running"))
|
||||
}
|
||||
}
|
||||
if v := getString(params, "name"); v != "" {
|
||||
listener.Name = v
|
||||
}
|
||||
if v := getString(params, "bind_host"); v != "" {
|
||||
listener.BindHost = v
|
||||
}
|
||||
if v := int(getFloat64(params, "bind_port")); v > 0 {
|
||||
listener.BindPort = v
|
||||
}
|
||||
if v := getString(params, "profile_id"); v != "" {
|
||||
listener.ProfileID = v
|
||||
}
|
||||
if v, ok := params["remark"]; ok {
|
||||
listener.Remark, _ = v.(string)
|
||||
}
|
||||
if cfgRaw, ok := params["config"]; ok && cfgRaw != nil {
|
||||
cfgBytes, _ := json.Marshal(cfgRaw)
|
||||
listener.ConfigJSON = string(cfgBytes)
|
||||
}
|
||||
if _, ok := params["callback_host"]; ok {
|
||||
pcfg := &c2.ListenerConfig{}
|
||||
raw := strings.TrimSpace(listener.ConfigJSON)
|
||||
if raw == "" {
|
||||
raw = "{}"
|
||||
}
|
||||
_ = json.Unmarshal([]byte(raw), pcfg)
|
||||
pcfg.CallbackHost = strings.TrimSpace(getString(params, "callback_host"))
|
||||
pcfg.ApplyDefaults()
|
||||
cfgBytes, err := json.Marshal(pcfg)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
listener.ConfigJSON = string(cfgBytes)
|
||||
}
|
||||
if err := m.DB().UpdateC2Listener(listener); err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
return makeC2Result(map[string]interface{}{"listener": listener}, nil)
|
||||
|
||||
case "start":
|
||||
listener, err := m.StartListener(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
return makeC2Result(map[string]interface{}{"listener": listener}, nil)
|
||||
|
||||
case "stop":
|
||||
err := m.StopListener(id)
|
||||
return makeC2Result(map[string]interface{}{"stopped": err == nil}, err)
|
||||
|
||||
case "delete":
|
||||
err := m.DeleteListener(id)
|
||||
return makeC2Result(map[string]interface{}{"deleted": err == nil}, err)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_session — 会话统一工具
|
||||
// ============================================================================
|
||||
|
||||
func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Session,
|
||||
Description: `C2 会话管理。通过 action 参数选择操作：
|
||||
- list: 列出会话（可按 listener_id/status/os/search/suspicious 过滤）
|
||||
- get: 获取会话详情及最近任务历史（需 session_id）
|
||||
- set_sleep: 设置心跳间隔（需 session_id）
|
||||
- kill: 下发 exit 任务让 implant 退出（需 session_id）
|
||||
- delete: 删除单个会话记录（需 session_id）
|
||||
- delete_batch: 批量删除会话（需 session_ids 数组）`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/set_sleep/kill/delete/delete_batch", "enum": []string{"list", "get", "set_sleep", "kill", "delete", "delete_batch"}},
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "会话 ID（get/set_sleep/kill/delete 需要）"},
|
||||
"session_ids": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "description": "会话 ID 列表（delete_batch）"},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "按监听器过滤（list）"},
|
||||
"status": map[string]interface{}{"type": "string", "description": "按状态过滤: active/sleeping/dead/killed（list）"},
|
||||
"os": map[string]interface{}{"type": "string", "description": "按 OS 过滤: linux/windows/darwin（list）"},
|
||||
"search": map[string]interface{}{"type": "string", "description": "模糊搜索 hostname/username/IP（list）"},
|
||||
"suspicious": map[string]interface{}{"type": "boolean", "description": "仅疑似误报：离线且 tcp_* / unknown / PID 0（list）"},
|
||||
"limit": map[string]interface{}{"type": "integer", "description": "返回数量上限（list）"},
|
||||
"sleep_seconds": map[string]interface{}{"type": "integer", "description": "心跳间隔秒数（set_sleep）"},
|
||||
"jitter_percent": map[string]interface{}{"type": "integer", "description": "抖动百分比 0-100（set_sleep）"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
id := getString(params, "session_id")
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
filter := database.ListC2SessionsFilter{
|
||||
ListenerID: getString(params, "listener_id"),
|
||||
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||
Status: getString(params, "status"),
|
||||
OS: getString(params, "os"),
|
||||
Search: getString(params, "search"),
|
||||
}
|
||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
||||
filter.Limit = limit
|
||||
}
|
||||
if v, ok := params["suspicious"].(bool); ok && v {
|
||||
filter.Suspicious = true
|
||||
}
|
||||
sessions, err := m.DB().ListC2SessionsForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"sessions": sessions, "count": len(sessions)}, err)
|
||||
|
||||
case "get":
|
||||
session, err := m.DB().GetC2Session(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if session == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("session not found"))
|
||||
}
|
||||
tasks, _ := m.DB().ListC2Tasks(database.ListC2TasksFilter{SessionID: id, Limit: 10})
|
||||
return makeC2Result(map[string]interface{}{"session": session, "tasks": tasks}, nil)
|
||||
|
||||
case "set_sleep":
|
||||
sleep := int(getFloat64(params, "sleep_seconds"))
|
||||
jitter := int(getFloat64(params, "jitter_percent"))
|
||||
task, err := m.SetSessionSleep(id, sleep, jitter)
|
||||
out := map[string]interface{}{
|
||||
"updated": err == nil,
|
||||
"sleep_seconds": sleep,
|
||||
"jitter_percent": jitter,
|
||||
}
|
||||
if task != nil {
|
||||
out["task_id"] = task.ID
|
||||
}
|
||||
return makeC2Result(out, err)
|
||||
|
||||
case "kill":
|
||||
task, err := m.EnqueueTask(c2.EnqueueTaskInput{
|
||||
SessionID: id,
|
||||
TaskType: c2.TaskTypeExit,
|
||||
Payload: map[string]interface{}{},
|
||||
Source: "ai",
|
||||
ConversationID: agent.ConversationIDFromContext(ctx),
|
||||
UserCtx: ctx,
|
||||
})
|
||||
return makeC2Result(map[string]interface{}{"task": task}, err)
|
||||
|
||||
case "delete":
|
||||
err := m.DB().DeleteC2Session(id)
|
||||
return makeC2Result(map[string]interface{}{"deleted": err == nil}, err)
|
||||
|
||||
case "delete_batch":
|
||||
rawIDs, _ := params["session_ids"].([]interface{})
|
||||
ids := make([]string, 0, len(rawIDs))
|
||||
for _, v := range rawIDs {
|
||||
if s, ok := v.(string); ok && strings.TrimSpace(s) != "" {
|
||||
ids = append(ids, strings.TrimSpace(s))
|
||||
}
|
||||
}
|
||||
n, err := m.DB().DeleteC2SessionsByIDs(ids)
|
||||
return makeC2Result(map[string]interface{}{"deleted": n}, err)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_task — 任务下发统一工具（合并所有 task 类型）
|
||||
// ============================================================================
|
||||
|
||||
func registerC2TaskTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Task,
|
||||
Description: `在 C2 会话上下发任务。所有任务类型通过 task_type 参数指定：
|
||||
- exec: 执行命令（需 command）
|
||||
- shell: 交互式命令，保持 cwd（需 command）
|
||||
- pwd/ps/screenshot/socks_stop: 无额外参数
|
||||
- cd/ls: 需 path
|
||||
- kill_proc: 需 pid
|
||||
- upload: 需 remote_path + file_id
|
||||
- download: 需 remote_path
|
||||
- port_fwd: 需 action(start/stop) + local_port + remote_host + remote_port
|
||||
- socks_start: 需 port（默认 1080）
|
||||
- load_assembly: 需 data(base64) 或 file_id，可选 args
|
||||
- persist: 可选 method(auto/cron/bashrc/launchagent/registry/schtasks)
|
||||
返回 task_id，用 c2_task_manage 的 wait/get_result 获取结果。`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "C2 会话 ID（s_xxx）"},
|
||||
"task_type": map[string]interface{}{"type": "string", "description": "任务类型", "enum": []string{"exec", "shell", "pwd", "cd", "ls", "ps", "kill_proc", "upload", "download", "screenshot", "port_fwd", "socks_start", "socks_stop", "load_assembly", "persist"}},
|
||||
"command": map[string]interface{}{"type": "string", "description": "命令（exec/shell）"},
|
||||
"path": map[string]interface{}{"type": "string", "description": "路径（cd/ls）"},
|
||||
"pid": map[string]interface{}{"type": "integer", "description": "进程 ID（kill_proc）"},
|
||||
"remote_path": map[string]interface{}{"type": "string", "description": "远程路径（upload/download）"},
|
||||
"file_id": map[string]interface{}{"type": "string", "description": "服务端文件 ID（upload/load_assembly）"},
|
||||
"data": map[string]interface{}{"type": "string", "description": "base64 数据（load_assembly）"},
|
||||
"args": map[string]interface{}{"type": "string", "description": "命令行参数（load_assembly）"},
|
||||
"action": map[string]interface{}{"type": "string", "description": "start/stop（port_fwd）"},
|
||||
"local_port": map[string]interface{}{"type": "integer", "description": "本地端口（port_fwd）"},
|
||||
"remote_host": map[string]interface{}{"type": "string", "description": "远程主机（port_fwd）"},
|
||||
"remote_port": map[string]interface{}{"type": "integer", "description": "远程端口（port_fwd）"},
|
||||
"port": map[string]interface{}{"type": "integer", "description": "SOCKS5 端口（socks_start），默认 1080"},
|
||||
"method": map[string]interface{}{"type": "string", "description": "持久化方法（persist）: auto/cron/bashrc/launchagent/registry/schtasks"},
|
||||
"timeout_seconds": map[string]interface{}{"type": "integer", "description": "超时秒数，默认 60"},
|
||||
},
|
||||
"required": []string{"session_id", "task_type"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
sessionID := getString(params, "session_id")
|
||||
taskTypeStr := getString(params, "task_type")
|
||||
taskType := c2.TaskType(taskTypeStr)
|
||||
timeout := getFloat64(params, "timeout_seconds")
|
||||
|
||||
payload := map[string]interface{}{"timeout_seconds": timeout}
|
||||
|
||||
switch taskType {
|
||||
case c2.TaskTypeExec, c2.TaskTypeShell:
|
||||
payload["command"] = getString(params, "command")
|
||||
case c2.TaskTypeCd, c2.TaskTypeLs:
|
||||
payload["path"] = getString(params, "path")
|
||||
case c2.TaskTypeKillProc:
|
||||
payload["pid"] = params["pid"]
|
||||
case c2.TaskTypeUpload:
|
||||
payload["remote_path"] = getString(params, "remote_path")
|
||||
payload["file_id"] = getString(params, "file_id")
|
||||
case c2.TaskTypeDownload:
|
||||
payload["remote_path"] = getString(params, "remote_path")
|
||||
case c2.TaskTypePortFwd:
|
||||
payload["action"] = getString(params, "action")
|
||||
payload["local_port"] = params["local_port"]
|
||||
payload["remote_host"] = getString(params, "remote_host")
|
||||
payload["remote_port"] = params["remote_port"]
|
||||
case c2.TaskTypeSocksStart:
|
||||
payload["port"] = params["port"]
|
||||
case c2.TaskTypeLoadAssembly:
|
||||
payload["data"] = getString(params, "data")
|
||||
payload["file_id"] = getString(params, "file_id")
|
||||
payload["args"] = getString(params, "args")
|
||||
case c2.TaskTypePersist:
|
||||
payload["method"] = getString(params, "method")
|
||||
case c2.TaskTypePwd, c2.TaskTypePs, c2.TaskTypeScreenshot, c2.TaskTypeSocksStop:
|
||||
// no extra params
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unsupported task_type: %s", taskTypeStr))
|
||||
}
|
||||
|
||||
input := c2.EnqueueTaskInput{
|
||||
SessionID: sessionID,
|
||||
TaskType: taskType,
|
||||
Payload: payload,
|
||||
Source: "ai",
|
||||
ConversationID: agent.ConversationIDFromContext(ctx),
|
||||
UserCtx: ctx,
|
||||
}
|
||||
task, err := m.EnqueueTask(input)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"task_id": task.ID, "status": task.Status}, nil)
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_task_manage — 任务管理工具（查询/等待/取消）
|
||||
// ============================================================================
|
||||
|
||||
func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2TaskManage,
|
||||
Description: `C2 任务管理。通过 action 参数选择操作：
|
||||
- get_result: 获取任务详情和结果（需 task_id）
|
||||
- wait: 阻塞等待任务完成并返回结果（需 task_id）
|
||||
- list: 列出任务（可按 session_id/status 过滤）
|
||||
- cancel: 取消排队中的任务（需 task_id）`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: get_result/wait/list/cancel", "enum": []string{"get_result", "wait", "list", "cancel"}},
|
||||
"task_id": map[string]interface{}{"type": "string", "description": "任务 ID（get_result/wait/cancel 需要）"},
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "按会话过滤（list）"},
|
||||
"status": map[string]interface{}{"type": "string", "description": "按状态过滤: queued/sent/running/success/failed/cancelled（list）"},
|
||||
"limit": map[string]interface{}{"type": "integer", "description": "返回数量上限（list）"},
|
||||
"timeout_seconds": map[string]interface{}{"type": "integer", "description": "等待超时秒数（wait），默认 60"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
|
||||
switch action {
|
||||
case "get_result":
|
||||
id := getString(params, "task_id")
|
||||
task, err := m.DB().GetC2Task(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if task == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("task not found"))
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"task": task}, nil)
|
||||
|
||||
case "wait":
|
||||
id := getString(params, "task_id")
|
||||
timeout := int(getFloat64(params, "timeout_seconds"))
|
||||
if timeout <= 0 {
|
||||
timeout = 60
|
||||
}
|
||||
deadline := time.Now().Add(time.Duration(timeout) * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
task, err := m.DB().GetC2Task(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if task == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("task not found"))
|
||||
}
|
||||
if task.Status == "success" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
return makeC2Result(map[string]interface{}{"task": task}, nil)
|
||||
}
|
||||
select {
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
case <-ctx.Done():
|
||||
return makeC2Result(nil, ctx.Err())
|
||||
}
|
||||
}
|
||||
return makeC2Result(nil, fmt.Errorf("timeout waiting for task completion"))
|
||||
|
||||
case "list":
|
||||
filter := database.ListC2TasksFilter{
|
||||
SessionID: getString(params, "session_id"),
|
||||
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||
Status: getString(params, "status"),
|
||||
}
|
||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
||||
filter.Limit = limit
|
||||
}
|
||||
tasks, err := m.DB().ListC2TasksForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"tasks": tasks, "count": len(tasks)}, err)
|
||||
|
||||
case "cancel":
|
||||
id := getString(params, "task_id")
|
||||
err := m.CancelTask(id)
|
||||
return makeC2Result(map[string]interface{}{"cancelled": err == nil}, err)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_payload — Payload 统一工具
|
||||
// ============================================================================
|
||||
|
||||
func registerC2PayloadTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListenPort int) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Payload,
|
||||
Description: fmt.Sprintf(`C2 Payload 生成。通过 action 参数选择操作：
|
||||
- oneliner: 生成单行 payload。kind 必须与监听器协议一致，否则会失败：
|
||||
• tcp_reverse：默认仅支持 build 加密 Beacon；若监听器 config.allow_legacy_shell=true，才可用 kind: bash, nc, nc_mkfifo, python, perl, powershell。
|
||||
• http_beacon / https_beacon / websocket：仅 HTTP(S) Beacon 轮询，oneliner 只能用 kind: curl_beacon（脚本内用 bash+curl，与「tcp 的 bash」不同）。curl_beacon 返回串末尾含「 &」用于把整个 bash -c 放后台；若用 exec/execute 同步执行，必须整段原样复制（含末尾 &）。若删掉 &，内部 while 死循环占满前台，调用会一直阻塞到超时/杀进程。
|
||||
• 公网部署 tcp_reverse 请用 build 生成加密 Beacon，勿开启 allow_legacy_shell。
|
||||
• 省略 kind 时，会按监听器类型自动选第一个兼容类型（HTTP 系默认为 curl_beacon）。
|
||||
- build: 交叉编译 beacon 二进制。支持 http_beacon / https_beacon / websocket / tcp_reverse（tcp_reverse 植入端回连后先发魔数 CSB1，再经 AES-GCM 解密且校验 ImplantToken 后才登记会话）。
|
||||
依赖的监听器 bind_port 须避开本服务 Web 端口 %d（配置 server.port，与 c2_listener 描述一致），否则 Beacon 无法正确回连。`, webListenPort),
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: oneliner/build", "enum": []string{"oneliner", "build"}},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID（必填）。oneliner 前请确认该监听器的 type，再选兼容的 kind"},
|
||||
"kind": map[string]interface{}{"type": "string", "description": "仅 action=oneliner 需要。tcp_reverse: bash|nc|nc_mkfifo|python|perl|powershell；http_beacon|https_beacon|websocket: 仅 curl_beacon"},
|
||||
"host": map[string]interface{}{"type": "string", "description": "oneliner/build 可选覆盖：非空则强制用作植入回连主机。留空时顺序为：监听器 callback_host（create/update 的 callback_host 参数写入）→ bind_host（0.0.0.0 时尝试本机对外 IP 探测）"},
|
||||
"os": map[string]interface{}{"type": "string", "description": "目标 OS（build）: linux/windows/darwin", "default": "linux"},
|
||||
"arch": map[string]interface{}{"type": "string", "description": "目标架构（build）: amd64/arm64/386/arm", "default": "amd64"},
|
||||
"sleep_seconds": map[string]interface{}{"type": "integer", "description": "默认心跳间隔（build）"},
|
||||
"jitter_percent": map[string]interface{}{"type": "integer", "description": "默认抖动百分比（build）"},
|
||||
},
|
||||
"required": []string{"action", "listener_id"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
listenerID := getString(params, "listener_id")
|
||||
|
||||
switch action {
|
||||
case "oneliner":
|
||||
listener, err := m.DB().GetC2Listener(listenerID)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if listener == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("listener not found"))
|
||||
}
|
||||
host := c2.ResolveBeaconDialHost(listener, getString(params, "host"), l, listenerID)
|
||||
kind := c2.OnelinerKind(getString(params, "kind"))
|
||||
if kind == "" {
|
||||
compatible := c2.OnelinerKindsForListener(listener.Type)
|
||||
if len(compatible) > 0 {
|
||||
kind = compatible[0]
|
||||
}
|
||||
}
|
||||
if !c2.IsOnelinerCompatible(listener.Type, kind) {
|
||||
compatible := c2.OnelinerKindsForListener(listener.Type)
|
||||
names := make([]string, len(compatible))
|
||||
for i, k := range compatible {
|
||||
names[i] = string(k)
|
||||
}
|
||||
return makeC2Result(nil, fmt.Errorf("监听器类型 %s 不支持 %s，兼容类型: %v", listener.Type, kind, names))
|
||||
}
|
||||
if err := c2.ValidateOnelinerForListener(listener, kind); err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
input := c2.OnelinerInput{
|
||||
Kind: kind,
|
||||
Host: host,
|
||||
Port: listener.BindPort,
|
||||
HTTPBaseURL: fmt.Sprintf("http://%s:%d", host, listener.BindPort),
|
||||
ImplantToken: listener.ImplantToken,
|
||||
}
|
||||
oneliner, err := c2.GenerateOneliner(input)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
out := map[string]interface{}{
|
||||
"oneliner": oneliner, "kind": input.Kind, "host": host, "port": listener.BindPort,
|
||||
}
|
||||
if kind == c2.OnelinerCurl {
|
||||
out["usage_note"] = "同步 exec/execute：整段原样执行（末尾须有「 &」）。去掉则 while 永不结束，工具会一直卡住。"
|
||||
}
|
||||
return makeC2Result(out, nil)
|
||||
|
||||
case "build":
|
||||
builder := c2.NewPayloadBuilder(m, l, "", "")
|
||||
input := c2.PayloadBuilderInput{
|
||||
ListenerID: listenerID,
|
||||
OS: getString(params, "os"),
|
||||
Arch: getString(params, "arch"),
|
||||
SleepSeconds: int(getFloat64(params, "sleep_seconds")),
|
||||
JitterPercent: int(getFloat64(params, "jitter_percent")),
|
||||
Host: strings.TrimSpace(getString(params, "host")),
|
||||
}
|
||||
result, err := builder.BuildBeacon(input)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = m.DB().RecordC2PayloadArtifact(filepath.Base(result.OutputPath), result.PayloadID, result.ListenerID, principal.UserID)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{
|
||||
"payload_id": result.PayloadID, "download_path": result.DownloadPath,
|
||||
"os": result.OS, "arch": result.Arch, "size_bytes": result.SizeBytes,
|
||||
}, nil)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_event — 事件查询工具
|
||||
// ============================================================================
|
||||
|
||||
func registerC2EventTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Event,
|
||||
Description: "获取 C2 事件（上线/掉线/任务/错误），支持按级别/类别/会话/任务/时间过滤",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"level": map[string]interface{}{"type": "string", "description": "级别过滤: info/warn/critical"},
|
||||
"category": map[string]interface{}{"type": "string", "description": "类别过滤: listener/session/task/payload/opsec"},
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "按会话过滤"},
|
||||
"task_id": map[string]interface{}{"type": "string", "description": "按任务过滤"},
|
||||
"since": map[string]interface{}{"type": "string", "description": "起始时间（RFC3339 格式，如 2025-01-01T00:00:00Z）"},
|
||||
"limit": map[string]interface{}{"type": "integer", "default": 50, "description": "返回数量"},
|
||||
},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
filter := database.ListC2EventsFilter{
|
||||
Level: getString(params, "level"),
|
||||
Category: getString(params, "category"),
|
||||
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||
SessionID: getString(params, "session_id"),
|
||||
TaskID: getString(params, "task_id"),
|
||||
Limit: int(getFloat64(params, "limit")),
|
||||
}
|
||||
if filter.Limit <= 0 {
|
||||
filter.Limit = 50
|
||||
}
|
||||
if since := getString(params, "since"); since != "" {
|
||||
if t, err := time.Parse(time.RFC3339, since); err == nil {
|
||||
filter.Since = &t
|
||||
}
|
||||
}
|
||||
events, err := m.DB().ListC2EventsForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"events": events, "count": len(events)}, err)
|
||||
})
|
||||
}
|
||||
|
||||
func c2ToolAccess(ctx context.Context) database.RBACListAccess {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return database.RBACListAccess{Scope: database.RBACScopeAssigned}
|
||||
}
|
||||
return database.RBACListAccess{UserID: principal.UserID, Scope: principal.ScopeFor("c2:read")}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_profile — Malleable Profile 管理工具（新增）
|
||||
// ============================================================================
|
||||
|
||||
func registerC2ProfileTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2Profile,
|
||||
Description: `C2 Malleable Profile 管理（控制 beacon 通信伪装）。通过 action 参数选择操作：
|
||||
- list: 列出所有 Profile
|
||||
- get: 获取 Profile 详情（需 profile_id）
|
||||
- create: 创建 Profile（需 name，可选 user_agent/uris/request_headers/response_headers/body_template/jitter_min_ms/jitter_max_ms）
|
||||
- update: 更新 Profile（需 profile_id）
|
||||
- delete: 删除 Profile（需 profile_id）`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/delete", "enum": []string{"list", "get", "create", "update", "delete"}},
|
||||
"profile_id": map[string]interface{}{"type": "string", "description": "Profile ID（get/update/delete 需要）"},
|
||||
"name": map[string]interface{}{"type": "string", "description": "Profile 名称"},
|
||||
"user_agent": map[string]interface{}{"type": "string", "description": "User-Agent 字符串"},
|
||||
"uris": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "description": "beacon 请求的 URI 列表"},
|
||||
"request_headers": map[string]interface{}{"type": "object", "description": "自定义请求头"},
|
||||
"response_headers": map[string]interface{}{"type": "object", "description": "自定义响应头"},
|
||||
"body_template": map[string]interface{}{"type": "string", "description": "响应体模板"},
|
||||
"jitter_min_ms": map[string]interface{}{"type": "integer", "description": "最小抖动（毫秒）"},
|
||||
"jitter_max_ms": map[string]interface{}{"type": "integer", "description": "最大抖动（毫秒）"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
id := getString(params, "profile_id")
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
profiles, err := m.DB().ListC2Profiles()
|
||||
return makeC2Result(map[string]interface{}{"profiles": profiles, "count": len(profiles)}, err)
|
||||
|
||||
case "get":
|
||||
profile, err := m.DB().GetC2Profile(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if profile == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("profile not found"))
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"profile": profile}, nil)
|
||||
|
||||
case "create":
|
||||
profile := &database.C2Profile{
|
||||
ID: "p_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14],
|
||||
Name: getString(params, "name"),
|
||||
UserAgent: getString(params, "user_agent"),
|
||||
BodyTemplate: getString(params, "body_template"),
|
||||
JitterMinMS: int(getFloat64(params, "jitter_min_ms")),
|
||||
JitterMaxMS: int(getFloat64(params, "jitter_max_ms")),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if uris, ok := params["uris"]; ok {
|
||||
if arr, ok := uris.([]interface{}); ok {
|
||||
for _, u := range arr {
|
||||
if s, ok := u.(string); ok {
|
||||
profile.URIs = append(profile.URIs, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if rh, ok := params["request_headers"]; ok {
|
||||
if m, ok := rh.(map[string]interface{}); ok {
|
||||
profile.RequestHeaders = make(map[string]string)
|
||||
for k, v := range m {
|
||||
profile.RequestHeaders[k], _ = v.(string)
|
||||
}
|
||||
}
|
||||
}
|
||||
if rh, ok := params["response_headers"]; ok {
|
||||
if m, ok := rh.(map[string]interface{}); ok {
|
||||
profile.ResponseHeaders = make(map[string]string)
|
||||
for k, v := range m {
|
||||
profile.ResponseHeaders[k], _ = v.(string)
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := m.DB().CreateC2Profile(profile); err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"profile": profile}, nil)
|
||||
|
||||
case "update":
|
||||
profile, err := m.DB().GetC2Profile(id)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if profile == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("profile not found"))
|
||||
}
|
||||
if v := getString(params, "name"); v != "" {
|
||||
profile.Name = v
|
||||
}
|
||||
if v := getString(params, "user_agent"); v != "" {
|
||||
profile.UserAgent = v
|
||||
}
|
||||
if v := getString(params, "body_template"); v != "" {
|
||||
profile.BodyTemplate = v
|
||||
}
|
||||
if v := int(getFloat64(params, "jitter_min_ms")); v > 0 {
|
||||
profile.JitterMinMS = v
|
||||
}
|
||||
if v := int(getFloat64(params, "jitter_max_ms")); v > 0 {
|
||||
profile.JitterMaxMS = v
|
||||
}
|
||||
if uris, ok := params["uris"]; ok {
|
||||
if arr, ok := uris.([]interface{}); ok {
|
||||
profile.URIs = nil
|
||||
for _, u := range arr {
|
||||
if s, ok := u.(string); ok {
|
||||
profile.URIs = append(profile.URIs, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if rh, ok := params["request_headers"]; ok {
|
||||
if mp, ok := rh.(map[string]interface{}); ok {
|
||||
profile.RequestHeaders = make(map[string]string)
|
||||
for k, v := range mp {
|
||||
profile.RequestHeaders[k], _ = v.(string)
|
||||
}
|
||||
}
|
||||
}
|
||||
if rh, ok := params["response_headers"]; ok {
|
||||
if mp, ok := rh.(map[string]interface{}); ok {
|
||||
profile.ResponseHeaders = make(map[string]string)
|
||||
for k, v := range mp {
|
||||
profile.ResponseHeaders[k], _ = v.(string)
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := m.DB().UpdateC2Profile(profile); err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{"profile": profile}, nil)
|
||||
|
||||
case "delete":
|
||||
err := m.DB().DeleteC2Profile(id)
|
||||
return makeC2Result(map[string]interface{}{"deleted": err == nil}, err)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_file — 文件管理工具（新增）
|
||||
// ============================================================================
|
||||
|
||||
func registerC2FileTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
s.RegisterTool(mcp.Tool{
|
||||
Name: builtin.ToolC2File,
|
||||
Description: `C2 文件管理。通过 action 参数选择操作：
|
||||
- list: 列出会话的文件传输记录（需 session_id）
|
||||
- get_result: 获取任务结果文件路径（截图等，需 task_id）`,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get_result", "enum": []string{"list", "get_result"}},
|
||||
"session_id": map[string]interface{}{"type": "string", "description": "会话 ID（list 需要）"},
|
||||
"task_id": map[string]interface{}{"type": "string", "description": "任务 ID（get_result 需要）"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
}, func(ctx context.Context, params map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
action := getString(params, "action")
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
sessionID := getString(params, "session_id")
|
||||
if sessionID == "" {
|
||||
return makeC2Result(nil, fmt.Errorf("session_id required"))
|
||||
}
|
||||
files, err := m.DB().ListC2FilesBySession(sessionID)
|
||||
return makeC2Result(map[string]interface{}{"files": files, "count": len(files)}, err)
|
||||
|
||||
case "get_result":
|
||||
taskID := getString(params, "task_id")
|
||||
task, err := m.DB().GetC2Task(taskID)
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if task == nil {
|
||||
return makeC2Result(nil, fmt.Errorf("task not found"))
|
||||
}
|
||||
if task.ResultBlobPath == "" {
|
||||
return makeC2Result(map[string]interface{}{"has_file": false, "task_id": taskID}, nil)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{
|
||||
"has_file": true,
|
||||
"task_id": taskID,
|
||||
"file_path": task.ResultBlobPath,
|
||||
}, nil)
|
||||
|
||||
default:
|
||||
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 工具函数
|
||||
// ============================================================================
|
||||
|
||||
func getString(params map[string]interface{}, key string) string {
|
||||
if v, ok := params[key]; ok {
|
||||
if s, ok := v.(string); ok {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func getFloat64(params map[string]interface{}, key string) float64 {
|
||||
if v, ok := params[key]; ok {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n
|
||||
case int:
|
||||
return float64(n)
|
||||
case string:
|
||||
if f, err := strconv.ParseFloat(n, 64); err == nil {
|
||||
return f
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -1,69 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestC2ListenerCreateInheritsConversationProject(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "c2-tools.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
user, err := db.CreateRBACUser("c2-agent", "C2 Agent", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
project, err := db.CreateProject(&database.Project{Name: "engagement"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
conversation, err := db.CreateConversation("project chat", database.ConversationCreateMeta{ProjectID: project.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{
|
||||
"c2:read": true, "c2:write": true,
|
||||
})
|
||||
ctx := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), conversation.ID), principal)
|
||||
server := mcp.NewServer(zap.NewNop())
|
||||
server.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||
registerC2Tools(server, c2.NewManager(db, zap.NewNop(), t.TempDir()), zap.NewNop(), 8080)
|
||||
|
||||
result, _, err := server.CallTool(ctx, builtin.ToolC2Listener, map[string]interface{}{
|
||||
"action": "create",
|
||||
"name": "tcp-reverse-2222",
|
||||
"type": "tcp_reverse",
|
||||
"bind_host": "0.0.0.0",
|
||||
"bind_port": 2222,
|
||||
})
|
||||
if err != nil || result == nil || result.IsError {
|
||||
t.Fatalf("create listener result=%#v err=%v text=%q", result, err, toolResultText(result))
|
||||
}
|
||||
|
||||
listeners, err := db.ListC2Listeners()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(listeners) != 1 {
|
||||
t.Fatalf("listener count=%d, want 1", len(listeners))
|
||||
}
|
||||
if listeners[0].ProjectID != project.ID {
|
||||
t.Fatalf("listener project_id=%q, want %q", listeners[0].ProjectID, project.ID)
|
||||
}
|
||||
}
|
||||
@@ -1,105 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestCORSMiddlewareAllowsSameOriginAndRejectsForeignOrigin(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(corsMiddleware(nil))
|
||||
router.GET("/test", func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
||||
|
||||
same := httptest.NewRequest(http.MethodGet, "http://app.example/test", nil)
|
||||
same.Host = "app.example"
|
||||
same.Header.Set("Origin", "http://app.example")
|
||||
sameW := httptest.NewRecorder()
|
||||
router.ServeHTTP(sameW, same)
|
||||
if sameW.Code != http.StatusNoContent || sameW.Header().Get("Access-Control-Allow-Origin") != "http://app.example" {
|
||||
t.Fatalf("same-origin response = %d, allow-origin=%q", sameW.Code, sameW.Header().Get("Access-Control-Allow-Origin"))
|
||||
}
|
||||
|
||||
foreign := httptest.NewRequest(http.MethodGet, "http://app.example/test", nil)
|
||||
foreign.Host = "app.example"
|
||||
foreign.Header.Set("Origin", "https://evil.example")
|
||||
foreignW := httptest.NewRecorder()
|
||||
router.ServeHTTP(foreignW, foreign)
|
||||
if foreignW.Code != http.StatusForbidden {
|
||||
t.Fatalf("foreign-origin response = %d, want %d", foreignW.Code, http.StatusForbidden)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCORSMiddlewareAllowsBrowserExtensionWithoutConfiguration(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(corsMiddleware(nil))
|
||||
router.POST("/api/auth/login", func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodOptions, "https://server.example/api/auth/login", nil)
|
||||
req.Host = "server.example"
|
||||
req.Header.Set("Origin", "chrome-extension://abcdefghijklmnopabcdefghijklmnop")
|
||||
req.Header.Set("Access-Control-Request-Method", http.MethodPost)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusNoContent {
|
||||
t.Fatalf("preflight response = %d, want %d", w.Code, http.StatusNoContent)
|
||||
}
|
||||
if got := w.Header().Get("Access-Control-Allow-Origin"); got != "chrome-extension://abcdefghijklmnopabcdefghijklmnop" {
|
||||
t.Fatalf("allow-origin = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCORSMiddlewareRejectsInvalidExtensionOrigins(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, origin := range []string{
|
||||
"chrome-extension://too-short",
|
||||
"chrome-extension://qrstuvwxyzabcdefqrstuvwxyzabcdef",
|
||||
"chrome-extension://abcdefghijklmnopabcdefghijklmnop:8443",
|
||||
"moz-extension://abcdefghijklmnopabcdefghijklmnop",
|
||||
} {
|
||||
t.Run(origin, func(t *testing.T) {
|
||||
router := gin.New()
|
||||
router.Use(corsMiddleware(nil))
|
||||
router.GET("/test", func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "https://server.example/test", nil)
|
||||
req.Host = "server.example"
|
||||
req.Header.Set("Origin", origin)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("response = %d, want %d", w.Code, http.StatusForbidden)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCORSMiddlewareRejectsUnsafeConfiguredEntries(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, configured := range []string{
|
||||
"*",
|
||||
"null",
|
||||
"https://trusted.example/extra",
|
||||
"https://trusted.example?trusted=true",
|
||||
} {
|
||||
t.Run(configured, func(t *testing.T) {
|
||||
router := gin.New()
|
||||
router.Use(corsMiddleware([]string{configured}))
|
||||
router.GET("/test", func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "https://server.example/test", nil)
|
||||
req.Host = "server.example"
|
||||
req.Header.Set("Origin", "https://trusted.example")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("response = %d, want %d", w.Code, http.StatusForbidden)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,213 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// peekedConn 在已预读首字节后仍将连接交给 net/http 或 crypto/tls。
|
||||
type peekedConn struct {
|
||||
net.Conn
|
||||
r *bufio.Reader
|
||||
}
|
||||
|
||||
func (c *peekedConn) Read(p []byte) (int, error) {
|
||||
return c.r.Read(p)
|
||||
}
|
||||
|
||||
// oneConnListener 供 http.Server.Serve 处理单条 TCP 连接（含 keep-alive）。
|
||||
type oneConnListener struct {
|
||||
conn net.Conn
|
||||
addr net.Addr
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (l *oneConnListener) Accept() (net.Conn, error) {
|
||||
var c net.Conn
|
||||
l.once.Do(func() {
|
||||
c = l.conn
|
||||
l.conn = nil
|
||||
})
|
||||
if c == nil {
|
||||
return nil, net.ErrClosed
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (l *oneConnListener) Close() error { return nil }
|
||||
func (l *oneConnListener) Addr() net.Addr { return l.addr }
|
||||
|
||||
// httpServerForTLSConn 从已有 Server 复制可服务字段，用于已握手 TLS 连接上的 HTTP 服务。
|
||||
// 不能复制整个 http.Server（内含 atomic/noCopy 字段）。
|
||||
func httpServerForTLSConn(src *http.Server) *http.Server {
|
||||
return &http.Server{
|
||||
Handler: src.Handler,
|
||||
DisableGeneralOptionsHandler: src.DisableGeneralOptionsHandler,
|
||||
ReadTimeout: src.ReadTimeout,
|
||||
ReadHeaderTimeout: src.ReadHeaderTimeout,
|
||||
WriteTimeout: src.WriteTimeout,
|
||||
IdleTimeout: src.IdleTimeout,
|
||||
MaxHeaderBytes: src.MaxHeaderBytes,
|
||||
ConnState: src.ConnState,
|
||||
ErrorLog: src.ErrorLog,
|
||||
BaseContext: src.BaseContext,
|
||||
ConnContext: src.ConnContext,
|
||||
}
|
||||
}
|
||||
|
||||
func isTLSHandshakeRecord(b byte) bool {
|
||||
return b == 0x16
|
||||
}
|
||||
|
||||
func newHTTPToHTTPSRedirectHandler(httpsPort int) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
host := r.Host
|
||||
if h, _, err := net.SplitHostPort(host); err == nil {
|
||||
host = h
|
||||
}
|
||||
var target string
|
||||
if httpsPort == 443 {
|
||||
target = fmt.Sprintf("https://%s%s", host, r.URL.RequestURI())
|
||||
} else {
|
||||
target = fmt.Sprintf("https://%s:%d%s", host, httpsPort, r.URL.RequestURI())
|
||||
}
|
||||
http.Redirect(w, r, target, http.StatusPermanentRedirect)
|
||||
})
|
||||
}
|
||||
|
||||
func portFromListenAddr(addr string) int {
|
||||
_, portStr, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return 443
|
||||
}
|
||||
p, err := strconv.Atoi(portStr)
|
||||
if err != nil || p <= 0 {
|
||||
return 443
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func ensureMainTLSConfigCerts(mode mainTLSMode, tlsConf *tls.Config, certFile, keyFile string) (*tls.Config, error) {
|
||||
if mode != mainTLSFromFiles {
|
||||
return tlsConf, nil
|
||||
}
|
||||
if tlsConf == nil {
|
||||
tlsConf = &tls.Config{MinVersion: tls.VersionTLS12}
|
||||
}
|
||||
if len(tlsConf.Certificates) > 0 {
|
||||
return tlsConf, nil
|
||||
}
|
||||
cert, err := tls.LoadX509KeyPair(certFile, keyFile)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tlsConf.Certificates = []tls.Certificate{cert}
|
||||
return tlsConf, nil
|
||||
}
|
||||
|
||||
type mainServerMux struct {
|
||||
ln net.Listener
|
||||
httpsSrv *http.Server
|
||||
redirectSrv *http.Server
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
func newMainServerMux(ln net.Listener, httpsSrv *http.Server, httpsPort int, logger *zap.Logger) *mainServerMux {
|
||||
return &mainServerMux{
|
||||
ln: ln,
|
||||
httpsSrv: httpsSrv,
|
||||
redirectSrv: &http.Server{Handler: newHTTPToHTTPSRedirectHandler(httpsPort), ReadHeaderTimeout: 10 * time.Second},
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mainServerMux) Serve() error {
|
||||
for {
|
||||
conn, err := m.ln.Accept()
|
||||
if err != nil {
|
||||
if errors.Is(err, net.ErrClosed) {
|
||||
return http.ErrServerClosed
|
||||
}
|
||||
return err
|
||||
}
|
||||
go m.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mainServerMux) handleConn(raw net.Conn) {
|
||||
if err := raw.SetReadDeadline(time.Now().Add(10 * time.Second)); err != nil {
|
||||
_ = raw.Close()
|
||||
return
|
||||
}
|
||||
br := bufio.NewReader(raw)
|
||||
b, err := br.Peek(1)
|
||||
if err != nil {
|
||||
_ = raw.Close()
|
||||
return
|
||||
}
|
||||
_ = raw.SetReadDeadline(time.Time{})
|
||||
|
||||
pc := &peekedConn{Conn: raw, r: br}
|
||||
ocl := &oneConnListener{conn: pc, addr: raw.LocalAddr()}
|
||||
|
||||
if isTLSHandshakeRecord(b[0]) {
|
||||
m.serveHTTPS(pc, raw.LocalAddr())
|
||||
return
|
||||
}
|
||||
if err := m.redirectSrv.Serve(ocl); err != nil && !errors.Is(err, net.ErrClosed) && !errors.Is(err, http.ErrServerClosed) {
|
||||
m.logger.Debug("HTTP 重定向连接处理结束", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// serveHTTPS 在已嗅探为 TLS 的连接上完成握手，再按 ALPN 走 HTTP/2 或 HTTP/1.1。
|
||||
// 不能对同一 http.Server 并发调用 Serve(TLSConfig!=nil)，否则握手/ALPN 会异常（浏览器 ERR_SSL_PROTOCOL_ERROR）。
|
||||
func (m *mainServerMux) serveHTTPS(pc *peekedConn, localAddr net.Addr) {
|
||||
tlsConn := tls.Server(pc, m.httpsSrv.TLSConfig)
|
||||
handCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
if err := tlsConn.HandshakeContext(handCtx); err != nil {
|
||||
m.logger.Debug("TLS 握手失败", zap.Error(err))
|
||||
_ = pc.Close()
|
||||
return
|
||||
}
|
||||
|
||||
srv := m.httpsSrv
|
||||
if srv.TLSNextProto != nil {
|
||||
proto := tlsConn.ConnectionState().NegotiatedProtocol
|
||||
if fn := srv.TLSNextProto[proto]; fn != nil {
|
||||
fn(srv, tlsConn, srv.Handler)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
plain := httpServerForTLSConn(srv)
|
||||
ocl := &oneConnListener{conn: tlsConn, addr: localAddr}
|
||||
if err := plain.Serve(ocl); err != nil && !errors.Is(err, net.ErrClosed) && !errors.Is(err, http.ErrServerClosed) {
|
||||
m.logger.Debug("HTTPS 连接处理结束", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mainServerMux) Shutdown(ctx context.Context) error {
|
||||
_ = m.ln.Close()
|
||||
var err1, err2 error
|
||||
if m.httpsSrv != nil {
|
||||
err1 = m.httpsSrv.Shutdown(ctx)
|
||||
}
|
||||
if m.redirectSrv != nil {
|
||||
err2 = m.redirectSrv.Shutdown(ctx)
|
||||
}
|
||||
if err1 != nil {
|
||||
return err1
|
||||
}
|
||||
return err2
|
||||
}
|
||||
@@ -1,150 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
)
|
||||
|
||||
func TestNewHTTPToHTTPSRedirectHandler(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
httpsPort int
|
||||
host string
|
||||
uri string
|
||||
wantTarget string
|
||||
}{
|
||||
{
|
||||
name: "non standard port",
|
||||
httpsPort: 8080,
|
||||
host: "127.0.0.1:8080",
|
||||
uri: "/login?next=/",
|
||||
wantTarget: "https://127.0.0.1:8080/login?next=/",
|
||||
},
|
||||
{
|
||||
name: "standard port",
|
||||
httpsPort: 443,
|
||||
host: "example.com:80",
|
||||
uri: "/",
|
||||
wantTarget: "https://example.com/",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
h := newHTTPToHTTPSRedirectHandler(tt.httpsPort)
|
||||
req := httptest.NewRequest(http.MethodGet, "http://"+tt.host+tt.uri, nil)
|
||||
req.Host = tt.host
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusPermanentRedirect {
|
||||
t.Fatalf("status = %d, want %d", rec.Code, http.StatusPermanentRedirect)
|
||||
}
|
||||
if got := rec.Header().Get("Location"); got != tt.wantTarget {
|
||||
t.Fatalf("Location = %q, want %q", got, tt.wantTarget)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsTLSHandshakeRecord(t *testing.T) {
|
||||
t.Parallel()
|
||||
if !isTLSHandshakeRecord(0x16) {
|
||||
t.Fatal("expected TLS handshake record")
|
||||
}
|
||||
if isTLSHandshakeRecord('G') {
|
||||
t.Fatal("GET should not be TLS")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerHTTPRedirectEnabled(t *testing.T) {
|
||||
t.Parallel()
|
||||
disabled := false
|
||||
enabled := true
|
||||
if config.ServerHTTPRedirectEnabled(nil) {
|
||||
t.Fatal("nil config should disable redirect")
|
||||
}
|
||||
if !config.ServerHTTPRedirectEnabled(&config.ServerConfig{TLSEnabled: true}) {
|
||||
t.Fatal("HTTPS without explicit flag should enable redirect")
|
||||
}
|
||||
if config.ServerHTTPRedirectEnabled(&config.ServerConfig{TLSEnabled: true, TLSHTTPRedirect: &disabled}) {
|
||||
t.Fatal("explicit false should disable redirect")
|
||||
}
|
||||
if !config.ServerHTTPRedirectEnabled(&config.ServerConfig{TLSEnabled: true, TLSHTTPRedirect: &enabled}) {
|
||||
t.Fatal("explicit true should enable redirect")
|
||||
}
|
||||
if config.ServerHTTPRedirectEnabled(&config.ServerConfig{}) {
|
||||
t.Fatal("plain HTTP should not redirect")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMainServerMuxHTTPRedirectAndHTTPS(t *testing.T) {
|
||||
cert, err := generateMainServerSelfSignedCert()
|
||||
if err != nil {
|
||||
t.Fatalf("generate cert: %v", err)
|
||||
}
|
||||
handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = io.WriteString(w, "ok")
|
||||
})
|
||||
srv := &http.Server{Handler: handler, TLSConfig: &tls.Config{
|
||||
MinVersion: tls.VersionTLS12,
|
||||
Certificates: []tls.Certificate{cert},
|
||||
}}
|
||||
if err := http2.ConfigureServer(srv, &http2.Server{}); err != nil {
|
||||
t.Fatalf("configure http2: %v", err)
|
||||
}
|
||||
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("listen: %v", err)
|
||||
}
|
||||
defer ln.Close()
|
||||
|
||||
mux := newMainServerMux(ln, srv, portFromListenAddr(ln.Addr().String()), nil)
|
||||
go func() { _ = mux.Serve() }()
|
||||
|
||||
client := &http.Client{
|
||||
Transport: &http.Transport{
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true, MinVersion: tls.VersionTLS12},
|
||||
},
|
||||
CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
}
|
||||
addr := ln.Addr().String()
|
||||
|
||||
httpResp, err := client.Get("http://" + addr + "/")
|
||||
if err != nil {
|
||||
t.Fatalf("http get: %v", err)
|
||||
}
|
||||
_ = httpResp.Body.Close()
|
||||
if httpResp.StatusCode != http.StatusPermanentRedirect {
|
||||
t.Fatalf("http status = %d, want %d", httpResp.StatusCode, http.StatusPermanentRedirect)
|
||||
}
|
||||
if got := httpResp.Header.Get("Location"); got != "https://127.0.0.1:"+strconv.Itoa(portFromListenAddr(addr))+"/" {
|
||||
t.Fatalf("Location = %q", got)
|
||||
}
|
||||
|
||||
httpsResp, err := client.Get("https://" + addr + "/")
|
||||
if err != nil {
|
||||
t.Fatalf("https get: %v", err)
|
||||
}
|
||||
defer httpsResp.Body.Close()
|
||||
if httpsResp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("https status = %d, want %d", httpsResp.StatusCode, http.StatusOK)
|
||||
}
|
||||
body, _ := io.ReadAll(httpsResp.Body)
|
||||
if string(body) != "ok" {
|
||||
t.Fatalf("body = %q, want ok", body)
|
||||
}
|
||||
}
|
||||
@@ -1,86 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"crypto/x509/pkix"
|
||||
"encoding/pem"
|
||||
"fmt"
|
||||
"math/big"
|
||||
"net"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
)
|
||||
|
||||
// mainTLSMode 主 Web 服务 TLS 启动方式。
|
||||
type mainTLSMode int
|
||||
|
||||
const (
|
||||
mainTLSOff mainTLSMode = iota
|
||||
mainTLSFromFiles
|
||||
mainTLSInMemorySelfSigned
|
||||
)
|
||||
|
||||
// prepareMainServerTLS 根据 server 配置决定主站是否启用 HTTPS（及 HTTP/2 协商）。
|
||||
// fromFiles：使用 tls_cert_path + tls_key_path，由 http.Server.ListenAndServeTLS 加载 PEM。
|
||||
// inMemory：tls_auto_self_sign 生成的自签证书，仅用于本地/测试。
|
||||
func prepareMainServerTLS(cfg *config.ServerConfig) (mode mainTLSMode, tlsConf *tls.Config, certFile, keyFile string, err error) {
|
||||
if cfg == nil || !config.MainWebUIUsesHTTPS(cfg) {
|
||||
return mainTLSOff, nil, "", "", nil
|
||||
}
|
||||
certFile = strings.TrimSpace(cfg.TLSCertPath)
|
||||
keyFile = strings.TrimSpace(cfg.TLSKeyPath)
|
||||
if certFile != "" && keyFile != "" {
|
||||
// 证书由 ListenAndServeTLS 从文件加载；此处仅提供最小 TLS 配置供 http2.ConfigureServer 合并 ALPN。
|
||||
return mainTLSFromFiles, &tls.Config{MinVersion: tls.VersionTLS12}, certFile, keyFile, nil
|
||||
}
|
||||
if cfg.TLSAutoSelfSign {
|
||||
cert, genErr := generateMainServerSelfSignedCert()
|
||||
if genErr != nil {
|
||||
return mainTLSOff, nil, "", "", fmt.Errorf("生成自签 TLS 证书: %w", genErr)
|
||||
}
|
||||
tlsConf = &tls.Config{
|
||||
MinVersion: tls.VersionTLS12,
|
||||
Certificates: []tls.Certificate{cert},
|
||||
}
|
||||
return mainTLSInMemorySelfSigned, tlsConf, "", "", nil
|
||||
}
|
||||
return mainTLSOff, nil, "", "", fmt.Errorf("server: 已启用 TLS（tls_enabled / tls_auto_self_sign / 证书路径），请设置 tls_cert_path 与 tls_key_path，或将 tls_auto_self_sign 设为 true（仅测试环境）")
|
||||
}
|
||||
|
||||
func generateMainServerSelfSignedCert() (tls.Certificate, error) {
|
||||
priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
serial, err := rand.Int(rand.Reader, big.NewInt(1<<62))
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
tmpl := &x509.Certificate{
|
||||
SerialNumber: serial,
|
||||
Subject: pkix.Name{CommonName: "CyberStrikeAI"},
|
||||
NotBefore: time.Now().Add(-1 * time.Hour),
|
||||
NotAfter: time.Now().Add(365 * 24 * time.Hour),
|
||||
KeyUsage: x509.KeyUsageDigitalSignature,
|
||||
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
|
||||
IPAddresses: []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")},
|
||||
DNSNames: []string{"localhost"},
|
||||
}
|
||||
der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &priv.PublicKey, priv)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
keyDER, err := x509.MarshalECPrivateKey(priv)
|
||||
if err != nil {
|
||||
return tls.Certificate{}, err
|
||||
}
|
||||
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
|
||||
keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
|
||||
return tls.X509KeyPair(certPEM, keyPEM)
|
||||
}
|
||||
@@ -1,407 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
)
|
||||
|
||||
func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string]interface{}) error {
|
||||
return func(ctx context.Context, toolName string, args map[string]interface{}) error {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return fmt.Errorf("missing authenticated principal")
|
||||
}
|
||||
require := func(permission string) error {
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
resource := func(permission, resourceType, argument string) error {
|
||||
if err := require(permission); err != nil {
|
||||
return err
|
||||
}
|
||||
id := mcpAuthorizationString(args, argument)
|
||||
if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||
}
|
||||
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, id); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
toolExecutionResource := func(permission string) error {
|
||||
if err := require(permission); err != nil {
|
||||
return err
|
||||
}
|
||||
id := mcpAuthorizationString(args, "execution_id")
|
||||
if id == "" || db == nil || !db.UserCanAccessToolExecution(principal.UserID, principal.ScopeFor(permission), id) {
|
||||
return fmt.Errorf("no access to tool execution %s", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
switch toolName {
|
||||
case builtin.ToolWebshellExec, builtin.ToolWebshellFileWrite:
|
||||
return resource("webshell:write", "webshell", "connection_id")
|
||||
case builtin.ToolWebshellFileList, builtin.ToolWebshellFileRead:
|
||||
return resource("webshell:read", "webshell", "connection_id")
|
||||
case builtin.ToolManageWebshellList:
|
||||
return require("webshell:read")
|
||||
case builtin.ToolManageWebshellAdd:
|
||||
return require("webshell:write")
|
||||
case builtin.ToolManageWebshellUpdate, builtin.ToolManageWebshellTest:
|
||||
return resource("webshell:write", "webshell", "connection_id")
|
||||
case builtin.ToolManageWebshellDelete:
|
||||
return resource("webshell:delete", "webshell", "connection_id")
|
||||
case builtin.ToolRecordVulnerability:
|
||||
if err := require("vulnerability:write"); err != nil {
|
||||
return err
|
||||
}
|
||||
conversationID := mcpAuthorizationString(args, "conversation_id")
|
||||
if conversationID == "" {
|
||||
conversationID = mcpAuthorizationConversationID(ctx)
|
||||
}
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("vulnerability:write"), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
return nil
|
||||
case builtin.ToolListVulnerabilities:
|
||||
if err := require("vulnerability:read"); err != nil {
|
||||
return err
|
||||
}
|
||||
conversationID := mcpAuthorizationConversationID(ctx)
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("vulnerability:read"), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
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:
|
||||
return authorizeProjectTool(ctx, principal, db, "project:read")
|
||||
case builtin.ToolListKnowledgeRiskTypes, builtin.ToolSearchKnowledgeBase:
|
||||
return require("knowledge:read")
|
||||
case builtin.ToolAnalyzeImage:
|
||||
return require("agent:execute")
|
||||
case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution:
|
||||
return toolExecutionResource("monitor:read")
|
||||
case builtin.ToolCancelToolExecution:
|
||||
return toolExecutionResource("monitor:write")
|
||||
case builtin.ToolBatchTaskList:
|
||||
return require("tasks:read")
|
||||
case builtin.ToolBatchTaskGet:
|
||||
return resource("tasks:read", "batch_task", "queue_id")
|
||||
case builtin.ToolBatchTaskCreate:
|
||||
if err := require("tasks:write"); err != nil {
|
||||
return err
|
||||
}
|
||||
if projectID := mcpAuthorizationString(args, "project_id"); projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("tasks:write"), "project", projectID)) {
|
||||
return fmt.Errorf("no access to project %s", projectID)
|
||||
}
|
||||
return nil
|
||||
case builtin.ToolBatchTaskDelete, builtin.ToolBatchTaskRemove:
|
||||
return resource("tasks:delete", "batch_task", "queue_id")
|
||||
case builtin.ToolBatchTaskStart, builtin.ToolBatchTaskRerun, builtin.ToolBatchTaskPause,
|
||||
builtin.ToolBatchTaskUpdateMetadata, builtin.ToolBatchTaskUpdateSchedule,
|
||||
builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate:
|
||||
return resource("tasks:write", "batch_task", "queue_id")
|
||||
case builtin.ToolC2Listener:
|
||||
return authorizeC2Action(ctx, principal, db, args, "c2_listener", "listener_id")
|
||||
case builtin.ToolC2Session, builtin.ToolC2Task, builtin.ToolC2File:
|
||||
if toolName == builtin.ToolC2File && mcpAuthorizationString(args, "action") == "get_result" {
|
||||
return authorizeC2Action(ctx, principal, db, args, "c2_task", "task_id")
|
||||
}
|
||||
return authorizeC2Action(ctx, principal, db, args, "c2_session", "session_id")
|
||||
case builtin.ToolC2TaskManage:
|
||||
return authorizeC2Action(ctx, principal, db, args, "c2_task", "task_id")
|
||||
case builtin.ToolC2Payload:
|
||||
return resource("c2:write", "c2_listener", "listener_id")
|
||||
case builtin.ToolC2Event:
|
||||
if id := mcpAuthorizationString(args, "session_id"); id != "" {
|
||||
return resource("c2:read", "c2_session", "session_id")
|
||||
}
|
||||
if id := mcpAuthorizationString(args, "task_id"); id != "" {
|
||||
return resource("c2:read", "c2_task", "task_id")
|
||||
}
|
||||
if filter := mcpEffectiveProjectFilter(ctx, db); filter != "" {
|
||||
return require("c2:read")
|
||||
}
|
||||
if principal.ScopeFor("c2:read") != database.RBACScopeAll {
|
||||
return fmt.Errorf("unfiltered C2 event list requires global scope")
|
||||
}
|
||||
return require("c2:read")
|
||||
case builtin.ToolC2Profile:
|
||||
// Profiles are process-global and do not yet have an owner. Writes are
|
||||
// therefore reserved for global scope; reads require c2:read.
|
||||
if mcpAuthorizationString(args, "action") == "list" || mcpAuthorizationString(args, "action") == "get" {
|
||||
return require("c2:read")
|
||||
}
|
||||
permission := "c2:write"
|
||||
if mcpAuthorizationString(args, "action") == "delete" {
|
||||
permission = "c2:delete"
|
||||
}
|
||||
if principal.ScopeFor(permission) != database.RBACScopeAll {
|
||||
return fmt.Errorf("C2 profile mutation requires global scope")
|
||||
}
|
||||
if mcpAuthorizationString(args, "action") == "delete" {
|
||||
return require("c2:delete")
|
||||
}
|
||||
return require("c2:write")
|
||||
default:
|
||||
if builtin.IsBuiltinTool(toolName) {
|
||||
return fmt.Errorf("no authorization policy registered for builtin tool %s", toolName)
|
||||
}
|
||||
if principal.HasPermission("agent:local-execute") {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("missing agent:local-execute")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func externalMCPToolAuthorizer() func(context.Context, string, map[string]interface{}) error {
|
||||
return func(ctx context.Context, toolName string, _ map[string]interface{}) error {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return fmt.Errorf("missing authenticated principal")
|
||||
}
|
||||
if !principal.HasPermission("mcp:external:execute") {
|
||||
return fmt.Errorf("missing permission mcp:external:execute")
|
||||
}
|
||||
if principal.ScopeFor("mcp:external:execute") != database.RBACScopeAll {
|
||||
return fmt.Errorf("external MCP invocation requires global scope")
|
||||
}
|
||||
if strings.TrimSpace(toolName) == "" {
|
||||
return fmt.Errorf("missing external tool name")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func authorizeC2Action(ctx context.Context, principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error {
|
||||
action := mcpAuthorizationString(args, "action")
|
||||
permission := "c2:write"
|
||||
if action == "list" || action == "get" || action == "get_result" || action == "wait" {
|
||||
permission = "c2:read"
|
||||
} else if action == "delete" || action == "delete_batch" {
|
||||
permission = "c2:delete"
|
||||
}
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
id := mcpAuthorizationString(args, argument)
|
||||
if action == "delete_batch" {
|
||||
ids := mcpAuthorizationStrings(args, argument+"s")
|
||||
if len(ids) == 0 {
|
||||
return fmt.Errorf("missing resource identifiers %ss", argument)
|
||||
}
|
||||
for _, candidate := range ids {
|
||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, candidate)
|
||||
}
|
||||
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, candidate); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if id == "" {
|
||||
if action == "create" {
|
||||
projectID := mcpAuthorizationString(args, "project_id")
|
||||
if projectID == "" {
|
||||
projectID = mcpEffectiveProjectFilter(ctx, db)
|
||||
if projectID == database.ProjectFilterUnbound {
|
||||
projectID = ""
|
||||
}
|
||||
}
|
||||
if projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID)) {
|
||||
return fmt.Errorf("no access to project %s", projectID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if action == "list" {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("missing resource identifier %s", argument)
|
||||
}
|
||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||
}
|
||||
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, id); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func authorizeMCPProjectResourceBoundary(ctx context.Context, db *database.DB, resourceType, resourceID string) error {
|
||||
filter := mcpEffectiveProjectFilter(ctx, db)
|
||||
if filter == "" || db == nil {
|
||||
return nil
|
||||
}
|
||||
projectID, ok, err := mcpResourceProjectID(db, resourceType, resourceID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if filter == database.ProjectFilterUnbound {
|
||||
if projectID != "" {
|
||||
return fmt.Errorf("resource %s %s belongs to project %s, current conversation is unbound", resourceType, resourceID, projectID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if projectID != filter {
|
||||
if projectID == "" {
|
||||
return fmt.Errorf("resource %s %s is unbound, current conversation project is %s", resourceType, resourceID, filter)
|
||||
}
|
||||
return fmt.Errorf("resource %s %s belongs to project %s, current conversation project is %s", resourceType, resourceID, projectID, filter)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mcpResourceProjectID(db *database.DB, resourceType, resourceID string) (string, bool, error) {
|
||||
switch resourceType {
|
||||
case "webshell":
|
||||
conn, err := db.GetWebshellConnection(resourceID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if conn == nil {
|
||||
return "", true, fmt.Errorf("webshell not found")
|
||||
}
|
||||
return strings.TrimSpace(conn.ProjectID), true, nil
|
||||
case "c2_listener":
|
||||
listener, err := db.GetC2Listener(resourceID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if listener == nil {
|
||||
return "", true, fmt.Errorf("listener not found")
|
||||
}
|
||||
return strings.TrimSpace(listener.ProjectID), true, nil
|
||||
case "c2_session":
|
||||
session, err := db.GetC2Session(resourceID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if session == nil {
|
||||
return "", true, fmt.Errorf("session not found")
|
||||
}
|
||||
return mcpResourceProjectID(db, "c2_listener", session.ListenerID)
|
||||
case "c2_task":
|
||||
task, err := db.GetC2Task(resourceID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if task == nil {
|
||||
return "", true, fmt.Errorf("task not found")
|
||||
}
|
||||
return mcpResourceProjectIDFromC2Session(db, task.SessionID)
|
||||
default:
|
||||
return "", false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func mcpResourceProjectIDFromC2Session(db *database.DB, sessionID string) (string, bool, error) {
|
||||
session, err := db.GetC2Session(sessionID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if session == nil {
|
||||
return "", true, fmt.Errorf("session not found")
|
||||
}
|
||||
return mcpResourceProjectID(db, "c2_listener", session.ListenerID)
|
||||
}
|
||||
|
||||
func mcpAuthorizationStrings(args map[string]interface{}, key string) []string {
|
||||
values := []string{}
|
||||
switch raw := args[key].(type) {
|
||||
case []string:
|
||||
for _, value := range raw {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
values = append(values, value)
|
||||
}
|
||||
}
|
||||
case []interface{}:
|
||||
for _, item := range raw {
|
||||
if value, ok := item.(string); ok {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
values = append(values, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
func authorizeProjectTool(ctx context.Context, principal authctx.Principal, db *database.DB, permission string) error {
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
conversationID := mcpAuthorizationConversationID(ctx)
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
projectID, err := db.GetConversationProjectID(conversationID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("no access to project: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(projectID) == "" {
|
||||
return fmt.Errorf("当前对话未绑定项目，无法使用项目黑板工具，请先在对话中选择项目或创建带项目的对话")
|
||||
}
|
||||
if !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) {
|
||||
return fmt.Errorf("no access to project %s", projectID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mcpAuthorizationConversationID(ctx context.Context) string {
|
||||
if id := strings.TrimSpace(agent.ConversationIDFromContext(ctx)); id != "" {
|
||||
return id
|
||||
}
|
||||
return strings.TrimSpace(mcp.MCPConversationIDFromContext(ctx))
|
||||
}
|
||||
|
||||
func mcpAuthorizationString(args map[string]interface{}, key string) string {
|
||||
value, _ := args[key].(string)
|
||||
return strings.TrimSpace(value)
|
||||
}
|
||||
@@ -1,261 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestMCPToolAuthorizerEnforcesPermissionAndResource(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-authz.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
user, err := db.CreateRBACUser("mcp-user", "MCP User", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, id := range []string{"ws_allowed", "ws_hidden"} {
|
||||
if err := db.CreateWebshellConnection(&database.WebShellConnection{ID: id, URL: "http://127.0.0.1/" + id, Type: "php", Method: "post", CmdParam: "cmd", CreatedAt: time.Now()}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "webshell", "ws_allowed"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"mcp:write": true, "webshell:write": true})
|
||||
ctx := authctx.WithPrincipal(context.Background(), principal)
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
if err := authorize(ctx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": "ws_allowed"}); err != nil {
|
||||
t.Fatalf("allowed resource denied: %v", err)
|
||||
}
|
||||
if err := authorize(ctx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": "ws_hidden"}); err == nil {
|
||||
t.Fatal("foreign webshell resource was allowed")
|
||||
}
|
||||
if err := authorize(ctx, builtin.ToolManageWebshellDelete, map[string]interface{}{"connection_id": "ws_allowed"}); err == nil {
|
||||
t.Fatal("delete without webshell:delete was allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPToolAuthorizerEnforcesConversationProjectBoundary(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-project-boundary.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
user, err := db.CreateRBACUser("boundary-user", "Boundary User", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
project, err := db.CreateProject(&database.Project{Name: "Project 123"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
projectConv, err := db.CreateConversation("project conversation", database.ConversationCreateMeta{ProjectID: project.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
unboundConv, err := db.CreateConversation("unbound conversation", database.ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
wsProject := database.WebShellConnection{ID: "ws_project", ProjectID: project.ID, URL: "http://127.0.0.1/project.php", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
wsUnbound := database.WebShellConnection{ID: "ws_unbound", URL: "http://127.0.0.1/unbound.php", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
if err := db.CreateWebshellConnection(&wsProject); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateWebshellConnection(&wsUnbound); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, id := range []string{wsProject.ID, wsUnbound.ID} {
|
||||
if err := db.AssignResourceToUser(user.ID, "webshell", id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
listener := &database.C2Listener{ID: "l_project", ProjectID: project.ID, Name: "project listener", Type: "tcp_reverse", BindHost: "127.0.0.1", BindPort: 5555, OwnerUserID: user.ID, CreatedAt: now}
|
||||
if err := db.CreateC2Listener(listener); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "c2_listener", listener.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
session := &database.C2Session{ID: "s_project", ListenerID: listener.ID, ImplantUUID: "implant-project", Status: "active", FirstSeenAt: now, LastCheckIn: now}
|
||||
if err := db.UpsertC2Session(session); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{
|
||||
"webshell:read": true, "webshell:write": true,
|
||||
"c2:read": true, "c2:write": true,
|
||||
})
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
unboundCtx := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), unboundConv.ID), principal)
|
||||
projectCtx := authctx.WithPrincipal(mcp.WithMCPProjectID(mcp.WithMCPConversationID(context.Background(), projectConv.ID), project.ID), principal)
|
||||
projectCtxFromConversationOnly := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), projectConv.ID), principal)
|
||||
|
||||
if err := authorize(unboundCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err == nil {
|
||||
t.Fatal("unbound conversation was allowed to use project-bound webshell")
|
||||
}
|
||||
if err := authorize(unboundCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsUnbound.ID}); err != nil {
|
||||
t.Fatalf("unbound webshell denied in unbound conversation: %v", err)
|
||||
}
|
||||
if err := authorize(projectCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err != nil {
|
||||
t.Fatalf("project webshell denied in project conversation: %v", err)
|
||||
}
|
||||
if err := authorize(projectCtxFromConversationOnly, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err != nil {
|
||||
t.Fatalf("project webshell denied when only conversation id is present: %v", err)
|
||||
}
|
||||
if err := authorize(projectCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsUnbound.ID}); err == nil {
|
||||
t.Fatal("project conversation was allowed to use unbound webshell by id")
|
||||
}
|
||||
if err := authorize(unboundCtx, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err == nil {
|
||||
t.Fatal("unbound conversation was allowed to use project-bound c2 session")
|
||||
}
|
||||
if err := authorize(projectCtx, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err != nil {
|
||||
t.Fatalf("project c2 session denied in project conversation: %v", err)
|
||||
}
|
||||
if err := authorize(projectCtxFromConversationOnly, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err != nil {
|
||||
t.Fatalf("project c2 session denied when only conversation id is present: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
permissions := map[string]bool{}
|
||||
for permission := range security.PermissionCatalog {
|
||||
permissions[permission] = true
|
||||
}
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("admin", "admin", database.RBACScopeAll, permissions))
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
args := map[string]interface{}{
|
||||
"action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x",
|
||||
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x", "execution_id": "x",
|
||||
}
|
||||
for _, toolName := range builtin.GetAllBuiltinTools() {
|
||||
err := authorize(ctx, toolName, args)
|
||||
if err != nil && strings.Contains(err.Error(), "no authorization policy registered") {
|
||||
t.Errorf("builtin tool %s has no explicit policy", toolName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPExecutionControlAuthorizationUsesExecutionScope(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-exec-authz.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
user, err := db.CreateRBACUser("exec-user", "Exec User", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||
ID: "exec-owned",
|
||||
ToolName: "lab::slow",
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
OwnerUserID: user.ID,
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||
ID: "exec-hidden",
|
||||
ToolName: "lab::slow",
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"monitor:read": true, "monitor:write": true})
|
||||
ctx := authctx.WithPrincipal(context.Background(), principal)
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
if err := authorize(ctx, builtin.ToolWaitToolExecution, map[string]interface{}{"execution_id": "exec-owned"}); err != nil {
|
||||
t.Fatalf("owned execution denied: %v", err)
|
||||
}
|
||||
if err := authorize(ctx, builtin.ToolCancelToolExecution, map[string]interface{}{"execution_id": "exec-hidden"}); err == nil {
|
||||
t.Fatal("foreign execution was allowed")
|
||||
}
|
||||
}
|
||||
|
||||
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}))
|
||||
if err := authorize(ctx, "server::tool", nil); err == nil {
|
||||
t.Fatal("agent:execute alone authorized an external MCP tool")
|
||||
}
|
||||
ctx = authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAll, map[string]bool{"mcp:external:execute": true}))
|
||||
if err := authorize(ctx, "server::tool", nil); err != nil {
|
||||
t.Fatalf("dedicated external MCP permission rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfiguredCommandToolRequiresLocalExecutePermission(t *testing.T) {
|
||||
authorize := mcpToolAuthorizer(nil)
|
||||
agentOnly := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
|
||||
if err := authorize(agentOnly, "nmap_scan", nil); err == nil {
|
||||
t.Fatal("agent:execute alone authorized a configured command tool")
|
||||
}
|
||||
local := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:local-execute": true}))
|
||||
if err := authorize(local, "nmap_scan", nil); err != nil {
|
||||
t.Fatalf("agent:local-execute rejected: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -1,59 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestStandaloneMCPPrefersUserRBACAndDisablesGlobalTokenByDefault(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-http-auth.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
auth := security.NewAuthManager(12)
|
||||
if _, err := auth.AttachRBACStore(db); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hash, err := security.HashPassword("admin-secret")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpdateRBACAdminPassword(hash); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
token, _, err := auth.Authenticate("admin", "admin-secret")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
server := mcp.NewServer(zap.NewNop())
|
||||
server.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||
a := &App{config: &config.Config{MCP: config.MCPConfig{AuthHeader: "X-MCP-Token", AuthHeaderValue: "static-secret"}}, auth: auth, mcpServer: server}
|
||||
body := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)
|
||||
|
||||
userReq := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body))
|
||||
userReq.Header.Set("Authorization", "Bearer "+token)
|
||||
userW := httptest.NewRecorder()
|
||||
a.mcpHandlerWithAuth(userW, userReq)
|
||||
if userW.Code != http.StatusOK {
|
||||
t.Fatalf("user bearer status = %d: %s", userW.Code, userW.Body.String())
|
||||
}
|
||||
|
||||
staticReq := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body))
|
||||
staticReq.Header.Set("X-MCP-Token", "static-secret")
|
||||
staticW := httptest.NewRecorder()
|
||||
a.mcpHandlerWithAuth(staticW, staticReq)
|
||||
if staticW.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("global static token status = %d, want 401", staticW.Code)
|
||||
}
|
||||
}
|
||||
@@ -1,26 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
)
|
||||
|
||||
func mcpEffectiveProjectFilter(ctx context.Context, db *database.DB) string {
|
||||
if projectID := strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx)); projectID != "" {
|
||||
return projectID
|
||||
}
|
||||
if conversationID := mcpAuthorizationConversationID(ctx); conversationID != "" {
|
||||
if db != nil {
|
||||
if projectID, err := db.GetConversationProjectID(conversationID); err == nil {
|
||||
if projectID = strings.TrimSpace(projectID); projectID != "" {
|
||||
return projectID
|
||||
}
|
||||
}
|
||||
}
|
||||
return database.ProjectFilterUnbound
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -1,389 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/project"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func projectIDFromConversation(db *database.DB, ctx context.Context) (string, error) {
|
||||
convID := agent.ConversationIDFromContext(ctx)
|
||||
if convID == "" {
|
||||
return "", fmt.Errorf("无法确定当前对话，请在对话上下文中使用项目事实工具")
|
||||
}
|
||||
pid, err := db.GetConversationProjectID(convID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(pid) == "" {
|
||||
return "", fmt.Errorf("当前对话未绑定项目，请先在对话中选择项目或创建带项目的对话")
|
||||
}
|
||||
return pid, nil
|
||||
}
|
||||
|
||||
func textResult(msg string, isErr bool) *mcp.ToolResult {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: msg}},
|
||||
IsError: isErr,
|
||||
}
|
||||
}
|
||||
|
||||
// registerProjectFactTools 注册项目黑板 MCP 工具。
|
||||
func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *config.Config, logger *zap.Logger) {
|
||||
if db == nil || cfg == nil || !cfg.Project.Enabled {
|
||||
if logger != nil {
|
||||
logger.Info("项目黑板工具未注册（未启用）")
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
upsertTool := mcp.Tool{
|
||||
Name: builtin.ToolUpsertProjectFact,
|
||||
Description: "写入或更新项目黑板事实，用于跨会话沉淀可复现上下文（非正式漏洞条目；可交付漏洞另用 record_vulnerability）。" +
|
||||
"边渗透边记录：每确认新认知（端口/入口/凭据/可利用点）后立即调用，同 fact_key 覆盖更新，勿等会话结束。" +
|
||||
"禁止仅写结论：summary 须含什么+在哪+如何验证；body 须含攻击链/请求响应/命令等复现细节。" +
|
||||
"发现类建议 fact_key 为 finding|chain|exploit|poc/<slug>，category 对应 finding|chain|exploit|poc，body 按攻击链模板填写。" +
|
||||
"环境类用 target|auth|infra|business/<slug>。同 fact_key 覆盖更新。需当前对话已绑定项目。",
|
||||
ShortDescription: "写入/更新项目事实（含攻击链 body）",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "项目内唯一 key：target/primary_domain、finding/sqli-login、exploit/upload-rce 等",
|
||||
},
|
||||
"category": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "target | auth | infra | business | finding | chain | exploit | poc | note",
|
||||
"enum": []string{"target", "auth", "infra", "business", "finding", "chain", "exploit", "poc", "note"},
|
||||
},
|
||||
"summary": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "索引用一行：结论 + 位置 + 触发/验证要点（勿仅写「存在 XSS」等空话）",
|
||||
},
|
||||
"body": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "完整可复现详情（仅 get_project_fact 返回）：须含攻击链步骤、原始 HTTP/命令、响应现象、证据与关联。" +
|
||||
"发现/利用类首次写入必填；环境类建议含来源证据。攻击链类可参考模板章节：结论、目标与入口、攻击链、Exploit/POC、关键证据、关联、备注。" +
|
||||
"更新已有 fact_key 时若省略或留空 body，将保留库中已有 body（可只改 summary）。",
|
||||
},
|
||||
"confidence": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "confirmed | tentative | deprecated",
|
||||
"enum": []string{"confirmed", "tentative", "deprecated"},
|
||||
},
|
||||
"pinned": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "是否优先出现在黑板索引",
|
||||
},
|
||||
"related_vulnerability_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "可选：关联的漏洞记录 ID",
|
||||
},
|
||||
"links": map[string]interface{}{
|
||||
"type": "array",
|
||||
"description": "可选：关系边（from → 当前 fact）。finding 至少 1 条 {from:target/*, type:discovered_on}；finding 上记录 exploit 用 {from:exploit/*, type:exploits}。省略保留已有边；传 [] 清空全部关系边。",
|
||||
"items": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"from": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "来源 fact_key：存储为 from → 当前 fact",
|
||||
},
|
||||
"type": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "depends_on | leads_to | enables | exploits | discovered_on | contains | part_of | supports",
|
||||
},
|
||||
"confidence": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "confirmed | tentative | deprecated",
|
||||
},
|
||||
},
|
||||
"required": []string{"from", "type"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": []string{"fact_key", "summary"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(upsertTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
factKey, _ := args["fact_key"].(string)
|
||||
summary, _ := args["summary"].(string)
|
||||
if strings.TrimSpace(factKey) == "" || strings.TrimSpace(summary) == "" {
|
||||
return textResult("错误: fact_key 与 summary 必填", true), nil
|
||||
}
|
||||
if len([]rune(summary)) > cfg.Project.FactSummaryMaxRunesEffective() {
|
||||
return textResult(fmt.Sprintf("错误: summary 过长（最多 %d 字）", cfg.Project.FactSummaryMaxRunesEffective()), true), nil
|
||||
}
|
||||
f := &database.ProjectFact{
|
||||
ProjectID: projectID,
|
||||
FactKey: factKey,
|
||||
Category: strArg(args, "category"),
|
||||
Summary: summary,
|
||||
Body: strArg(args, "body"),
|
||||
Confidence: strArg(args, "confidence"),
|
||||
Pinned: boolArg(args, "pinned"),
|
||||
RelatedVulnerabilityID: strArg(args, "related_vulnerability_id"),
|
||||
}
|
||||
if convID := agent.ConversationIDFromContext(ctx); convID != "" {
|
||||
f.SourceConversationID = convID
|
||||
}
|
||||
created, err := db.UpsertProjectFact(f)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
if _, hasLinks := args["links"]; hasLinks {
|
||||
linkInputs, err := project.ParseFactLinkInputs(args["links"])
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
convID := agent.ConversationIDFromContext(ctx)
|
||||
if err := project.PersistFactLinksFromParsed(db, projectID, created.FactKey, convID, linkInputs, true); err != nil {
|
||||
return textResult("错误: 保存关系边失败: "+err.Error(), true), nil
|
||||
}
|
||||
created, _ = db.GetProjectFactByKey(projectID, created.FactKey)
|
||||
} else if parsed := project.ParseLinksFromBody(created.Body); len(parsed) > 0 {
|
||||
if err := project.PersistFactIncomingLinks(db, projectID, created.FactKey, parsed, true); err != nil {
|
||||
return textResult("错误: 从 body 解析边失败: "+err.Error(), true), nil
|
||||
}
|
||||
created, _ = db.GetProjectFactByKey(projectID, created.FactKey)
|
||||
}
|
||||
msg := fmt.Sprintf("事实已保存。\nfact_key: %s\nid: %s\nconfidence: %s", created.FactKey, created.ID, created.Confidence)
|
||||
if in, _ := db.ListIncomingProjectFactEdges(projectID, created.FactKey); len(in) > 0 {
|
||||
msg += "\n关系边: " + project.FormatFactLinksText(in)
|
||||
}
|
||||
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
|
||||
msg += warn
|
||||
}
|
||||
return textResult(msg, false), nil
|
||||
})
|
||||
|
||||
getTool := mcp.Tool{
|
||||
Name: builtin.ToolGetProjectFact,
|
||||
Description: "按 fact_key 获取项目事实完整 body 与元数据。摘要不足时必须调用本工具，禁止臆造细节。",
|
||||
ShortDescription: "按 key 获取事实详情",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string", "description": "事实 key"},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(getTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if key == "" {
|
||||
return textResult("错误: fact_key 必填", true), nil
|
||||
}
|
||||
f, err := db.GetProjectFactByKey(projectID, key)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
msg := fmt.Sprintf("fact_key: %s\ncategory: %s\nconfidence: %s\nsummary: %s\nupdated_at: %s",
|
||||
f.FactKey, f.Category, f.Confidence, f.Summary, f.UpdatedAt.Format("2006-01-02 15:04:05"))
|
||||
if f.RelatedVulnerabilityID != "" {
|
||||
msg += fmt.Sprintf("\nrelated_vulnerability_id: %s", f.RelatedVulnerabilityID)
|
||||
}
|
||||
if f.SourceConversationID != "" {
|
||||
msg += fmt.Sprintf("\nsource_conversation_id: %s", f.SourceConversationID)
|
||||
}
|
||||
if in, _ := db.ListIncomingProjectFactEdges(projectID, f.FactKey); len(in) > 0 {
|
||||
msg += "\n关系边（from → 本 fact）:\n"
|
||||
for _, e := range in {
|
||||
msg += fmt.Sprintf("- %s ← %s (%s)\n", e.EdgeType, e.SourceFactKey, e.Confidence)
|
||||
}
|
||||
}
|
||||
if out, _ := db.ListOutgoingProjectFactEdges(projectID, f.FactKey); len(out) > 0 {
|
||||
msg += "指向其他事实:\n"
|
||||
for _, e := range out {
|
||||
msg += fmt.Sprintf("- %s → %s (%s)\n", e.EdgeType, e.TargetFactKey, e.Confidence)
|
||||
}
|
||||
}
|
||||
msg += "\n\n--- body ---\n" + f.Body
|
||||
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
|
||||
msg += warn
|
||||
}
|
||||
return textResult(msg, false), nil
|
||||
})
|
||||
|
||||
listTool := mcp.Tool{
|
||||
Name: builtin.ToolListProjectFacts,
|
||||
Description: "列出当前项目的事实（分页）。",
|
||||
ShortDescription: "列出项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"category": map[string]interface{}{"type": "string"},
|
||||
"confidence": map[string]interface{}{"type": "string"},
|
||||
"limit": map[string]interface{}{"type": "integer"},
|
||||
"offset": map[string]interface{}{"type": "integer"},
|
||||
},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(listTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
limit := intArg(args, "limit", 50)
|
||||
offset := intArg(args, "offset", 0)
|
||||
filter := database.ProjectFactListFilter{
|
||||
Category: strArg(args, "category"),
|
||||
Confidence: strArg(args, "confidence"),
|
||||
}
|
||||
list, err := db.ListProjectFacts(projectID, filter, limit, offset)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("共 %d 条（limit=%d offset=%d）:\n", len(list), limit, offset))
|
||||
for _, f := range list {
|
||||
b.WriteString(fmt.Sprintf("- [%s] %s — %s (%s)\n", f.FactKey, f.Category, f.Summary, f.Confidence))
|
||||
}
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
|
||||
searchTool := mcp.Tool{
|
||||
Name: builtin.ToolSearchProjectFacts,
|
||||
Description: "按关键词搜索项目事实（summary/body/fact_key）。",
|
||||
ShortDescription: "搜索项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"query": map[string]interface{}{"type": "string"},
|
||||
"limit": map[string]interface{}{"type": "integer"},
|
||||
"offset": map[string]interface{}{"type": "integer"},
|
||||
},
|
||||
"required": []string{"query"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(searchTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
q := strings.TrimSpace(strArg(args, "query"))
|
||||
if q == "" {
|
||||
return textResult("错误: query 必填", true), nil
|
||||
}
|
||||
list, err := db.ListProjectFacts(projectID, database.ProjectFactListFilter{Search: q}, intArg(args, "limit", 30), intArg(args, "offset", 0))
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("搜索 \"%s\" 命中 %d 条:\n", q, len(list)))
|
||||
for _, f := range list {
|
||||
b.WriteString(fmt.Sprintf("- [%s] %s — %s\n", f.FactKey, f.Category, f.Summary))
|
||||
}
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
|
||||
deprecateTool := mcp.Tool{
|
||||
Name: builtin.ToolDeprecateProjectFact,
|
||||
Description: "将事实标记为 deprecated，从黑板索引中排除。",
|
||||
ShortDescription: "废弃项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(deprecateTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if err := db.DeprecateProjectFact(projectID, key); err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
return textResult("事实已标记为 deprecated: "+key, false), nil
|
||||
})
|
||||
|
||||
restoreTool := mcp.Tool{
|
||||
Name: builtin.ToolRestoreProjectFact,
|
||||
Description: "将已废弃（deprecated）的事实恢复为 tentative 或 confirmed，重新参与黑板索引。",
|
||||
ShortDescription: "恢复已废弃的项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string"},
|
||||
"confidence": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "恢复后的置信度：tentative（默认）或 confirmed",
|
||||
"enum": []string{"tentative", "confirmed"},
|
||||
},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(restoreTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if key == "" {
|
||||
return textResult("错误: fact_key 必填", true), nil
|
||||
}
|
||||
conf := strArg(args, "confidence")
|
||||
if err := db.RestoreProjectFact(projectID, key, conf); err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
if conf == "" {
|
||||
conf = "tentative"
|
||||
}
|
||||
return textResult(fmt.Sprintf("事实已恢复为 %s: %s", conf, key), false), nil
|
||||
})
|
||||
|
||||
if logger != nil {
|
||||
logger.Debug("项目黑板 MCP 工具注册成功")
|
||||
}
|
||||
}
|
||||
|
||||
func strArg(args map[string]interface{}, key string) string {
|
||||
if v, ok := args[key].(string); ok {
|
||||
return v
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func boolArg(args map[string]interface{}, key string) bool {
|
||||
if v, ok := args[key].(bool); ok {
|
||||
return v
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func intArg(args map[string]interface{}, key string, def int) int {
|
||||
switch v := args[key].(type) {
|
||||
case float64:
|
||||
return int(v)
|
||||
case int:
|
||||
return v
|
||||
case int64:
|
||||
return int(v)
|
||||
default:
|
||||
return def
|
||||
}
|
||||
}
|
||||
@@ -1,13 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/vision"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func registerVisionTools(mcpServer *mcp.Server, cfg *config.Config, logger *zap.Logger) {
|
||||
vision.RegisterAnalyzeImageTool(mcpServer, cfg, logger)
|
||||
}
|
||||
@@ -1,466 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func conversationIDFromToolCtx(ctx context.Context) string {
|
||||
if id := agent.ConversationIDFromContext(ctx); id != "" {
|
||||
return id
|
||||
}
|
||||
return mcp.MCPConversationIDFromContext(ctx)
|
||||
}
|
||||
|
||||
// canAccessVulnerability 校验当前对话是否有权查看该漏洞（默认项目隔离，未绑项目则仅本会话）。
|
||||
func canAccessVulnerability(vuln *database.Vulnerability, convID, projectID string) bool {
|
||||
if vuln == nil || convID == "" {
|
||||
return false
|
||||
}
|
||||
if projectID != "" {
|
||||
if strings.TrimSpace(vuln.ProjectID) == projectID {
|
||||
return true
|
||||
}
|
||||
// 历史记录：写入时尚未绑定 project_id，但属于同一会话
|
||||
if strings.TrimSpace(vuln.ProjectID) == "" && vuln.ConversationID == convID {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
return vuln.ConversationID == convID
|
||||
}
|
||||
|
||||
func buildVulnerabilityListFilter(db *database.DB, ctx context.Context, args map[string]interface{}) (database.VulnerabilityListFilter, string, error) {
|
||||
convID := conversationIDFromToolCtx(ctx)
|
||||
if convID == "" {
|
||||
return database.VulnerabilityListFilter{}, "", fmt.Errorf("无法确定当前对话，请在对话上下文中使用漏洞查询工具")
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, err := db.GetConversationProjectID(convID); err == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
scope := strings.TrimSpace(strArg(args, "scope"))
|
||||
if scope == "" {
|
||||
if projectID != "" {
|
||||
scope = "project"
|
||||
} else {
|
||||
scope = "conversation"
|
||||
}
|
||||
}
|
||||
|
||||
filter := database.VulnerabilityListFilter{
|
||||
Severity: strings.TrimSpace(strArg(args, "severity")),
|
||||
Status: strings.TrimSpace(strArg(args, "status")),
|
||||
}
|
||||
if q := strings.TrimSpace(strArg(args, "q")); q != "" {
|
||||
filter.Search = q
|
||||
} else {
|
||||
filter.Search = strings.TrimSpace(strArg(args, "search"))
|
||||
}
|
||||
|
||||
var scopeLabel string
|
||||
switch scope {
|
||||
case "project":
|
||||
if projectID == "" {
|
||||
return filter, "", fmt.Errorf("当前对话未绑定项目，无法按项目列出漏洞；请使用 scope=conversation，或先在对话中绑定项目")
|
||||
}
|
||||
filter.ProjectID = projectID
|
||||
scopeLabel = fmt.Sprintf("项目 %s", projectID)
|
||||
case "conversation":
|
||||
filter.ConversationID = convID
|
||||
scopeLabel = fmt.Sprintf("会话 %s", convID)
|
||||
default:
|
||||
return filter, "", fmt.Errorf("scope 仅支持 project 或 conversation，当前值: %s", scope)
|
||||
}
|
||||
return filter, scopeLabel, nil
|
||||
}
|
||||
|
||||
func formatVulnerabilityListItem(v *database.Vulnerability) string {
|
||||
line := fmt.Sprintf("- id=%s | %s | %s | %s", v.ID, v.Severity, v.Status, v.Title)
|
||||
if v.Type != "" {
|
||||
line += fmt.Sprintf(" | type=%s", v.Type)
|
||||
}
|
||||
if v.Target != "" {
|
||||
line += fmt.Sprintf(" | target=%s", truncateRunes(v.Target, 80))
|
||||
}
|
||||
return line
|
||||
}
|
||||
|
||||
func formatVulnerabilityDetail(v *database.Vulnerability) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("漏洞ID: %s\n", v.ID))
|
||||
b.WriteString(fmt.Sprintf("标题: %s\n", v.Title))
|
||||
b.WriteString(fmt.Sprintf("严重程度: %s\n", v.Severity))
|
||||
b.WriteString(fmt.Sprintf("状态: %s\n", v.Status))
|
||||
if v.Type != "" {
|
||||
b.WriteString(fmt.Sprintf("类型: %s\n", v.Type))
|
||||
}
|
||||
if v.Target != "" {
|
||||
b.WriteString(fmt.Sprintf("目标: %s\n", v.Target))
|
||||
}
|
||||
if v.ProjectID != "" {
|
||||
b.WriteString(fmt.Sprintf("项目ID: %s\n", v.ProjectID))
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("会话ID: %s\n", v.ConversationID))
|
||||
if !v.CreatedAt.IsZero() {
|
||||
b.WriteString(fmt.Sprintf("创建时间: %s\n", v.CreatedAt.Format("2006-01-02 15:04:05")))
|
||||
}
|
||||
if v.Description != "" {
|
||||
b.WriteString("\n--- 描述 ---\n")
|
||||
b.WriteString(v.Description)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Preconditions != "" {
|
||||
b.WriteString("\n--- 前置条件 ---\n")
|
||||
b.WriteString(v.Preconditions)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.ReproSteps != "" {
|
||||
b.WriteString("\n--- 复现步骤 ---\n")
|
||||
b.WriteString(v.ReproSteps)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Evidence != "" {
|
||||
b.WriteString("\n--- 证据 / POC ---\n")
|
||||
b.WriteString(v.Evidence)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Impact != "" {
|
||||
b.WriteString("\n--- 影响 ---\n")
|
||||
b.WriteString(v.Impact)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Recommendation != "" {
|
||||
b.WriteString("\n--- 修复建议 ---\n")
|
||||
b.WriteString(v.Recommendation)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.RetestNotes != "" {
|
||||
b.WriteString("\n--- 复测方式 ---\n")
|
||||
b.WriteString(v.RetestNotes)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func missingVulnerabilityReproFields(args map[string]interface{}) []string {
|
||||
required := []struct {
|
||||
key string
|
||||
label string
|
||||
}{
|
||||
{"target", "target（受影响的 URL/IP/服务/接口）"},
|
||||
{"vulnerability_type", "vulnerability_type（漏洞类型）"},
|
||||
{"description", "description（漏洞摘要与触发点）"},
|
||||
{"reproduction_steps", "reproduction_steps（可逐步执行的复现步骤）"},
|
||||
{"evidence", "evidence（POC、原始请求/响应、命令输出或截图/日志证据）"},
|
||||
{"impact", "impact（确认后的实际影响）"},
|
||||
{"recommendation", "recommendation（修复建议）"},
|
||||
}
|
||||
missing := make([]string, 0)
|
||||
for _, item := range required {
|
||||
if strings.TrimSpace(strArg(args, item.key)) == "" {
|
||||
missing = append(missing, item.label)
|
||||
}
|
||||
}
|
||||
return missing
|
||||
}
|
||||
|
||||
func truncateRunes(s string, max int) string {
|
||||
r := []rune(s)
|
||||
if len(r) <= max {
|
||||
return s
|
||||
}
|
||||
return string(r[:max]) + "…"
|
||||
}
|
||||
|
||||
// registerVulnerabilityTools 注册漏洞记录与查询 MCP 工具。
|
||||
func registerVulnerabilityTools(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
registerRecordVulnerabilityTool(mcpServer, db, logger)
|
||||
registerListVulnerabilitiesTool(mcpServer, db, logger)
|
||||
registerGetVulnerabilityTool(mcpServer, db, logger)
|
||||
if logger != nil {
|
||||
logger.Debug("漏洞 MCP 工具注册成功", zap.Strings("tools", []string{
|
||||
builtin.ToolRecordVulnerability,
|
||||
builtin.ToolListVulnerabilities,
|
||||
builtin.ToolGetVulnerability,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
func registerRecordVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolRecordVulnerability,
|
||||
Description: "记录发现的漏洞详情到漏洞管理系统。必须按“仅看本记录即可复现”的标准填写：目标、漏洞类型、触发点、复现步骤、证据/POC、实际影响和修复建议；前置条件与复测方式为推荐填写项。边渗透边记录：每验证出一条可复现漏洞后立即调用，勿等会话结束。记录前可先 list_vulnerabilities 避免重复。",
|
||||
ShortDescription: "记录可复现的漏洞详情到漏洞管理系统",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"title": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞标题（必需）。建议格式：<资产/接口> 存在 <漏洞类型>，例如“/api/login 存在 SQL 注入”。",
|
||||
},
|
||||
"description": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞摘要与触发点（必需）：说明哪个功能/参数/入口存在问题、为什么可被利用。不要只写结论。",
|
||||
},
|
||||
"severity": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞严重程度：critical（严重）、high（高）、medium（中）、low（低）、info（信息）",
|
||||
"enum": []string{"critical", "high", "medium", "low", "info"},
|
||||
},
|
||||
"vulnerability_type": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞类型，如：SQL注入、XSS、CSRF、命令注入等（必需）",
|
||||
},
|
||||
"target": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "受影响的目标（必需）：尽量精确到 URL、IP:端口、服务名、接口路径和参数名。",
|
||||
},
|
||||
"preconditions": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "前置条件（推荐填写）：登录状态、权限、账号、Header/Cookie、特定数据、网络位置、环境/版本等；无前置条件可写“无”。",
|
||||
},
|
||||
"reproduction_steps": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "复现步骤（必需）：按 1/2/3 编号，写清入口、参数、payload、执行命令、观察点。应让未参与对话的人照做即可复现。",
|
||||
},
|
||||
"evidence": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "证据 / POC（必需）：原始 HTTP 请求/响应、curl/工具命令、截图文字说明、日志、DNSLog/回连记录、数据库结果、文件路径、时间戳等。优先放最小可验证证据。",
|
||||
},
|
||||
"impact": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞影响说明（必需）：结合已验证事实说明可造成什么后果，避免泛泛而谈。",
|
||||
},
|
||||
"recommendation": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "修复建议（必需）：给出针对该触发点/参数/组件的具体修复和复测建议。",
|
||||
},
|
||||
"retest_notes": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "复测方式（推荐填写）：修复后如何验证漏洞已关闭，包括应返回的状态码、错误信息或访问控制结果。",
|
||||
},
|
||||
},
|
||||
"required": []string{"title", "description", "severity", "vulnerability_type", "target", "reproduction_steps", "evidence", "impact", "recommendation"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
conversationID := strings.TrimSpace(strArg(args, "conversation_id"))
|
||||
if conversationID == "" {
|
||||
conversationID = conversationIDFromToolCtx(ctx)
|
||||
}
|
||||
if conversationID == "" {
|
||||
return textResult("错误: conversation_id 未设置。这是系统错误，请重试。", true), nil
|
||||
}
|
||||
|
||||
title := strings.TrimSpace(strArg(args, "title"))
|
||||
if title == "" {
|
||||
return textResult("错误: title 参数必需且不能为空", true), nil
|
||||
}
|
||||
|
||||
severity := strings.TrimSpace(strArg(args, "severity"))
|
||||
if severity == "" {
|
||||
return textResult("错误: severity 参数必需且不能为空", true), nil
|
||||
}
|
||||
|
||||
validSeverities := map[string]bool{
|
||||
"critical": true, "high": true, "medium": true, "low": true, "info": true,
|
||||
}
|
||||
if !validSeverities[severity] {
|
||||
return textResult(fmt.Sprintf("错误: severity 必须是 critical、high、medium、low 或 info 之一，当前值: %s", severity), true), nil
|
||||
}
|
||||
if missing := missingVulnerabilityReproFields(args); len(missing) > 0 {
|
||||
return textResult("错误: 漏洞记录缺少必填信息，请补充后再记录：\n- "+strings.Join(missing, "\n- ")+"\n\n必填项用于确保单条记录可独立复现；前置条件和复测方式为推荐填写项。", true), nil
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, perr := db.GetConversationProjectID(conversationID); perr == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: conversationID,
|
||||
ProjectID: projectID,
|
||||
Title: title,
|
||||
Description: strArg(args, "description"),
|
||||
Severity: severity,
|
||||
Status: "open",
|
||||
Type: strArg(args, "vulnerability_type"),
|
||||
Target: strArg(args, "target"),
|
||||
Preconditions: strArg(args, "preconditions"),
|
||||
ReproSteps: strArg(args, "reproduction_steps"),
|
||||
Evidence: strArg(args, "evidence"),
|
||||
Impact: strArg(args, "impact"),
|
||||
Recommendation: strArg(args, "recommendation"),
|
||||
RetestNotes: strArg(args, "retest_notes"),
|
||||
}
|
||||
|
||||
created, err := db.CreateVulnerability(vuln)
|
||||
if err != nil {
|
||||
if logger != nil {
|
||||
logger.Error("记录漏洞失败", zap.Error(err))
|
||||
}
|
||||
return textResult(fmt.Sprintf("记录漏洞失败: %v", err), true), nil
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = db.SetResourceOwner("vulnerability", created.ID, principal.UserID)
|
||||
_ = db.AssignResourceToUser(principal.UserID, "vulnerability", created.ID)
|
||||
}
|
||||
db.NotifyVulnerabilityCreated(created)
|
||||
|
||||
if logger != nil {
|
||||
logger.Info("漏洞记录成功",
|
||||
zap.String("id", created.ID),
|
||||
zap.String("title", created.Title),
|
||||
zap.String("severity", created.Severity),
|
||||
zap.String("conversation_id", conversationID),
|
||||
)
|
||||
}
|
||||
|
||||
return textResult(fmt.Sprintf("漏洞已成功记录！\n\n漏洞ID: %s\n标题: %s\n严重程度: %s\n状态: %s\n\n可使用 get_vulnerability(id) 查看详情，或 list_vulnerabilities 查看列表。",
|
||||
created.ID, created.Title, created.Severity, created.Status), false), nil
|
||||
})
|
||||
}
|
||||
|
||||
func registerListVulnerabilitiesTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolListVulnerabilities,
|
||||
Description: "列出当前授权范围内的漏洞（摘要）。默认：对话已绑定项目时列出该项目下全部漏洞；未绑项目时仅列出当前会话漏洞。可用 scope=conversation 仅看本会话。记录新漏洞前建议先调用以避免重复。",
|
||||
ShortDescription: "列出漏洞（默认当前项目）",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"scope": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "范围：project（默认，需绑定项目）| conversation（仅当前会话）",
|
||||
"enum": []string{"project", "conversation"},
|
||||
},
|
||||
"severity": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "按严重程度筛选：critical、high、medium、low、info",
|
||||
"enum": []string{"critical", "high", "medium", "low", "info"},
|
||||
},
|
||||
"status": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "按状态筛选：open、confirmed、fixed、false_positive、ignored",
|
||||
"enum": []string{"open", "confirmed", "fixed", "false_positive", "ignored"},
|
||||
},
|
||||
"q": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "关键词搜索（标题、描述、类型、目标等）",
|
||||
},
|
||||
"limit": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "返回条数上限，默认 30，最大 100",
|
||||
},
|
||||
"offset": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "分页偏移，默认 0",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
filter, scopeLabel, err := buildVulnerabilityListFilter(db, ctx, args)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
|
||||
limit := intArg(args, "limit", 30)
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 30
|
||||
}
|
||||
offset := intArg(args, "offset", 0)
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
|
||||
total, err := db.CountVulnerabilities(filter)
|
||||
if err != nil {
|
||||
if logger != nil {
|
||||
logger.Warn("统计漏洞失败", zap.Error(err))
|
||||
}
|
||||
total = 0
|
||||
}
|
||||
|
||||
list, err := db.ListVulnerabilities(limit, offset, filter)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("范围: %s\n总计: %d | 本页: %d 条 (limit=%d offset=%d)\n\n", scopeLabel, total, len(list), limit, offset))
|
||||
if len(list) == 0 {
|
||||
b.WriteString("（暂无漏洞记录）\n")
|
||||
} else {
|
||||
for _, v := range list {
|
||||
b.WriteString(formatVulnerabilityListItem(v))
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if total > offset+len(list) {
|
||||
b.WriteString(fmt.Sprintf("\n（还有更多，可增大 offset 或使用 q/severity/status 筛选）\n"))
|
||||
}
|
||||
}
|
||||
b.WriteString("\n需要 POC 与完整字段请对具体 id 调用 get_vulnerability。")
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
}
|
||||
|
||||
func registerGetVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolGetVulnerability,
|
||||
Description: "按漏洞 ID 获取完整详情（含 POC、影响、修复建议）。仅能访问当前项目或当前会话下的漏洞（与 list_vulnerabilities 授权范围一致）。",
|
||||
ShortDescription: "按 ID 获取漏洞详情",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞 ID（list_vulnerabilities 返回的 id）",
|
||||
},
|
||||
},
|
||||
"required": []string{"id"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
convID := conversationIDFromToolCtx(ctx)
|
||||
if convID == "" {
|
||||
return textResult("错误: 无法确定当前对话，请在对话上下文中使用本工具", true), nil
|
||||
}
|
||||
|
||||
id := strings.TrimSpace(strArg(args, "id"))
|
||||
if id == "" {
|
||||
return textResult("错误: id 必填", true), nil
|
||||
}
|
||||
|
||||
vuln, err := db.GetVulnerability(id)
|
||||
if err != nil {
|
||||
return textResult("错误: 漏洞不存在或查询失败", true), nil
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, perr := db.GetConversationProjectID(convID); perr == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
if !canAccessVulnerability(vuln, convID, projectID) {
|
||||
return textResult("错误: 无权访问该漏洞（仅可查看当前项目或当前会话下的记录）", true), nil
|
||||
}
|
||||
|
||||
return textResult(formatVulnerabilityDetail(vuln), false), nil
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user