Add files via upload

This commit is contained in:
公明
2026-08-18 20:34:53 +08:00
committed by GitHub
parent 98c380ed5f
commit 5aa2c382b0
94 changed files with 27142 additions and 0 deletions
+2256
View File
File diff suppressed because it is too large Load Diff
+511
View File
@@ -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
}
+201
View File
@@ -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)
}
}
+228
View File
@@ -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 "" },
}
}
+104
View File
@@ -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 子系统已关闭")
}
}
+919
View File
@@ -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 / onelinerlist/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 加密 BeaconAES-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": "监听器 IDget/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": "会话 IDget/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/killedlist"},
"os": map[string]interface{}{"type": "string", "description": "按 OS 过滤: linux/windows/darwinlist"},
"search": map[string]interface{}{"type": "string", "description": "模糊搜索 hostname/username/IPlist"},
"suspicious": map[string]interface{}{"type": "boolean", "description": "仅疑似误报:离线且 tcp_* / unknown / PID 0list"},
"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-100set_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 会话 IDs_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": "进程 IDkill_proc"},
"remote_path": map[string]interface{}{"type": "string", "description": "远程路径(upload/download"},
"file_id": map[string]interface{}{"type": "string", "description": "服务端文件 IDupload/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/stopport_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": "任务 IDget_result/wait/cancel 需要)"},
"session_id": map[string]interface{}{"type": "string", "description": "按会话过滤(list"},
"status": map[string]interface{}{"type": "string", "description": "按状态过滤: queued/sent/running/success/failed/cancelledlist"},
"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_reversetcp_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|powershellhttp_beacon|https_beacon|websocket: 仅 curl_beacon"},
"host": map[string]interface{}{"type": "string", "description": "oneliner/build 可选覆盖:非空则强制用作植入回连主机。留空时顺序为:监听器 callback_hostcreate/update 的 callback_host 参数写入)→ bind_host0.0.0.0 时尝试本机对外 IP 探测)"},
"os": map[string]interface{}{"type": "string", "description": "目标 OSbuild: 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 IDget/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": "会话 IDlist 需要)"},
"task_id": map[string]interface{}{"type": "string", "description": "任务 IDget_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
}
+69
View File
@@ -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)
}
}
+105
View File
@@ -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)
}
})
}
}
+213
View File
@@ -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)
}
}
+86
View File
@@ -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。
// inMemorytls_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: 已启用 TLStls_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)
}
+407
View File
@@ -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)
}
+261
View File
@@ -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)
}
}
+59
View File
@@ -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)
}
}
+26
View File
@@ -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 ""
}
+389
View File
@@ -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|pocbody 按攻击链模板填写。" +
"环境类用 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": "项目内唯一 keytarget/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
}
}
+13
View File
@@ -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)
}
+466
View File
@@ -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", "evidencePOC、原始请求/响应、命令输出或截图/日志证据)"},
{"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": "漏洞 IDlist_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
+449
View File
@@ -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)
}
}
+167
View File
@@ -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
}
+222
View File
@@ -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()
}
+75
View File
@@ -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)
}
}
}
+631
View File
@@ -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
+30
View File
@@ -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)
}
+56
View File
@@ -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
+486
View File
@@ -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
}
+75
View File
@@ -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
}
+106
View File
@@ -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
+102
View File
@@ -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)
}
}
+122
View File
@@ -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)
}
}
+132
View File
@@ -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)
}
}
+125
View File
@@ -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
}
+104
View File
@@ -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
}
+635
View File
@@ -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{}
}
+112
View File
@@ -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
}
+410
View File
@@ -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)
}
}
+82
View File
@@ -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)
}
}
+125
View File
@@ -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
}
+93
View File
@@ -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
+727
View File
@@ -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")
}
}
+174
View File
@@ -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
}
+89
View File
@@ -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")
}
}
+93
View File
@@ -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
}
+142
View File
@@ -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
}
+33
View File
@@ -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
}
+547
View File
@@ -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
}
+215
View File
@@ -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)
}
}
+178
View File
@@ -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
}
+468
View File
@@ -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
}
+286
View File
@@ -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)
}
}
+125
View File
@@ -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
}
+60
View File
@@ -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, "")
}
+76
View File
@@ -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)
}
}
+99
View File
@@ -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
}
+56
View File
@@ -0,0 +1,56 @@
package project
import "strings"
// FactIndexSectionHeading 黑板索引可读标题行前缀(块内保留,供 Agent 阅读)。
const FactIndexSectionHeading = "## 项目黑板索引"
// FactIndexSectionStartMarker / EndMarkerHTML 注释边界,供程序化替换;对模型无指令语义。
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
}
+154
View File
@@ -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)
}
}
+256
View File
@@ -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
}
+68
View File
@@ -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))
}
}
+407
View File
@@ -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()
}
+96
View File
@@ -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 写入解析后的 linksparsed 为 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 写入出边(图连线等低层 APIbody 同步请用 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
}
+296
View File
@@ -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")
}
}
+231
View File
@@ -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)
})
}
+161
View File
@@ -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)
}
}
+23
View File
@@ -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)
}
+135
View File
@@ -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>
- linksupsert 参数: [{ "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)
}
+42
View File
@@ -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")
}
}
+99
View File
@@ -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\nscope_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("\nscope_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
}
+40
View File
@@ -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)
}
}
+21
View File
@@ -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
}
+26
View File
@@ -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_imagequestion 写明「只输出验证码字符」;识别失败则刷新验证码后重新保存再识;复杂滑块/行为验证码勿指望单次识图成功。\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())
}
+69
View File
@@ -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)
}
+58
View File
@@ -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)
}
}
+67
View File
@@ -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()
}
+76
View File
@@ -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))
}
}
+108
View File
@@ -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()
}
+73
View File
@@ -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
}