From 9f38fda15fae817b832a2165351e25a7b66fa70e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sat, 15 Aug 2026 02:12:18 +0800 Subject: [PATCH] Add files via upload --- internal/attackchain/builder.go | 952 ++++++++++++++++++++ internal/attackchain/promote_project.go | 203 +++++ internal/attackchain/truncate.go | 248 +++++ internal/attackchain/truncate_test.go | 63 ++ internal/project/blackboard.go | 99 ++ internal/project/blackboard_refresh.go | 56 ++ internal/project/blackboard_refresh_test.go | 154 ++++ internal/project/fact_body_links.go | 256 ++++++ internal/project/fact_body_links_test.go | 68 ++ internal/project/fact_edges.go | 407 +++++++++ internal/project/fact_edges_apply.go | 96 ++ internal/project/fact_edges_test.go | 296 ++++++ internal/project/fact_index_links.go | 231 +++++ internal/project/fact_index_links_test.go | 161 ++++ internal/project/fact_recording_prompt.go | 23 + internal/project/fact_template.go | 135 +++ internal/project/fact_template_test.go | 42 + internal/project/scope_block.go | 99 ++ internal/project/scope_block_test.go | 40 + internal/project/stats.go | 21 + internal/project/vision_image_prompt.go | 26 + internal/project/workspace.go | 69 ++ internal/project/workspace_test.go | 58 ++ internal/vision/client.go | 134 +++ internal/vision/client_test.go | 12 + internal/vision/path.go | 72 ++ internal/vision/path_test.go | 52 ++ internal/vision/preprocess.go | 212 +++++ internal/vision/preprocess_test.go | 109 +++ internal/vision/tool.go | 125 +++ 30 files changed, 4519 insertions(+) create mode 100644 internal/attackchain/builder.go create mode 100644 internal/attackchain/promote_project.go create mode 100644 internal/attackchain/truncate.go create mode 100644 internal/attackchain/truncate_test.go create mode 100644 internal/project/blackboard.go create mode 100644 internal/project/blackboard_refresh.go create mode 100644 internal/project/blackboard_refresh_test.go create mode 100644 internal/project/fact_body_links.go create mode 100644 internal/project/fact_body_links_test.go create mode 100644 internal/project/fact_edges.go create mode 100644 internal/project/fact_edges_apply.go create mode 100644 internal/project/fact_edges_test.go create mode 100644 internal/project/fact_index_links.go create mode 100644 internal/project/fact_index_links_test.go create mode 100644 internal/project/fact_recording_prompt.go create mode 100644 internal/project/fact_template.go create mode 100644 internal/project/fact_template_test.go create mode 100644 internal/project/scope_block.go create mode 100644 internal/project/scope_block_test.go create mode 100644 internal/project/stats.go create mode 100644 internal/project/vision_image_prompt.go create mode 100644 internal/project/workspace.go create mode 100644 internal/project/workspace_test.go create mode 100644 internal/vision/client.go create mode 100644 internal/vision/client_test.go create mode 100644 internal/vision/path.go create mode 100644 internal/vision/path_test.go create mode 100644 internal/vision/preprocess.go create mode 100644 internal/vision/preprocess_test.go create mode 100644 internal/vision/tool.go diff --git a/internal/attackchain/builder.go b/internal/attackchain/builder.go new file mode 100644 index 00000000..f257f5d9 --- /dev/null +++ b/internal/attackchain/builder.go @@ -0,0 +1,952 @@ +package attackchain + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/openai" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// Builder 攻击链构建器 +type Builder struct { + db *database.DB + logger *zap.Logger + openAIClient *openai.Client + openAIConfig *config.OpenAIConfig + tokenCounter agent.TokenCounter + maxTokens int // 最大tokens限制，默认100000 +} + +// Node 攻击链节点（使用database包的类型） +type Node = database.AttackChainNode + +// Edge 攻击链边（使用database包的类型） +type Edge = database.AttackChainEdge + +// Chain 完整的攻击链 +type Chain struct { + Nodes []Node `json:"nodes"` + Edges []Edge `json:"edges"` +} + +// NewBuilder 创建新的攻击链构建器 +func NewBuilder(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.Logger) *Builder { + transport := &http.Transport{ + MaxIdleConns: 100, + MaxIdleConnsPerHost: 10, + IdleConnTimeout: 90 * time.Second, + } + httpClient := &http.Client{Timeout: 5 * time.Minute, Transport: transport} + + // 优先使用配置文件中的统一 Token 上限（config.yaml -> openai.max_total_tokens） + maxTokens := 0 + if openAIConfig != nil && openAIConfig.MaxTotalTokens > 0 { + maxTokens = openAIConfig.MaxTotalTokens + } else if openAIConfig != nil { + // 如果未显式配置 max_total_tokens，则根据模型设置一个合理的默认值 + model := strings.ToLower(openAIConfig.Model) + if strings.Contains(model, "gpt-4") { + maxTokens = 128000 // gpt-4通常支持128k + } else if strings.Contains(model, "gpt-3.5") { + maxTokens = 16000 // gpt-3.5-turbo通常支持16k + } else if strings.Contains(model, "deepseek") { + maxTokens = 131072 // deepseek-chat通常支持131k + } else { + maxTokens = 100000 // 兜底默认值 + } + } else { + // 没有 OpenAI 配置时使用兜底值，避免为 0 + maxTokens = 100000 + } + + return &Builder{ + db: db, + logger: logger, + openAIClient: openai.NewClient(openAIConfig, httpClient, logger), + openAIConfig: openAIConfig, + tokenCounter: agent.NewTikTokenCounter(), + maxTokens: maxTokens, + } +} + +// BuildChainFromConversation 从对话构建攻击链（单次 LLM 调用；输入为当前任务轮次的 last_react 轨迹，与继续对话续跑范围一致）。 +func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID string) (*Chain, error) { + b.logger.Info("开始构建攻击链（简化版本）", zap.String("conversationId", conversationID)) + + // 0. 首先检查是否有实际的工具执行记录 + messages, err := b.db.GetMessages(conversationID) + if err != nil { + return nil, fmt.Errorf("获取对话消息失败: %w", err) + } + + if len(messages) == 0 { + b.logger.Info("对话中没有数据", zap.String("conversationId", conversationID)) + return &Chain{Nodes: []Node{}, Edges: []Edge{}}, nil + } + + // 检查是否有实际的工具执行：assistant 的 mcp_execution_ids，或过程详情中的 tool_call/tool_result + //（多代理下若 MCP 未返回 execution_id，IDs 可能为空，但工具已通过 Eino 执行并写入 process_details） + hasToolExecutions := false + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "assistant") { + if len(messages[i].MCPExecutionIDs) > 0 { + hasToolExecutions = true + break + } + } + } + if !hasToolExecutions { + if pdOK, err := b.db.ConversationHasToolProcessDetails(conversationID); err != nil { + b.logger.Warn("查询过程详情判定工具执行失败", zap.Error(err)) + } else if pdOK { + hasToolExecutions = true + } + } + + // 检查任务是否被取消（通过检查最后一条assistant消息内容或process_details） + taskCancelled := false + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "assistant") { + content := strings.ToLower(messages[i].Content) + if strings.Contains(content, "取消") || strings.Contains(content, "cancelled") { + taskCancelled = true + } + break + } + } + + // 如果任务被取消且没有实际工具执行，返回空攻击链 + if taskCancelled && !hasToolExecutions { + b.logger.Info("任务已取消且没有实际工具执行，返回空攻击链", + zap.String("conversationId", conversationID), + zap.Bool("taskCancelled", taskCancelled), + zap.Bool("hasToolExecutions", hasToolExecutions)) + return &Chain{Nodes: []Node{}, Edges: []Edge{}}, nil + } + + // 如果没有实际工具执行，也返回空攻击链（避免AI编造） + if !hasToolExecutions { + b.logger.Info("没有实际工具执行记录，返回空攻击链", + zap.String("conversationId", conversationID)) + return &Chain{Nodes: []Node{}, Edges: []Edge{}}, nil + } + + // 1. 优先尝试从数据库获取保存的最后一轮ReAct输入和输出 + reactInputJSON, modelOutput, err := b.db.GetAgentTrace(conversationID) + if err != nil { + b.logger.Warn("获取保存的ReAct数据失败，将使用消息历史构建", zap.Error(err)) + // 继续使用原来的逻辑 + reactInputJSON = "" + modelOutput = "" + } + + // var userInput string + var reactInputFinal string + var dataSource string // 记录数据来源 + + // 优先使用落库的代理轨迹（与继续对话 loadHistoryFromAgentTrace 同源），并裁剪为「当前任务轮次」 + if reactInputJSON != "" { + trimmedJSON := agent.ExtractLastUserTurnTraceJSON(reactInputJSON) + hash := sha256.Sum256([]byte(trimmedJSON)) + reactInputHash := hex.EncodeToString(hash[:])[:16] + + var messageCount int + if msgs, parseErr := agent.ParseTraceMessages(trimmedJSON); parseErr == nil { + messageCount = len(msgs) + msgs = agent.MergeAssistantTraceOutput(msgs, modelOutput) + reactInputFinal = b.formatAgentTraceFromChatMessages(msgs) + } else { + b.logger.Warn("解析代理轨迹失败，回退原始 JSON 格式化", zap.Error(parseErr)) + reactInputFinal = b.formatAgentTraceInputFromJSON(trimmedJSON) + if strings.TrimSpace(modelOutput) != "" { + reactInputFinal += "\n\n## 助手结论（last_react_output）\n\n" + modelOutput + } + } + + dataSource = "last_user_turn_agent_trace" + b.logger.Info("使用当前任务轮次代理轨迹构建攻击链（与续跑上下文范围一致）", + zap.String("conversationId", conversationID), + zap.String("dataSource", dataSource), + zap.Int("traceInputSizeBeforeTrim", len(reactInputJSON)), + zap.Int("traceInputSizeAfterTrim", len(trimmedJSON)), + zap.Int("messageCount", messageCount), + zap.String("reactInputHash", reactInputHash), + zap.Int("modelOutputSize", len(modelOutput))) + } else { + // 2. 如果没有保存的ReAct数据，从对话消息构建 + dataSource = "messages_table" + b.logger.Info("从消息历史构建ReAct数据", + zap.String("conversationId", conversationID), + zap.String("dataSource", dataSource), + zap.Int("messageCount", len(messages))) + + // 提取用户输入（最后一条user消息） + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "user") { + // userInput = messages[i].Content + break + } + } + + // 提取最后一轮ReAct的输入（历史消息+当前用户输入） + reactInputFinal = b.buildAgentTraceInput(messages) + + // 提取大模型最后的输出（最后一条assistant消息） + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "assistant") { + modelOutput = messages[i].Content + break + } + } + } + + // 多代理：保存的轨迹列可能仅为首轮用户消息，不含工具轨迹；补充最后一轮助手的过程详情（与单代理完整轨迹对齐） + hasMCPOnAssistant := false + var lastAssistantID string + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "assistant") { + lastAssistantID = messages[i].ID + if len(messages[i].MCPExecutionIDs) > 0 { + hasMCPOnAssistant = true + } + break + } + } + if lastAssistantID != "" { + pdHasTools, _ := b.db.ConversationHasToolProcessDetails(conversationID) + if pdHasTools && !(hasMCPOnAssistant && reactInputContainsToolTrace(reactInputJSON)) { + detailsMap, err := b.db.GetProcessDetailsByConversation(conversationID) + if err != nil { + b.logger.Warn("加载过程详情用于攻击链失败", zap.Error(err)) + } else if dets := detailsMap[lastAssistantID]; len(dets) > 0 { + extra := b.formatProcessDetailsForAttackChain(dets) + if strings.TrimSpace(extra) != "" { + reactInputFinal = reactInputFinal + "\n\n## 执行过程与工具记录（含多代理编排与子任务）\n\n" + extra + b.logger.Info("攻击链输入已补充过程详情", + zap.String("conversationId", conversationID), + zap.String("messageId", lastAssistantID), + zap.Int("detailEvents", len(dets))) + } + } + } + } + + // 3. 按 token 预算压缩输入，再构建 prompt（避免超出模型上下文） + reactInputFinal, modelOutput, _ = b.fitAttackChainPayload(reactInputFinal, modelOutput) + + // 4. 构建 prompt 并单次调用大模型（助手结论已并入轨迹时不再重复传入） + promptAssistantOut := modelOutput + if reactInputJSON != "" { + promptAssistantOut = "" + } + prompt := b.buildSimplePrompt(reactInputFinal, promptAssistantOut) + // fmt.Println(prompt) + // 6. 调用AI生成攻击链（一次性，不做任何处理） + chainJSON, err := b.callAIForChainGeneration(ctx, prompt) + if err != nil { + return nil, fmt.Errorf("AI生成失败: %w", err) + } + + // 7. 解析JSON并生成节点/边ID（前端需要有效的ID） + chainData, err := b.parseChainJSON(chainJSON) + if err != nil { + // 如果解析失败，返回空链，让前端处理错误 + b.logger.Warn("解析攻击链JSON失败", zap.Error(err), zap.String("raw_json", chainJSON)) + return &Chain{ + Nodes: []Node{}, + Edges: []Edge{}, + }, nil + } + + b.logger.Info("攻击链构建完成", + zap.String("conversationId", conversationID), + zap.String("dataSource", dataSource), + zap.Int("nodes", len(chainData.Nodes)), + zap.Int("edges", len(chainData.Edges))) + + // 保存到数据库（供后续加载使用） + if err := b.saveChain(conversationID, chainData.Nodes, chainData.Edges); err != nil { + b.logger.Warn("保存攻击链到数据库失败", zap.Error(err)) + // 即使保存失败，也返回数据给前端 + } + + // 直接返回，不做任何处理和校验 + return chainData, nil +} + +// reactInputContainsToolTrace 判断保存的 ReAct JSON 是否包含可解析的工具调用轨迹（单代理完整保存时为 true）。 +func reactInputContainsToolTrace(reactInputJSON string) bool { + s := strings.TrimSpace(reactInputJSON) + if s == "" { + return false + } + return strings.Contains(s, "tool_calls") || + strings.Contains(s, "tool_call_id") || + strings.Contains(s, `"role":"tool"`) || + strings.Contains(s, `"role": "tool"`) +} + +// formatProcessDetailsForAttackChain 将最后一轮助手的过程详情格式化为攻击链分析的输入（覆盖多代理下 last_react_input 不完整的情况）。 +func (b *Builder) formatProcessDetailsForAttackChain(details []database.ProcessDetail) string { + if len(details) == 0 { + return "" + } + var sb strings.Builder + for _, d := range details { + // 目标：以主 agent（编排器）视角输出整轮迭代 + // - 保留：编排器工具调用/结果、对子代理的 task 调度、子代理最终回复（不含推理） + // - 丢弃：thinking/planning/progress 等噪声、子代理的工具细节与推理过程 + if d.EventType == "progress" || d.EventType == "thinking" || d.EventType == "reasoning_chain" || d.EventType == "planning" { + continue + } + + // 解析 data（JSON string），用于识别 einoRole / toolName 等 + var dataMap map[string]interface{} + if strings.TrimSpace(d.Data) != "" { + _ = json.Unmarshal([]byte(d.Data), &dataMap) + } + einoRole := "" + if v, ok := dataMap["einoRole"]; ok { + einoRole = strings.ToLower(strings.TrimSpace(fmt.Sprint(v))) + } + toolName := "" + if v, ok := dataMap["toolName"]; ok { + toolName = strings.TrimSpace(fmt.Sprint(v)) + } + + // 1) 编排器的工具调用/结果：保留（这是“主 agent 调了什么工具”） + if (d.EventType == "tool_call" || d.EventType == "tool_result" || d.EventType == "tool_calls_detected" || d.EventType == "iteration") && einoRole == "orchestrator" { + sb.WriteString("[") + sb.WriteString(d.EventType) + sb.WriteString("] ") + sb.WriteString(strings.TrimSpace(d.Message)) + sb.WriteString("\n") + if strings.TrimSpace(d.Data) != "" { + sb.WriteString(d.Data) + sb.WriteString("\n") + } + sb.WriteString("\n") + continue + } + + // 2) 子代理调度：tool_call(toolName=="task") 代表编排器把子任务派发出去；保留（只需任务，不要子代理推理） + if d.EventType == "tool_call" && strings.EqualFold(toolName, "task") { + sb.WriteString("[dispatch_subagent_task] ") + sb.WriteString(strings.TrimSpace(d.Message)) + sb.WriteString("\n") + if strings.TrimSpace(d.Data) != "" { + sb.WriteString(d.Data) + sb.WriteString("\n") + } + sb.WriteString("\n") + continue + } + + // 3) 子代理最终回复：保留（只保留最终输出，不保留分析过程） + if d.EventType == "eino_agent_reply" && einoRole == "sub" { + sb.WriteString("[subagent_final_reply] ") + sb.WriteString(strings.TrimSpace(d.Message)) + sb.WriteString("\n") + // data 里含 einoAgent 等元信息，保留有助于追踪“哪个子代理说的” + if strings.TrimSpace(d.Data) != "" { + sb.WriteString(d.Data) + sb.WriteString("\n") + } + sb.WriteString("\n") + continue + } + + // 其他事件默认丢弃，避免把子代理工具细节/推理塞进 prompt，偏离“主 agent 一轮迭代”的视角。 + } + return strings.TrimSpace(sb.String()) +} + +// buildAgentTraceInput 构建最后一轮 ReAct 的输入（从最后一条 user 消息起，不含更早轮次）。 +func (b *Builder) buildAgentTraceInput(messages []database.Message) string { + start := 0 + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "user") { + start = i + break + } + } + var builder strings.Builder + for _, msg := range messages[start:] { + builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", msg.Role, msg.Content)) + } + return builder.String() +} + +// extractUserInputFromReActInput 从保存的ReAct输入（JSON格式的messages数组）中提取最后一条用户输入 +// func (b *Builder) extractUserInputFromReActInput(reactInputJSON string) string { +// // reactInputJSON是JSON格式的ChatMessage数组，需要解析 +// var messages []map[string]interface{} +// if err := json.Unmarshal([]byte(reactInputJSON), &messages); err != nil { +// b.logger.Warn("解析ReAct输入JSON失败", zap.Error(err)) +// return "" +// } + +// // 从后往前查找最后一条user消息 +// for i := len(messages) - 1; i >= 0; i-- { +// if role, ok := messages[i]["role"].(string); ok && strings.EqualFold(role, "user") { +// if content, ok := messages[i]["content"].(string); ok { +// return content +// } +// } +// } + +// return "" +// } + +// formatAgentTraceInputFromJSON 将 JSON 轨迹转为可读文本（会先按当前任务轮次裁剪）。 +func (b *Builder) formatAgentTraceInputFromJSON(reactInputJSON string) string { + trimmed := agent.ExtractLastUserTurnTraceJSON(reactInputJSON) + msgs, err := agent.ParseTraceMessages(trimmed) + if err != nil { + b.logger.Warn("解析ReAct输入JSON失败", zap.Error(err)) + return trimmed + } + return b.formatAgentTraceFromChatMessages(msgs) +} + +// formatAgentTraceFromChatMessages 将代理消息带格式化为攻击链分析输入（与续跑轨迹字段一致）。 +func (b *Builder) formatAgentTraceFromChatMessages(msgs []agent.ChatMessage) string { + var builder strings.Builder + for _, msg := range msgs { + role := msg.Role + content := msg.Content + + if strings.EqualFold(role, "assistant") && len(msg.ToolCalls) > 0 { + if content != "" { + builder.WriteString(fmt.Sprintf("[%s]: %s\n", role, content)) + } + builder.WriteString(fmt.Sprintf("[%s] 工具调用 (%d个):\n", role, len(msg.ToolCalls))) + for i, tc := range msg.ToolCalls { + args := "" + if tc.Function.Arguments != nil { + if b, err := json.Marshal(tc.Function.Arguments); err == nil { + args = string(b) + } + } + builder.WriteString(fmt.Sprintf(" [工具调用 %d]\n", i+1)) + builder.WriteString(fmt.Sprintf(" ID: %s\n", tc.ID)) + builder.WriteString(fmt.Sprintf(" 工具名称: %s\n", tc.Function.Name)) + builder.WriteString(fmt.Sprintf(" 参数: %s\n", args)) + } + builder.WriteString("\n") + continue + } + + if strings.EqualFold(role, "tool") { + if msg.ToolCallID != "" { + builder.WriteString(fmt.Sprintf("[%s] (tool_call_id: %s):\n%s\n\n", role, msg.ToolCallID, content)) + } else { + builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", role, content)) + } + continue + } + + builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", role, content)) + } + return builder.String() +} + +// buildSimplePrompt 构建简化的prompt +func (b *Builder) buildSimplePrompt(reactInput, modelOutput string) string { + return fmt.Sprintf(`你是专业的安全测试分析师和攻击链构建专家。你的任务是根据**当前任务轮次**的对话记录和工具执行结果，一次性输出攻击链 JSON（不要分多轮追问）。 + +## 输入范围（与「继续对话」续跑一致） +- 下方「ReAct 轨迹」仅包含**最后一次用户提问之后**的消息与工具结果（last_react 当前任务轮次），不含更早的用户提问轮次。 +- 「助手结论」为同轮任务的最终输出摘要（last_react_output）；节点须与轨迹中的实际工具执行一致，严禁编造。 + +## 核心目标 + +构建一个能够讲述完整攻击故事的攻击链让学习者能够： +1. 理解渗透测试的完整流程和思维逻辑（从目标识别到漏洞发现的每一步） +2. 学习如何从失败中获取线索并调整策略 +3. 掌握工具使用的实际效果和局限性 +4. 理解漏洞发现和利用的因果关系 + +**关键原则**：完整性优先。必须包含所有有意义的工具执行和关键步骤，不要为了控制节点数量而遗漏重要信息。 + +## 构建流程（按此顺序思考） + +### 第一步：理解上下文 +仔细分析ReAct输入中的工具调用序列和大模型输出，识别： +- 测试目标（IP、域名、URL等） +- 实际执行的工具和参数 +- 工具返回的关键信息（成功结果、错误信息、超时等） +- AI的分析和决策过程 + +### 第二步：提取关键节点 +从工具执行记录中提取有意义的节点，**确保不遗漏任何关键步骤**： +- **target节点**：每个独立的测试目标创建一个target节点 +- **action节点**：每个有意义的工具执行创建一个action节点（包括提供线索的失败、成功的信息收集、漏洞验证等） +- **vulnerability节点**：每个真实确认的漏洞创建一个vulnerability节点 +- **完整性检查**：对照ReAct输入中的工具调用序列，确保每个有意义的工具执行都被包含在攻击链中 + +### 第三步：构建逻辑关系（树状结构） +**重要：必须构建树状结构，而不是简单的线性链。** +按照因果关系连接节点，形成树状图（因为是单agent执行，所以可以不按照时间顺序）： +- **分支结构**：一个节点可以有多个后续节点（例如：端口扫描发现多个端口后，可以同时进行多个不同的测试） +- **汇聚结构**：多个节点可以指向同一个节点（例如：多个不同的测试都发现了同一个漏洞） +- 识别哪些action是基于前面action的结果而执行的 +- 识别哪些vulnerability是由哪些action发现的 +- 识别失败节点如何为后续成功提供线索 +- **避免线性链**：不要将所有节点连成一条线，应该根据实际的并行测试和分支探索构建树状结构 + +### 第四步：优化和精简 +- **完整性检查**：确保所有有意义的工具执行都被包含，不要遗漏关键步骤 +- **合并规则**：只合并真正相似或重复的action节点（如多次相同工具的相似调用） +- **删除规则**：只删除完全无价值的失败节点（完全无输出、纯系统错误、重复的相同失败） +- **重要提醒**：宁可保留更多节点，也不要遗漏关键步骤。攻击链必须完整展现渗透测试过程 +- 确保攻击链逻辑连贯，能够讲述完整故事 + +## 节点类型详解 + +### target（目标节点） +- **用途**：标识测试目标 +- **创建规则**：每个独立目标（不同IP/域名）创建一个target节点 +- **多目标处理**：不同目标的节点不相互连接，各自形成独立的子图 +- **metadata.target**：精确记录目标标识（IP地址、域名、URL等） + +### action（行动节点） +- **用途**：记录工具执行和AI分析结果 +- **标签规则**： + * 15-25个汉字，动宾结构 + * 成功节点：描述执行结果（如"扫描端口发现80/443/8080"、"目录扫描发现/admin路径"） + * 失败节点：描述失败原因（如"尝试SQL注入（被WAF拦截）"、"端口扫描超时（目标不可达）"） +- **ai_analysis要求**： + * 成功节点：总结工具执行的关键发现，说明这些发现的意义 + * 失败节点：必须说明失败原因、获得的线索、这些线索如何指引后续行动 + * 不超过150字，要具体、有信息量 +- **findings要求**： + * 提取工具返回结果中的关键信息点 + * 每个finding应该是独立的、有价值的信息片段 + * 成功节点：列出关键发现（如["80端口开放", "443端口开放", "HTTP服务为Apache 2.4"]） + * 失败节点：列出失败线索（如["WAF拦截", "返回403", "检测到Cloudflare"]） +- **status标记**： + * 成功节点：不设置或设为"success" + * 提供线索的失败节点：必须设为"failed_insight" +- **risk_score**：始终为0（action节点不评估风险） + +### vulnerability（漏洞节点） +- **用途**：记录真实确认的安全漏洞 +- **创建规则**： + * 必须是真实确认的漏洞，不是所有发现都是漏洞 + * 需要明确的漏洞证据（如SQL注入返回数据库错误、XSS成功执行等） +- **risk_score规则**： + * critical（90-100）：可导致系统完全沦陷（RCE、SQL注入导致数据泄露等） + * high（80-89）：可导致敏感信息泄露或权限提升 + * medium（60-79）：存在安全风险但影响有限 + * low（40-59）：轻微安全问题 +- **metadata要求**： + * vulnerability_type：漏洞类型（SQL注入、XSS、RCE等） + * description：详细描述漏洞位置、原理、影响 + * severity：critical/high/medium/low + * location：精确的漏洞位置（URL、参数、文件路径等） + +## 节点过滤和合并规则 + +### 必须保留的失败节点 +以下失败情况必须创建节点，因为它们提供了有价值的线索： +- 工具返回明确的错误信息（权限错误、连接拒绝、认证失败等） +- 超时或连接失败（可能表明防火墙、网络隔离等） +- WAF/防火墙拦截（返回403、406等，表明存在防护机制） +- 工具未安装或配置错误（但执行了调用） +- 目标不可达（DNS解析失败、网络不通等） + +### 应该删除的失败节点 +以下情况不应创建节点： +- 完全无输出的工具调用 +- 纯系统错误（与目标无关，如本地环境问题） +- 重复的相同失败（多次相同错误只保留第一次） + +### 节点合并规则 +以下情况应合并节点： +- 同一工具的多次相似调用（如多次nmap扫描不同端口范围，合并为一个"端口扫描"节点） +- 同一目标的多个相似探测（如多个目录扫描工具，合并为一个"目录扫描"节点） + +### 节点数量控制 +- **完整性优先**：必须包含所有有意义的工具执行和关键步骤，不要为了控制数量而删除重要节点 +- **建议范围**：单目标通常8-15个节点，但如果实际执行步骤较多，可以适当增加（最多20个节点） +- **优先保留**：关键成功步骤、提供线索的失败、发现的漏洞、重要的信息收集步骤 +- **可以合并**：同一工具的多次相似调用（如多次nmap扫描不同端口范围，合并为一个"端口扫描"节点） +- **可以删除**：完全无输出的工具调用、纯系统错误、重复的相同失败（多次相同错误只保留第一次） +- **重要原则**：宁可节点稍多，也不要遗漏关键步骤。攻击链必须能够完整展现渗透测试的完整过程 + +## 边的类型和权重 + +### 边的类型 +- **leads_to**：表示"导致"或"引导到"，用于action→action、target→action + * 例如：端口扫描 → 目录扫描（因为发现了80端口，所以进行目录扫描） +- **discovers**：表示"发现"，**专门用于action→vulnerability** + * 例如：SQL注入测试 → SQL注入漏洞 + * **重要**：所有action→vulnerability的边都必须使用discovers类型，即使多个action都指向同一个vulnerability，也应该统一使用discovers +- **enables**：表示"使能"或"促成"，**仅用于vulnerability→vulnerability、action→action（当后续行动依赖前面结果时）** + * 例如：信息泄露漏洞 → 权限提升漏洞（通过信息泄露获得的信息促成了权限提升） + * **重要**：enables不能用于action→vulnerability，action→vulnerability必须使用discovers + +### 边的权重 +- **权重1-2**：弱关联（如初步探测到进一步探测） +- **权重3-4**：中等关联（如发现端口到服务识别） +- **权重5-7**：强关联（如发现漏洞、关键信息泄露） +- **权重8-10**：极强关联（如漏洞利用成功、权限提升） + +### DAG结构要求（有向无环图） +**关键：必须确保生成的是真正的DAG（有向无环图），不能有任何循环。** + +- **节点编号规则**：节点id从"node_1"开始递增（node_1, node_2, node_3...） +- **边的方向规则**：所有边的source节点id必须严格小于target节点id（source < target），这是确保无环的关键 + * 例如：node_1 → node_2 ✓（正确） + * 例如：node_2 → node_1 ✗（错误，会形成环） + * 例如：node_3 → node_5 ✓（正确） +- **无环验证**：在输出JSON前，必须检查所有边，确保没有任何一条边的source >= target +- **无孤立节点**：确保每个节点至少有一条边连接（除了可能的根节点） +- **DAG结构特点**： + * 一个节点可以有多个后续节点（分支），例如：node_2（端口扫描）可以同时连接到node_3、node_4、node_5等多个节点 + * 多个节点可以汇聚到一个节点（汇聚），例如：node_3、node_4、node_5都指向node_6（漏洞节点） + * 避免将所有节点连成一条线，应该根据实际的并行测试和分支探索构建DAG结构 +- **拓扑排序验证**：如果按照节点id从小到大排序，所有边都应该从左指向右（从上指向下），这样就能保证无环 + +## 攻击链逻辑连贯性要求 + +构建的攻击链应该能够回答以下问题： +1. **起点**：测试从哪里开始？（target节点） +2. **探索过程**：如何逐步收集信息？（action节点序列） +3. **失败与调整**：遇到障碍时如何调整策略？（failed_insight节点） +4. **关键发现**：发现了哪些重要信息？（action的findings） +5. **漏洞确认**：如何确认漏洞存在？（action→vulnerability） +6. **攻击路径**：完整的攻击路径是什么？（从target到vulnerability的路径） + +## 当前任务 ReAct 轨迹（含工具执行；助手结论见轨迹末尾 assistant） + +%s +%s + +## 输出格式 + +严格按照以下JSON格式输出，不要添加任何其他文字： + +**重要：示例展示的是树状结构，注意node_2（端口扫描）同时连接到多个后续节点（node_3、node_4），形成分支结构。** + +{ + "nodes": [ + { + "id": "node_1", + "type": "target", + "label": "测试目标: example.com", + "risk_score": 40, + "metadata": { + "target": "example.com" + } + }, + { + "id": "node_2", + "type": "action", + "label": "扫描端口发现80/443/8080", + "risk_score": 0, + "metadata": { + "tool_name": "nmap", + "tool_intent": "端口扫描", + "ai_analysis": "使用nmap对目标进行端口扫描，发现80、443、8080端口开放。80端口运行HTTP服务，443端口运行HTTPS服务，8080端口可能为管理后台。这些开放端口为后续Web应用测试提供了入口。", + "findings": ["80端口开放", "443端口开放", "8080端口开放", "HTTP服务为Apache 2.4"] + } + }, + { + "id": "node_3", + "type": "action", + "label": "目录扫描发现/admin后台", + "risk_score": 0, + "metadata": { + "tool_name": "dirsearch", + "tool_intent": "目录扫描", + "ai_analysis": "使用dirsearch对目标进行目录扫描，发现/admin目录存在且可访问。该目录可能为管理后台，是重要的测试目标。", + "findings": ["/admin目录存在", "返回200状态码", "疑似管理后台"] + } + }, + { + "id": "node_4", + "type": "action", + "label": "识别Web服务为Apache 2.4", + "risk_score": 0, + "metadata": { + "tool_name": "whatweb", + "tool_intent": "Web服务识别", + "ai_analysis": "识别出目标运行Apache 2.4服务器，这为后续的漏洞测试提供了重要信息。", + "findings": ["Apache 2.4", "PHP版本信息"] + } + }, + { + "id": "node_5", + "type": "action", + "label": "尝试SQL注入（被WAF拦截）", + "risk_score": 0, + "metadata": { + "tool_name": "sqlmap", + "tool_intent": "SQL注入检测", + "ai_analysis": "对/login.php进行SQL注入测试时被WAF拦截，返回403错误。错误信息显示检测到Cloudflare防护。这表明目标部署了WAF，需要调整测试策略。", + "findings": ["WAF拦截", "返回403", "检测到Cloudflare", "目标部署WAF"], + "status": "failed_insight" + } + }, + { + "id": "node_6", + "type": "vulnerability", + "label": "SQL注入漏洞", + "risk_score": 85, + "metadata": { + "vulnerability_type": "SQL注入", + "description": "在/admin/login.php的username参数发现SQL注入漏洞，可通过注入payload绕过登录验证，直接获取管理员权限。漏洞返回数据库错误信息，确认存在注入点。", + "severity": "high", + "location": "/admin/login.php?username=" + } + } + ], + "edges": [ + { + "source": "node_1", + "target": "node_2", + "type": "leads_to", + "weight": 3 + }, + { + "source": "node_2", + "target": "node_3", + "type": "leads_to", + "weight": 4 + }, + { + "source": "node_2", + "target": "node_4", + "type": "leads_to", + "weight": 3 + }, + { + "source": "node_3", + "target": "node_5", + "type": "leads_to", + "weight": 4 + }, + { + "source": "node_5", + "target": "node_6", + "type": "discovers", + "weight": 7 + } + ] +} + +## 重要提醒 + +1. **严禁杜撰**：只使用ReAct输入中实际执行的工具和实际返回的结果。如无实际数据，返回空的nodes和edges数组。 +2. **DAG结构必须**：必须构建真正的DAG（有向无环图），不能有任何循环。所有边的source节点id必须严格小于target节点id（source < target）。 +3. **拓扑顺序**：节点应该按照逻辑顺序编号，target节点通常是node_1，后续的action节点按执行顺序递增，vulnerability节点在最后。 +4. **完整性优先**：必须包含所有有意义的工具执行和关键步骤，不要为了控制节点数量而删除重要节点。攻击链必须能够完整展现从目标识别到漏洞发现的完整过程。 +5. **逻辑连贯**：确保攻击链能够讲述一个完整、连贯的渗透测试故事，包括所有关键步骤和决策点。 +6. **教育价值**：优先保留有教育意义的节点，帮助学习者理解渗透测试思维和完整流程。 +7. **准确性**：所有节点信息必须基于实际数据，不要推测或假设。 +8. **完整性检查**：确保每个节点都有必要的metadata字段，每条边都有正确的source和target，没有孤立节点，没有循环。 +9. **不要过度精简**：如果实际执行步骤较多，可以适当增加节点数量（最多20个），确保不遗漏关键步骤。 +10. **输出前验证**：在输出JSON前，必须验证所有边都满足source < target的条件，确保DAG结构正确。 + +现在开始分析并构建攻击链：`, reactInput, assistantOutSection(modelOutput)) +} + +func assistantOutSection(modelOutput string) string { + modelOutput = strings.TrimSpace(modelOutput) + if modelOutput == "" { + return "" + } + return "\n## 助手结论（补充）\n\n" + modelOutput + "\n" +} + +// saveChain 保存攻击链到数据库 +func (b *Builder) saveChain(conversationID string, nodes []Node, edges []Edge) error { + // 先删除旧的攻击链数据 + if err := b.db.DeleteAttackChain(conversationID); err != nil { + b.logger.Warn("删除旧攻击链失败", zap.Error(err)) + } + + for _, node := range nodes { + metadataJSON, _ := json.Marshal(node.Metadata) + if err := b.db.SaveAttackChainNode(conversationID, node.ID, node.Type, node.Label, "", string(metadataJSON), node.RiskScore); err != nil { + b.logger.Warn("保存攻击链节点失败", zap.String("nodeId", node.ID), zap.Error(err)) + } + } + + // 保存边 + for _, edge := range edges { + if err := b.db.SaveAttackChainEdge(conversationID, edge.ID, edge.Source, edge.Target, edge.Type, edge.Weight); err != nil { + b.logger.Warn("保存攻击链边失败", zap.String("edgeId", edge.ID), zap.Error(err)) + } + } + + return nil +} + +// LoadChainFromDatabase 从数据库加载攻击链 +func (b *Builder) LoadChainFromDatabase(conversationID string) (*Chain, error) { + nodes, err := b.db.LoadAttackChainNodes(conversationID) + if err != nil { + return nil, fmt.Errorf("加载攻击链节点失败: %w", err) + } + + edges, err := b.db.LoadAttackChainEdges(conversationID) + if err != nil { + return nil, fmt.Errorf("加载攻击链边失败: %w", err) + } + + return &Chain{ + Nodes: nodes, + Edges: edges, + }, nil +} + +// callAIForChainGeneration 调用AI生成攻击链 +func (b *Builder) callAIForChainGeneration(ctx context.Context, prompt string) (string, error) { + requestBody := map[string]interface{}{ + "model": b.openAIConfig.Model, + "messages": []map[string]interface{}{ + { + "role": "system", + "content": "你是一个专业的安全测试分析师，擅长构建攻击链图。请严格按照JSON格式返回攻击链数据。", + }, + { + "role": "user", + "content": prompt, + }, + }, + "temperature": 0.3, + "max_completion_tokens": attackChainMaxCompletionTokens(b.maxTokens), + } + + var apiResponse struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + + if b.openAIClient == nil { + return "", fmt.Errorf("OpenAI客户端未初始化") + } + if err := b.openAIClient.ChatCompletion(ctx, requestBody, &apiResponse); err != nil { + var apiErr *openai.APIError + if errors.As(err, &apiErr) { + bodyStr := strings.ToLower(apiErr.Body) + if strings.Contains(bodyStr, "context") || strings.Contains(bodyStr, "length") || strings.Contains(bodyStr, "too long") { + return "", fmt.Errorf("context length exceeded") + } + } else if strings.Contains(strings.ToLower(err.Error()), "context") || strings.Contains(strings.ToLower(err.Error()), "length") { + return "", fmt.Errorf("context length exceeded") + } + return "", fmt.Errorf("请求失败: %w", err) + } + + if len(apiResponse.Choices) == 0 { + return "", fmt.Errorf("API未返回有效响应") + } + + content := strings.TrimSpace(apiResponse.Choices[0].Message.Content) + // 尝试提取JSON（可能包含markdown代码块） + content = strings.TrimPrefix(content, "```json") + content = strings.TrimPrefix(content, "```") + content = strings.TrimSuffix(content, "```") + content = strings.TrimSpace(content) + + return content, nil +} + +// ChainJSON 攻击链JSON结构 +type ChainJSON struct { + Nodes []struct { + ID string `json:"id"` + Type string `json:"type"` + Label string `json:"label"` + RiskScore int `json:"risk_score"` + Metadata map[string]interface{} `json:"metadata"` + } `json:"nodes"` + Edges []struct { + Source string `json:"source"` + Target string `json:"target"` + Type string `json:"type"` + Weight int `json:"weight"` + } `json:"edges"` +} + +// parseChainJSON 解析攻击链JSON +func (b *Builder) parseChainJSON(chainJSON string) (*Chain, error) { + var chainData ChainJSON + if err := json.Unmarshal([]byte(chainJSON), &chainData); err != nil { + return nil, fmt.Errorf("解析JSON失败: %w", err) + } + + // 创建节点ID映射（AI返回的ID -> 新的UUID） + nodeIDMap := make(map[string]string) + + // 转换为Chain结构 + nodes := make([]Node, 0, len(chainData.Nodes)) + for _, n := range chainData.Nodes { + // 生成新的UUID节点ID + newNodeID := fmt.Sprintf("node_%s", uuid.New().String()) + nodeIDMap[n.ID] = newNodeID + + node := Node{ + ID: newNodeID, + Type: n.Type, + Label: n.Label, + RiskScore: n.RiskScore, + Metadata: n.Metadata, + } + if node.Metadata == nil { + node.Metadata = make(map[string]interface{}) + } + nodes = append(nodes, node) + } + + // 转换边 + edges := make([]Edge, 0, len(chainData.Edges)) + for _, e := range chainData.Edges { + sourceID, ok := nodeIDMap[e.Source] + if !ok { + continue + } + targetID, ok := nodeIDMap[e.Target] + if !ok { + continue + } + + // 生成边的ID（前端需要） + edgeID := fmt.Sprintf("edge_%s", uuid.New().String()) + + edges = append(edges, Edge{ + ID: edgeID, + Source: sourceID, + Target: targetID, + Type: e.Type, + Weight: e.Weight, + }) + } + + return &Chain{ + Nodes: nodes, + Edges: edges, + }, nil +} + +// 以下所有方法已不再使用，已删除以简化代码 diff --git a/internal/attackchain/promote_project.go b/internal/attackchain/promote_project.go new file mode 100644 index 00000000..d8a9cd80 --- /dev/null +++ b/internal/attackchain/promote_project.go @@ -0,0 +1,203 @@ +package attackchain + +import ( + "fmt" + "regexp" + "strings" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/project" + + "github.com/google/uuid" +) + +var promoteSlugSanitizer = regexp.MustCompile(`[^a-z0-9._/-]+`) + +// PromoteToProjectResult 攻击链沉淀结果。 +type PromoteToProjectResult struct { + FactsCreated int `json:"facts_created"` + FactsUpdated int `json:"facts_updated"` + EdgesCreated int `json:"edges_created"` + FactKeys []string `json:"fact_keys"` + Graph *database.ProjectFactGraph `json:"graph,omitempty"` +} + +// PromoteToProject 将对话攻击链沉淀为项目事实与边。 +func PromoteToProject(db *database.DB, projectID, conversationID string) (*PromoteToProjectResult, error) { + if db == nil { + return nil, fmt.Errorf("database 未初始化") + } + projectID = strings.TrimSpace(projectID) + conversationID = strings.TrimSpace(conversationID) + if projectID == "" || conversationID == "" { + return nil, fmt.Errorf("project_id 与 conversation_id 必填") + } + if _, err := db.GetProject(projectID); err != nil { + return nil, fmt.Errorf("项目不存在") + } + conv, err := db.GetConversation(conversationID) + if err != nil { + return nil, fmt.Errorf("对话不存在") + } + if pid := strings.TrimSpace(conv.ProjectID); pid != "" && pid != projectID { + return nil, fmt.Errorf("对话已绑定其他项目") + } + + nodes, err := db.LoadAttackChainNodes(conversationID) + if err != nil { + return nil, err + } + edges, err := db.LoadAttackChainEdges(conversationID) + if err != nil { + return nil, err + } + if len(nodes) == 0 { + return nil, fmt.Errorf("该对话尚无攻击链，请先在对话中生成攻击链") + } + + res := &PromoteToProjectResult{} + nodeToKey := make(map[string]string, len(nodes)) + usedKeys := map[string]int{} + + for _, node := range nodes { + key := allocatePromoteFactKey(node, usedKeys) + nodeToKey[node.ID] = key + category := mapPromoteNodeCategory(node.Type) + existing, getErr := db.GetProjectFactByKey(projectID, key) + f := &database.ProjectFact{ + ProjectID: projectID, + FactKey: key, + Category: category, + Summary: strings.TrimSpace(node.Label), + Body: formatPromotedFactBody(node, conversationID), + Confidence: "tentative", + SourceConversationID: conversationID, + } + if getErr == nil && existing != nil { + f.ID = existing.ID + f.CreatedAt = existing.CreatedAt + if strings.TrimSpace(f.Summary) == "" { + f.Summary = existing.Summary + } + if _, err := db.UpsertProjectFact(f); err != nil { + return nil, err + } + res.FactsUpdated++ + } else { + if _, err := db.UpsertProjectFact(f); err != nil { + return nil, err + } + res.FactsCreated++ + } + res.FactKeys = append(res.FactKeys, key) + } + + for _, edge := range edges { + srcKey, ok1 := nodeToKey[edge.Source] + tgtKey, ok2 := nodeToKey[edge.Target] + if !ok1 || !ok2 || srcKey == tgtKey { + continue + } + edgeType := mapPromoteEdgeType(edge.Type) + incoming, _ := db.ListIncomingProjectFactEdges(projectID, tgtKey) + merged := project.MergeLinkFromInputsUnique(promoteFromEdgeInputsFromDB(incoming), []database.ProjectFactEdgeFromInput{{From: srcKey, Type: edgeType}}) + if err := db.ReplaceIncomingProjectFactEdges(projectID, tgtKey, merged); err != nil { + return nil, err + } + res.EdgesCreated++ + if fact, err := db.GetProjectFactByKey(projectID, tgtKey); err == nil { + in, _ := db.ListIncomingProjectFactEdges(projectID, tgtKey) + fact.Body = project.SyncBodyLinksSection(fact.Body, in) + _, _ = db.UpsertProjectFact(fact) + } + } + + graph, _ := project.BuildProjectFactGraph(db, projectID, "full", true) + res.Graph = graph + return res, nil +} + +func promoteFromEdgeInputsFromDB(edges []*database.ProjectFactEdge) []database.ProjectFactEdgeFromInput { + out := make([]database.ProjectFactEdgeFromInput, 0, len(edges)) + for _, e := range edges { + out = append(out, database.ProjectFactEdgeFromInput{From: e.SourceFactKey, Type: e.EdgeType, Confidence: e.Confidence}) + } + return out +} + +func mapPromoteNodeCategory(nodeType string) string { + switch strings.ToLower(strings.TrimSpace(nodeType)) { + case "target": + return project.FactCategoryTarget + case "vulnerability": + return project.FactCategoryFinding + case "action": + return project.FactCategoryChain + default: + return project.FactCategoryNote + } +} + +func mapPromoteEdgeType(t string) string { + switch strings.ToLower(strings.TrimSpace(t)) { + case "discovers", "discovered_on", "targets": + return "discovered_on" + case "exploits": + return "exploits" + case "enables": + return "enables" + case "depends_on": + return "depends_on" + default: + return "leads_to" + } +} + +func allocatePromoteFactKey(node Node, used map[string]int) string { + prefix := "chain/" + switch strings.ToLower(strings.TrimSpace(node.Type)) { + case "target": + prefix = "target/" + case "vulnerability": + prefix = "finding/" + case "action": + prefix = "chain/" + } + base := promoteSlugify(node.Label) + if base == "" { + base = promoteSlugify(node.ID) + } + if base == "" { + base = uuid.New().String()[:8] + } + key := prefix + base + if n, ok := used[key]; ok { + n++ + used[key] = n + key = fmt.Sprintf("%s-%d", key, n) + } else { + used[key] = 1 + } + return key +} + +func promoteSlugify(s string) string { + s = strings.ToLower(strings.TrimSpace(s)) + s = strings.NewReplacer(" ", "-", "—", "-", "–", "-", "/", "-").Replace(s) + s = promoteSlugSanitizer.ReplaceAllString(s, "-") + s = strings.Trim(s, "-") + if len(s) > 64 { + s = s[:64] + } + return s +} + +func formatPromotedFactBody(node Node, conversationID string) string { + var b strings.Builder + b.WriteString("## 来源\n") + b.WriteString(fmt.Sprintf("- 对话攻击链沉淀\n- source_conversation_id: %s\n- node_id: %s\n- node_type: %s\n\n", conversationID, node.ID, node.Type)) + b.WriteString("## 摘要\n") + b.WriteString(strings.TrimSpace(node.Label)) + b.WriteString("\n\n## 关联\n- 结构化关系边（自动同步）:\n （见项目攻击路径图）\n") + return b.String() +} diff --git a/internal/attackchain/truncate.go b/internal/attackchain/truncate.go new file mode 100644 index 00000000..ba379b3b --- /dev/null +++ b/internal/attackchain/truncate.go @@ -0,0 +1,248 @@ +package attackchain + +import ( + "strings" + "unicode/utf8" + + "go.uber.org/zap" +) + +const ( + attackChainTruncationMarker = "\n\n...[攻击链输入已截断 / attack chain input truncated]...\n\n" + attackChainSystemReserve = 256 + attackChainSafetyReserve = 2048 +) + +// attackChainMaxCompletionTokens 为攻击链 JSON 输出预留的 completion token 上限。 +func attackChainMaxCompletionTokens(maxTotal int) int { + const capTokens = 16384 + if maxTotal <= 0 { + return 8192 + } + v := maxTotal / 8 + if v < 4096 { + v = 4096 + } + if v > capTokens { + v = capTokens + } + return v +} + +func (b *Builder) modelName() string { + if b.openAIConfig != nil && b.openAIConfig.Model != "" { + return b.openAIConfig.Model + } + return "gpt-4" +} + +func (b *Builder) countTokens(text string) int { + if text == "" { + return 0 + } + n, err := b.tokenCounter.Count(b.modelName(), text) + if err != nil { + return utf8.RuneCountInString(text) / 4 + } + return n +} + +// attackChainPayloadTokenBudget 计算 reactInput + modelOutput 可用的 token 预算。 +func (b *Builder) attackChainPayloadTokenBudget() int { + maxTotal := b.maxTokens + if maxTotal <= 0 { + maxTotal = 100000 + } + templateTok := b.countTokens(b.buildSimplePrompt("", "")) + completion := attackChainMaxCompletionTokens(maxTotal) + reserve := templateTok + attackChainSystemReserve + completion + attackChainSafetyReserve + budget := maxTotal - reserve + minBudget := maxTotal * 35 / 100 + if budget < minBudget { + budget = minBudget + } + if budget < 4096 { + budget = 4096 + } + return budget +} + +// fitAttackChainPayload 在构建最终 prompt 前压缩 ReAct 轨迹与模型输出，避免超出模型上下文。 +func (b *Builder) fitAttackChainPayload(reactInput, modelOutput string) (string, string, bool) { + budget := b.attackChainPayloadTokenBudget() + modelBudget := budget * 15 / 100 + if modelBudget < 512 { + modelBudget = 512 + } + reactBudget := budget - modelBudget + + origReactTok := b.countTokens(reactInput) + origModelTok := b.countTokens(modelOutput) + truncated := false + + outModel := modelOutput + if origModelTok > modelBudget { + outModel = truncateTextByTokens(b, modelOutput, modelBudget) + truncated = true + } + + outReact := reactInput + perToolLimits := []int{12000, 6000, 3000, 1500, 800} + for _, lim := range perToolLimits { + compact := compactFormattedToolBodies(outReact, lim) + if compact != outReact { + outReact = compact + truncated = true + } + if b.countTokens(outReact) <= reactBudget { + break + } + } + + if b.countTokens(outReact) > reactBudget { + outReact = truncateTextByTokens(b, outReact, reactBudget) + truncated = true + } + + if truncated { + b.logger.Info("攻击链输入已按 token 预算截断", + zap.Int("maxTotalTokens", b.maxTokens), + zap.Int("payloadBudget", budget), + zap.Int("reactBudget", reactBudget), + zap.Int("modelBudget", modelBudget), + zap.Int("reactInputTokensBefore", origReactTok), + zap.Int("reactInputTokensAfter", b.countTokens(outReact)), + zap.Int("modelOutputTokensBefore", origModelTok), + zap.Int("modelOutputTokensAfter", b.countTokens(outModel)), + zap.Int("maxCompletionTokens", attackChainMaxCompletionTokens(b.maxTokens)), + ) + } + + return outReact, outModel, truncated +} + +// compactFormattedToolBodies 缩短格式化 trace 中 [tool] 消息的正文，保留工具头与调用 ID。 +func compactFormattedToolBodies(s string, maxRunesPerBody int) string { + if maxRunesPerBody <= 0 || s == "" { + return s + } + const marker = "[tool]" + var out strings.Builder + remaining := s + changed := false + for { + idx := strings.Index(remaining, marker) + if idx < 0 { + out.WriteString(remaining) + break + } + out.WriteString(remaining[:idx]) + remaining = remaining[idx:] + nl := strings.IndexByte(remaining, '\n') + if nl < 0 { + out.WriteString(remaining) + break + } + header := remaining[:nl+1] + remaining = remaining[nl+1:] + bodyEnd := strings.Index(remaining, "\n\n[") + var body, rest string + if bodyEnd < 0 { + body = remaining + rest = "" + } else { + body = remaining[:bodyEnd] + rest = remaining[bodyEnd:] + } + if runeLen(body) > maxRunesPerBody { + body = truncateRunesWithNotice(body, maxRunesPerBody) + changed = true + } + out.WriteString(header) + out.WriteString(body) + remaining = rest + if rest == "" { + break + } + } + if !changed { + return s + } + return out.String() +} + +func truncateTextByTokens(b *Builder, text string, maxTokens int) string { + if maxTokens <= 0 || text == "" { + return "" + } + if b.countTokens(text) <= maxTokens { + return text + } + markerTok := b.countTokens(attackChainTruncationMarker) + usable := maxTokens - markerTok + if usable < 256 { + usable = maxTokens / 2 + } + headBudget := usable * 60 / 100 + tailBudget := usable - headBudget + head := takeTokensFromStart(b, text, headBudget) + tail := takeTokensFromEnd(b, text, tailBudget) + return head + attackChainTruncationMarker + tail +} + +func takeTokensFromStart(b *Builder, text string, maxTokens int) string { + rs := []rune(text) + if len(rs) == 0 || maxTokens <= 0 { + return "" + } + lo, hi := 0, len(rs) + for lo < hi { + mid := (lo + hi + 1) / 2 + if b.countTokens(string(rs[:mid])) <= maxTokens { + lo = mid + } else { + hi = mid - 1 + } + } + return string(rs[:lo]) +} + +func takeTokensFromEnd(b *Builder, text string, maxTokens int) string { + rs := []rune(text) + if len(rs) == 0 || maxTokens <= 0 { + return "" + } + lo, hi := 0, len(rs) + for lo < hi { + mid := (lo + hi) / 2 + if b.countTokens(string(rs[mid:])) <= maxTokens { + hi = mid + } else { + lo = mid + 1 + } + } + return string(rs[lo:]) +} + +func truncateRunesWithNotice(s string, maxRunes int) string { + rs := []rune(s) + if len(rs) <= maxRunes { + return s + } + const notice = "\n...[工具输出已截断 / tool output truncated]...\n" + noticeRunes := []rune(notice) + keep := maxRunes - len(noticeRunes) + if keep < 200 { + keep = maxRunes * 2 / 3 + } + if keep < 1 { + return notice + } + head := keep * 70 / 100 + tail := keep - head + return string(rs[:head]) + notice + string(rs[len(rs)-tail:]) +} + +func runeLen(s string) int { + return len([]rune(s)) +} diff --git a/internal/attackchain/truncate_test.go b/internal/attackchain/truncate_test.go new file mode 100644 index 00000000..2cb4563c --- /dev/null +++ b/internal/attackchain/truncate_test.go @@ -0,0 +1,63 @@ +package attackchain + +import ( + "strings" + "testing" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +func testBuilder(maxTotal int) *Builder { + return &Builder{ + logger: zap.NewNop(), + openAIConfig: &config.OpenAIConfig{Model: "gpt-4"}, + tokenCounter: agent.NewTikTokenCounter(), + maxTokens: maxTotal, + } +} + +func TestCompactFormattedToolBodies(t *testing.T) { + long := strings.Repeat("x", 20000) + in := "[user]: hi\n\n[tool] (tool_call_id: abc):\n" + long + "\n\n[assistant]: done\n" + out := compactFormattedToolBodies(in, 500) + if strings.Contains(out, strings.Repeat("x", 10000)) { + t.Fatal("expected tool body to be truncated") + } + if !strings.Contains(out, "[user]: hi") { + t.Fatal("expected user header preserved") + } + if !strings.Contains(out, "[assistant]: done") { + t.Fatal("expected assistant header preserved") + } +} + +func TestFitAttackChainPayloadWithinBudget(t *testing.T) { + b := testBuilder(32000) + react := strings.Repeat("scan ", 50000) + model := strings.Repeat("result ", 10000) + r, m, truncated := b.fitAttackChainPayload(react, model) + if !truncated { + t.Fatal("expected truncation for large payload") + } + prompt := b.buildSimplePrompt(r, m) + total := b.countTokens(prompt) + attackChainMaxCompletionTokens(b.maxTokens) + attackChainSystemReserve + if total > b.maxTokens+attackChainSafetyReserve { + t.Fatalf("prompt still too large: estimated %d > max %d", total, b.maxTokens) + } + _ = m +} + +func TestAttackChainMaxCompletionTokens(t *testing.T) { + if got := attackChainMaxCompletionTokens(120000); got != 15000 && got != 16384 { + // 120000/8 = 15000 + if got < 4096 || got > 16384 { + t.Fatalf("unexpected completion cap: %d", got) + } + } + if got := attackChainMaxCompletionTokens(0); got != 8192 { + t.Fatalf("expected default 8192, got %d", got) + } +} diff --git a/internal/project/blackboard.go b/internal/project/blackboard.go new file mode 100644 index 00000000..d1e2aec9 --- /dev/null +++ b/internal/project/blackboard.go @@ -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 +} diff --git a/internal/project/blackboard_refresh.go b/internal/project/blackboard_refresh.go new file mode 100644 index 00000000..6a494727 --- /dev/null +++ b/internal/project/blackboard_refresh.go @@ -0,0 +1,56 @@ +package project + +import "strings" + +// FactIndexSectionHeading 黑板索引可读标题行前缀（块内保留，供 Agent 阅读）。 +const FactIndexSectionHeading = "## 项目黑板索引" + +// FactIndexSectionStartMarker / EndMarker：HTML 注释边界，供程序化替换；对模型无指令语义。 +const ( + FactIndexSectionStartMarker = "" + FactIndexSectionEndMarker = "" +) + +// 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 +} diff --git a/internal/project/blackboard_refresh_test.go b/internal/project/blackboard_refresh_test.go new file mode 100644 index 00000000..31e9db4d --- /dev/null +++ b/internal/project/blackboard_refresh_test.go @@ -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) + } +} diff --git a/internal/project/fact_body_links.go b/internal/project/fact_body_links.go new file mode 100644 index 00000000..8c0bd39c --- /dev/null +++ b/internal/project/fact_body_links.go @@ -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 +} diff --git a/internal/project/fact_body_links_test.go b/internal/project/fact_body_links_test.go new file mode 100644 index 00000000..1b5daa95 --- /dev/null +++ b/internal/project/fact_body_links_test.go @@ -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)) + } +} diff --git a/internal/project/fact_edges.go b/internal/project/fact_edges.go new file mode 100644 index 00000000..d9d15795 --- /dev/null +++ b/internal/project/fact_edges.go @@ -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:） + 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() +} diff --git a/internal/project/fact_edges_apply.go b/internal/project/fact_edges_apply.go new file mode 100644 index 00000000..870861e4 --- /dev/null +++ b/internal/project/fact_edges_apply.go @@ -0,0 +1,96 @@ +package project + +import ( + "cyberstrike-ai/internal/database" +) + +// ApplyFactOutgoingLinks 替换某事实的出边（links 为 nil 时不修改）。 +func ApplyFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput) error { + if links == nil { + return nil + } + return db.ReplaceOutgoingProjectFactEdges(projectID, sourceFactKey, sourceConversationID, links) +} + +// ResolveFactLinkInputs 合并 links 数组与 links_text 文本（数组优先）。 +func ResolveFactLinkInputs(links []database.ProjectFactEdgeFromInput, linksText string) ([]database.ProjectFactEdgeFromInput, error) { + if len(links) > 0 { + return links, nil + } + return ParseFactLinksText(linksText) +} + +// ApplyFactIncomingLinks 替换某事实的入边（links 为 nil 时不修改）。 +func ApplyFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput) error { + if links == nil { + return nil + } + return db.ReplaceIncomingProjectFactEdges(projectID, targetFactKey, links) +} + +// PersistFactIncomingLinks 写入入边并可选同步当前事实 body「关联」段。 +func PersistFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput, syncBody bool) error { + if links == nil { + return nil + } + if err := ApplyFactIncomingLinks(db, projectID, targetFactKey, links); err != nil { + return err + } + if !syncBody { + return nil + } + f, err := db.GetProjectFactByKey(projectID, targetFactKey) + if err != nil { + return nil + } + in, err := db.ListIncomingProjectFactEdges(projectID, targetFactKey) + if err != nil { + return err + } + f.Body = SyncBodyLinksSection(f.Body, in) + _, err = db.UpsertProjectFact(f) + return err +} + +// PersistFactLinksFromParsed 写入解析后的 links（parsed 为 nil 表示不修改）。 +func PersistFactLinksFromParsed(db *database.DB, projectID, factKey, sourceConversationID string, parsed *ParsedFactLinks, syncBody bool) error { + if parsed == nil || parsed.Incoming == nil { + return nil + } + return PersistFactIncomingLinks(db, projectID, factKey, parsed.Incoming, syncBody) +} + +// PersistFactOutgoingLinks 写入出边（图连线等低层 API；body 同步请用 PersistFactIncomingLinks）。 +func PersistFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput, syncBody bool) error { + if links == nil { + return nil + } + return ApplyFactOutgoingLinks(db, projectID, sourceFactKey, sourceConversationID, links) +} + +// LinkCountMap 项目内各 fact 的入/出边计数。 +type LinkCountMap map[string]LinkCounts + +// LinkCounts 单 fact 的入/出边数。 +type LinkCounts struct { + Outgoing int `json:"outgoing"` + Incoming int `json:"incoming"` +} + +// LoadProjectFactLinkCounts 批量加载边计数。 +func LoadProjectFactLinkCounts(db *database.DB, projectID string) (LinkCountMap, error) { + edges, err := db.ListProjectFactEdgesByProject(projectID) + if err != nil { + return nil, err + } + m := LinkCountMap{} + for _, e := range edges { + c := m[e.SourceFactKey] + c.Outgoing++ + m[e.SourceFactKey] = c + c = m[e.TargetFactKey] + c.Incoming++ + m[e.TargetFactKey] = c + } + return m, nil +} diff --git a/internal/project/fact_edges_test.go b/internal/project/fact_edges_test.go new file mode 100644 index 00000000..2e4b3775 --- /dev/null +++ b/internal/project/fact_edges_test.go @@ -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") + } +} diff --git a/internal/project/fact_index_links.go b/internal/project/fact_index_links.go new file mode 100644 index 00000000..32732894 --- /dev/null +++ b/internal/project/fact_index_links.go @@ -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) + }) +} diff --git a/internal/project/fact_index_links_test.go b/internal/project/fact_index_links_test.go new file mode 100644 index 00000000..a5794b9d --- /dev/null +++ b/internal/project/fact_index_links_test.go @@ -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) + } +} diff --git a/internal/project/fact_recording_prompt.go b/internal/project/fact_recording_prompt.go new file mode 100644 index 00000000..7d986a46 --- /dev/null +++ b/internal/project/fact_recording_prompt.go @@ -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) +} diff --git a/internal/project/fact_template.go b/internal/project/fact_template.go new file mode 100644 index 00000000..c94c819c --- /dev/null +++ b/internal/project/fact_template.go @@ -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 = `## 结论（可验证，一句话） +<勿仅写「存在漏洞」；写明类型 + 位置 + 触发条件> + +## 目标与入口 +- 目标: +- 入口: <路径 / 接口 / 参数> +- 前置条件: <匿名 / 角色 / Cookie / 其他依赖> + +## 攻击链（逐步可复现） +1. <侦察/发现> +2. <利用/触发> +3. <影响证明（读文件、RCE 回显、越权数据等）> + +## Exploit / POC +### 请求 +` + "```http\n HTTP/1.1\nHost: ...\n...\n\n\n```" + ` + +### 响应 / 现象 +<关键响应片段、状态码、差异点> + +### 命令 / 脚本（如有） +` + "```bash\n\n```" + ` + +## 关键证据 +- <工具输出摘要 / 截图路径 / 会话或消息 ID> + +## 关联 +- related_vulnerability_id: <可选，对应 record_vulnerability 的 id> +- links（upsert 参数）: [{ "from": "", "type": "discovered_on|..." }]（from → 当前 fact） +- 依赖事实（body 可读镜像）: + +## 备注与不确定性 +<待验证假设、环境差异、绕过尝试记录>` + +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) +} diff --git a/internal/project/fact_template_test.go b/internal/project/fact_template_test.go new file mode 100644 index 00000000..172bc0b6 --- /dev/null +++ b/internal/project/fact_template_test.go @@ -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") + } +} \ No newline at end of file diff --git a/internal/project/scope_block.go b/internal/project/scope_block.go new file mode 100644 index 00000000..e52cf1ea --- /dev/null +++ b/internal/project/scope_block.go @@ -0,0 +1,99 @@ +package project + +import ( + "encoding/json" + "fmt" + "strings" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" +) + +// projectScopePayload 解析 projects.scope_json（约定字段，可扩展）。 +type projectScopePayload struct { + Targets []string `json:"targets"` + Exclude []string `json:"exclude"` + Notes string `json:"notes"` +} + +// BuildScopeBlock 将项目 scope_json 格式化为 Agent 可读的授权范围块。 +func BuildScopeBlock(proj *database.Project) string { + if proj == nil { + return "" + } + raw := strings.TrimSpace(proj.ScopeJSON) + if raw == "" { + return "" + } + + var payload projectScopePayload + if err := json.Unmarshal([]byte(raw), &payload); err != nil { + return fmt.Sprintf("## 项目测试范围（project: %s）\n（scope_json 非合法 JSON，请人工核对配置）\n```\n%s\n```\n"+ + "仅对明确授权目标执行测试；超出范围须停止并说明。\n", proj.Name, truncateRunes(raw, 800)) + } + + var b strings.Builder + b.WriteString(fmt.Sprintf("## 项目测试范围（project: %s, id: %s）\n", proj.Name, proj.ID)) + b.WriteString("以下为授权边界，**必须遵守**：仅测试列出的 targets，避开 exclude，不得擅自扩大范围。\n") + + if len(payload.Targets) > 0 { + b.WriteString("\n**允许测试（targets）**：\n") + for _, t := range payload.Targets { + t = strings.TrimSpace(t) + if t != "" { + b.WriteString("- " + t + "\n") + } + } + } + if len(payload.Exclude) > 0 { + b.WriteString("\n**明确排除（exclude）**：\n") + for _, t := range payload.Exclude { + t = strings.TrimSpace(t) + if t != "" { + b.WriteString("- " + t + "\n") + } + } + } + if n := strings.TrimSpace(payload.Notes); n != "" { + b.WriteString("\n**说明（notes）**：\n" + n + "\n") + } + if len(payload.Targets) == 0 && len(payload.Exclude) == 0 && strings.TrimSpace(payload.Notes) == "" { + b.WriteString("\n（scope_json 已配置但未识别 targets/exclude/notes 字段，原始内容供参考）\n```json\n") + b.WriteString(truncateRunes(raw, 1200)) + b.WriteString("\n```\n") + } + b.WriteString("\n若目标不在 targets 内或命中 exclude，不得主动扫描/利用；需用户明确扩大授权后再继续。\n") + return b.String() +} + +func truncateRunes(s string, max int) string { + r := []rune(s) + if len(r) <= max { + return s + } + return string(r[:max]) + "…" +} + +// BuildProjectBlackboardBlock 组合测试范围 + 事实黑板索引。 +func BuildProjectBlackboardBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) { + projectID = strings.TrimSpace(projectID) + if projectID == "" { + return "", nil + } + proj, err := db.GetProject(projectID) + if err != nil { + return "", err + } + parts := []string{} + if scope := strings.TrimSpace(BuildScopeBlock(proj)); scope != "" { + parts = append(parts, scope) + } + index, err := BuildFactIndexBlock(db, projectID, cfg) + if err != nil { + return "", err + } + if strings.TrimSpace(index) != "" { + parts = append(parts, index) + } + return strings.Join(parts, "\n\n"), nil +} diff --git a/internal/project/scope_block_test.go b/internal/project/scope_block_test.go new file mode 100644 index 00000000..11a5a264 --- /dev/null +++ b/internal/project/scope_block_test.go @@ -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) + } +} diff --git a/internal/project/stats.go b/internal/project/stats.go new file mode 100644 index 00000000..b6e1d1b3 --- /dev/null +++ b/internal/project/stats.go @@ -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 +} diff --git a/internal/project/vision_image_prompt.go b/internal/project/vision_image_prompt.go new file mode 100644 index 00000000..12e901fb --- /dev/null +++ b/internal/project/vision_image_prompt.go @@ -0,0 +1,26 @@ +package project + +import "strings" + +// VisionImageSectionMarker 图片分析 section 标题（与 AppendVisionImageAnalysisIfReady 注入一致）。 +const VisionImageSectionMarker = "## 图片分析" + +// VisionImageAnalysisSection 单/多代理共用的图片分析提示（analyze_image；上下文仅保留文字摘要）。 +func VisionImageAnalysisSection() string { + var b strings.Builder + b.WriteString(VisionImageSectionMarker) + b.WriteString("\n\n") + b.WriteString("- 遇到图片文件（截图、验证码、登录页、报告配图）时，若存在工具 analyze_image，请传入服务器上的文件路径进行分析。\n") + b.WriteString("- 不要对二进制图片使用 read_file 指望理解内容；用户消息中「📎 xxx.png: /path」即为可传给 analyze_image 的路径。\n") + b.WriteString("- 验证码类：若已从页面或接口保存为本地图片（如 captcha.png），用 analyze_image，question 写明「只输出验证码字符」；识别失败则刷新验证码后重新保存再识；复杂滑块/行为验证码勿指望单次识图成功。\n") + b.WriteString("- 委派子代理时，若子任务含验证码/截图识读，在 task description 中写明图片路径与期望输出格式。\n") + return b.String() +} + +// AppendVisionImageAnalysisIfReady 仅在 vision.enabled 且 model 已配置时追加图片分析提示。 +func AppendVisionImageAnalysisIfReady(base string, visionReady bool) string { + if !visionReady { + return base + } + return AppendSystemPromptBlock(base, VisionImageAnalysisSection()) +} diff --git a/internal/project/workspace.go b/internal/project/workspace.go new file mode 100644 index 00000000..55a9137d --- /dev/null +++ b/internal/project/workspace.go @@ -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//; otherwise conversations//. +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) +} diff --git a/internal/project/workspace_test.go b/internal/project/workspace_test.go new file mode 100644 index 00000000..dd62b162 --- /dev/null +++ b/internal/project/workspace_test.go @@ -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) + } +} diff --git a/internal/vision/client.go b/internal/vision/client.go new file mode 100644 index 00000000..721aa18d --- /dev/null +++ b/internal/vision/client.go @@ -0,0 +1,134 @@ +package vision + +import ( + "context" + "encoding/base64" + "fmt" + "net" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/openai" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" + "github.com/cloudwego/eino/schema" +) + +// Client 调用独立 Vision ChatModel（单次 Generate）。 +type Client struct { + cfg config.VisionConfig + mainOA config.OpenAIConfig +} + +// NewClient 构造视觉客户端。 +func NewClient(visionCfg config.VisionConfig, mainOpenAI config.OpenAIConfig) *Client { + return &Client{cfg: visionCfg, mainOA: mainOpenAI} +} + +// Analyze 将图片字节送入 VL 模型并返回文本描述。 +func (c *Client) Analyze(ctx context.Context, img ImagePayload, question string) (string, error) { + if len(img.Bytes) == 0 { + return "", fmt.Errorf("empty image payload") + } + mime := strings.TrimSpace(img.MIMEType) + if mime == "" { + mime = "image/jpeg" + } + oa := c.cfg.OpenAICfgEffective(c.mainOA) + if strings.TrimSpace(oa.APIKey) == "" { + return "", fmt.Errorf("vision API key is empty (set vision.api_key or openai.api_key)") + } + if strings.TrimSpace(oa.Model) == "" { + return "", fmt.Errorf("vision model is empty") + } + + timeout := time.Duration(c.cfg.TimeoutSecondsEffective()) * time.Second + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + httpClient := &http.Client{ + Timeout: timeout + 15*time.Second, + Transport: &http.Transport{ + DialContext: (&net.Dialer{ + Timeout: 60 * time.Second, + KeepAlive: 60 * time.Second, + }).DialContext, + ResponseHeaderTimeout: timeout + 10*time.Second, + }, + } + httpClient = openai.NewEinoHTTPClient(&oa, httpClient) + + maxCompletionTokens := oa.MaxCompletionTokensEffective() + modelCfg := &einoopenai.ChatModelConfig{ + APIKey: oa.APIKey, + BaseURL: strings.TrimSuffix(oa.BaseURL, "/"), + Model: oa.Model, + HTTPClient: httpClient, + MaxCompletionTokens: &maxCompletionTokens, + } + chatModel, err := einoopenai.NewChatModel(ctx, modelCfg) + if err != nil { + return "", fmt.Errorf("vision chat model: %w", err) + } + + b64 := base64.StdEncoding.EncodeToString(img.Bytes) + detail := schema.ImageURLDetailLow + switch c.cfg.DetailEffective() { + case "high": + detail = schema.ImageURLDetailHigh + case "auto": + detail = schema.ImageURLDetailAuto + } + + prompt := buildVisionPrompt(question) + userMsg := &schema.Message{ + Role: schema.User, + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: prompt}, + { + Type: schema.ChatMessagePartTypeImageURL, + Image: &schema.MessageInputImage{ + MessagePartCommon: schema.MessagePartCommon{ + Base64Data: &b64, + MIMEType: mime, + }, + Detail: detail, + }, + }, + }, + } + + resp, err := chatModel.Generate(ctx, []*schema.Message{userMsg}) + if err != nil { + return "", fmt.Errorf("vision generate: %w", err) + } + if resp == nil || strings.TrimSpace(resp.Content) == "" { + return "", fmt.Errorf("vision model returned empty content") + } + return strings.TrimSpace(resp.Content), nil +} + +func buildVisionPrompt(question string) string { + q := strings.TrimSpace(question) + if q == "" { + q = "请对图片做通用描述，侧重授权安全测试场景（可见文本、表单、按钮、验证码、错误信息、技术栈线索）。" + } + extra := "" + if looksLikeCaptchaQuestion(q) { + extra = "\n若为验证码：仅输出你辨认出的字符序列，不要空格、标点、解释；看不清则明确说无法识别。" + } + return `你是授权安全测试助手。请根据图片回答用户问题，只描述你能从图中确认的内容，不要编造。 +用户问题：` + q + extra +} + +func looksLikeCaptchaQuestion(q string) bool { + s := strings.ToLower(q) + for _, kw := range []string{"验证码", "captcha", "verification code", "verify code", "vcode", "图形码"} { + if strings.Contains(s, kw) { + return true + } + } + return strings.Contains(s, "只输出") && (strings.Contains(s, "字符") || strings.Contains(s, "character")) +} diff --git a/internal/vision/client_test.go b/internal/vision/client_test.go new file mode 100644 index 00000000..101aa943 --- /dev/null +++ b/internal/vision/client_test.go @@ -0,0 +1,12 @@ +package vision + +import "testing" + +func TestLooksLikeCaptchaQuestion(t *testing.T) { + if !looksLikeCaptchaQuestion("识别验证码，只输出字符") { + t.Fatal("expected captcha hint") + } + if looksLikeCaptchaQuestion("描述登录页布局") { + t.Fatal("expected non-captcha") + } +} diff --git a/internal/vision/path.go b/internal/vision/path.go new file mode 100644 index 00000000..3d9756ed --- /dev/null +++ b/internal/vision/path.go @@ -0,0 +1,72 @@ +package vision + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +var allowedImageExt = map[string]struct{}{ + ".png": {}, ".jpg": {}, ".jpeg": {}, ".webp": {}, ".gif": {}, + ".bmp": {}, ".tif": {}, ".tiff": {}, +} + +// ResolveImagePath 解析并校验可读图片路径（支持任意目录；仍校验扩展名与常规文件）。 +func ResolveImagePath(path string, cwd string) (string, error) { + p := strings.TrimSpace(path) + if p == "" { + return "", fmt.Errorf("path is empty") + } + cwdTrim := strings.TrimSpace(cwd) + if cwdTrim == "" { + var err error + cwdTrim, err = os.Getwd() + if err != nil { + return "", fmt.Errorf("getwd: %w", err) + } + } + cwdAbs, err := filepath.Abs(filepath.Clean(cwdTrim)) + if err != nil { + return "", err + } + + var candidate string + if filepath.IsAbs(p) { + candidate = filepath.Clean(p) + } else { + candidate = filepath.Clean(filepath.Join(cwdAbs, p)) + } + resolved := normalizeAbsPath(candidate) + if resolved == "" { + return "", fmt.Errorf("invalid path") + } + + ext := strings.ToLower(filepath.Ext(resolved)) + if _, ok := allowedImageExt[ext]; !ok { + return "", fmt.Errorf("unsupported image extension %q", ext) + } + + st, err := os.Stat(resolved) + if err != nil { + return "", fmt.Errorf("stat: %w", err) + } + if st.IsDir() { + return "", fmt.Errorf("not a regular file") + } + if st.Size() > 0 && st.Size() > 1<<30 { + return "", fmt.Errorf("file too large on disk") + } + return resolved, nil +} + +func normalizeAbsPath(p string) string { + abs, err := filepath.Abs(filepath.Clean(p)) + if err != nil { + return "" + } + if link, err := filepath.EvalSymlinks(abs); err == nil { + return link + } + return abs +} diff --git a/internal/vision/path_test.go b/internal/vision/path_test.go new file mode 100644 index 00000000..b38206bf --- /dev/null +++ b/internal/vision/path_test.go @@ -0,0 +1,52 @@ +package vision + +import ( + "os" + "path/filepath" + "testing" +) + +func TestResolveImagePath_underCWD(t *testing.T) { + dir := t.TempDir() + img := filepath.Join(dir, "shot.png") + if err := os.WriteFile(img, []byte{0x89, 0x50, 0x4e, 0x47}, 0o644); err != nil { + t.Fatal(err) + } + got, err := ResolveImagePath(img, dir) + if err != nil { + t.Fatal(err) + } + want := normalizeAbsPath(img) + if got != want { + t.Fatalf("got %q want %q", got, want) + } +} + +func TestResolveImagePath_absoluteOutsideCWD(t *testing.T) { + dir := t.TempDir() + cwd := t.TempDir() + img := filepath.Join(dir, "remote.png") + if err := os.WriteFile(img, []byte{0x89, 0x50, 0x4e, 0x47}, 0o644); err != nil { + t.Fatal(err) + } + got, err := ResolveImagePath(img, cwd) + if err != nil { + t.Fatalf("expected absolute path outside cwd to be allowed: %v", err) + } + want := normalizeAbsPath(img) + if got != want { + t.Fatalf("got %q want %q", got, want) + } +} + +func TestResolveImagePath_rejectsNonImageExt(t *testing.T) { + dir := t.TempDir() + f := filepath.Join(dir, "notes.txt") + if err := os.WriteFile(f, []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + _, err := ResolveImagePath(f, dir) + if err == nil { + t.Fatal("expected error for non-image extension") + } +} diff --git a/internal/vision/preprocess.go b/internal/vision/preprocess.go new file mode 100644 index 00000000..860dab63 --- /dev/null +++ b/internal/vision/preprocess.go @@ -0,0 +1,212 @@ +package vision + +import ( + "bytes" + "fmt" + "image" + "os" + "strings" + + "github.com/disintegration/imaging" +) + +// ImagePayload 送入 VL API 的图片字节与 MIME。 +type ImagePayload struct { + Bytes []byte + MIMEType string +} + +// PreprocessMeta 记录缩放与编码结果，供工具输出与排障。 +type PreprocessMeta struct { + OriginalPath string + OriginalBytes int64 + OriginalWidth int + OriginalHeight int + OutputWidth int + OutputHeight int + OutputBytes int + OutputMIMEType string + JPEGQuality int // 0 表示未 JPEG 重编码（原图直传） + PreprocessMode string // passthrough | jpeg +} + +// PreprocessOptions 图片预处理参数。 +type PreprocessOptions struct { + MaxImageBytes int64 + MaxDimension int + JPEGQuality int + MaxPayloadBytes int64 + SkipPreprocessBelowBytes int64 // 0 = 始终压缩；>0 时小图+尺寸合规可直传 +} + +// PreprocessImageFile 读取图片；大图或超尺寸走 imaging 缩放+JPEG，否则可原图直传。 +func PreprocessImageFile(path string, opt PreprocessOptions) (ImagePayload, PreprocessMeta, error) { + var meta PreprocessMeta + meta.OriginalPath = path + + st, err := os.Stat(path) + if err != nil { + return ImagePayload{}, meta, err + } + meta.OriginalBytes = st.Size() + if opt.MaxImageBytes > 0 && st.Size() > opt.MaxImageBytes { + return ImagePayload{}, meta, fmt.Errorf("file size %d exceeds max_image_bytes %d", st.Size(), opt.MaxImageBytes) + } + + cfgW, cfgH, format, err := imageDimensions(path) + if err != nil { + return ImagePayload{}, meta, err + } + meta.OriginalWidth = cfgW + meta.OriginalHeight = cfgH + + maxDim := opt.MaxDimension + if maxDim <= 0 { + maxDim = 2048 + } + maxPayload := opt.MaxPayloadBytes + if maxPayload <= 0 { + maxPayload = 512 * 1024 + } + + if payload, meta, ok, err := tryPassthrough(path, st.Size(), cfgW, cfgH, format, opt, maxDim, maxPayload); ok { + return payload, meta, err + } + + return compressWithImaging(path, opt, maxDim, maxPayload, meta) +} + +func tryPassthrough(path string, size int64, w, h int, format string, opt PreprocessOptions, maxDim int, maxPayload int64) (ImagePayload, PreprocessMeta, bool, error) { + var meta PreprocessMeta + meta.OriginalPath = path + meta.OriginalBytes = size + meta.OriginalWidth = w + meta.OriginalHeight = h + + threshold := opt.SkipPreprocessBelowBytes + if threshold <= 0 { + return ImagePayload{}, meta, false, nil + } + if size > threshold { + return ImagePayload{}, meta, false, nil + } + longEdge := w + if h > longEdge { + longEdge = h + } + if longEdge > maxDim { + return ImagePayload{}, meta, false, nil + } + if size > maxPayload { + return ImagePayload{}, meta, false, nil + } + + raw, err := os.ReadFile(path) + if err != nil { + return ImagePayload{}, meta, false, err + } + mime := mimeFromImageFormat(format) + if mime == "" { + return ImagePayload{}, meta, false, nil + } + + meta.OutputWidth = w + meta.OutputHeight = h + meta.OutputBytes = len(raw) + meta.OutputMIMEType = mime + meta.PreprocessMode = "passthrough" + return ImagePayload{Bytes: raw, MIMEType: mime}, meta, true, nil +} + +func compressWithImaging(path string, opt PreprocessOptions, maxDim int, maxPayload int64, meta PreprocessMeta) (ImagePayload, PreprocessMeta, error) { + src, err := imaging.Open(path) + if err != nil { + return ImagePayload{}, meta, fmt.Errorf("open image: %w", err) + } + bounds := src.Bounds() + meta.OriginalWidth = bounds.Dx() + meta.OriginalHeight = bounds.Dy() + + dst := imaging.Fit(src, maxDim, maxDim, imaging.Lanczos) + outBounds := dst.Bounds() + meta.OutputWidth = outBounds.Dx() + meta.OutputHeight = outBounds.Dy() + + quality := opt.JPEGQuality + if quality <= 0 || quality > 100 { + quality = 82 + } + + dim := maxDim + for attempt := 0; attempt < 6; attempt++ { + if attempt > 0 { + dim = int(float64(dim) * 0.85) + if dim < 256 { + dim = 256 + } + dst = imaging.Fit(src, dim, dim, imaging.Lanczos) + outBounds = dst.Bounds() + meta.OutputWidth = outBounds.Dx() + meta.OutputHeight = outBounds.Dy() + } + q := quality + for q >= 60 { + var buf bytes.Buffer + if err := imaging.Encode(&buf, dst, imaging.JPEG, imaging.JPEGQuality(q)); err != nil { + return ImagePayload{}, meta, fmt.Errorf("encode jpeg: %w", err) + } + if int64(buf.Len()) <= maxPayload { + meta.JPEGQuality = q + meta.OutputBytes = buf.Len() + meta.OutputMIMEType = "image/jpeg" + meta.PreprocessMode = "jpeg" + return ImagePayload{Bytes: buf.Bytes(), MIMEType: "image/jpeg"}, meta, nil + } + q -= 5 + } + quality = 75 + } + return ImagePayload{}, meta, fmt.Errorf("could not compress image under max_payload_bytes %d", maxPayload) +} + +func imageDimensions(path string) (w, h int, format string, err error) { + f, err := os.Open(path) + if err != nil { + return 0, 0, "", err + } + defer f.Close() + cfg, format, err := image.DecodeConfig(f) + if err != nil { + return 0, 0, "", fmt.Errorf("decode image config: %w", err) + } + return cfg.Width, cfg.Height, format, nil +} + +func mimeFromImageFormat(format string) string { + switch strings.ToLower(strings.TrimSpace(format)) { + case "jpeg", "jpg": + return "image/jpeg" + case "png": + return "image/png" + case "gif": + return "image/gif" + case "webp": + return "image/webp" + case "bmp": + return "image/bmp" + case "tiff": + return "image/tiff" + default: + return "" + } +} + +// DecodeImageConfig 用于测试：确认文件可被解码。 +func DecodeImageConfig(path string) (image.Config, string, error) { + f, err := os.Open(path) + if err != nil { + return image.Config{}, "", err + } + defer f.Close() + return image.DecodeConfig(f) +} diff --git a/internal/vision/preprocess_test.go b/internal/vision/preprocess_test.go new file mode 100644 index 00000000..a9b9e068 --- /dev/null +++ b/internal/vision/preprocess_test.go @@ -0,0 +1,109 @@ +package vision + +import ( + "image" + "image/color" + "image/png" + "os" + "path/filepath" + "testing" + + "github.com/disintegration/imaging" +) + +func TestPreprocessImageFile_scalesAndLimitsPayload(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "big.png") + img := imaging.New(3000, 2000, color.White) + if err := imaging.Save(img, path); err != nil { + t.Fatal(err) + } + + out, meta, err := PreprocessImageFile(path, PreprocessOptions{ + MaxImageBytes: 10 * 1024 * 1024, + MaxDimension: 1024, + JPEGQuality: 85, + MaxPayloadBytes: 600 * 1024, + SkipPreprocessBelowBytes: 0, + }) + if err != nil { + t.Fatal(err) + } + if len(out.Bytes) == 0 { + t.Fatal("empty output") + } + if meta.PreprocessMode != "jpeg" { + t.Fatalf("mode: %s", meta.PreprocessMode) + } + if meta.OutputWidth > 1024 || meta.OutputHeight > 1024 { + t.Fatalf("expected fit within 1024, got %dx%d", meta.OutputWidth, meta.OutputHeight) + } + if int64(len(out.Bytes)) > 600*1024 { + t.Fatalf("payload %d exceeds max", len(out.Bytes)) + } +} + +func TestPreprocessImageFile_passthroughSmallPNG(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "small.png") + if err := imaging.Save(imaging.New(400, 300, color.White), path); err != nil { + t.Fatal(err) + } + + out, meta, err := PreprocessImageFile(path, PreprocessOptions{ + MaxImageBytes: 5 * 1024 * 1024, + MaxDimension: 2048, + MaxPayloadBytes: 512 * 1024, + SkipPreprocessBelowBytes: 2 * 1024 * 1024, + }) + if err != nil { + t.Fatal(err) + } + if meta.PreprocessMode != "passthrough" { + t.Fatalf("expected passthrough, got %s", meta.PreprocessMode) + } + if out.MIMEType != "image/png" { + t.Fatalf("mime: %s", out.MIMEType) + } + if meta.OutputWidth != 400 || meta.OutputHeight != 300 { + t.Fatalf("dims: %dx%d", meta.OutputWidth, meta.OutputHeight) + } +} + +func TestPreprocessImageFile_passthroughDisabled(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "small.png") + if err := imaging.Save(imaging.New(100, 100, color.White), path); err != nil { + t.Fatal(err) + } + + _, meta, err := PreprocessImageFile(path, PreprocessOptions{ + MaxDimension: 2048, + MaxPayloadBytes: 512 * 1024, + SkipPreprocessBelowBytes: 0, + }) + if err != nil { + t.Fatal(err) + } + if meta.PreprocessMode != "jpeg" { + t.Fatalf("expected jpeg compress, got %s", meta.PreprocessMode) + } +} + +func TestPreprocessImageFile_rejectsOversizeFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "tiny.png") + f, err := os.Create(path) + if err != nil { + t.Fatal(err) + } + if err := png.Encode(f, image.NewRGBA(image.Rect(0, 0, 2, 2))); err != nil { + t.Fatal(err) + } + f.Close() + + _, _, err = PreprocessImageFile(path, PreprocessOptions{MaxImageBytes: 1}) + if err == nil { + t.Fatal("expected error when file exceeds max_image_bytes") + } +} diff --git a/internal/vision/tool.go b/internal/vision/tool.go new file mode 100644 index 00000000..d31ca928 --- /dev/null +++ b/internal/vision/tool.go @@ -0,0 +1,125 @@ +package vision + +import ( + "context" + "fmt" + "os" + "strings" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" + + "go.uber.org/zap" +) + +// RegisterAnalyzeImageTool 在 vision.enabled 且 model 已配置时注册 MCP 工具 analyze_image。 +func RegisterAnalyzeImageTool(mcpServer *mcp.Server, cfg *config.Config, logger *zap.Logger) { + if mcpServer == nil || cfg == nil { + return + } + if !cfg.Vision.Ready() { + if cfg.Vision.Enabled && logger != nil { + logger.Warn("vision.enabled 但 vision.model 为空，跳过注册 analyze_image") + } + return + } + + cwd, err := os.Getwd() + if err != nil { + if logger != nil { + logger.Warn("vision: getwd failed, skip analyze_image", zap.Error(err)) + } + return + } + + preOpt := PreprocessOptions{ + MaxImageBytes: cfg.Vision.MaxImageBytesEffective(), + MaxDimension: cfg.Vision.MaxDimensionEffective(), + JPEGQuality: cfg.Vision.JPEGQualityEffective(), + MaxPayloadBytes: cfg.Vision.MaxPayloadBytesEffective(), + SkipPreprocessBelowBytes: cfg.Vision.SkipPreprocessBelowBytesEffective(), + } + client := NewClient(cfg.Vision, cfg.OpenAI) + + tool := mcp.Tool{ + Name: builtin.ToolAnalyzeImage, + Description: "分析服务器上的本地图片并返回文字描述（验证码、UI 元素、报错、架构图要点等）。" + + "输入为文件路径（如用户上传的 chat_uploads 路径或工具截图路径）。" + + "输出仅为文本，不含图片数据。不要对二进制图片使用 read_file 指望理解内容。", + ShortDescription: "分析本地图片并返回文字描述（验证码/UI/报错等）", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "图片绝对路径或相对于进程工作目录的路径", + }, + "question": map[string]interface{}{ + "type": "string", + "description": "可选：希望模型重点回答的问题。验证码图建议：只输出验证码字符，不要空格和解释", + }, + }, + "required": []string{"path"}, + }, + } + + handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + path, _ := args["path"].(string) + question, _ := args["question"].(string) + + abs, err := ResolveImagePath(path, cwd) + if err != nil { + return textResult(fmt.Sprintf("路径校验失败: %v", err), true), nil + } + + img, meta, err := PreprocessImageFile(abs, preOpt) + if err != nil { + return textResult(fmt.Sprintf("图片预处理失败: %v", err), true), nil + } + + summary, err := client.Analyze(ctx, img, question) + if err != nil { + return textResult(fmt.Sprintf("视觉模型调用失败: %v", err), true), nil + } + + body := formatAnalysisResult(abs, meta, summary) + return textResult(body, false), nil + } + + mcpServer.RegisterTool(tool, handler) + if logger != nil { + logger.Debug("vision: analyze_image 工具已注册", zap.String("model", cfg.Vision.Model)) + } +} + +func textResult(text string, isError bool) *mcp.ToolResult { + return &mcp.ToolResult{ + Content: []mcp.Content{{Type: "text", Text: text}}, + IsError: isError, + } +} + +func formatAnalysisResult(path string, meta PreprocessMeta, summary string) string { + var b strings.Builder + b.WriteString("## Image analysis\n") + b.WriteString("- **path**: ") + b.WriteString(path) + b.WriteString("\n") + switch meta.PreprocessMode { + case "passthrough": + b.WriteString(fmt.Sprintf("- **preprocess**: passthrough %dx%d, %s, %dKB (original %dKB)\n\n", + meta.OutputWidth, meta.OutputHeight, meta.OutputMIMEType, + (meta.OutputBytes+1023)/1024, (meta.OriginalBytes+1023)/1024)) + default: + b.WriteString(fmt.Sprintf("- **preprocess**: %dx%d → %dx%d, jpeg q=%d, %dKB (original %dKB)\n\n", + meta.OriginalWidth, meta.OriginalHeight, + meta.OutputWidth, meta.OutputHeight, + meta.JPEGQuality, (meta.OutputBytes+1023)/1024, + (meta.OriginalBytes+1023)/1024)) + } + b.WriteString("### Summary\n") + b.WriteString(strings.TrimSpace(summary)) + b.WriteString("\n") + return b.String() +}