Add files via upload

2026-06-10 16:23:54 +02:00 · 2025-11-17 02:43:36 +08:00
parent 71b6c8aa2d
commit 9f862ce721
10 changed files with 2646 additions and 12 deletions
@@ -0,0 +1,168 @@
+package database
+
+import (
+	"database/sql"
+	"encoding/json"
+	"fmt"
+
+	"go.uber.org/zap"
+)
+
+// AttackChainNode 攻击链节点
+type AttackChainNode struct {
+	ID              string                 `json:"id"`
+	Type            string                 `json:"type"` // tool, vulnerability, target, exploit
+	Label           string                 `json:"label"`
+	ToolExecutionID string                 `json:"tool_execution_id,omitempty"`
+	Metadata        map[string]interface{} `json:"metadata"`
+	RiskScore       int                    `json:"risk_score"`
+}
+
+// AttackChainEdge 攻击链边
+type AttackChainEdge struct {
+	ID     string `json:"id"`
+	Source string `json:"source"`
+	Target string `json:"target"`
+	Type   string `json:"type"` // leads_to, exploits, enables, depends_on
+	Weight int    `json:"weight"`
+}
+
+// SaveAttackChainNode 保存攻击链节点
+func (db *DB) SaveAttackChainNode(conversationID, nodeID, nodeType, nodeName, toolExecutionID, metadata string, riskScore int) error {
+	var toolExecID sql.NullString
+	if toolExecutionID != "" {
+		toolExecID = sql.NullString{String: toolExecutionID, Valid: true}
+	}
+
+	var metadataJSON sql.NullString
+	if metadata != "" {
+		metadataJSON = sql.NullString{String: metadata, Valid: true}
+	}
+
+	query := `
+		INSERT OR REPLACE INTO attack_chain_nodes 
+		(id, conversation_id, node_type, node_name, tool_execution_id, metadata, risk_score, created_at)
+		VALUES (?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
+	`
+
+	_, err := db.Exec(query, nodeID, conversationID, nodeType, nodeName, toolExecID, metadataJSON, riskScore)
+	if err != nil {
+		db.logger.Error("保存攻击链节点失败", zap.Error(err), zap.String("nodeId", nodeID))
+		return err
+	}
+
+	return nil
+}
+
+// SaveAttackChainEdge 保存攻击链边
+func (db *DB) SaveAttackChainEdge(conversationID, edgeID, sourceNodeID, targetNodeID, edgeType string, weight int) error {
+	query := `
+		INSERT OR REPLACE INTO attack_chain_edges 
+		(id, conversation_id, source_node_id, target_node_id, edge_type, weight, created_at)
+		VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
+	`
+
+	_, err := db.Exec(query, edgeID, conversationID, sourceNodeID, targetNodeID, edgeType, weight)
+	if err != nil {
+		db.logger.Error("保存攻击链边失败", zap.Error(err), zap.String("edgeId", edgeID))
+		return err
+	}
+
+	return nil
+}
+
+// LoadAttackChainNodes 加载攻击链节点
+func (db *DB) LoadAttackChainNodes(conversationID string) ([]AttackChainNode, error) {
+	query := `
+		SELECT id, node_type, node_name, tool_execution_id, metadata, risk_score
+		FROM attack_chain_nodes
+		WHERE conversation_id = ?
+		ORDER BY created_at ASC
+	`
+
+	rows, err := db.Query(query, conversationID)
+	if err != nil {
+		return nil, fmt.Errorf("查询攻击链节点失败: %w", err)
+	}
+	defer rows.Close()
+
+	var nodes []AttackChainNode
+	for rows.Next() {
+		var node AttackChainNode
+		var toolExecID sql.NullString
+		var metadataJSON sql.NullString
+
+		err := rows.Scan(&node.ID, &node.Type, &node.Label, &toolExecID, &metadataJSON, &node.RiskScore)
+		if err != nil {
+			db.logger.Warn("扫描攻击链节点失败", zap.Error(err))
+			continue
+		}
+
+		if toolExecID.Valid {
+			node.ToolExecutionID = toolExecID.String
+		}
+
+		if metadataJSON.Valid && metadataJSON.String != "" {
+			if err := json.Unmarshal([]byte(metadataJSON.String), &node.Metadata); err != nil {
+				db.logger.Warn("解析节点元数据失败", zap.Error(err))
+				node.Metadata = make(map[string]interface{})
+			}
+		} else {
+			node.Metadata = make(map[string]interface{})
+		}
+
+		nodes = append(nodes, node)
+	}
+
+	return nodes, nil
+}
+
+// LoadAttackChainEdges 加载攻击链边
+func (db *DB) LoadAttackChainEdges(conversationID string) ([]AttackChainEdge, error) {
+	query := `
+		SELECT id, source_node_id, target_node_id, edge_type, weight
+		FROM attack_chain_edges
+		WHERE conversation_id = ?
+		ORDER BY created_at ASC
+	`
+
+	rows, err := db.Query(query, conversationID)
+	if err != nil {
+		return nil, fmt.Errorf("查询攻击链边失败: %w", err)
+	}
+	defer rows.Close()
+
+	var edges []AttackChainEdge
+	for rows.Next() {
+		var edge AttackChainEdge
+
+		err := rows.Scan(&edge.ID, &edge.Source, &edge.Target, &edge.Type, &edge.Weight)
+		if err != nil {
+			db.logger.Warn("扫描攻击链边失败", zap.Error(err))
+			continue
+		}
+
+		edges = append(edges, edge)
+	}
+
+	return edges, nil
+}
+
+// DeleteAttackChain 删除对话的攻击链数据
+func (db *DB) DeleteAttackChain(conversationID string) error {
+	// 先删除边（因为有外键约束）
+	_, err := db.Exec("DELETE FROM attack_chain_edges WHERE conversation_id = ?", conversationID)
+	if err != nil {
+		db.logger.Warn("删除攻击链边失败", zap.Error(err))
+	}
+
+	// 再删除节点
+	_, err = db.Exec("DELETE FROM attack_chain_nodes WHERE conversation_id = ?", conversationID)
+	if err != nil {
+		db.logger.Error("删除攻击链节点失败", zap.Error(err), zap.String("conversationId", conversationID))
+		return err
+	}
+
+	return nil
+}
+
@@ -101,6 +101,36 @@ func (db *DB) initTables() error {
 		updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
 	);`

+	// 创建攻击链节点表
+	createAttackChainNodesTable := `
+	CREATE TABLE IF NOT EXISTS attack_chain_nodes (
+		id TEXT PRIMARY KEY,
+		conversation_id TEXT NOT NULL,
+		node_type TEXT NOT NULL,
+		node_name TEXT NOT NULL,
+		tool_execution_id TEXT,
+		metadata TEXT,
+		risk_score INTEGER DEFAULT 0,
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE,
+		FOREIGN KEY (tool_execution_id) REFERENCES tool_executions(id) ON DELETE SET NULL
+	);`
+
+	// 创建攻击链边表
+	createAttackChainEdgesTable := `
+	CREATE TABLE IF NOT EXISTS attack_chain_edges (
+		id TEXT PRIMARY KEY,
+		conversation_id TEXT NOT NULL,
+		source_node_id TEXT NOT NULL,
+		target_node_id TEXT NOT NULL,
+		edge_type TEXT NOT NULL,
+		weight INTEGER DEFAULT 1,
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE,
+		FOREIGN KEY (source_node_id) REFERENCES attack_chain_nodes(id) ON DELETE CASCADE,
+		FOREIGN KEY (target_node_id) REFERENCES attack_chain_nodes(id) ON DELETE CASCADE
+	);`
+
 	// 创建索引
 	createIndexes := `
 	CREATE INDEX IF NOT EXISTS idx_messages_conversation_id ON messages(conversation_id);
@@ -110,6 +140,10 @@ func (db *DB) initTables() error {
 	CREATE INDEX IF NOT EXISTS idx_tool_executions_tool_name ON tool_executions(tool_name);
 	CREATE INDEX IF NOT EXISTS idx_tool_executions_start_time ON tool_executions(start_time);
 	CREATE INDEX IF NOT EXISTS idx_tool_executions_status ON tool_executions(status);
+	CREATE INDEX IF NOT EXISTS idx_chain_nodes_conversation ON attack_chain_nodes(conversation_id);
+	CREATE INDEX IF NOT EXISTS idx_chain_edges_conversation ON attack_chain_edges(conversation_id);
+	CREATE INDEX IF NOT EXISTS idx_chain_edges_source ON attack_chain_edges(source_node_id);
+	CREATE INDEX IF NOT EXISTS idx_chain_edges_target ON attack_chain_edges(target_node_id);
 	`

 	if _, err := db.Exec(createConversationsTable); err != nil {
@@ -132,6 +166,14 @@ func (db *DB) initTables() error {
 		return fmt.Errorf("创建tool_stats表失败: %w", err)
 	}

+	if _, err := db.Exec(createAttackChainNodesTable); err != nil {
+		return fmt.Errorf("创建attack_chain_nodes表失败: %w", err)
+	}
+
+	if _, err := db.Exec(createAttackChainEdgesTable); err != nil {
+		return fmt.Errorf("创建attack_chain_edges表失败: %w", err)
+	}
+
 	if _, err := db.Exec(createIndexes); err != nil {
 		return fmt.Errorf("创建索引失败: %w", err)
 	}