Add files via upload

2026-03-31 16:20:28 +02:00 · 2025-12-25 22:12:41 +08:00
parent 99cf5e78a9
commit 025704cbf7
11 changed files with 1683 additions and 28 deletions
--- a/internal/database/database.go
+++ b/internal/database/database.go
@@ -170,6 +170,25 @@ func (db *DB) initTables() error {
 		UNIQUE(conversation_id, group_id)
 	);`

+	// 创建漏洞表
+	createVulnerabilitiesTable := `
+	CREATE TABLE IF NOT EXISTS vulnerabilities (
+		id TEXT PRIMARY KEY,
+		conversation_id TEXT NOT NULL,
+		title TEXT NOT NULL,
+		description TEXT,
+		severity TEXT NOT NULL,
+		status TEXT NOT NULL DEFAULT 'open',
+		vulnerability_type TEXT,
+		target TEXT,
+		proof TEXT,
+		impact TEXT,
+		recommendation TEXT,
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
+	);`
+
 	// 创建索引
 	createIndexes := `
 	CREATE INDEX IF NOT EXISTS idx_messages_conversation_id ON messages(conversation_id);
@@ -189,6 +208,10 @@ func (db *DB) initTables() error {
 	CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_conversation ON conversation_group_mappings(conversation_id);
 	CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_group ON conversation_group_mappings(group_id);
 	CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned);
+	CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id);
+	CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity);
+	CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status);
+	CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at);
 	`

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

+	if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
+		return fmt.Errorf("创建vulnerabilities表失败: %w", err)
+	}
+
 	// 为已有表添加新字段（如果不存在）- 必须在创建索引之前
 	if err := db.migrateConversationsTable(); err != nil {
 		db.logger.Warn("迁移conversations表失败", zap.Error(err))
--- a/internal/database/vulnerability.go
+++ b/internal/database/vulnerability.go
@@ -0,0 +1,246 @@
+package database
+
+import (
+	"database/sql"
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"go.uber.org/zap"
+)
+
+// Vulnerability 漏洞
+type Vulnerability struct {
+	ID              string    `json:"id"`
+	ConversationID string    `json:"conversation_id"`
+	Title           string    `json:"title"`
+	Description     string    `json:"description"`
+	Severity        string    `json:"severity"` // critical, high, medium, low, info
+	Status          string    `json:"status"`   // open, confirmed, fixed, false_positive
+	Type            string    `json:"type"`
+	Target          string    `json:"target"`
+	Proof           string    `json:"proof"`
+	Impact          string    `json:"impact"`
+	Recommendation  string    `json:"recommendation"`
+	CreatedAt       time.Time `json:"created_at"`
+	UpdatedAt       time.Time `json:"updated_at"`
+}
+
+// CreateVulnerability 创建漏洞
+func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
+	if vuln.ID == "" {
+		vuln.ID = uuid.New().String()
+	}
+	if vuln.Status == "" {
+		vuln.Status = "open"
+	}
+	now := time.Now()
+	if vuln.CreatedAt.IsZero() {
+		vuln.CreatedAt = now
+	}
+	vuln.UpdatedAt = now
+
+	query := `
+		INSERT INTO vulnerabilities (
+			id, conversation_id, title, description, severity, status,
+			vulnerability_type, target, proof, impact, recommendation,
+			created_at, updated_at
+		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+	`
+
+	_, err := db.Exec(
+		query,
+		vuln.ID, vuln.ConversationID, vuln.Title, vuln.Description,
+		vuln.Severity, vuln.Status, vuln.Type, vuln.Target,
+		vuln.Proof, vuln.Impact, vuln.Recommendation,
+		vuln.CreatedAt, vuln.UpdatedAt,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("创建漏洞失败: %w", err)
+	}
+
+	return vuln, nil
+}
+
+// GetVulnerability 获取漏洞
+func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
+	var vuln Vulnerability
+	query := `
+		SELECT id, conversation_id, title, description, severity, status,
+		       vulnerability_type, target, proof, impact, recommendation,
+		       created_at, updated_at
+		FROM vulnerabilities
+		WHERE id = ?
+	`
+
+	err := db.QueryRow(query, id).Scan(
+		&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description,
+		&vuln.Severity, &vuln.Status, &vuln.Type, &vuln.Target,
+		&vuln.Proof, &vuln.Impact, &vuln.Recommendation,
+		&vuln.CreatedAt, &vuln.UpdatedAt,
+	)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			return nil, fmt.Errorf("漏洞不存在")
+		}
+		return nil, fmt.Errorf("获取漏洞失败: %w", err)
+	}
+
+	return &vuln, nil
+}
+
+// ListVulnerabilities 列出漏洞
+func (db *DB) ListVulnerabilities(limit, offset int, conversationID, severity, status string) ([]*Vulnerability, error) {
+	query := `
+		SELECT id, conversation_id, title, description, severity, status,
+		       vulnerability_type, target, proof, impact, recommendation,
+		       created_at, updated_at
+		FROM vulnerabilities
+		WHERE 1=1
+	`
+	args := []interface{}{}
+
+	if conversationID != "" {
+		query += " AND conversation_id = ?"
+		args = append(args, conversationID)
+	}
+	if severity != "" {
+		query += " AND severity = ?"
+		args = append(args, severity)
+	}
+	if status != "" {
+		query += " AND status = ?"
+		args = append(args, status)
+	}
+
+	query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
+	args = append(args, limit, offset)
+
+	rows, err := db.Query(query, args...)
+	if err != nil {
+		return nil, fmt.Errorf("查询漏洞列表失败: %w", err)
+	}
+	defer rows.Close()
+
+	var vulnerabilities []*Vulnerability
+	for rows.Next() {
+		var vuln Vulnerability
+		err := rows.Scan(
+			&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description,
+			&vuln.Severity, &vuln.Status, &vuln.Type, &vuln.Target,
+			&vuln.Proof, &vuln.Impact, &vuln.Recommendation,
+			&vuln.CreatedAt, &vuln.UpdatedAt,
+		)
+		if err != nil {
+			db.logger.Warn("扫描漏洞记录失败", zap.Error(err))
+			continue
+		}
+		vulnerabilities = append(vulnerabilities, &vuln)
+	}
+
+	return vulnerabilities, nil
+}
+
+// UpdateVulnerability 更新漏洞
+func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
+	vuln.UpdatedAt = time.Now()
+
+	query := `
+		UPDATE vulnerabilities
+		SET title = ?, description = ?, severity = ?, status = ?,
+		    vulnerability_type = ?, target = ?, proof = ?, impact = ?,
+		    recommendation = ?, updated_at = ?
+		WHERE id = ?
+	`
+
+	_, err := db.Exec(
+		query,
+		vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
+		vuln.Type, vuln.Target, vuln.Proof, vuln.Impact,
+		vuln.Recommendation, vuln.UpdatedAt, id,
+	)
+	if err != nil {
+		return fmt.Errorf("更新漏洞失败: %w", err)
+	}
+
+	return nil
+}
+
+// DeleteVulnerability 删除漏洞
+func (db *DB) DeleteVulnerability(id string) error {
+	_, err := db.Exec("DELETE FROM vulnerabilities WHERE id = ?", id)
+	if err != nil {
+		return fmt.Errorf("删除漏洞失败: %w", err)
+	}
+	return nil
+}
+
+// GetVulnerabilityStats 获取漏洞统计
+func (db *DB) GetVulnerabilityStats(conversationID string) (map[string]interface{}, error) {
+	stats := make(map[string]interface{})
+
+	// 总漏洞数
+	var totalCount int
+	query := "SELECT COUNT(*) FROM vulnerabilities"
+	args := []interface{}{}
+	if conversationID != "" {
+		query += " WHERE conversation_id = ?"
+		args = append(args, conversationID)
+	}
+	err := db.QueryRow(query, args...).Scan(&totalCount)
+	if err != nil {
+		return nil, fmt.Errorf("获取总漏洞数失败: %w", err)
+	}
+	stats["total"] = totalCount
+
+	// 按严重程度统计
+	severityQuery := "SELECT severity, COUNT(*) FROM vulnerabilities"
+	if conversationID != "" {
+		severityQuery += " WHERE conversation_id = ?"
+	}
+	severityQuery += " GROUP BY severity"
+
+	rows, err := db.Query(severityQuery, args...)
+	if err != nil {
+		return nil, fmt.Errorf("获取严重程度统计失败: %w", err)
+	}
+	defer rows.Close()
+
+	severityStats := make(map[string]int)
+	for rows.Next() {
+		var severity string
+		var count int
+		if err := rows.Scan(&severity, &count); err != nil {
+			continue
+		}
+		severityStats[severity] = count
+	}
+	stats["by_severity"] = severityStats
+
+	// 按状态统计
+	statusQuery := "SELECT status, COUNT(*) FROM vulnerabilities"
+	if conversationID != "" {
+		statusQuery += " WHERE conversation_id = ?"
+	}
+	statusQuery += " GROUP BY status"
+
+	rows, err = db.Query(statusQuery, args...)
+	if err != nil {
+		return nil, fmt.Errorf("获取状态统计失败: %w", err)
+	}
+	defer rows.Close()
+
+	statusStats := make(map[string]int)
+	for rows.Next() {
+		var status string
+		var count int
+		if err := rows.Scan(&status, &count); err != nil {
+			continue
+		}
+		statusStats[status] = count
+	}
+	stats["by_status"] = statusStats
+
+	return stats, nil
+}
+