From 386ec7b835b2c94ed187acefbbde31b7387cd238 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:06:02 +0800 Subject: [PATCH] Add files via upload --- internal/app/app.go | 11 + internal/app/vulnerability_tools.go | 1 + internal/database/database.go | 34 +++ internal/database/vulnerability.go | 1 - internal/database/vulnerability_alert.go | 215 ++++++++++++++++++ internal/database/vulnerability_alert_test.go | 88 +++++++ 6 files changed, 349 insertions(+), 1 deletion(-) create mode 100644 internal/database/vulnerability_alert.go create mode 100644 internal/database/vulnerability_alert_test.go diff --git a/internal/app/app.go b/internal/app/app.go index 9ce10592..aa56adb3 100644 --- a/internal/app/app.go +++ b/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) diff --git a/internal/app/vulnerability_tools.go b/internal/app/vulnerability_tools.go index 70575b28..f3dd7bfe 100644 --- a/internal/app/vulnerability_tools.go +++ b/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("漏洞记录成功", diff --git a/internal/database/database.go b/internal/database/database.go index 20754537..57fa34fc 100644 --- a/internal/database/database.go +++ b/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, diff --git a/internal/database/vulnerability.go b/internal/database/vulnerability.go index 51d2a3db..664b2ba8 100644 --- a/internal/database/vulnerability.go +++ b/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 } diff --git a/internal/database/vulnerability_alert.go b/internal/database/vulnerability_alert.go new file mode 100644 index 00000000..23a5b760 --- /dev/null +++ b/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<