Add files via upload

internal/app/app.go

@@ -66,6 +66,7 @@ type App struct {
 	slackCancel        context.CancelFunc        // Slack Socket Mode 取消函数
 	discordCancel      context.CancelFunc        // Discord Gateway 取消函数
 	qqCancel           context.CancelFunc        // QQ WebSocket 取消函数
+	alertCancel        context.CancelFunc        // 漏洞提醒持久化投递 worker
 	c2Manager          *c2.Manager               // C2 管理器（未启用 C2 时为 nil）
 	c2Watchdog         *c2.SessionWatchdog       // C2 会话看门狗
 	c2WatchdogCancel   context.CancelFunc        // 看门狗取消函数
@@ -416,6 +417,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
 	auditHandler := handler.NewAuditHandler(db, auditSvc, log.Logger)
 	robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger)
 	robotHandler.SetAudit(auditSvc)
+	db.SetVulnerabilityCreatedHook(robotHandler.NotifyNewVulnerability)
 	openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, conversationHandler, agentHandler)
 
 	// 创建 App 实例（部分字段稍后填充）
@@ -444,6 +446,9 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
 	}
 	// 飞书/钉钉长连接（无需公网），启用时在后台启动；后续前端应用配置时会通过 RestartRobotConnections 重启
 	app.startRobotConnections()
+	alertCtx, alertCancel := context.WithCancel(context.Background())
+	app.alertCancel = alertCancel
+	go robotHandler.RunVulnerabilityAlertWorker(alertCtx)
 
 	// 设置漏洞工具注册器（内置工具，必须设置）
 	vulnerabilityRegistrar := func() error {
@@ -705,6 +710,10 @@ func (a *App) Shutdown() {
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
 	_ = einoobserve.ShutdownOtel(shutdownCtx)
 	shutdownCancel()
+	if a.alertCancel != nil {
+		a.alertCancel()
+		a.alertCancel = nil
+	}
 
 	// 停止钉钉/飞书长连接
 	a.robotMu.Lock()
@@ -1194,6 +1203,8 @@ func setupRoutes(
 		protected.DELETE("/vulnerabilities/batch", vulnerabilityHandler.BatchDeleteVulnerabilities)
 		protected.GET("/vulnerabilities/filter-options", vulnerabilityHandler.GetVulnerabilityFilterOptions)
 		protected.GET("/vulnerabilities/stats", vulnerabilityHandler.GetVulnerabilityStats)
+		protected.GET("/vulnerability-alerts/subscription", vulnerabilityHandler.GetMyAlertSubscription)
+		protected.PUT("/vulnerability-alerts/subscription", vulnerabilityHandler.UpdateMyAlertSubscription)
 		protected.GET("/vulnerabilities/:id", vulnerabilityHandler.GetVulnerability)
 		protected.POST("/vulnerabilities", vulnerabilityHandler.CreateVulnerability)
 		protected.PUT("/vulnerabilities/:id", vulnerabilityHandler.UpdateVulnerability)

internal/app/vulnerability_tools.go

@@ -318,6 +318,7 @@ func registerRecordVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, log
 			_ = db.SetResourceOwner("vulnerability", created.ID, principal.UserID)
 			_ = db.AssignResourceToUser(principal.UserID, "vulnerability", created.ID)
 		}
+		db.NotifyVulnerabilityCreated(created)
 
 		if logger != nil {
 			logger.Info("漏洞记录成功",

internal/database/database.go

@@ -58,6 +58,7 @@ type DB struct {
 	checkpointDone           chan struct{}
 	closeOnce                sync.Once
 	closeErr                 error
+	vulnerabilityCreatedHook func(*Vulnerability)
 }
 
 // startPassiveCheckpointLoop 启动后台 PASSIVE checkpoint 循环。
@@ -404,6 +405,33 @@ func (db *DB) initTables() error {
 		FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL
 	);`
 
+	createVulnerabilityAlertSubscriptionsTable := `
+	CREATE TABLE IF NOT EXISTS vulnerability_alert_subscriptions (
+		user_id TEXT PRIMARY KEY,
+		enabled INTEGER NOT NULL DEFAULT 0,
+		min_severity TEXT NOT NULL DEFAULT 'high',
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE
+	);`
+	createVulnerabilityAlertDeliveriesTable := `
+	CREATE TABLE IF NOT EXISTS vulnerability_alert_deliveries (
+		id INTEGER PRIMARY KEY AUTOINCREMENT,
+		vulnerability_id TEXT NOT NULL,
+		user_id TEXT NOT NULL,
+		platform TEXT NOT NULL,
+		external_user_id TEXT NOT NULL,
+		status TEXT NOT NULL DEFAULT 'pending',
+		attempts INTEGER NOT NULL DEFAULT 0,
+		next_attempt_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		last_error TEXT NOT NULL DEFAULT '',
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		UNIQUE(vulnerability_id, platform, external_user_id),
+		FOREIGN KEY (vulnerability_id) REFERENCES vulnerabilities(id) ON DELETE CASCADE,
+		FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE
+	);`
+
 	// 创建批量任务队列表
 	createBatchTaskQueuesTable := `
 	CREATE TABLE IF NOT EXISTS batch_task_queues (
@@ -794,6 +822,12 @@ func (db *DB) initTables() error {
 	if err := db.initRBACTables(); err != nil {
 		return fmt.Errorf("创建RBAC表失败: %w", err)
 	}
+	if _, err := db.Exec(createVulnerabilityAlertSubscriptionsTable); err != nil {
+		return fmt.Errorf("创建漏洞提醒订阅表失败: %w", err)
+	}
+	if _, err := db.Exec(createVulnerabilityAlertDeliveriesTable); err != nil {
+		return fmt.Errorf("创建漏洞提醒投递表失败: %w", err)
+	}
 
 	for tableName, ddl := range map[string]string{
 		"workflow_definitions": createWorkflowDefinitionsTable,

internal/database/vulnerability.go

@@ -191,7 +191,6 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
 	if err != nil {
 		return nil, fmt.Errorf("创建漏洞失败: %w", err)
 	}
-
 	return vuln, nil
 }
 

internal/database/vulnerability_alert.go

@@ -0,0 +1,215 @@
+package database
+
+import (
+	"database/sql"
+	"fmt"
+	"strings"
+	"time"
+)
+
+// VulnerabilityAlertSubscription is the single source of truth shared by Web
+// settings and robot commands. Alerts are opt-in and user scoped.
+type VulnerabilityAlertSubscription struct {
+	UserID      string    `json:"user_id"`
+	Enabled     bool      `json:"enabled"`
+	MinSeverity string    `json:"min_severity"`
+	CreatedAt   time.Time `json:"created_at"`
+	UpdatedAt   time.Time `json:"updated_at"`
+}
+
+type VulnerabilityAlertRecipient struct {
+	UserID         string
+	Platform       string
+	ExternalUserID string
+}
+
+type VulnerabilityAlertDelivery struct {
+	ID             int64
+	Vulnerability  *Vulnerability
+	UserID         string
+	Platform       string
+	ExternalUserID string
+	Attempts       int
+}
+
+var vulnerabilitySeverityRank = map[string]int{
+	"info": 0, "low": 1, "medium": 2, "high": 3, "critical": 4,
+}
+
+func NormalizeVulnerabilityAlertSeverity(value string) (string, error) {
+	value = strings.ToLower(strings.TrimSpace(value))
+	if value == "" {
+		value = "high"
+	}
+	if _, ok := vulnerabilitySeverityRank[value]; !ok {
+		return "", fmt.Errorf("invalid minimum severity %q", value)
+	}
+	return value, nil
+}
+
+func (db *DB) GetVulnerabilityAlertSubscription(userID string) (*VulnerabilityAlertSubscription, error) {
+	userID = strings.TrimSpace(userID)
+	var sub VulnerabilityAlertSubscription
+	var enabled int
+	var createdAt, updatedAt string
+	err := db.QueryRow(`SELECT user_id, enabled, min_severity, created_at, updated_at
+		FROM vulnerability_alert_subscriptions WHERE user_id = ?`, userID).
+		Scan(&sub.UserID, &enabled, &sub.MinSeverity, &createdAt, &updatedAt)
+	if err == sql.ErrNoRows {
+		now := time.Now()
+		return &VulnerabilityAlertSubscription{UserID: userID, MinSeverity: "high", CreatedAt: now, UpdatedAt: now}, nil
+	}
+	if err != nil {
+		return nil, err
+	}
+	sub.Enabled = enabled != 0
+	sub.CreatedAt = parseDBTime(createdAt)
+	sub.UpdatedAt = parseDBTime(updatedAt)
+	return &sub, nil
+}
+
+func (db *DB) UpsertVulnerabilityAlertSubscription(userID string, enabled bool, minSeverity string) (*VulnerabilityAlertSubscription, error) {
+	userID = strings.TrimSpace(userID)
+	if userID == "" {
+		return nil, fmt.Errorf("user id is required")
+	}
+	severity, err := NormalizeVulnerabilityAlertSeverity(minSeverity)
+	if err != nil {
+		return nil, err
+	}
+	now := time.Now()
+	_, err = db.Exec(`INSERT INTO vulnerability_alert_subscriptions
+		(user_id, enabled, min_severity, created_at, updated_at) VALUES (?, ?, ?, ?, ?)
+		ON CONFLICT(user_id) DO UPDATE SET enabled = excluded.enabled,
+		min_severity = excluded.min_severity, updated_at = excluded.updated_at`,
+		userID, boolToInt(enabled), severity, now, now)
+	if err != nil {
+		return nil, err
+	}
+	return db.GetVulnerabilityAlertSubscription(userID)
+}
+
+// ListVulnerabilityAlertRecipients applies the same RBAC ownership/assignment
+// boundaries as the vulnerability list, then expands only enabled robot bindings.
+func (db *DB) ListVulnerabilityAlertRecipients(vuln *Vulnerability) ([]VulnerabilityAlertRecipient, error) {
+	if vuln == nil {
+		return nil, nil
+	}
+	rank, ok := vulnerabilitySeverityRank[strings.ToLower(strings.TrimSpace(vuln.Severity))]
+	if !ok {
+		return nil, nil
+	}
+	rows, err := db.Query(`
+		SELECT DISTINCT s.user_id, b.platform, b.external_user_id, s.min_severity
+		FROM vulnerability_alert_subscriptions s
+		JOIN rbac_users u ON u.id = s.user_id AND u.enabled = 1
+		JOIN robot_user_bindings b ON b.rbac_user_id = s.user_id AND b.enabled = 1
+		WHERE s.enabled = 1 AND (
+			EXISTS (SELECT 1 FROM vulnerabilities v WHERE v.id = ? AND v.owner_user_id = s.user_id)
+			OR EXISTS (SELECT 1 FROM rbac_resource_assignments ra WHERE ra.user_id = s.user_id AND ra.resource_type = 'vulnerability' AND ra.resource_id = ?)
+			OR (? <> '' AND (EXISTS (SELECT 1 FROM projects p WHERE p.id = ? AND p.owner_user_id = s.user_id)
+				OR EXISTS (SELECT 1 FROM rbac_resource_assignments pra WHERE pra.user_id = s.user_id AND pra.resource_type = 'project' AND pra.resource_id = ?)))
+			OR (? <> '' AND (EXISTS (SELECT 1 FROM conversations c WHERE c.id = ? AND c.owner_user_id = s.user_id)
+				OR EXISTS (SELECT 1 FROM rbac_resource_assignments cra WHERE cra.user_id = s.user_id AND cra.resource_type = 'conversation' AND cra.resource_id = ?)))
+		)`, vuln.ID, vuln.ID, vuln.ProjectID, vuln.ProjectID, vuln.ProjectID,
+		vuln.ConversationID, vuln.ConversationID, vuln.ConversationID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	out := make([]VulnerabilityAlertRecipient, 0)
+	for rows.Next() {
+		var recipient VulnerabilityAlertRecipient
+		var minimum string
+		if err := rows.Scan(&recipient.UserID, &recipient.Platform, &recipient.ExternalUserID, &minimum); err != nil {
+			return nil, err
+		}
+		if rank >= vulnerabilitySeverityRank[minimum] {
+			out = append(out, recipient)
+		}
+	}
+	return out, rows.Err()
+}
+
+func (db *DB) SetVulnerabilityCreatedHook(hook func(*Vulnerability)) {
+	db.vulnerabilityCreatedHook = hook
+}
+
+// NotifyVulnerabilityCreated must be called after resource ownership has been
+// committed. Delivery runs asynchronously and never delays the write path.
+func (db *DB) NotifyVulnerabilityCreated(vulnerability *Vulnerability) {
+	if db == nil || vulnerability == nil || db.vulnerabilityCreatedHook == nil {
+		return
+	}
+	created := *vulnerability
+	go db.vulnerabilityCreatedHook(&created)
+}
+
+func (db *DB) EnqueueVulnerabilityAlertDeliveries(vulnerabilityID string, recipients []VulnerabilityAlertRecipient) error {
+	now := time.Now()
+	tx, err := db.Begin()
+	if err != nil {
+		return err
+	}
+	defer func() { _ = tx.Rollback() }()
+	for _, r := range recipients {
+		if _, err := tx.Exec(`INSERT INTO vulnerability_alert_deliveries
+			(vulnerability_id, user_id, platform, external_user_id, status, attempts, next_attempt_at, created_at, updated_at)
+			VALUES (?, ?, ?, ?, 'pending', 0, ?, ?, ?)
+			ON CONFLICT(vulnerability_id, platform, external_user_id) DO NOTHING`,
+			vulnerabilityID, r.UserID, r.Platform, r.ExternalUserID, now, now, now); err != nil {
+			return err
+		}
+	}
+	return tx.Commit()
+}
+
+func (db *DB) ListDueVulnerabilityAlertDeliveries(limit int) ([]VulnerabilityAlertDelivery, error) {
+	if limit <= 0 || limit > 100 {
+		limit = 50
+	}
+	rows, err := db.Query(`SELECT d.id, d.user_id, d.platform, d.external_user_id, d.attempts,
+		v.id, COALESCE(v.conversation_id,''), COALESCE(v.project_id,''), v.title, COALESCE(v.description,''),
+		v.severity, v.status, COALESCE(v.vulnerability_type,''), COALESCE(v.target,''),
+		COALESCE(v.impact,''), COALESCE(v.recommendation,''), v.created_at, v.updated_at
+		FROM vulnerability_alert_deliveries d JOIN vulnerabilities v ON v.id = d.vulnerability_id
+		WHERE d.status IN ('pending','retry') AND d.next_attempt_at <= ?
+		ORDER BY d.next_attempt_at, d.id LIMIT ?`, time.Now(), limit)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var out []VulnerabilityAlertDelivery
+	for rows.Next() {
+		var d VulnerabilityAlertDelivery
+		v := &Vulnerability{}
+		if err := rows.Scan(&d.ID, &d.UserID, &d.Platform, &d.ExternalUserID, &d.Attempts,
+			&v.ID, &v.ConversationID, &v.ProjectID, &v.Title, &v.Description, &v.Severity, &v.Status,
+			&v.Type, &v.Target, &v.Impact, &v.Recommendation, &v.CreatedAt, &v.UpdatedAt); err != nil {
+			return nil, err
+		}
+		d.Vulnerability = v
+		out = append(out, d)
+	}
+	return out, rows.Err()
+}
+
+func (db *DB) MarkVulnerabilityAlertDeliverySent(id int64) error {
+	_, err := db.Exec(`UPDATE vulnerability_alert_deliveries SET status='sent', attempts=attempts+1, last_error='', updated_at=? WHERE id=?`, time.Now(), id)
+	return err
+}
+
+func (db *DB) MarkVulnerabilityAlertDeliveryFailed(id int64, attempts int, sendErr error) error {
+	status := "retry"
+	if attempts >= 5 {
+		status = "failed"
+	}
+	delay := time.Minute * time.Duration(1<<min(attempts, 6))
+	message := ""
+	if sendErr != nil {
+		message = sendErr.Error()
+	}
+	_, err := db.Exec(`UPDATE vulnerability_alert_deliveries SET status=?, attempts=?, next_attempt_at=?, last_error=?, updated_at=? WHERE id=?`,
+		status, attempts, time.Now().Add(delay), message, time.Now(), id)
+	return err
+}

internal/database/vulnerability_alert_test.go

@@ -0,0 +1,88 @@
+package database_test
+
+import (
+	"testing"
+	"time"
+
+	"cyberstrike-ai/internal/database"
+	"cyberstrike-ai/internal/security"
+
+	"go.uber.org/zap"
+)
+
+func TestVulnerabilityAlertSubscriptionIsOptInAndRBACScoped(t *testing.T) {
+	db, err := database.NewDB(t.TempDir()+"/alerts.db", zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = db.Close() })
+	if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
+		t.Fatal(err)
+	}
+	owner, _ := db.CreateRBACUser("alert-owner", "Owner", "hash", true, []string{database.RBACSystemRoleOperator})
+	other, _ := db.CreateRBACUser("alert-other", "Other", "hash", true, []string{database.RBACSystemRoleOperator})
+
+	for i, pair := range []struct {
+		user     *database.RBACUser
+		external string
+	}{{owner, "owner"}, {other, "other"}} {
+		code := "code-" + string(rune('a'+i))
+		if err := db.CreateRobotBindingCode(pair.user.ID, code, time.Now().Add(time.Minute)); err != nil {
+			t.Fatal(err)
+		}
+		if _, err := db.ConsumeRobotBindingCode("wecom", pair.external, code); err != nil {
+			t.Fatal(err)
+		}
+	}
+	defaultSub, err := db.GetVulnerabilityAlertSubscription(owner.ID)
+	if err != nil || defaultSub.Enabled || defaultSub.MinSeverity != "high" {
+		t.Fatalf("unsafe default: %#v %v", defaultSub, err)
+	}
+	if _, err := db.UpsertVulnerabilityAlertSubscription(owner.ID, true, "high"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.UpsertVulnerabilityAlertSubscription(other.ID, true, "low"); err != nil {
+		t.Fatal(err)
+	}
+
+	vuln, err := db.CreateVulnerability(&database.Vulnerability{Title: "owned", Severity: "high"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := db.SetResourceOwner("vulnerability", vuln.ID, owner.ID); err != nil {
+		t.Fatal(err)
+	}
+	recipients, err := db.ListVulnerabilityAlertRecipients(vuln)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(recipients) != 1 || recipients[0].UserID != owner.ID || recipients[0].ExternalUserID != "owner" {
+		t.Fatalf("alert escaped RBAC boundary: %#v", recipients)
+	}
+	if err := db.EnqueueVulnerabilityAlertDeliveries(vuln.ID, recipients); err != nil {
+		t.Fatal(err)
+	}
+	if err := db.EnqueueVulnerabilityAlertDeliveries(vuln.ID, recipients); err != nil {
+		t.Fatal(err)
+	}
+	deliveries, err := db.ListDueVulnerabilityAlertDeliveries(10)
+	if err != nil || len(deliveries) != 1 {
+		t.Fatalf("outbox is not durable/deduplicated: %#v %v", deliveries, err)
+	}
+	if err := db.MarkVulnerabilityAlertDeliverySent(deliveries[0].ID); err != nil {
+		t.Fatal(err)
+	}
+	deliveries, _ = db.ListDueVulnerabilityAlertDeliveries(10)
+	if len(deliveries) != 0 {
+		t.Fatalf("sent delivery remained due: %#v", deliveries)
+	}
+
+	vuln.Severity = "medium"
+	recipients, err = db.ListVulnerabilityAlertRecipients(vuln)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(recipients) != 0 {
+		t.Fatalf("minimum severity was ignored: %#v", recipients)
+	}
+}