From 74a4ec7be3cc15bcb47c68809c6bfc39c58fabbc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=85=AC=E6=98=8E?=
<83812544+Ed1s0nZ@users.noreply.github.com>
Date: Wed, 24 Dec 2025 01:16:10 +0800
Subject: [PATCH] Add files via upload
---
internal/attackchain/builder.go | 1633 +++----------------------------
web/static/css/style.css | 205 +++-
web/static/js/chat.js | 728 ++++++++++++--
web/templates/index.html | 16 +-
4 files changed, 939 insertions(+), 1643 deletions(-)
diff --git a/internal/attackchain/builder.go b/internal/attackchain/builder.go
index aba0dde0..25b25627 100644
--- a/internal/attackchain/builder.go
+++ b/internal/attackchain/builder.go
@@ -6,7 +6,6 @@ import (
"errors"
"fmt"
"net/http"
- "sort"
"strings"
"time"
@@ -82,111 +81,148 @@ func NewBuilder(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.
}
}
-// BuildChainFromConversation 从对话构建攻击链（一次性生成整个图）
+// BuildChainFromConversation 从对话构建攻击链（简化版本：用户输入+最后一轮ReAct输入+大模型输出）
func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID string) (*Chain, error) {
- b.logger.Info("开始构建攻击链（一次性生成）", zap.String("conversationId", conversationID))
+ b.logger.Info("开始构建攻击链（简化版本）", zap.String("conversationId", conversationID))
- // 1. 获取对话消息和工具执行记录
+ // 1. 获取对话消息
messages, err := b.db.GetMessages(conversationID)
if err != nil {
return nil, fmt.Errorf("获取对话消息失败: %w", err)
}
- executions, err := b.getToolExecutionsByConversation(conversationID)
- if err != nil {
- return nil, fmt.Errorf("获取工具执行记录失败: %w", err)
- }
-
- // 获取过程详情
- processDetailsMap, err := b.db.GetProcessDetailsByConversation(conversationID)
- if err != nil {
- b.logger.Warn("获取过程详情失败", zap.Error(err))
- processDetailsMap = make(map[string][]database.ProcessDetail)
- }
-
- if len(executions) == 0 && len(messages) == 0 {
+ if len(messages) == 0 {
b.logger.Info("对话中没有数据", zap.String("conversationId", conversationID))
return &Chain{Nodes: []Node{}, Edges: []Edge{}}, nil
}
- // 2. 准备上下文数据
- contextData, err := b.prepareContextData(messages, executions, processDetailsMap)
- if err != nil {
- return nil, fmt.Errorf("准备上下文数据失败: %w", err)
+ // 2. 提取用户输入（最后一条user消息）
+ var userInput string
+ for i := len(messages) - 1; i >= 0; i-- {
+ if strings.EqualFold(messages[i].Role, "user") {
+ userInput = messages[i].Content
+ break
+ }
}
- // 3. 一次性生成攻击链（带重试和压缩机制）
- chain, err := b.generateChainWithRetry(ctx, contextData, 5)
- if err != nil {
- return nil, fmt.Errorf("生成攻击链失败: %w", err)
+ // 3. 提取最后一轮ReAct的输入（历史消息+当前用户输入）
+ // 最后一轮ReAct的输入 = 所有历史消息（包括当前用户输入）
+ reactInput := b.buildReActInput(messages)
+
+ // 4. 提取大模型最后的输出（最后一条assistant消息）
+ var modelOutput string
+ for i := len(messages) - 1; i >= 0; i-- {
+ if strings.EqualFold(messages[i].Role, "assistant") {
+ modelOutput = messages[i].Content
+ break
+ }
}
- // 4. 保存到数据库
- if err := b.saveChain(conversationID, chain.Nodes, chain.Edges); err != nil {
- b.logger.Warn("保存攻击链失败", zap.Error(err))
- // 不返回错误，继续返回结果
+ // 5. 构建简化的prompt，一次性传递给大模型
+ prompt := b.buildSimplePrompt(userInput, reactInput, modelOutput)
+
+ // 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, nil) // executions为nil，因为我们不再使用tool_execution_id
+ 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.Int("nodes", len(chain.Nodes)),
- zap.Int("edges", len(chain.Edges)))
+ zap.Int("nodes", len(chainData.Nodes)),
+ zap.Int("edges", len(chainData.Edges)))
- return chain, nil
+ // 保存到数据库（供后续加载使用）
+ if err := b.saveChain(conversationID, chainData.Nodes, chainData.Edges); err != nil {
+ b.logger.Warn("保存攻击链到数据库失败", zap.Error(err))
+ // 即使保存失败，也返回数据给前端
+ }
+
+ // 直接返回，不做任何处理和校验
+ return chainData, nil
}
-// getToolExecutionsByConversation 获取对话的工具执行记录
-func (b *Builder) getToolExecutionsByConversation(conversationID string) ([]*mcp.ToolExecution, error) {
- // 通过conversation_id关联messages，再通过mcp_execution_ids关联tool_executions
- // 简化实现：直接查询所有工具执行记录，然后过滤（实际应该优化查询）
- allExecutions, err := b.db.LoadToolExecutions()
- if err != nil {
- return nil, err
- }
-
- // 获取对话的消息，提取mcp_execution_ids
- messages, err := b.db.GetMessages(conversationID)
- if err != nil {
- return nil, err
- }
-
- // 收集所有execution IDs
- executionIDSet := make(map[string]bool)
+// buildReActInput 构建最后一轮ReAct的输入（历史消息+当前用户输入）
+func (b *Builder) buildReActInput(messages []database.Message) string {
+ var builder strings.Builder
for _, msg := range messages {
- if len(msg.MCPExecutionIDs) > 0 {
- for _, id := range msg.MCPExecutionIDs {
- executionIDSet[id] = true
- }
- }
+ builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", msg.Role, msg.Content))
}
-
- // 过滤执行记录
- var filteredExecutions []*mcp.ToolExecution
- for _, exec := range allExecutions {
- if executionIDSet[exec.ID] {
- filteredExecutions = append(filteredExecutions, exec)
- }
- }
-
- // 按时间排序
- sort.Slice(filteredExecutions, func(i, j int) bool {
- return filteredExecutions[i].StartTime.Before(filteredExecutions[j].StartTime)
- })
-
- return filteredExecutions, nil
+ return builder.String()
}
-// saveChain 保存攻击链到数据库
+// buildSimplePrompt 构建简化的prompt
+func (b *Builder) buildSimplePrompt(userInput, reactInput, modelOutput string) string {
+ return fmt.Sprintf(`你是一个专业的安全测试分析师。请根据以下信息生成攻击链图。
+
+## 用户输入
+%s
+
+## 最后一轮ReAct的输入（历史对话上下文）
+%s
+
+## 大模型最后的输出
+%s
+
+## 任务要求
+
+请根据上述信息，生成一个清晰的攻击链图。攻击链应该包含：
+1. **target（目标）**：从用户输入中提取的测试目标
+2. **action（行动）**：从ReAct输入和模型输出中提取的关键测试步骤
+3. **vulnerability（漏洞）**：从模型输出中提取的发现的漏洞
+
+## 输出格式
+
+请以JSON格式返回攻击链，格式如下：
+{
+ "nodes": [
+ {
+ "id": "node_1",
+ "type": "target|action|vulnerability",
+ "label": "节点标签",
+ "risk_score": 0-100,
+ "metadata": {
+ "target": "目标（target节点）",
+ "tool_name": "工具名称（action节点）",
+ "description": "描述（vulnerability节点）"
+ }
+ }
+ ],
+ "edges": [
+ {
+ "source": "node_1",
+ "target": "node_2",
+ "type": "leads_to|discovers|enables",
+ "weight": 1-5
+ }
+ ]
+}
+
+只返回JSON，不要包含其他解释文字。`, userInput, reactInput, modelOutput)
+}
+
+// saveChain 保存攻击链到数据库（简化版本，移除tool_execution_id）
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))
}
- // 保存节点
+ // 保存节点（不保存tool_execution_id）
for _, node := range nodes {
metadataJSON, _ := json.Marshal(node.Metadata)
- if err := b.db.SaveAttackChainNode(conversationID, node.ID, node.Type, node.Label, node.ToolExecutionID, string(metadataJSON), node.RiskScore); err != nil {
+ 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))
}
}
@@ -219,1035 +255,6 @@ func (b *Builder) LoadChainFromDatabase(conversationID string) (*Chain, error) {
}, nil
}
-// ContextData 上下文数据（用于一次性生成攻击链）
-type ContextData struct {
- Messages []database.Message `json:"messages"`
- Executions []*mcp.ToolExecution `json:"executions"`
- ProcessDetails map[string][]database.ProcessDetail `json:"process_details"`
- SummarizedItems map[string]string `json:"summarized_items"` // 已总结的项目（key: 原始ID, value: 总结内容）
-}
-
-// prepareContextData 准备上下文数据
-func (b *Builder) prepareContextData(messages []database.Message, executions []*mcp.ToolExecution, processDetails map[string][]database.ProcessDetail) (*ContextData, error) {
- return &ContextData{
- Messages: messages,
- Executions: executions,
- ProcessDetails: processDetails,
- SummarizedItems: make(map[string]string),
- }, nil
-}
-
-// generateChainWithRetry 生成攻击链（带重试和压缩机制）
-func (b *Builder) generateChainWithRetry(ctx context.Context, contextData *ContextData, maxRetries int) (*Chain, error) {
- // 在第一次尝试前，先检查tokens并压缩（如果需要）
- totalTokens, err := b.countPromptTokens(contextData)
- if err == nil && totalTokens > b.maxTokens {
- b.logger.Info("检测到tokens超过限制，提前压缩",
- zap.Int("totalTokens", totalTokens),
- zap.Int("maxTokens", b.maxTokens))
- if err := b.compressContextData(ctx, contextData); err != nil {
- return nil, fmt.Errorf("压缩上下文失败: %w", err)
- }
- }
-
- for attempt := 0; attempt < maxRetries; attempt++ {
- b.logger.Info("尝试生成攻击链",
- zap.Int("attempt", attempt+1),
- zap.Int("maxRetries", maxRetries))
-
- // 构建提示词
- prompt, err := b.buildChainGenerationPrompt(contextData)
- if err != nil {
- return nil, fmt.Errorf("构建提示词失败: %w", err)
- }
-
- // 调用AI生成攻击链
- chainJSON, err := b.callAIForChainGeneration(ctx, prompt)
- if err != nil {
- // 检查是否是上下文过长错误
- if strings.Contains(err.Error(), "context length") || strings.Contains(err.Error(), "too long") || strings.Contains(err.Error(), "context length exceeded") {
- b.logger.Warn("上下文过长，尝试压缩",
- zap.Int("attempt", attempt+1),
- zap.Error(err))
-
- // 使用分片压缩
- if err := b.compressContextData(ctx, contextData); err != nil {
- return nil, fmt.Errorf("压缩上下文失败: %w", err)
- }
-
- // 重试
- continue
- }
-
- return nil, fmt.Errorf("AI生成失败: %w", err)
- }
-
- // 解析JSON（传入executions用于ID映射）
- chain, err := b.parseChainJSON(chainJSON, contextData.Executions)
- if err != nil {
- return nil, fmt.Errorf("解析攻击链JSON失败: %w", err)
- }
-
- return chain, nil
- }
-
- return nil, fmt.Errorf("生成攻击链失败：超过最大重试次数 %d", maxRetries)
-}
-
-// buildChainGenerationPrompt 构建攻击链生成提示词
-func (b *Builder) buildChainGenerationPrompt(contextData *ContextData) (string, error) {
- var promptBuilder strings.Builder
-
- promptBuilder.WriteString(`你是一个专业的安全测试分析师。请根据以下对话和工具执行记录，生成清晰、有教育意义的攻击链图。
-
-## 核心原则
-
-**目标：让不懂渗透测试的同学可以通过这个攻击链路学习到知识，而不是无数个节点看花眼。**
-**即便某些工具执行或漏洞挖掘没有成功，只要它们提供了关键线索、错误提示或下一步思路，也要被保留下来。**
-
-**关键要求：**
-1. **节点标签必须简洁明了**：每个节点标签控制在15-25个汉字以内，使用简洁的动宾结构（如"扫描端口"、"发现SQL注入"、"验证漏洞"），避免冗长描述
-2. **严格控制节点数量**：优先保留关键步骤，避免生成过多细碎节点。理想情况下，单个目标的攻击链应控制在8-15个节点以内
-3. **确保DAG结构**：生成的图必须是有向无环图（DAG），不允许出现循环。边的方向必须符合时间顺序和逻辑关系（从早期步骤指向后期步骤）
-4. **层次清晰**：攻击链应该呈现清晰的层次结构：目标 → 信息收集 → 漏洞发现 → 漏洞利用 → 后续行动
-
-## 任务要求
-
-1. **节点类型（简化，只保留3种）**：
- - **target（目标）**：从用户输入中提取测试目标（IP、域名、URL等）
- - **重要：如果对话中测试了多个不同的目标（如先测试A网页，后测试B网页），必须为每个不同的目标创建独立的target节点**
- - 每个target节点只关联属于它的action节点（通过工具执行参数中的目标来判断）
- - 不同目标的action节点之间**不应该**建立关联关系
- - **action（行动）**：**工具执行 + AI分析结果 = 一个action节点**
- - 将每个工具执行和AI对该工具结果的分析合并为一个action节点
- - **节点标签必须简洁**：控制在15-25个汉字，使用动宾结构，突出关键信息
- - 好的示例："扫描端口发现22/80/443"、"SQL注入验证成功"、"WAF拦截暴露厂商"
- - 避免的示例："使用Nmap工具对目标192.168.1.1进行了全面的端口扫描，发现了22、80、443等多个端口开放"（过于冗长）
- - 默认关注成功的执行；但如果执行失败却提供了有价值的线索（错误信息、资产指纹、下一步建议等），也要保留，记为"带线索的失败"行动
- - **重要：action节点必须关联到正确的target节点（通过工具执行参数判断目标）**
- - **vulnerability（漏洞）**：从工具执行结果和AI分析中提取的**真实漏洞**（不是所有发现都是漏洞）。若验证失败但能明确表明某个漏洞利用方向不可行，可作为行动节点的线索描述，而不是漏洞节点。
-
-2. **过滤规则（重要！）**：
- - **默认忽略**彻底无效的信息：完全没有输出、没有任何线索的失败执行仍需过滤
- - **必须保留**下列失败执行：
- - 错误信息里包含了潜在线索、受限条件、可复现的报错
- - 虽未找到漏洞，但收集到了资产信息、技术栈或后续测试方向
- - 用户特别关注的失败尝试
- - **保留策略**：只要行动节点能给后续测试提供启发，就保留；否则忽略
-
-3. **建立清晰的关联关系（确保DAG结构）**：
- - target → action：目标指向属于它的所有行动（通过工具执行参数判断目标）
- - action → action：行动之间的逻辑顺序（按时间顺序，但只连接有逻辑关系的）
- - **重要：只连接属于同一目标的action节点，不同目标的action节点之间不应该连接**
- - **必须确保无环**：只能从早期步骤指向后期步骤，不能形成循环（如A→B→C→A）
- - 优先连接直接相关的步骤，避免过度连接导致图过于复杂
- - action → vulnerability：行动发现的漏洞
- - vulnerability → vulnerability：漏洞间的因果关系（如SQL注入 → 信息泄露）
- - **重要：只连接属于同一目标的漏洞，不同目标的漏洞之间不应该连接**
- - **必须确保无环**：漏洞间的因果关系也必须是单向的，不能形成循环
-
-4. **节点属性**：
- - 每个节点需要：id, type, label, risk_score, metadata
- - action节点需要：
- - tool_name: 工具名称
- - tool_intent: 工具调用意图（如"端口扫描"、"漏洞扫描"）
- - ai_analysis: AI对工具结果的分析总结（简洁，不超过100字，失败节点需解释线索价值）
- - findings: 关键发现（列表）
- - status: "success" | "failed_insight"（失败但有价值的线索）
- - hints: ["下一步建议1", "限制条件2"]（失败节点可提供的线索列表）
- - vulnerability节点需要：type, description, severity, location
-
-## 对话数据
-
-`)
-
- // 添加消息
- promptBuilder.WriteString("\n### 对话消息\n\n")
- for i, msg := range contextData.Messages {
- promptBuilder.WriteString(fmt.Sprintf("消息%d [%s]:\n", i+1, msg.Role))
-
- isUserMessage := strings.EqualFold(msg.Role, "user")
- // 用户输入必须原样提供给攻击链模型
- if isUserMessage {
- promptBuilder.WriteString(fmt.Sprintf("%s\n\n", msg.Content))
- } else if summary, ok := contextData.SummarizedItems[msg.ID]; ok {
- promptBuilder.WriteString(fmt.Sprintf("[已总结] %s\n\n", summary))
- } else {
- content := msg.Content
- if len(content) > 5000 {
- content = content[:5000] + "..."
- }
- promptBuilder.WriteString(fmt.Sprintf("%s\n\n", content))
- }
-
- // 添加过程详情
- if details, ok := contextData.ProcessDetails[msg.ID]; ok {
- for _, detail := range details {
- if detail.EventType == "thinking" {
- thinkingText := detail.Message
- if summary, ok := contextData.SummarizedItems[detail.ID]; ok {
- thinkingText = "[已总结] " + summary
- } else if len(thinkingText) > 2000 {
- thinkingText = thinkingText[:2000] + "..."
- }
- promptBuilder.WriteString(fmt.Sprintf("思考过程: %s\n", thinkingText))
- }
- }
- }
- promptBuilder.WriteString("\n")
- }
-
- // 添加工具执行记录（关联对应的AI回复）
- promptBuilder.WriteString("\n### 工具执行记录（包含对应的AI分析）\n\n")
-
- // 构建工具执行ID到消息的映射（找到工具执行后AI的回复）
- execToMessageMap := b.buildExecutionToMessageMap(contextData)
-
- for i, exec := range contextData.Executions {
- // 检查是否是错误/失败的执行
- isError := exec.Error != "" || (exec.Result != nil && exec.Result.IsError)
-
- statusText := "成功"
- if isError {
- statusText = "失败（可能包含线索）"
- }
-
- promptBuilder.WriteString(fmt.Sprintf("执行%d [%s] (ID: %s) - 状态: %s\n", i+1, exec.ToolName, exec.ID, statusText))
- promptBuilder.WriteString(fmt.Sprintf("参数: %s\n", b.formatArguments(exec.Arguments)))
-
- if isError && exec.Error != "" {
- promptBuilder.WriteString(fmt.Sprintf("错误信息: %s\n", exec.Error))
- }
-
- // 检查是否已总结
- var resultText string
- if exec.Result != nil {
- for _, content := range exec.Result.Content {
- if content.Type == "text" {
- resultText += content.Text + "\n"
- }
- }
- }
-
- // 检查结果是否为空或无效
- if strings.TrimSpace(resultText) == "" {
- if isError {
- promptBuilder.WriteString("工具执行结果: [失败但未返回正文]\n")
- } else {
- promptBuilder.WriteString("工具执行结果: **已忽略（结果为空）**\n\n")
- continue
- }
- } else {
- if summary, ok := contextData.SummarizedItems[exec.ID]; ok {
- promptBuilder.WriteString(fmt.Sprintf("工具执行结果: [已总结] %s\n", summary))
- } else {
- if len(resultText) > 5000 {
- resultText = resultText[:5000] + "..."
- }
- promptBuilder.WriteString(fmt.Sprintf("工具执行结果: %s\n", resultText))
- }
- }
-
- // 添加对应的AI分析（工具执行后AI的回复）
- if aiMessage, ok := execToMessageMap[exec.ID]; ok {
- aiContent := aiMessage.Content
- if len(aiContent) > 2000 {
- aiContent = aiContent[:2000] + "..."
- }
- promptBuilder.WriteString(fmt.Sprintf("AI分析: %s\n", aiContent))
- }
-
- promptBuilder.WriteString("\n")
- }
-
- promptBuilder.WriteString(`
-
-## 输出格式
-
-请以JSON格式返回攻击链，格式如下：
-
-{
- "nodes": [
- {
- "id": "node_1",
- "type": "target|action|vulnerability",
- "label": "节点标签（清晰、简洁，action节点要描述"做了什么"和"发现了什么"）",
- "risk_score": 0-100,
- "tool_execution_id": "执行记录的真实ID（action节点必须使用上面执行记录中的ID字段）",
- "metadata": {
- "target": "目标（target节点）",
- "tool_name": "工具名称（action节点）",
- "tool_intent": "工具调用意图（action节点，如"端口扫描"、"漏洞扫描"）",
- "ai_analysis": "AI对工具结果的分析总结（action节点，不超过100字）",
- "findings": ["发现1", "发现2"]（action节点，关键发现列表）,
- "vulnerability_type": "漏洞类型（vulnerability节点）",
- "description": "描述（vulnerability节点）",
- "severity": "critical|high|medium|low（vulnerability节点）",
- "location": "漏洞位置（vulnerability节点）"
- }
- }
- ],
- "edges": [
- {
- "source": "node_1",
- "target": "node_2",
- "type": "leads_to|discovers|enables",
- "weight": 1-5
- }
- ]
-}
-
-## 重要要求
-
-1. **节点合并和标签优化**：
- - 每个工具执行和对应的AI分析必须合并为一个action节点
- - **action节点的label必须简洁**：控制在15-25个汉字，使用动宾结构
- - 好的示例："扫描端口发现22/80/443"、"验证SQL注入成功"、"WAF拦截暴露厂商"
- - 避免冗长描述，关键信息放在metadata中详细说明
- - 若为失败但有线索的行动，请在metadata.status中标记为"failed_insight"，并在findings/hints里写清线索价值
-
-2. **过滤无效节点**：
- - **必须忽略**没有任何输出、没有线索的失败执行
- - **必须保留**失败但提供关键线索的执行，确保metadata里解释清楚
- - 只保留对学习或溯源有帮助的节点
-
-3. **简化结构**：
- - 只创建target、action、vulnerability三种节点
- - 不要创建discovery、decision等节点
- - 让攻击链清晰、有教育意义
-
-4. **关联关系（确保DAG结构）**：
- - target → action：目标指向属于它的所有行动（通过工具执行参数判断目标）
- - action → action：按时间顺序连接，但只连接有逻辑关系的
- - **重要：只连接属于同一目标的action节点，不同目标的action节点之间不应该连接**
- - **必须确保无环**：只能从早期步骤指向后期步骤，不能形成循环
- - 优先连接直接相关的步骤，避免过度连接
- - action → vulnerability：行动发现的漏洞
- - vulnerability → vulnerability：漏洞间的因果关系
- - **重要：只连接属于同一目标的漏洞，不同目标的漏洞之间不应该连接**
- - **必须确保无环**：漏洞间的因果关系也必须是单向的
-
-5. **多目标处理（重要！）**：
- - 如果对话中测试了多个不同的目标（如先测试A网页，后测试B网页），必须：
- - 为每个不同的目标创建独立的target节点
- - 每个target节点只关联属于它的action和vulnerability节点
- - 不同目标的节点之间**不应该**建立任何关联关系
- - 这样会形成多个独立的攻击链分支，每个分支对应一个测试目标
-
-6. **节点数量控制和合并策略**：
- - **严格控制节点数量**：单个目标的攻击链理想情况下应控制在8-15个节点以内
- - 如果节点太多（>20个），优先保留最重要的节点，合并或删除次要节点
- - 合并相似的action节点（如同一工具的连续调用，如果结果相似）
- - 对于同一类型的多个发现，考虑合并为一个节点（如"发现多个开放端口"而不是为每个端口创建节点）
-
-7. **DAG结构验证**：
- - 生成后必须检查：确保图中不存在循环（即不存在路径A→B→...→A）
- - 边的方向必须符合时间顺序：早期步骤指向后期步骤
- - 如果发现循环，必须断开形成循环的边，保留最重要的连接
-
-只返回JSON，不要包含其他解释文字。`)
-
- return promptBuilder.String(), nil
-}
-
-// buildExecutionToMessageMap 构建工具执行ID到AI消息的映射
-// 找到每个工具执行后AI的回复消息
-func (b *Builder) buildExecutionToMessageMap(contextData *ContextData) map[string]database.Message {
- execToMessageMap := make(map[string]database.Message)
-
- // 遍历消息，找到包含工具执行ID的消息（通常是assistant消息）
- for _, msg := range contextData.Messages {
- if msg.Role != "assistant" {
- continue
- }
-
- // 检查消息中是否引用了工具执行ID
- // 通常工具执行后，AI会在回复中引用这些执行ID
- for _, execID := range msg.MCPExecutionIDs {
- // 找到对应的工具执行
- for _, exec := range contextData.Executions {
- if exec.ID == execID {
- // 如果这个执行还没有关联的消息，或者当前消息时间更晚，则更新
- if existingMsg, exists := execToMessageMap[execID]; !exists || msg.CreatedAt.After(existingMsg.CreatedAt) {
- execToMessageMap[execID] = msg
- }
- break
- }
- }
- }
- }
-
- // 如果通过MCPExecutionIDs找不到，尝试按时间顺序匹配
- // 找到每个工具执行后最近的assistant消息
- for _, exec := range contextData.Executions {
- if _, exists := execToMessageMap[exec.ID]; exists {
- continue
- }
-
- // 找到执行时间之后最近的assistant消息
- var closestMsg *database.Message
- for i := range contextData.Messages {
- msg := &contextData.Messages[i]
- if msg.Role == "assistant" && msg.CreatedAt.After(exec.StartTime) {
- if closestMsg == nil || msg.CreatedAt.Before(closestMsg.CreatedAt) {
- closestMsg = msg
- }
- }
- }
-
- if closestMsg != nil {
- execToMessageMap[exec.ID] = *closestMsg
- }
- }
-
- return execToMessageMap
-}
-
-// formatArguments 格式化工具参数
-func (b *Builder) formatArguments(args map[string]interface{}) string {
- if args == nil {
- return "{}"
- }
- jsonData, _ := json.Marshal(args)
- return string(jsonData)
-}
-
-// countPromptTokens 计算prompt的总tokens数
-func (b *Builder) countPromptTokens(contextData *ContextData) (int, error) {
- prompt, err := b.buildChainGenerationPrompt(contextData)
- if err != nil {
- return 0, fmt.Errorf("构建提示词失败: %w", err)
- }
-
- if b.tokenCounter == nil || b.openAIConfig == nil {
- // 如果没有token计数器或配置，使用简单的估算（4个字符=1个token）
- return len(prompt) / 4, nil
- }
-
- model := b.openAIConfig.Model
- if model == "" {
- model = "gpt-4" // 默认模型
- }
-
- count, err := b.tokenCounter.Count(model, prompt)
- if err != nil {
- // 如果计算失败，使用估算
- return len(prompt) / 4, nil
- }
- return count, nil
-}
-
-// compressContextData 使用分片压缩方式压缩上下文数据
-func (b *Builder) compressContextData(ctx context.Context, contextData *ContextData) error {
- // 计算当前tokens
- totalTokens, err := b.countPromptTokens(contextData)
- if err != nil {
- return fmt.Errorf("计算tokens失败: %w", err)
- }
-
- b.logger.Info("开始压缩上下文",
- zap.Int("totalTokens", totalTokens),
- zap.Int("maxTokens", b.maxTokens))
-
- // 如果tokens在限制内，不需要压缩
- if totalTokens <= b.maxTokens {
- return nil
- }
-
- // 计算需要分成多少份
- numChunks := (totalTokens + b.maxTokens - 1) / b.maxTokens // 向上取整
- if numChunks < 2 {
- numChunks = 2 // 至少分成2份
- }
-
- b.logger.Info("将上下文分成多个片段进行压缩",
- zap.Int("totalTokens", totalTokens),
- zap.Int("maxTokens", b.maxTokens),
- zap.Int("numChunks", numChunks))
-
- // 按时间顺序将数据分成多个片段
- chunks, err := b.splitContextDataByTime(contextData, numChunks)
- if err != nil {
- return fmt.Errorf("分割上下文数据失败: %w", err)
- }
-
- // 对每个片段进行摘要
- summaries := make([]string, 0, len(chunks))
- for i, chunk := range chunks {
- b.logger.Info("压缩片段",
- zap.Int("chunkIndex", i+1),
- zap.Int("totalChunks", len(chunks)),
- zap.Int("chunkSize", len(chunk.Messages)+len(chunk.Executions)))
-
- summary, err := b.summarizeContextChunk(ctx, chunk)
- if err != nil {
- // 检查是否是认证错误
- if strings.Contains(err.Error(), "Authentication") || strings.Contains(err.Error(), "api key") || strings.Contains(err.Error(), "invalid") {
- return fmt.Errorf("压缩片段%d失败（API认证错误，请检查OpenAI配置）: %w", i+1, err)
- }
- return fmt.Errorf("压缩片段%d失败: %w", i+1, err)
- }
- summaries = append(summaries, summary)
- }
-
- // 将摘要合并到contextData中
- // 保留用户消息，清空其他数据，用摘要替换
- var userMessages []database.Message
- for _, msg := range contextData.Messages {
- if strings.EqualFold(msg.Role, "user") {
- userMessages = append(userMessages, msg)
- }
- }
-
- // 清空非用户消息和执行记录
- contextData.Messages = userMessages
- contextData.Executions = []*mcp.ToolExecution{}
- contextData.ProcessDetails = make(map[string][]database.ProcessDetail)
-
- // 创建一个综合摘要消息
- combinedSummary := strings.Join(summaries, "\n\n---\n\n")
- summaryMsg := database.Message{
- ID: uuid.New().String(),
- Role: "assistant",
- Content: fmt.Sprintf("[上下文摘要 - 包含%d个片段的压缩内容]\n\n%s", len(summaries), combinedSummary),
- CreatedAt: time.Now(),
- }
- contextData.Messages = append(contextData.Messages, summaryMsg)
-
- // 检查压缩后的tokens
- compressedTokens, err := b.countPromptTokens(contextData)
- if err != nil {
- return fmt.Errorf("计算压缩后tokens失败: %w", err)
- }
-
- b.logger.Info("压缩完成",
- zap.Int("originalTokens", totalTokens),
- zap.Int("compressedTokens", compressedTokens),
- zap.Int("reduction", totalTokens-compressedTokens))
-
- // 如果压缩后仍然超过限制，递归压缩
- if compressedTokens > b.maxTokens {
- b.logger.Info("压缩后仍然超过限制，继续递归压缩",
- zap.Int("compressedTokens", compressedTokens),
- zap.Int("maxTokens", b.maxTokens))
- return b.compressContextData(ctx, contextData)
- }
-
- return nil
-}
-
-// ContextChunk 上下文数据片段
-type ContextChunk struct {
- Messages []database.Message
- Executions []*mcp.ToolExecution
- ProcessDetails map[string][]database.ProcessDetail
-}
-
-// splitContextDataByTime 按时间顺序将上下文数据分成多个片段
-func (b *Builder) splitContextDataByTime(contextData *ContextData, numChunks int) ([]*ContextChunk, error) {
- if numChunks <= 0 {
- return nil, fmt.Errorf("片段数量必须大于0")
- }
-
- // 收集所有带时间戳的项目
- type timeItem struct {
- time time.Time
- itemType string // "message", "execution", "thinking"
- message *database.Message
- execution *mcp.ToolExecution
- processDetail *database.ProcessDetail
- }
-
- var items []timeItem
-
- // 添加消息（跳过已总结的）
- for i := range contextData.Messages {
- msg := &contextData.Messages[i]
- if _, alreadySummarized := contextData.SummarizedItems[msg.ID]; alreadySummarized {
- continue
- }
- items = append(items, timeItem{
- time: msg.CreatedAt,
- itemType: "message",
- message: msg,
- })
- }
-
- // 添加工具执行（跳过已总结的）
- for _, exec := range contextData.Executions {
- if _, alreadySummarized := contextData.SummarizedItems[exec.ID]; alreadySummarized {
- continue
- }
- items = append(items, timeItem{
- time: exec.StartTime,
- itemType: "execution",
- execution: exec,
- })
- }
-
- // 添加思考过程（跳过已总结的）
- for _, details := range contextData.ProcessDetails {
- for i := range details {
- detail := &details[i]
- if detail.EventType == "thinking" {
- if _, alreadySummarized := contextData.SummarizedItems[detail.ID]; alreadySummarized {
- continue
- }
- items = append(items, timeItem{
- time: detail.CreatedAt,
- itemType: "thinking",
- processDetail: detail,
- })
- }
- }
- }
-
- if len(items) == 0 {
- return nil, fmt.Errorf("没有可分割的数据")
- }
-
- // 按时间排序
- sort.Slice(items, func(i, j int) bool {
- return items[i].time.Before(items[j].time)
- })
-
- // 计算每个片段的大小
- chunkSize := (len(items) + numChunks - 1) / numChunks // 向上取整
-
- // 创建片段
- chunks := make([]*ContextChunk, 0, numChunks)
- for i := 0; i < len(items); i += chunkSize {
- end := i + chunkSize
- if end > len(items) {
- end = len(items)
- }
-
- chunk := &ContextChunk{
- Messages: []database.Message{},
- Executions: []*mcp.ToolExecution{},
- ProcessDetails: make(map[string][]database.ProcessDetail),
- }
-
- for j := i; j < end; j++ {
- item := items[j]
- switch item.itemType {
- case "message":
- chunk.Messages = append(chunk.Messages, *item.message)
- case "execution":
- chunk.Executions = append(chunk.Executions, item.execution)
- case "thinking":
- if item.processDetail != nil {
- msgID := item.processDetail.MessageID
- chunk.ProcessDetails[msgID] = append(chunk.ProcessDetails[msgID], *item.processDetail)
- }
- }
- }
-
- chunks = append(chunks, chunk)
- }
-
- return chunks, nil
-}
-
-// getModelMaxContextLength 获取模型的最大上下文长度
-func (b *Builder) getModelMaxContextLength() int {
- if b.openAIConfig == nil {
- return 131072 // 默认值
- }
- model := strings.ToLower(b.openAIConfig.Model)
- if strings.Contains(model, "gpt-4") {
- return 128000
- } else if strings.Contains(model, "gpt-3.5") {
- return 16000
- } else if strings.Contains(model, "deepseek") {
- return 131072
- }
- return 131072 // 默认值
-}
-
-// summarizeContextChunk 总结一个上下文片段
-func (b *Builder) summarizeContextChunk(ctx context.Context, chunk *ContextChunk) (string, error) {
- // 先构建内容
- content, err := b.buildChunkContent(chunk)
- if err != nil {
- return "", err
- }
-
- // 使用AI总结
- promptTemplate := `请详细总结以下安全测试对话片段的关键信息。虽然需要压缩内容，但必须保留所有重要的技术细节和上下文信息，确保后续攻击链生成时能够准确理解整个测试过程。
-
-**必须详细保留的内容：**
-1. **所有工具执行记录**：
- - 工具名称、执行参数、执行结果（包括成功和失败）
- - 失败执行的错误信息、状态码、响应头等关键线索
- - 工具输出的关键数据（端口、服务版本、漏洞信息等）
- - 每个工具执行的时间顺序和上下文关系
-
-2. **所有发现的漏洞和潜在安全问题**：
- - 漏洞类型、严重程度、位置、利用方式
- - 验证过程和结果
- - 漏洞之间的关联关系
-
-3. **所有测试目标和资产信息**：
- - IP地址、域名、URL、端口等
- - 发现的服务、技术栈、版本信息
- - 资产之间的关联关系
-
-4. **所有测试步骤和决策过程**：
- - 每个测试步骤的详细描述（做了什么、为什么做、结果如何）
- - AI的分析思路和决策依据
- - 失败尝试的原因和从中获得的线索
-
-5. **所有关键发现和线索**：
- - 成功发现的详细信息
- - 失败但提供线索的尝试（错误信息、限制条件、下一步建议等）
- - 收集到的任何有价值的信息（凭据、令牌、配置信息等）
-
-**总结要求：**
-- 用结构化的方式组织信息，按时间顺序或逻辑顺序排列
-- 对于每个工具执行，必须包含：工具名、目标、参数、结果/错误、关键发现
-- 对于每个漏洞，必须包含：类型、位置、严重程度、验证结果
-- 保留所有技术细节，不要过度简化
-- 确保后续AI能够根据这个摘要完整重建攻击链
-
-对话片段：
-%s
-
-请给出详细且结构化的技术摘要（建议1000-2000字，确保信息完整）：`
-
- // 检查prompt tokens，如果超过限制，需要进一步压缩内容
- maxContextLength := b.getModelMaxContextLength()
- maxPromptTokens := maxContextLength - 2000 // 留出空间给响应和系统消息
-
- // 尝试构建完整prompt并检查tokens
- fullPrompt := fmt.Sprintf(promptTemplate, content)
- promptTokens, err := b.countTextTokens(fullPrompt)
- if err != nil {
- // 如果计算失败，使用估算
- promptTokens = len(fullPrompt) / 4
- }
-
- // 如果prompt太大，需要进一步压缩内容
- if promptTokens > maxPromptTokens {
- b.logger.Warn("片段内容过大，需要进一步压缩",
- zap.Int("promptTokens", promptTokens),
- zap.Int("maxPromptTokens", maxPromptTokens))
-
- // 递归压缩：将chunk进一步分割
- compressedContent, err := b.compressLargeChunk(ctx, chunk, maxPromptTokens)
- if err != nil {
- return "", fmt.Errorf("压缩大片段失败: %w", err)
- }
- content = compressedContent
- }
-
- prompt := fmt.Sprintf(promptTemplate, content)
-
- // 检查配置
- if b.openAIConfig == nil {
- return "", fmt.Errorf("OpenAI配置未初始化")
- }
- if b.openAIConfig.APIKey == "" {
- return "", fmt.Errorf("OpenAI API Key未配置")
- }
- if b.openAIConfig.Model == "" {
- return "", fmt.Errorf("OpenAI Model未配置")
- }
-
- // 直接调用AI API进行总结
- requestBody := map[string]interface{}{
- "model": b.openAIConfig.Model,
- "messages": []map[string]interface{}{
- {
- "role": "system",
- "content": `你是一个资深的安全测试分析师和渗透测试专家，拥有丰富的实战经验。你的任务是总结安全测试对话片段，这些摘要将用于后续构建完整的攻击链图。
-
-**你的专业背景：**
-- 精通各种安全测试工具（Nmap、SQLMap、Burp Suite、Metasploit等）的使用和结果分析
-- 熟悉常见漏洞类型（SQL注入、XSS、文件上传、命令执行、目录遍历等）的识别和验证
-- 理解攻击链的构建逻辑：从信息收集 → 漏洞发现 → 漏洞利用 → 权限提升 → 横向移动
-- 能够识别失败尝试中的有价值线索（错误信息、状态码、WAF指纹、技术栈信息等）
-
-**你的总结原则：**
-1. **完整性优先**：虽然需要压缩，但必须保留所有技术细节，确保后续AI能够完整重建攻击链
-2. **结构化组织**：按时间顺序或逻辑顺序组织信息，让信息易于理解和追踪
-3. **技术精准**：使用准确的技术术语，保留具体的数值、版本号、端口号、URL等关键数据
-4. **上下文关联**：保留测试步骤之间的因果关系和逻辑关联
-5. **失败价值**：即使是失败的尝试，只要提供了线索（错误信息、限制条件、下一步建议），也要详细记录
-
-**你需要特别关注的信息类型：**
-- 工具执行：工具名、目标、参数、完整结果（包括错误和失败）
-- 漏洞发现：类型、位置、严重程度、验证方法、利用结果
-- 资产信息：IP、域名、端口、服务版本、技术栈
-- 测试策略：为什么选择这个工具、为什么测试这个目标、发现了什么线索
-- 关键数据：凭据、令牌、配置信息、敏感文件内容
-
-请用专业、详细、结构化的中文进行总结，确保信息完整且易于后续处理。`,
- },
- {
- "role": "user",
- "content": prompt,
- },
- },
- "temperature": 0.3,
- "max_tokens": 4000, // 增加摘要长度，以容纳更详细的内容
- }
-
- 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 {
- return "", fmt.Errorf("请求失败: %w", err)
- }
-
- if len(apiResponse.Choices) == 0 {
- return "", fmt.Errorf("API未返回有效响应")
- }
-
- return strings.TrimSpace(apiResponse.Choices[0].Message.Content), nil
-}
-
-// compressLongestItem 压缩最长的子节点（保留作为备用方法）
-func (b *Builder) compressLongestItem(ctx context.Context, contextData *ContextData) error {
- // 使用新的分片压缩方法
- return b.compressContextData(ctx, contextData)
-}
-
-// summarizeContent 总结内容
-func (b *Builder) summarizeContent(ctx context.Context, contentType, content string) (string, error) {
- var prompt string
- switch contentType {
- case "message":
- prompt = fmt.Sprintf(`请总结以下AI回复的关键信息，保留所有重要的安全发现、漏洞信息和测试结果。用简洁的中文总结，不超过500字。
-
-AI回复：
-%s
-
-总结：`, content)
- case "execution":
- prompt = fmt.Sprintf(`请总结以下工具执行结果的关键信息，保留所有发现的漏洞、重要发现和测试结果。用简洁的中文总结，不超过500字。
-
-工具执行结果：
-%s
-
-总结：`, content)
- case "thinking":
- prompt = fmt.Sprintf(`请总结以下AI思考过程的关键决策和思路，保留所有重要的决策点和测试策略。用简洁的中文总结，不超过300字。
-
-思考过程：
-%s
-
-总结：`, content)
- default:
- return "", fmt.Errorf("未知的内容类型: %s", contentType)
- }
-
- requestBody := map[string]interface{}{
- "model": b.openAIConfig.Model,
- "messages": []map[string]interface{}{
- {
- "role": "system",
- "content": `你是一个资深的安全测试分析师和渗透测试专家，拥有丰富的实战经验。你的任务是总结安全测试过程中的关键信息，这些摘要将用于构建攻击链图。
-
-**你的专业背景：**
-- 精通各种安全测试工具的使用和结果分析（Nmap、SQLMap、Burp Suite、Metasploit、Nuclei等）
-- 熟悉常见漏洞类型的识别和验证（SQL注入、XSS、文件上传、命令执行、目录遍历、SSRF等）
-- 理解攻击链的构建逻辑和测试流程
-- 能够识别失败尝试中的有价值线索
-
-**你的总结原则：**
-1. **保留技术细节**：保留所有重要的技术信息，包括工具名、参数、结果、错误信息、状态码等
-2. **突出关键发现**：重点记录发现的漏洞、安全问题、资产信息、凭据等
-3. **记录失败线索**：即使是失败的尝试，如果提供了错误信息、限制条件或下一步建议，也要详细记录
-4. **保持准确性**：使用准确的技术术语，保留具体的数值、版本号、端口号等关键数据
-5. **结构化表达**：用清晰、有条理的方式组织信息
-
-**根据内容类型，你需要特别关注：**
-- **AI回复**：提取安全发现、漏洞信息、测试结果、分析思路、决策依据
-- **工具执行**：记录工具名、目标、参数、完整结果（成功或失败）、关键发现、错误信息
-- **思考过程**：提取关键决策点、测试策略、分析思路、下一步计划
-
-请用专业、准确、简洁的中文进行总结，确保信息完整且易于理解。`,
- },
- {
- "role": "user",
- "content": prompt,
- },
- },
- "temperature": 0.3,
- "max_tokens": 1000,
- }
-
- 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 {
- return "", fmt.Errorf("请求失败: %w", err)
- }
-
- if len(apiResponse.Choices) == 0 {
- return "", fmt.Errorf("API未返回有效响应")
- }
-
- return strings.TrimSpace(apiResponse.Choices[0].Message.Content), nil
-}
-
-// buildChunkContent 构建chunk的文本内容
-func (b *Builder) buildChunkContent(chunk *ContextChunk) (string, error) {
- var contentBuilder strings.Builder
-
- // 添加消息
- for _, msg := range chunk.Messages {
- if strings.EqualFold(msg.Role, "user") {
- contentBuilder.WriteString(fmt.Sprintf("用户消息: %s\n\n", msg.Content))
- } else {
- contentBuilder.WriteString(fmt.Sprintf("AI回复: %s\n\n", msg.Content))
- }
- }
-
- // 添加工具执行
- for _, exec := range chunk.Executions {
- contentBuilder.WriteString(fmt.Sprintf("工具执行 [%s] (ID: %s):\n", exec.ToolName, exec.ID))
- contentBuilder.WriteString(fmt.Sprintf("参数: %s\n", b.formatArguments(exec.Arguments)))
-
- if exec.Error != "" {
- contentBuilder.WriteString(fmt.Sprintf("错误: %s\n", exec.Error))
- }
-
- if exec.Result != nil {
- var resultText string
- for _, content := range exec.Result.Content {
- if content.Type == "text" {
- resultText += content.Text + "\n"
- }
- }
- if resultText != "" {
- // 如果结果太长，截断
- if len(resultText) > 10000 {
- resultText = resultText[:10000] + "\n... [内容已截断]"
- }
- contentBuilder.WriteString(fmt.Sprintf("结果: %s\n", resultText))
- }
- }
- contentBuilder.WriteString("\n")
- }
-
- // 添加思考过程
- for _, details := range chunk.ProcessDetails {
- for _, detail := range details {
- if detail.EventType == "thinking" {
- thinkingText := detail.Message
- // 如果思考过程太长，截断
- if len(thinkingText) > 5000 {
- thinkingText = thinkingText[:5000] + "\n... [内容已截断]"
- }
- contentBuilder.WriteString(fmt.Sprintf("思考过程: %s\n\n", thinkingText))
- }
- }
- }
-
- content := contentBuilder.String()
- if content == "" {
- return "", fmt.Errorf("片段内容为空")
- }
- return content, nil
-}
-
-// compressLargeChunk 压缩过大的chunk（递归分割）
-func (b *Builder) compressLargeChunk(ctx context.Context, chunk *ContextChunk, maxTokens int) (string, error) {
- // 将chunk进一步分割成更小的子chunk
- // 简单策略：按消息和执行数量平均分割
- totalItems := len(chunk.Messages) + len(chunk.Executions)
- if totalItems <= 1 {
- // 如果只有一个项目，直接截断内容
- content, _ := b.buildChunkContent(chunk)
- if len(content) > maxTokens*4 { // 粗略估算：1 token ≈ 4字符
- content = content[:maxTokens*4] + "\n... [内容过大，已截断]"
- }
- return content, nil
- }
-
- // 分成2个子chunk
- mid := totalItems / 2
- subChunk1 := &ContextChunk{
- Messages: []database.Message{},
- Executions: []*mcp.ToolExecution{},
- ProcessDetails: make(map[string][]database.ProcessDetail),
- }
- subChunk2 := &ContextChunk{
- Messages: []database.Message{},
- Executions: []*mcp.ToolExecution{},
- ProcessDetails: make(map[string][]database.ProcessDetail),
- }
-
- // 分配消息
- for i, msg := range chunk.Messages {
- if i < mid {
- subChunk1.Messages = append(subChunk1.Messages, msg)
- } else {
- subChunk2.Messages = append(subChunk2.Messages, msg)
- }
- }
-
- // 分配执行
- execStart := len(chunk.Messages)
- for i, exec := range chunk.Executions {
- if execStart+i < mid {
- subChunk1.Executions = append(subChunk1.Executions, exec)
- } else {
- subChunk2.Executions = append(subChunk2.Executions, exec)
- }
- }
-
- // 递归压缩子chunk
- summary1, err := b.summarizeContextChunk(ctx, subChunk1)
- if err != nil {
- return "", fmt.Errorf("压缩子chunk1失败: %w", err)
- }
-
- summary2, err := b.summarizeContextChunk(ctx, subChunk2)
- if err != nil {
- return "", fmt.Errorf("压缩子chunk2失败: %w", err)
- }
-
- // 合并摘要
- return fmt.Sprintf("片段1摘要：\n%s\n\n---\n\n片段2摘要：\n%s", summary1, summary2), nil
-}
-
-// countTextTokens 计算文本的tokens数
-func (b *Builder) countTextTokens(text string) (int, error) {
- if b.tokenCounter == nil || b.openAIConfig == nil {
- return len(text) / 4, nil
- }
-
- model := b.openAIConfig.Model
- if model == "" {
- model = "gpt-4"
- }
-
- count, err := b.tokenCounter.Count(model, text)
- if err != nil {
- return len(text) / 4, nil
- }
- return count, nil
-}
-
// callAIForChainGeneration 调用AI生成攻击链
func (b *Builder) callAIForChainGeneration(ctx context.Context, prompt string) (string, error) {
requestBody := map[string]interface{}{
@@ -1329,33 +336,12 @@ func (b *Builder) parseChainJSON(chainJSON string, executions []*mcp.ToolExecuti
return nil, fmt.Errorf("解析JSON失败: %w", err)
}
- // 创建execution ID映射（AI可能返回简单的索引或ID，需要映射到真实的execution ID）
- executionMap := make(map[string]string) // AI返回的ID -> 真实execution ID
- for i, exec := range executions {
- // 支持多种可能的AI返回格式
- executionMap[fmt.Sprintf("exec_%d", i+1)] = exec.ID
- executionMap[fmt.Sprintf("execution_%d", i+1)] = exec.ID
- executionMap[exec.ID] = exec.ID // 如果AI直接返回真实ID
- executionMap[fmt.Sprintf("tool_%d", i+1)] = exec.ID // AI可能用tool_1格式
- executionMap[fmt.Sprintf("执行%d", i+1)] = exec.ID // 中文格式
- executionMap[fmt.Sprintf("执行_%d", i+1)] = exec.ID
- }
-
// 创建节点ID映射（AI返回的ID -> 新的UUID）
nodeIDMap := make(map[string]string)
- // 转换为Chain结构，并过滤无效节点
+ // 转换为Chain结构
nodes := make([]Node, 0, len(chainData.Nodes))
for _, n := range chainData.Nodes {
- // 过滤无效节点
- if b.shouldFilterNode(n, executions) {
- b.logger.Info("过滤无效节点",
- zap.String("nodeID", n.ID),
- zap.String("nodeType", n.Type),
- zap.String("label", n.Label))
- continue
- }
-
// 生成新的UUID节点ID
newNodeID := fmt.Sprintf("node_%s", uuid.New().String())
nodeIDMap[n.ID] = newNodeID
@@ -1370,420 +356,37 @@ func (b *Builder) parseChainJSON(chainJSON string, executions []*mcp.ToolExecuti
if node.Metadata == nil {
node.Metadata = make(map[string]interface{})
}
-
- // 处理tool_execution_id：如果是action或vulnerability节点，需要映射到真实的execution ID
- if n.ToolExecutionID != "" {
- if realExecID, ok := executionMap[n.ToolExecutionID]; ok {
- node.ToolExecutionID = realExecID
- } else {
- // 检查是否是真实的execution ID（UUID格式）
- // 如果是，直接使用；如果不是，尝试从节点ID推断
- if len(n.ToolExecutionID) > 20 { // UUID通常很长
- node.ToolExecutionID = n.ToolExecutionID
- } else {
- // 可能是简单的ID，尝试从节点ID推断
- if realExecID, ok := executionMap[n.ID]; ok {
- node.ToolExecutionID = realExecID
- } else {
- b.logger.Warn("无法映射tool_execution_id",
- zap.String("nodeID", n.ID),
- zap.String("toolExecutionID", n.ToolExecutionID))
- // 对于action节点，如果没有有效的execution ID，清空它（避免外键约束失败）
- if n.Type == "action" {
- node.ToolExecutionID = ""
- }
- }
- }
- }
- } else if n.Type == "action" || n.Type == "vulnerability" {
- // 如果AI没有提供tool_execution_id，尝试从节点ID推断
- // 例如：tool_1 -> 查找exec_1
- if realExecID, ok := executionMap[n.ID]; ok {
- node.ToolExecutionID = realExecID
- } else {
- b.logger.Warn("action/vulnerability节点缺少tool_execution_id",
- zap.String("nodeID", n.ID),
- zap.String("nodeType", n.Type))
- }
- }
-
nodes = append(nodes, node)
}
- // 转换边，更新source和target为新的节点ID
+ // 转换边
edges := make([]Edge, 0, len(chainData.Edges))
for _, e := range chainData.Edges {
sourceID, ok := nodeIDMap[e.Source]
if !ok {
- b.logger.Warn("边的源节点ID未找到", zap.String("source", e.Source))
continue
}
-
targetID, ok := nodeIDMap[e.Target]
if !ok {
- b.logger.Warn("边的目标节点ID未找到", zap.String("target", e.Target))
continue
}
- edge := Edge{
- ID: fmt.Sprintf("edge_%s", uuid.New().String()),
+ // 生成边的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,
- }
- edges = append(edges, edge)
+ })
}
- // 过滤掉指向已删除节点的边
- filteredEdges := make([]Edge, 0, len(edges))
- for _, edge := range edges {
- // 检查source和target节点是否都存在
- sourceExists := false
- targetExists := false
- for _, node := range nodes {
- if node.ID == edge.Source {
- sourceExists = true
- }
- if node.ID == edge.Target {
- targetExists = true
- }
- }
-
- if sourceExists && targetExists {
- filteredEdges = append(filteredEdges, edge)
- } else {
- b.logger.Warn("过滤无效边",
- zap.String("edgeID", edge.ID),
- zap.String("source", edge.Source),
- zap.String("target", edge.Target),
- zap.Bool("sourceExists", sourceExists),
- zap.Bool("targetExists", targetExists))
- }
- }
-
- // 验证和优化DAG结构
- optimizedEdges := b.optimizeDAGStructure(nodes, filteredEdges)
-
return &Chain{
Nodes: nodes,
- Edges: optimizedEdges,
+ Edges: edges,
}, nil
}
-// shouldFilterNode 判断是否应该过滤掉这个节点
-func (b *Builder) shouldFilterNode(n struct {
- ID string `json:"id"`
- Type string `json:"type"`
- Label string `json:"label"`
- RiskScore int `json:"risk_score"`
- ToolExecutionID string `json:"tool_execution_id,omitempty"`
- Metadata map[string]interface{} `json:"metadata"`
-}, executions []*mcp.ToolExecution) bool {
- // 只允许target、action、vulnerability三种节点类型
- if n.Type != "target" && n.Type != "action" && n.Type != "vulnerability" {
- return true
- }
-
- // 检查节点标签是否为空或无效
- if strings.TrimSpace(n.Label) == "" {
- return true
- }
-
- // 对于vulnerability节点，即使没有tool_execution_id也应该保留（漏洞可能不是直接来自工具执行）
- if n.Type == "vulnerability" {
- // 只要标签有意义就保留
- return false
- }
-
- // 对于target节点，只要标签有意义就保留
- if n.Type == "target" {
- return false
- }
-
- // 对于action节点，进行更宽松的检查
- if n.Type == "action" {
- // 如果executions为空（可能是压缩后的场景），只要标签有意义就保留
- if len(executions) == 0 {
- // 压缩场景下，只要标签不是明显无效就保留
- labelLower := strings.ToLower(n.Label)
- // 只过滤明显无效的标签
- invalidKeywords := []string{"空节点", "无效节点", "empty node", "invalid node"}
- for _, keyword := range invalidKeywords {
- if strings.Contains(labelLower, keyword) {
- return true
- }
- }
- return false
- }
-
- // 如果有tool_execution_id，尝试查找对应的工具执行
- if n.ToolExecutionID != "" {
- var exec *mcp.ToolExecution
- for _, e := range executions {
- if e.ID == n.ToolExecutionID {
- exec = e
- break
- }
- }
-
- if exec != nil {
- // 找到了对应的工具执行，检查是否有效
- // 检查工具执行是否错误或失败
- if exec.Error != "" || (exec.Result != nil && exec.Result.IsError) {
- // 失败但有线索的应该保留
- if !hasInsightfulFailure(n.Metadata) {
- // 即使没有明确标记为有线索，如果标签描述了具体内容，也保留
- labelLower := strings.ToLower(n.Label)
- // 如果标签包含具体的技术信息（端口、服务、漏洞等），说明有价值
- valuableKeywords := []string{"端口", "服务", "漏洞", "扫描", "发现", "获取", "验证", "port", "service", "vulnerability", "scan", "found", "discover"}
- hasValuableInfo := false
- for _, keyword := range valuableKeywords {
- if strings.Contains(labelLower, keyword) {
- hasValuableInfo = true
- break
- }
- }
- if !hasValuableInfo {
- return true
- }
- }
- }
-
- // 检查工具执行结果是否为空
- if exec.Result == nil || len(exec.Result.Content) == 0 {
- // 结果为空，但如果有线索或标签有意义，也保留
- if !hasInsightfulFailure(n.Metadata) {
- labelLower := strings.ToLower(n.Label)
- valuableKeywords := []string{"端口", "服务", "漏洞", "扫描", "发现", "获取", "验证", "port", "service", "vulnerability", "scan", "found", "discover"}
- hasValuableInfo := false
- for _, keyword := range valuableKeywords {
- if strings.Contains(labelLower, keyword) {
- hasValuableInfo = true
- break
- }
- }
- if !hasValuableInfo {
- return true
- }
- }
- } else {
- // 检查结果文本是否为空
- var resultText string
- for _, content := range exec.Result.Content {
- if content.Type == "text" {
- resultText += content.Text
- }
- }
- if strings.TrimSpace(resultText) == "" {
- // 结果文本为空，但如果有线索或标签有意义，也保留
- if !hasInsightfulFailure(n.Metadata) {
- labelLower := strings.ToLower(n.Label)
- valuableKeywords := []string{"端口", "服务", "漏洞", "扫描", "发现", "获取", "验证", "port", "service", "vulnerability", "scan", "found", "discover"}
- hasValuableInfo := false
- for _, keyword := range valuableKeywords {
- if strings.Contains(labelLower, keyword) {
- hasValuableInfo = true
- break
- }
- }
- if !hasValuableInfo {
- return true
- }
- }
- }
- }
- } else {
- // 找不到对应的工具执行，但可能是压缩后的场景
- // 只要标签有意义就保留（不要因为找不到execution就过滤掉）
- labelLower := strings.ToLower(n.Label)
- invalidKeywords := []string{"空节点", "无效节点", "empty node", "invalid node"}
- for _, keyword := range invalidKeywords {
- if strings.Contains(labelLower, keyword) {
- return true
- }
- }
- // 标签有意义，保留
- return false
- }
- } else {
- // 没有tool_execution_id，但可能是压缩后的场景或AI生成的节点
- // 只要标签有意义就保留
- labelLower := strings.ToLower(n.Label)
- invalidKeywords := []string{"空节点", "无效节点", "empty node", "invalid node"}
- for _, keyword := range invalidKeywords {
- if strings.Contains(labelLower, keyword) {
- return true
- }
- }
- // 标签有意义，保留
- return false
- }
- }
-
- // 默认保留（已经通过了所有检查）
- return false
-}
-
-// optimizeDAGStructure 优化DAG结构，检测并修复循环
-func (b *Builder) optimizeDAGStructure(nodes []Node, edges []Edge) []Edge {
- // 构建邻接表
- adjList := make(map[string][]string) // nodeID -> []targetNodeIDs
- nodeSet := make(map[string]bool)
- for _, node := range nodes {
- nodeSet[node.ID] = true
- }
-
- // 构建邻接表
- for _, edge := range edges {
- if nodeSet[edge.Source] && nodeSet[edge.Target] {
- adjList[edge.Source] = append(adjList[edge.Source], edge.Target)
- }
- }
-
- // 检测循环
- cycles := b.detectCycles(adjList, nodeSet)
- if len(cycles) == 0 {
- // 没有循环，直接返回
- return edges
- }
-
- b.logger.Warn("检测到攻击链中存在循环，正在修复",
- zap.Int("cycleCount", len(cycles)))
-
- // 构建边映射，方便删除
- edgeMap := make(map[string]Edge) // "source:target" -> Edge
- for _, edge := range edges {
- key := edge.Source + ":" + edge.Target
- edgeMap[key] = edge
- }
-
- // 删除形成循环的边（保留权重较低的边，通常权重低的边重要性较低）
- removedEdges := make(map[string]bool)
- for _, cycle := range cycles {
- // 找到循环中权重最低的边并删除
- minWeight := 999
- var edgeToRemove string
- for i := 0; i < len(cycle); i++ {
- source := cycle[i]
- target := cycle[(i+1)%len(cycle)]
- key := source + ":" + target
- if edge, exists := edgeMap[key]; exists {
- if edge.Weight < minWeight {
- minWeight = edge.Weight
- edgeToRemove = key
- }
- }
- }
- if edgeToRemove != "" {
- removedEdges[edgeToRemove] = true
- b.logger.Info("删除形成循环的边",
- zap.String("edge", edgeToRemove),
- zap.Int("weight", minWeight))
- }
- }
-
- // 重新构建边列表，排除已删除的边
- optimizedEdges := make([]Edge, 0, len(edges))
- for _, edge := range edges {
- key := edge.Source + ":" + edge.Target
- if !removedEdges[key] {
- optimizedEdges = append(optimizedEdges, edge)
- }
- }
-
- // 再次验证（递归处理，最多3次）
- if len(optimizedEdges) < len(edges) {
- // 重新构建邻接表
- newAdjList := make(map[string][]string)
- for _, edge := range optimizedEdges {
- newAdjList[edge.Source] = append(newAdjList[edge.Source], edge.Target)
- }
- newCycles := b.detectCycles(newAdjList, nodeSet)
- if len(newCycles) > 0 && len(removedEdges) < 10 { // 防止无限循环
- // 递归优化
- return b.optimizeDAGStructure(nodes, optimizedEdges)
- }
- }
-
- return optimizedEdges
-}
-
-// detectCycles 检测图中的循环（使用DFS）
-func (b *Builder) detectCycles(adjList map[string][]string, nodeSet map[string]bool) [][]string {
- visited := make(map[string]bool)
- recStack := make(map[string]bool)
- cycles := make([][]string, 0)
-
- var dfs func(node string, path []string) bool
- dfs = func(node string, path []string) bool {
- if !nodeSet[node] {
- return false
- }
-
- visited[node] = true
- recStack[node] = true
- currentPath := append(path, node)
-
- for _, neighbor := range adjList[node] {
- if !visited[neighbor] {
- if dfs(neighbor, currentPath) {
- return true
- }
- } else if recStack[neighbor] {
- // 找到循环：从neighbor到node的路径
- cycleStart := -1
- for i, n := range currentPath {
- if n == neighbor {
- cycleStart = i
- break
- }
- }
- if cycleStart >= 0 {
- cycle := append(currentPath[cycleStart:], neighbor)
- cycles = append(cycles, cycle)
- }
- return true
- }
- }
-
- recStack[node] = false
- return false
- }
-
- // 对所有节点进行DFS
- for node := range nodeSet {
- if !visited[node] {
- dfs(node, []string{})
- }
- }
-
- return cycles
-}
-
-func hasInsightfulFailure(metadata map[string]interface{}) bool {
- if metadata == nil {
- return false
- }
-
- if status, ok := metadata["status"].(string); ok {
- normalized := strings.ToLower(strings.TrimSpace(status))
- if normalized == "failed_insight" || normalized == "failed_clue" || normalized == "failed_with_hint" {
- return true
- }
- }
-
- if hint, ok := metadata["hint"].(string); ok && strings.TrimSpace(hint) != "" {
- return true
- }
-
- if hints, ok := metadata["hints"].([]interface{}); ok && len(hints) > 0 {
- return true
- }
-
- if insight, ok := metadata["insight"].(string); ok && strings.TrimSpace(insight) != "" {
- return true
- }
-
- return false
-}
+// 以下所有方法已不再使用，已删除以简化代码
diff --git a/web/static/css/style.css b/web/static/css/style.css
index 8bf60144..8c7682ff 100644
--- a/web/static/css/style.css
+++ b/web/static/css/style.css
@@ -1640,18 +1640,23 @@ header {
display: flex;
justify-content: space-between;
align-items: center;
- padding: 20px 24px;
+ padding: 24px 28px;
border-bottom: 1px solid var(--border-color);
- background: var(--bg-primary);
- border-top-left-radius: 12px;
- border-top-right-radius: 12px;
+ background: linear-gradient(135deg, #ffffff 0%, #f8f9fa 100%);
+ border-top-left-radius: 16px;
+ border-top-right-radius: 16px;
+ box-shadow: 0 2px 8px rgba(0, 0, 0, 0.04);
}
.modal-header h2 {
margin: 0;
- font-size: 1.25rem;
+ font-size: 1.5rem;
font-weight: 600;
color: var(--text-primary);
+ background: linear-gradient(135deg, var(--accent-color) 0%, #0052cc 100%);
+ -webkit-background-clip: text;
+ -webkit-text-fill-color: transparent;
+ background-clip: text;
}
.modal-close {
@@ -3732,6 +3737,9 @@ header {
max-height: 90vh;
display: flex;
flex-direction: column;
+ border-radius: 16px;
+ overflow: hidden;
+ box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
}
.attack-chain-body {
@@ -3740,64 +3748,188 @@ header {
flex: 1;
overflow: hidden;
padding: 0;
+ background: linear-gradient(135deg, #f8f9fa 0%, #ffffff 100%);
}
.attack-chain-controls {
- padding: 16px;
+ padding: 20px 24px;
border-bottom: 1px solid var(--border-color);
display: flex;
justify-content: space-between;
align-items: center;
flex-wrap: wrap;
- gap: 16px;
- background: var(--bg-secondary);
+ gap: 20px;
+ background: linear-gradient(135deg, #ffffff 0%, #f8f9fa 100%);
+ box-shadow: 0 2px 8px rgba(0, 0, 0, 0.04);
}
.attack-chain-info {
- font-size: 0.875rem;
- color: var(--text-secondary);
+ font-size: 0.9375rem;
+ color: var(--text-primary);
+ font-weight: 500;
+ padding: 8px 16px;
+ background: linear-gradient(135deg, rgba(0, 102, 255, 0.08) 0%, rgba(0, 102, 255, 0.04) 100%);
+ border: 1px solid rgba(0, 102, 255, 0.15);
+ border-radius: 12px;
+ box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
}
.attack-chain-legend {
display: flex;
- gap: 16px;
+ gap: 20px;
flex-wrap: wrap;
+ padding: 12px 20px;
+ background: var(--bg-primary);
+ border-radius: 12px;
+ border: 1px solid var(--border-color);
+ box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
}
.legend-item {
display: flex;
align-items: center;
- gap: 6px;
+ gap: 8px;
font-size: 0.875rem;
+ font-weight: 500;
+ color: var(--text-primary);
+ padding: 4px 0;
}
.legend-color {
- width: 16px;
- height: 16px;
- border-radius: 4px;
- border: 1px solid var(--border-color);
+ width: 20px;
+ height: 20px;
+ border-radius: 6px;
+ border: 2px solid rgba(255, 255, 255, 0.8);
+ box-shadow: 0 2px 4px rgba(0, 0, 0, 0.15);
+ transition: transform 0.2s ease;
+}
+
+.legend-item:hover .legend-color {
+ transform: scale(1.1);
}
.attack-chain-container {
flex: 1;
min-height: 0;
- background: var(--bg-primary);
- border: 1px solid var(--border-color);
+ background: linear-gradient(135deg, #fafbfc 0%, #f5f7fa 100%);
+ border: none;
position: relative;
+ overflow: hidden;
+}
+
+.attack-chain-container::before {
+ content: '';
+ position: absolute;
+ top: 0;
+ left: 0;
+ right: 0;
+ bottom: 0;
+ background:
+ radial-gradient(circle at 20% 30%, rgba(0, 102, 255, 0.03) 0%, transparent 50%),
+ radial-gradient(circle at 80% 70%, rgba(0, 102, 255, 0.03) 0%, transparent 50%);
+ pointer-events: none;
+ z-index: 0;
+}
+
+/* 攻击链筛选器样式 */
+.attack-chain-filters {
+ display: flex;
+ gap: 12px;
+ align-items: center;
+ flex-wrap: wrap;
+ padding: 12px 16px;
+ background: var(--bg-primary);
+ border-radius: 12px;
+ border: 1px solid var(--border-color);
+ box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
+}
+
+.attack-chain-filters input[type="text"] {
+ padding: 10px 16px;
+ border: 2px solid var(--border-color);
+ border-radius: 8px;
+ font-size: 0.875rem;
+ min-width: 220px;
+ background: var(--bg-primary);
+ color: var(--text-primary);
+ transition: all 0.2s ease;
+ box-shadow: 0 1px 2px rgba(0, 0, 0, 0.05);
+}
+
+.attack-chain-filters input[type="text"]:focus {
+ outline: none;
+ border-color: var(--accent-color);
+ box-shadow: 0 0 0 3px rgba(0, 102, 255, 0.1);
+}
+
+.attack-chain-filters input[type="text"]::placeholder {
+ color: var(--text-muted);
+}
+
+.attack-chain-filters select {
+ padding: 10px 16px;
+ border: 2px solid var(--border-color);
+ border-radius: 8px;
+ font-size: 0.875rem;
+ background: var(--bg-primary);
+ color: var(--text-primary);
+ cursor: pointer;
+ transition: all 0.2s ease;
+ box-shadow: 0 1px 2px rgba(0, 0, 0, 0.05);
+ appearance: none;
+ background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='12' height='12' viewBox='0 0 12 12'%3E%3Cpath fill='%23666' d='M6 9L1 4h10z'/%3E%3C/svg%3E");
+ background-repeat: no-repeat;
+ background-position: right 12px center;
+ padding-right: 40px;
+}
+
+.attack-chain-filters select:focus {
+ outline: none;
+ border-color: var(--accent-color);
+ box-shadow: 0 0 0 3px rgba(0, 102, 255, 0.1);
+}
+
+.attack-chain-filters select:hover {
+ border-color: var(--accent-color);
+}
+
+.attack-chain-filters button.btn-secondary {
+ padding: 10px 18px;
+ font-size: 0.875rem;
+ font-weight: 500;
+ border: 2px solid var(--border-color);
+ border-radius: 8px;
+ background: var(--bg-primary);
+ color: var(--text-primary);
+ cursor: pointer;
+ transition: all 0.2s ease;
+ box-shadow: 0 1px 2px rgba(0, 0, 0, 0.05);
+}
+
+.attack-chain-filters button.btn-secondary:hover {
+ background: var(--bg-tertiary);
+ border-color: var(--accent-color);
+ color: var(--accent-color);
+ transform: translateY(-1px);
+ box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}
.attack-chain-details {
- padding: 16px;
+ padding: 20px 24px;
border-top: 1px solid var(--border-color);
- background: var(--bg-secondary);
+ background: linear-gradient(135deg, #ffffff 0%, #f8f9fa 100%);
max-height: 200px;
overflow-y: auto;
+ box-shadow: 0 -2px 8px rgba(0, 0, 0, 0.04);
}
.attack-chain-details h3 {
- margin: 0 0 12px 0;
- font-size: 1rem;
+ margin: 0 0 16px 0;
+ font-size: 1.125rem;
+ font-weight: 600;
color: var(--text-primary);
+ padding-bottom: 12px;
+ border-bottom: 2px solid var(--border-color);
}
.node-detail-item {
@@ -3829,8 +3961,35 @@ header {
.modal-header-actions {
display: flex;
- gap: 8px;
+ gap: 10px;
align-items: center;
+ flex-wrap: wrap;
+}
+
+.attack-chain-action-btn {
+ padding: 10px 18px;
+ font-size: 0.875rem;
+ font-weight: 500;
+ border-radius: 8px;
+ transition: all 0.2s ease;
+ box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
+ border: 2px solid transparent;
+}
+
+.attack-chain-action-btn:hover {
+ transform: translateY(-2px);
+ box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
+}
+
+.attack-chain-action-btn.btn-primary {
+ background: linear-gradient(135deg, #0066ff 0%, #0052cc 100%);
+ color: white;
+ border-color: #0066ff;
+}
+
+.attack-chain-action-btn.btn-primary:hover {
+ background: linear-gradient(135deg, #0052cc 0%, #0040a3 100%);
+ border-color: #0052cc;
}
.loading-spinner {
diff --git a/web/static/js/chat.js b/web/static/js/chat.js
index 436af6ce..690f3705 100644
--- a/web/static/js/chat.js
+++ b/web/static/js/chat.js
@@ -1703,9 +1703,33 @@ function renderAttackChain(chainData) {
// 准备Cytoscape数据
const elements = [];
- // 添加节点，并预计算文字颜色和边框颜色
+ // 添加节点，并预计算文字颜色和边框颜色，同时为类型标签准备数据
chainData.nodes.forEach(node => {
const riskScore = node.risk_score || 0;
+ const nodeType = node.type || '';
+
+ // 根据节点类型设置类型标签文本和标识符（使用更现代的设计）
+ let typeLabel = '';
+ let typeBadge = '';
+ let typeColor = '';
+ if (nodeType === 'target') {
+ typeLabel = '目标';
+ typeBadge = '○'; // 使用空心圆，更现代
+ typeColor = '#1976d2'; // 蓝色
+ } else if (nodeType === 'action') {
+ typeLabel = '行动';
+ typeBadge = '▷'; // 使用更简洁的三角形
+ typeColor = '#f57c00'; // 橙色
+ } else if (nodeType === 'vulnerability') {
+ typeLabel = '漏洞';
+ typeBadge = '◇'; // 使用空心菱形，更精致
+ typeColor = '#d32f2f'; // 红色
+ } else {
+ typeLabel = nodeType;
+ typeBadge = '•';
+ typeColor = '#666';
+ }
+
// 根据风险分数计算文字颜色和边框颜色
let textColor, borderColor, textOutlineWidth, textOutlineColor;
if (riskScore >= 80) {
@@ -1734,11 +1758,19 @@ function renderAttackChain(chainData) {
textOutlineColor = '#fff';
}
+ // 构建带类型标签的显示文本：使用现代极简的设计风格
+ // 类型标签显示在顶部，使用简洁的格式，通过间距自然分隔
+ const displayLabel = typeBadge + ' ' + typeLabel + '\n\n' + node.label;
+
elements.push({
data: {
id: node.id,
- label: node.label,
- type: node.type,
+ label: displayLabel, // 使用包含类型标签的标签
+ originalLabel: node.label, // 保存原始标签用于搜索
+ type: nodeType,
+ typeLabel: typeLabel, // 保存类型标签文本
+ typeBadge: typeBadge, // 保存类型标识符
+ typeColor: typeColor, // 保存类型颜色
riskScore: riskScore,
toolExecutionId: node.tool_execution_id || '',
metadata: node.metadata || {},
@@ -1750,17 +1782,29 @@ function renderAttackChain(chainData) {
});
});
- // 添加边
+ // 添加边（只添加源节点和目标节点都存在的边）
+ const nodeIds = new Set(chainData.nodes.map(node => node.id));
chainData.edges.forEach(edge => {
- elements.push({
- data: {
- id: edge.id,
+ // 验证源节点和目标节点是否存在
+ if (nodeIds.has(edge.source) && nodeIds.has(edge.target)) {
+ elements.push({
+ data: {
+ id: edge.id,
+ source: edge.source,
+ target: edge.target,
+ type: edge.type || 'leads_to',
+ weight: edge.weight || 1
+ }
+ });
+ } else {
+ console.warn('跳过无效的边：源节点或目标节点不存在', {
+ edgeId: edge.id,
source: edge.source,
target: edge.target,
- type: edge.type || 'leads_to',
- weight: edge.weight || 1
- }
- });
+ sourceExists: nodeIds.has(edge.source),
+ targetExists: nodeIds.has(edge.target)
+ });
+ }
});
// 初始化Cytoscape
@@ -1772,37 +1816,108 @@ function renderAttackChain(chainData) {
selector: 'node',
style: {
'label': 'data(label)',
- // 统一节点大小，减少布局混乱（根据复杂度调整）
- 'width': isComplexGraph ? 70 : 'mapData(riskScore, 0, 100, 50, 80)',
- 'height': isComplexGraph ? 70 : 'mapData(riskScore, 0, 100, 50, 80)',
- 'shape': function(ele) {
+ // 增大节点尺寸，使其更加醒目和美观
+ // 根据节点类型调整大小，target节点更大（增加高度以容纳类型标签）
+ 'width': function(ele) {
const type = ele.data('type');
- if (type === 'vulnerability') return 'diamond';
- if (type === 'action') return 'round-rectangle';
- if (type === 'target') return 'star';
- return 'ellipse';
+ if (type === 'target') return isComplexGraph ? 240 : 260;
+ return isComplexGraph ? 220 : 240;
+ },
+ 'height': function(ele) {
+ const type = ele.data('type');
+ if (type === 'target') return isComplexGraph ? 115 : 125;
+ return isComplexGraph ? 110 : 120;
+ },
+ 'shape': function(ele) {
+ // 所有节点都使用圆角矩形，参考图片风格
+ return 'round-rectangle';
},
'background-color': function(ele) {
+ const type = ele.data('type');
const riskScore = ele.data('riskScore') || 0;
- if (riskScore >= 80) return '#ff4444'; // 红色
- if (riskScore >= 60) return '#ff8800'; // 橙色
- if (riskScore >= 40) return '#ffbb00'; // 黄色
- return '#88cc00'; // 绿色
+
+ // target节点使用更深的蓝色背景，增强对比度
+ if (type === 'target') {
+ return '#bbdefb'; // 更深的浅蓝色
+ }
+
+ // 其他节点根据风险分数，使用更饱和的背景色，增加视觉层次
+ if (riskScore >= 80) return '#ffcdd2'; // 更饱和的浅红色
+ if (riskScore >= 60) return '#ffe0b2'; // 更饱和的浅橙色
+ if (riskScore >= 40) return '#fff9c4'; // 更饱和的浅黄色
+ return '#dcedc8'; // 更饱和的浅绿色
},
- // 使用预计算的颜色数据
- 'color': 'data(textColor)',
- 'font-size': isComplexGraph ? '11px' : '13px', // 复杂图使用更小字体
- 'font-weight': 'bold',
+ // 根据节点类型和风险分数设置文字颜色
+ // 注意：由于标签包含类型标签和内容，颜色适用于所有文本
+ 'color': function(ele) {
+ const type = ele.data('type');
+ const riskScore = ele.data('riskScore') || 0;
+
+ if (type === 'target') {
+ return '#1976d2'; // 深蓝色文字
+ }
+
+ if (riskScore >= 80) return '#c62828'; // 深红色
+ if (riskScore >= 60) return '#e65100'; // 深橙色
+ if (riskScore >= 40) return '#f57f17'; // 深黄色
+ return '#558b2f'; // 深绿色
+ },
+ 'font-size': function(ele) {
+ // 由于标签包含类型标签和内容，使用合适的字体大小
+ const type = ele.data('type');
+ if (type === 'target') return isComplexGraph ? '16px' : '18px';
+ return isComplexGraph ? '15px' : '17px';
+ },
+ 'font-weight': '600',
+ 'font-family': '-apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif',
'text-valign': 'center',
'text-halign': 'center',
'text-wrap': 'wrap',
- 'text-max-width': isComplexGraph ? '90px' : '110px', // 复杂图限制文本宽度
- 'text-overflow-wrap': 'anywhere', // 允许在任何位置换行
- 'border-width': 2,
- 'border-color': 'data(borderColor)',
- 'overlay-padding': '4px',
- 'text-outline-width': 'data(textOutlineWidth)',
- 'text-outline-color': 'data(textOutlineColor)'
+ 'text-max-width': function(ele) {
+ const type = ele.data('type');
+ if (type === 'target') return isComplexGraph ? '220px' : '240px';
+ return isComplexGraph ? '200px' : '220px';
+ },
+ 'text-overflow-wrap': 'anywhere',
+ 'text-margin-y': 3, // 调整垂直边距以适应多行文本
+ 'padding': '14px', // 增加内边距，使节点内容更有呼吸感和现代感
+ // 根据节点类型设置边框样式，使用更粗的边框增强视觉效果
+ 'border-width': function(ele) {
+ const type = ele.data('type');
+ if (type === 'target') return 5;
+ return 4;
+ },
+ 'border-radius': '12px', // 增加圆角半径，使节点更圆润美观
+ 'border-color': function(ele) {
+ const type = ele.data('type');
+ const riskScore = ele.data('riskScore') || 0;
+
+ if (type === 'target') {
+ return '#1976d2'; // 蓝色边框
+ }
+
+ if (riskScore >= 80) return '#d32f2f'; // 红色边框
+ if (riskScore >= 60) return '#f57c00'; // 橙色边框
+ if (riskScore >= 40) return '#fbc02d'; // 黄色边框
+ return '#689f38'; // 绿色边框
+ },
+ 'border-style': function(ele) {
+ const type = ele.data('type');
+ // target和vulnerability使用实线，action可以使用虚线
+ if (type === 'action') return 'solid';
+ return 'solid';
+ },
+ 'overlay-padding': '12px',
+ // 移除文字轮廓，使用纯色文字
+ 'text-outline-width': 0,
+ // 增强阴影效果，使节点更立体更有层次感
+ // 增强阴影效果，使节点更立体更有层次感（使用更柔和的阴影）
+ 'shadow-blur': 20,
+ 'shadow-opacity': 0.25,
+ 'shadow-offset-x': 2,
+ 'shadow-offset-y': 6,
+ 'shadow-color': 'rgba(0, 0, 0, 0.15)',
+ 'background-opacity': 1
}
},
{
@@ -1811,35 +1926,74 @@ function renderAttackChain(chainData) {
'width': 'mapData(weight, 1, 5, 1.5, 3)',
'line-color': function(ele) {
const type = ele.data('type');
- if (type === 'discovers') return '#3498db'; // 浅蓝：action发现vulnerability
- if (type === 'targets') return '#0066ff'; // 蓝色：target指向action
- if (type === 'enables') return '#e74c3c'; // 深红：vulnerability间的因果关系
- if (type === 'leads_to') return '#666'; // 灰色：action之间的逻辑顺序
- return '#999';
+ // 参考图片风格，使用不同颜色和样式
+ if (type === 'discovers') return '#42a5f5'; // 浅蓝色：action发现vulnerability
+ if (type === 'targets') return '#1976d2'; // 深蓝色：target指向action（虚线）
+ if (type === 'enables') return '#e53935'; // 红色：vulnerability间的因果关系
+ if (type === 'leads_to') return '#616161'; // 灰色：action之间的逻辑顺序
+ return '#9e9e9e';
},
'target-arrow-color': function(ele) {
const type = ele.data('type');
- if (type === 'discovers') return '#3498db';
- if (type === 'targets') return '#0066ff';
- if (type === 'enables') return '#e74c3c';
- if (type === 'leads_to') return '#666';
- return '#999';
+ if (type === 'discovers') return '#42a5f5';
+ if (type === 'targets') return '#1976d2';
+ if (type === 'enables') return '#e53935';
+ if (type === 'leads_to') return '#616161';
+ return '#9e9e9e';
},
'target-arrow-shape': 'triangle',
'target-arrow-size': 8,
- // 对于复杂图，使用straight样式减少交叉；简单图使用bezier更美观
- 'curve-style': isComplexGraph ? 'straight' : 'bezier',
- 'control-point-step-size': isComplexGraph ? 40 : 60, // bezier控制点间距
- 'control-point-distance': isComplexGraph ? 30 : 50, // bezier控制点距离
- 'opacity': isComplexGraph ? 0.5 : 0.7, // 复杂图降低不透明度，减少视觉混乱
- 'line-style': 'solid'
+ // 使用bezier曲线，更美观
+ 'curve-style': 'bezier',
+ 'control-point-step-size': 50,
+ 'control-point-distance': 40,
+ 'opacity': 0.7,
+ // 根据边类型设置线条样式：targets使用虚线，其他使用实线
+ 'line-style': function(ele) {
+ const type = ele.data('type');
+ if (type === 'targets') return 'dashed'; // target相关的边使用虚线
+ return 'solid';
+ },
+ 'line-dash-pattern': function(ele) {
+ const type = ele.data('type');
+ if (type === 'targets') return [8, 4]; // 虚线模式
+ return [];
+ },
+ // 添加边的阴影效果（浅色主题使用浅阴影）
+ 'shadow-blur': 3,
+ 'shadow-opacity': 0.1,
+ 'shadow-offset-x': 1,
+ 'shadow-offset-y': 1,
+ 'shadow-color': '#000000'
}
},
{
selector: 'node:selected',
style: {
- 'border-width': 4,
- 'border-color': '#0066ff'
+ 'border-width': 5,
+ 'border-color': '#0066ff',
+ 'shadow-blur': 16,
+ 'shadow-opacity': 0.6,
+ 'shadow-offset-x': 4,
+ 'shadow-offset-y': 5,
+ 'shadow-color': '#0066ff',
+ 'z-index': 999,
+ 'opacity': 1,
+ 'overlay-opacity': 0.1,
+ 'overlay-color': '#0066ff'
+ }
+ },
+ {
+ selector: 'node:hover',
+ style: {
+ 'border-width': 5,
+ 'shadow-blur': 14,
+ 'shadow-opacity': 0.5,
+ 'shadow-offset-x': 3,
+ 'shadow-offset-y': 4,
+ 'z-index': 998,
+ 'overlay-opacity': 0.05,
+ 'overlay-color': '#333333'
}
}
],
@@ -1861,21 +2015,59 @@ function renderAttackChain(chainData) {
try {
cytoscape.use(cytoscapeDagre);
layoutName = 'dagre';
- // 根据图的复杂度调整布局参数，优化可读性
+
+ // 动态计算布局参数，基于容器尺寸和节点数量
+ const containerWidth = container ? container.offsetWidth : 1200;
+ const containerHeight = container ? container.offsetHeight : 800;
+
+ // 计算平均节点宽度（考虑不同类型节点的平均尺寸）
+ const avgNodeWidth = isComplexGraph ? 230 : 250; // 基于新的节点尺寸
+ const avgNodeHeight = isComplexGraph ? 97.5 : 107.5;
+
+ // 计算图的层级深度（估算）
+ const estimatedDepth = Math.ceil(Math.log2(Math.max(nodeCount, 2))) + 1;
+
+ // 动态计算节点水平间距：基于容器宽度和节点数量
+ // 目标：使用容器宽度的85-90%，让图充分展开
+ const maxLevelWidth = Math.max(1, Math.ceil(nodeCount / estimatedDepth));
+ const targetGraphWidth = containerWidth * 0.88; // 使用88%的容器宽度
+ const minNodeSep = avgNodeWidth * 0.6; // 最小间距为节点宽度的60%
+ const calculatedNodeSep = Math.max(
+ minNodeSep,
+ Math.min(
+ (targetGraphWidth - avgNodeWidth * maxLevelWidth) / Math.max(1, maxLevelWidth - 1),
+ avgNodeWidth * 1.5 // 最大间距不超过节点宽度的1.5倍
+ )
+ );
+
+ // 动态计算层级间距：基于容器高度和层级数
+ const targetGraphHeight = containerHeight * 0.85;
+ const calculatedRankSep = Math.max(
+ avgNodeHeight * 1.2, // 最小为节点高度的1.2倍
+ Math.min(
+ targetGraphHeight / Math.max(estimatedDepth - 1, 1),
+ avgNodeHeight * 2.5 // 最大不超过节点高度的2.5倍
+ )
+ );
+
+ // 边间距：基于节点间距的合理比例
+ const calculatedEdgeSep = Math.max(30, calculatedNodeSep * 0.25);
+
+ // 根据图的复杂度调整布局参数，优化可读性和空间利用率
layoutOptions = {
name: 'dagre',
rankDir: 'TB', // 从上到下
- spacingFactor: isComplexGraph ? 3.0 : 2.5, // 增加整体间距
- nodeSep: isComplexGraph ? 100 : 80, // 增加节点间距，提高可读性
- edgeSep: isComplexGraph ? 50 : 40, // 增加边间距
- rankSep: isComplexGraph ? 140 : 120, // 增加层级间距，让层次更清晰
+ spacingFactor: 1.0, // 使用1.0，因为我们已经动态计算了具体间距
+ nodeSep: Math.round(calculatedNodeSep), // 动态计算的节点间距
+ edgeSep: Math.round(calculatedEdgeSep), // 动态计算的边间距
+ rankSep: Math.round(calculatedRankSep), // 动态计算的层级间距
nodeDimensionsIncludeLabels: true, // 考虑标签大小
animate: false,
- padding: 50, // 增加边距
+ padding: Math.min(40, containerWidth * 0.02), // 动态边距，不超过容器宽度的2%
// 优化边的路由，减少交叉
edgeRouting: 'polyline',
- // 对齐方式：居中对齐，让图更整齐
- align: 'UL' // 上左对齐
+ // 对齐方式：使用上左对齐，然后手动居中
+ align: 'UL' // 上左对齐（dagre不支持'C'）
};
} catch (e) {
console.warn('dagre布局注册失败，使用默认布局:', e);
@@ -1884,16 +2076,324 @@ function renderAttackChain(chainData) {
console.warn('dagre布局插件未加载，使用默认布局');
}
- // 应用布局
- attackChainCytoscape.layout(layoutOptions).run();
+ // 应用布局，等待布局完成后再平衡和居中
+ const layout = attackChainCytoscape.layout(layoutOptions);
+ layout.one('layoutstop', () => {
+ // 布局完成后，先平衡分支，再居中显示
+ setTimeout(() => {
+ balanceBranches();
+ setTimeout(() => {
+ centerAttackChain();
+ }, 50);
+ }, 100);
+ });
+ layout.run();
- // 布局完成后，调整视图以适应所有节点
- // 使用更大的padding，让图更易读
- attackChainCytoscape.fit(undefined, 60); // 60px padding
+ // 平衡分支分布的函数 - 使分支在根节点左右平均分布
+ function balanceBranches() {
+ try {
+ if (!attackChainCytoscape) {
+ return;
+ }
+
+ // 动态计算节点间距，基于容器尺寸
+ const container = attackChainCytoscape.container();
+ const containerWidth = container ? container.offsetWidth : 1200;
+ const avgNodeWidth = isComplexGraph ? 230 : 250;
+ const estimatedDepth = Math.ceil(Math.log2(Math.max(nodeCount, 2))) + 1;
+ const maxLevelWidth = Math.max(1, Math.ceil(nodeCount / estimatedDepth));
+ const targetGraphWidth = containerWidth * 0.88;
+ const minNodeSep = avgNodeWidth * 0.6;
+ const spacing = Math.max(
+ minNodeSep,
+ Math.min(
+ (targetGraphWidth - avgNodeWidth * maxLevelWidth) / Math.max(1, maxLevelWidth - 1),
+ avgNodeWidth * 1.5
+ )
+ );
+
+ // 找到target节点作为根节点
+ const targetNodes = attackChainCytoscape.nodes().filter(node => {
+ return node.data('type') === 'target';
+ });
+
+ if (targetNodes.length === 0) {
+ return; // 没有target节点，无法平衡
+ }
+
+ const rootNode = targetNodes[0];
+ const rootPos = rootNode.position();
+ const rootX = rootPos.x;
+ const rootY = rootPos.y;
+
+ // 构建图的邻接表
+ const edges = attackChainCytoscape.edges();
+ const childrenMap = new Map();
+
+ edges.forEach(edge => {
+ const { source, target, valid } = getEdgeNodes(edge);
+ if (valid) {
+ const sourceId = source.id();
+ const targetId = target.id();
+
+ if (!childrenMap.has(sourceId)) {
+ childrenMap.set(sourceId, []);
+ }
+ childrenMap.get(sourceId).push(targetId);
+ }
+ });
+
+ // 计算每个节点的子树宽度（递归）
+ const subtreeWidth = new Map();
+ function calculateSubtreeWidth(nodeId) {
+ if (subtreeWidth.has(nodeId)) {
+ return subtreeWidth.get(nodeId);
+ }
+
+ const children = childrenMap.get(nodeId) || [];
+ if (children.length === 0) {
+ subtreeWidth.set(nodeId, 0);
+ return 0;
+ }
+
+ // 计算所有子树的宽度总和
+ let totalWidth = 0;
+ children.forEach(childId => {
+ totalWidth += calculateSubtreeWidth(childId);
+ });
+
+ // 使用动态计算的间距
+ const width = Math.max(totalWidth + (children.length - 1) * spacing, spacing);
+
+ subtreeWidth.set(nodeId, width);
+ return width;
+ }
+
+ // 计算所有子树宽度
+ const nodes = attackChainCytoscape.nodes();
+ nodes.forEach(node => {
+ calculateSubtreeWidth(node.id());
+ });
+
+ // 获取根节点的直接子节点
+ const rootChildren = childrenMap.get(rootNode.id()) || [];
+
+ if (rootChildren.length === 0) {
+ return; // 没有子节点
+ }
+
+ // 将子节点分成左右两组
+ const childWidths = rootChildren.map(childId => ({
+ id: childId,
+ width: subtreeWidth.get(childId) || 100
+ })).sort((a, b) => b.width - a.width);
+
+ const leftGroup = [];
+ const rightGroup = [];
+ let leftTotal = 0;
+ let rightTotal = 0;
+
+ // 贪心分配：将较大的子树交替分配到左右
+ childWidths.forEach(child => {
+ if (leftTotal <= rightTotal) {
+ leftGroup.push(child);
+ leftTotal += child.width;
+ } else {
+ rightGroup.push(child);
+ rightTotal += child.width;
+ }
+ });
+
+ // 计算左右两侧需要的总宽度（使用动态计算的间距）
+ const leftTotalWidth = leftGroup.length > 0 ? leftTotal + (leftGroup.length - 1) * spacing : 0;
+ const rightTotalWidth = rightGroup.length > 0 ? rightTotal + (rightGroup.length - 1) * spacing : 0;
+ // 根据容器宽度动态调整，充分利用水平空间
+ // 使用更大的宽度系数，让图充分利用容器空间（使用88%的容器宽度以匹配布局算法）
+ const maxSideWidth = Math.max(leftTotalWidth, rightTotalWidth);
+ const targetWidth = Math.max(maxSideWidth * 1.2, containerWidth * 0.88); // 使用88%的容器宽度以匹配布局
+ const maxWidth = Math.max(targetWidth, avgNodeWidth * 2);
+
+ // 递归调整子树位置
+ function adjustSubtree(nodeId, centerX, availableWidth) {
+ const node = attackChainCytoscape.getElementById(nodeId);
+ if (!node) return;
+
+ const currentPos = node.position();
+ const children = childrenMap.get(nodeId) || [];
+
+ if (children.length === 0) {
+ // 叶子节点
+ node.position({
+ x: centerX,
+ y: currentPos.y
+ });
+ return;
+ }
+
+ // 计算子节点的宽度
+ const childWidths = children.map(childId => ({
+ id: childId,
+ width: subtreeWidth.get(childId) || 100
+ }));
+
+ const totalChildWidth = childWidths.reduce((sum, c) => sum + c.width, 0);
+ const totalSpacing = (children.length - 1) * spacing;
+ const neededWidth = totalChildWidth + totalSpacing;
+
+ // 如果需要的宽度超过可用宽度，按比例缩放
+ const scale = neededWidth > availableWidth ? availableWidth / neededWidth : 1;
+ const scaledWidth = neededWidth * scale;
+
+ // 分配子节点位置
+ let currentOffset = -scaledWidth / 2;
+ childWidths.forEach((child, index) => {
+ const childWidth = child.width * scale;
+ const childCenterX = centerX + currentOffset + childWidth / 2;
+
+ adjustSubtree(child.id, childCenterX, childWidth);
+ currentOffset += childWidth + spacing * scale;
+ });
+
+ // 调整当前节点到子节点的中心
+ const childPositions = children.map(childId => {
+ const childNode = attackChainCytoscape.getElementById(childId);
+ return childNode ? childNode.position().x : centerX;
+ });
+ const childrenCenterX = childPositions.reduce((sum, x) => sum + x, 0) / childPositions.length;
+
+ node.position({
+ x: childrenCenterX,
+ y: currentPos.y
+ });
+ }
+
+ // 调整左侧子树
+ let leftOffset = -maxWidth / 2;
+ leftGroup.forEach((child, index) => {
+ const childWidth = child.width;
+ const childCenterX = rootX + leftOffset + childWidth / 2;
+ adjustSubtree(child.id, childCenterX, childWidth);
+ leftOffset += childWidth + spacing;
+ });
+
+ // 调整右侧子树
+ let rightOffset = maxWidth / 2;
+ rightGroup.forEach((child, index) => {
+ const childWidth = child.width;
+ const childCenterX = rootX + rightOffset - childWidth / 2;
+ adjustSubtree(child.id, childCenterX, childWidth);
+ rightOffset -= (childWidth + spacing);
+ });
+
+ // 重新计算根节点的中心位置：基于所有直接子节点的实际位置
+ const rootChildrenPositions = rootChildren.map(childId => {
+ const childNode = attackChainCytoscape.getElementById(childId);
+ return childNode ? childNode.position().x : rootX;
+ });
+
+ if (rootChildrenPositions.length > 0) {
+ // 计算所有子节点的平均 x 位置作为根节点的中心位置
+ const childrenCenterX = rootChildrenPositions.reduce((sum, x) => sum + x, 0) / rootChildrenPositions.length;
+ rootNode.position({
+ x: childrenCenterX,
+ y: rootY
+ });
+ } else {
+ // 如果没有子节点，保持原位置
+ rootNode.position({
+ x: rootX,
+ y: rootY
+ });
+ }
+
+ } catch (error) {
+ console.warn('平衡分支时出错:', error);
+ }
+ }
- // 如果图太复杂，稍微缩小视图以便看到全貌
- if (isComplexGraph && nodeCount > 20) {
- attackChainCytoscape.zoom(0.85);
+ // 居中攻击链的函数
+ function centerAttackChain() {
+ try {
+ if (!attackChainCytoscape) {
+ return;
+ }
+
+ const container = attackChainCytoscape.container();
+ if (!container) {
+ return;
+ }
+
+ const containerWidth = container.offsetWidth;
+ const containerHeight = container.offsetHeight;
+
+ if (containerWidth === 0 || containerHeight === 0) {
+ // 如果容器尺寸为0，延迟重试
+ setTimeout(centerAttackChain, 100);
+ return;
+ }
+
+ // 先fit以适应所有节点，使用更小的边距以更好地填充空间
+ attackChainCytoscape.fit(undefined, 60);
+
+ // 等待fit完成，然后根据图的宽度调整缩放，并整体居中
+ setTimeout(() => {
+ const extent = attackChainCytoscape.extent();
+ if (!extent || typeof extent.x1 === 'undefined' || typeof extent.x2 === 'undefined' ||
+ typeof extent.y1 === 'undefined' || typeof extent.y2 === 'undefined') {
+ return;
+ }
+
+ // 根据图的宽度和容器宽度，调整缩放以更好地利用水平空间
+ const graphWidth = extent.x2 - extent.x1;
+ const graphHeight = extent.y2 - extent.y1;
+ const availableWidth = containerWidth * 0.88; // 使用88%的容器宽度（与布局算法一致）
+ const availableHeight = containerHeight * 0.85; // 使用85%的容器高度
+ const currentZoom = attackChainCytoscape.zoom();
+
+ // 计算基于宽度和高度的缩放比例，选择较小的以适配
+ const widthScale = graphWidth > 0 ? availableWidth / (graphWidth * currentZoom) : 1;
+ const heightScale = graphHeight > 0 ? availableHeight / (graphHeight * currentZoom) : 1;
+ const scale = Math.min(widthScale, heightScale);
+
+ if (graphWidth > 0 && scale > 1 && scale < 1.4) {
+ // 如果图在当前缩放下太窄，稍微放大以填充空间，但不要过度放大
+ attackChainCytoscape.zoom(currentZoom * scale);
+ }
+
+ // 如果图太复杂，稍微缩小视图
+ if (isComplexGraph && nodeCount > 20) {
+ attackChainCytoscape.zoom(attackChainCytoscape.zoom() * 0.9);
+ }
+
+ // 计算图的中心点（在图形坐标系中）
+ const graphCenterX = (extent.x1 + extent.x2) / 2;
+ const graphCenterY = (extent.y1 + extent.y2) / 2;
+
+ // 获取当前的缩放和平移
+ const zoom = attackChainCytoscape.zoom();
+ const pan = attackChainCytoscape.pan();
+
+ // 计算图中心在当前视图中的位置
+ const graphCenterViewX = graphCenterX * zoom + pan.x;
+ const graphCenterViewY = graphCenterY * zoom + pan.y;
+
+ // 目标位置：容器中心
+ const desiredViewX = containerWidth / 2;
+ const desiredViewY = containerHeight / 2;
+
+ // 计算需要平移的距离
+ const deltaX = desiredViewX - graphCenterViewX;
+ const deltaY = desiredViewY - graphCenterViewY;
+
+ // 应用新的平移，使整个图居中（包括所有分支）
+ attackChainCytoscape.pan({
+ x: pan.x + deltaX,
+ y: pan.y + deltaY
+ });
+ }, 150);
+ } catch (error) {
+ console.warn('居中图表时出错:', error);
+ }
}
// 添加点击事件
@@ -1902,21 +2402,30 @@ function renderAttackChain(chainData) {
showNodeDetails(node.data());
});
- // 添加悬停效果
- attackChainCytoscape.on('mouseover', 'node', function(evt) {
- const node = evt.target;
- node.style('opacity', 0.8);
- });
-
- attackChainCytoscape.on('mouseout', 'node', function(evt) {
- const node = evt.target;
- node.style('opacity', 1);
- });
+ // 悬停渐变效果已移除
// 保存原始数据用于过滤
window.attackChainOriginalData = chainData;
}
+// 安全地获取边的源节点和目标节点
+function getEdgeNodes(edge) {
+ try {
+ const source = edge.source();
+ const target = edge.target();
+
+ // 检查源节点和目标节点是否存在
+ if (!source || !target || source.length === 0 || target.length === 0) {
+ return { source: null, target: null, valid: false };
+ }
+
+ return { source: source, target: target, valid: true };
+ } catch (error) {
+ console.warn('获取边的节点时出错:', error, edge.id());
+ return { source: null, target: null, valid: false };
+ }
+}
+
// 过滤攻击链节点（按搜索关键词）
function filterAttackChainNodes(searchText) {
if (!attackChainCytoscape || !window.attackChainOriginalData) {
@@ -1935,7 +2444,9 @@ function filterAttackChainNodes(searchText) {
// 过滤节点
attackChainCytoscape.nodes().forEach(node => {
- const label = (node.data('label') || '').toLowerCase();
+ // 使用原始标签进行搜索，不包含类型标签
+ const originalLabel = node.data('originalLabel') || node.data('label') || '';
+ const label = originalLabel.toLowerCase();
const type = (node.data('type') || '').toLowerCase();
const matches = label.includes(searchLower) || type.includes(searchLower);
@@ -1951,8 +2462,14 @@ function filterAttackChainNodes(searchText) {
// 隐藏没有可见源节点或目标节点的边
attackChainCytoscape.edges().forEach(edge => {
- const sourceVisible = edge.source().style('display') !== 'none';
- const targetVisible = edge.target().style('display') !== 'none';
+ const { source, target, valid } = getEdgeNodes(edge);
+ if (!valid) {
+ edge.style('display', 'none');
+ return;
+ }
+
+ const sourceVisible = source.style('display') !== 'none';
+ const targetVisible = target.style('display') !== 'none';
if (sourceVisible && targetVisible) {
edge.style('display', 'element');
} else {
@@ -1990,8 +2507,14 @@ function filterAttackChainByType(type) {
// 隐藏没有可见源节点或目标节点的边
attackChainCytoscape.edges().forEach(edge => {
- const sourceVisible = edge.source().style('display') !== 'none';
- const targetVisible = edge.target().style('display') !== 'none';
+ const { source, target, valid } = getEdgeNodes(edge);
+ if (!valid) {
+ edge.style('display', 'none');
+ return;
+ }
+
+ const sourceVisible = source.style('display') !== 'none';
+ const targetVisible = target.style('display') !== 'none';
if (sourceVisible && targetVisible) {
edge.style('display', 'element');
} else {
@@ -2039,8 +2562,14 @@ function filterAttackChainByRisk(riskLevel) {
// 隐藏没有可见源节点或目标节点的边
attackChainCytoscape.edges().forEach(edge => {
- const sourceVisible = edge.source().style('display') !== 'none';
- const targetVisible = edge.target().style('display') !== 'none';
+ const { source, target, valid } = getEdgeNodes(edge);
+ if (!valid) {
+ edge.style('display', 'none');
+ return;
+ }
+
+ const sourceVisible = source.style('display') !== 'none';
+ const targetVisible = target.style('display') !== 'none';
if (sourceVisible && targetVisible) {
edge.style('display', 'element');
} else {
@@ -2102,7 +2631,7 @@ function showNodeDetails(nodeData) {
类型: ${getNodeTypeLabel(nodeData.type)}
- 标签: ${escapeHtml(nodeData.label)}
+ 标签: ${escapeHtml(nodeData.originalLabel || nodeData.label)}
风险评分: ${nodeData.riskScore}/100
@@ -2444,12 +2973,15 @@ function exportAttackChain(format) {
// 也考虑边的范围
edges.forEach(edge => {
- const sourcePos = edge.source().position();
- const targetPos = edge.target().position();
- minX = Math.min(minX, sourcePos.x, targetPos.x);
- minY = Math.min(minY, sourcePos.y, targetPos.y);
- maxX = Math.max(maxX, sourcePos.x, targetPos.x);
- maxY = Math.max(maxY, sourcePos.y, targetPos.y);
+ const { source, target, valid } = getEdgeNodes(edge);
+ if (valid) {
+ const sourcePos = source.position();
+ const targetPos = target.position();
+ minX = Math.min(minX, sourcePos.x, targetPos.x);
+ minY = Math.min(minY, sourcePos.y, targetPos.y);
+ maxX = Math.max(maxX, sourcePos.x, targetPos.x);
+ maxY = Math.max(maxY, sourcePos.y, targetPos.y);
+ }
});
// 添加边距
@@ -2507,8 +3039,13 @@ function exportAttackChain(format) {
// 添加边（先绘制，这样节点会在上面）
edges.forEach(edge => {
- const sourcePos = edge.source().position();
- const targetPos = edge.target().position();
+ const { source, target, valid } = getEdgeNodes(edge);
+ if (!valid) {
+ return; // 跳过无效的边
+ }
+
+ const sourcePos = source.position();
+ const targetPos = target.position();
const edgeData = edge.data();
const edgeType = edgeData.type || 'leads_to';
@@ -2607,7 +3144,8 @@ function exportAttackChain(format) {
svg.appendChild(shapeElement);
// 添加文本标签（使用文本描边提高可读性）
- const label = (nodeData.label || nodeData.id || '').toString();
+ // 使用原始标签，不包含类型标签前缀
+ const label = (nodeData.originalLabel || nodeData.label || nodeData.id || '').toString();
const maxLength = 15;
// 创建文本组，包含描边和填充
diff --git a/web/templates/index.html b/web/templates/index.html
index cd22844e..d1c750c0 100644
--- a/web/templates/index.html
+++ b/web/templates/index.html
@@ -636,16 +636,16 @@