diff --git a/internal/handler/robot.go b/internal/handler/robot.go index a4cc1fce..45f0ec36 100644 --- a/internal/handler/robot.go +++ b/internal/handler/robot.go @@ -65,6 +65,7 @@ const ( robotCmdDoctor = "诊断" robotCmdConfirm = "确认" robotCmdCancel = "取消" + robotCmdVulnAlerts = "漏洞提醒" robotBindingCodeTTL = 5 * time.Minute ) @@ -88,6 +89,7 @@ type RobotHandler struct { runningCancels map[string]context.CancelFunc // key: "platform_userID", 用于停止命令中断任务 wecomReplay map[string]time.Time pendingConfirmations map[string]robotPendingConfirmation + alertWake chan struct{} audit *audit.Service } @@ -104,6 +106,7 @@ func NewRobotHandler(cfg *config.Config, db *database.DB, agentHandler *AgentHan runningCancels: make(map[string]context.CancelFunc), wecomReplay: make(map[string]time.Time), pendingConfirmations: make(map[string]robotPendingConfirmation), + alertWake: make(chan struct{}, 1), } } @@ -490,6 +493,9 @@ func (h *RobotHandler) cmdHelp(platform, userID string) string { if can("agent:execute") { b.WriteString("\n【模式 Mode】\n· 模式 / modes — 列出对话模式与当前选择\n· 模式 <名称> / mode — 切换对话模式\n· 停止 / stop — 中断当前任务\n") } + if can("vulnerability:read") { + b.WriteString("\n【漏洞提醒 Vulnerability alerts】\n· 漏洞提醒 — 查看订阅状态\n· 漏洞提醒 开启 / vuln alerts on — 开启提醒\n· 漏洞提醒 仅严重|高危以上|中危以上 / vuln alerts critical|high|medium — 设置最低级别\n· 漏洞提醒 关闭 / vuln alerts off — 关闭提醒\n") + } b.WriteString("\n【诊断 Diagnostics】\n") b.WriteString("· 权限 / permissions — 查看当前业务权限\n") if can("config:read") { @@ -1110,6 +1116,9 @@ func robotCommandPermission(text string) (string, bool) { return "config:read", true case text == robotCmdProjects || text == robotCmdProjectsList || text == "projects": return "project:read", true + case text == robotCmdVulnAlerts || strings.HasPrefix(text, robotCmdVulnAlerts+" "), + text == "vuln alerts" || strings.HasPrefix(text, "vuln alerts "): + return "vulnerability:read", true case text == robotCmdUnbindProject || text == "unbind project", strings.HasPrefix(text, robotCmdNewProject+" "), strings.HasPrefix(text, "new project "), strings.HasPrefix(text, robotCmdBindProject+" "), strings.HasPrefix(text, "bind project "): @@ -1253,6 +1262,12 @@ func (h *RobotHandler) handleRobotCommand(platform, userID, text string) (string } } switch { + case text == robotCmdVulnAlerts || text == "vuln alerts": + return h.cmdVulnerabilityAlerts(platform, userID, ""), true + case strings.HasPrefix(text, robotCmdVulnAlerts+" "): + return h.cmdVulnerabilityAlerts(platform, userID, strings.TrimSpace(text[len(robotCmdVulnAlerts)+1:])), true + case strings.HasPrefix(text, "vuln alerts "): + return h.cmdVulnerabilityAlerts(platform, userID, strings.TrimSpace(text[len("vuln alerts "):])), true case text == robotCmdHelp || text == "help" || text == "?" || text == "?": return h.cmdHelp(platform, userID), true case text == robotCmdIdentity || text == "whoami": diff --git a/internal/handler/vulnerability.go b/internal/handler/vulnerability.go index d52babf3..ab1eef22 100644 --- a/internal/handler/vulnerability.go +++ b/internal/handler/vulnerability.go @@ -101,6 +101,7 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) { _ = h.db.SetResourceOwner("vulnerability", created.ID, session.UserID) _ = h.db.AssignResourceToUser(session.UserID, "vulnerability", created.ID) } + h.db.NotifyVulnerabilityCreated(created) if h.audit != nil { h.audit.RecordOK(c, "vulnerability", "create", "创建漏洞记录", "vulnerability", created.ID, map[string]interface{}{ diff --git a/internal/handler/vulnerability_alert.go b/internal/handler/vulnerability_alert.go new file mode 100644 index 00000000..941874c9 --- /dev/null +++ b/internal/handler/vulnerability_alert.go @@ -0,0 +1,227 @@ +package handler + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/robot" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type updateVulnerabilityAlertRequest struct { + Enabled bool `json:"enabled"` + MinSeverity string `json:"min_severity"` +} + +func (h *VulnerabilityHandler) GetMyAlertSubscription(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + sub, err := h.db.GetVulnerabilityAlertSubscription(session.UserID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + bindings, _ := h.db.ListRobotUserBindings(session.UserID) + deliveryReady := false + for _, binding := range bindings { + if binding.Enabled && robot.SupportsProactive(binding.Platform) { + deliveryReady = true + break + } + } + c.JSON(http.StatusOK, gin.H{"subscription": sub, "bindings": bindings, "delivery_ready": deliveryReady}) +} + +func (h *VulnerabilityHandler) UpdateMyAlertSubscription(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + var req updateVulnerabilityAlertRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + sub, err := h.db.UpsertVulnerabilityAlertSubscription(session.UserID, req.Enabled, req.MinSeverity) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "vulnerability", "alert_subscription_update", "更新漏洞提醒订阅", "user", session.UserID, map[string]interface{}{ + "enabled": sub.Enabled, "min_severity": sub.MinSeverity, + }) + } + c.JSON(http.StatusOK, sub) +} + +func vulnerabilityAlertSeverityLabel(value string) string { + switch value { + case "critical": + return "严重" + case "high": + return "高危" + case "medium": + return "中危" + case "low": + return "低危" + default: + return "信息" + } +} + +func formatVulnerabilityRobotAlert(v *database.Vulnerability) string { + return fmt.Sprintf("🚨 新漏洞提醒\n\n标题:%s\n严重程度:%s(%s)\n目标:%s\n影响:%s\n修复建议:%s\n漏洞 ID:%s\n\n请核实目标版本与实际暴露情况。", + strings.TrimSpace(v.Title), vulnerabilityAlertSeverityLabel(v.Severity), v.Severity, + fallbackAlertText(v.Target), fallbackAlertText(v.Impact), fallbackAlertText(v.Recommendation), v.ID) +} + +func fallbackAlertText(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "未填写" + } + if len([]rune(value)) > 300 { + return string([]rune(value)[:300]) + "…" + } + return value +} + +// NotifyNewVulnerability pushes an alert to every eligible bound robot identity. +// A failure on one platform never blocks vulnerability creation or other recipients. +func (h *RobotHandler) NotifyNewVulnerability(v *database.Vulnerability) { + if h == nil || h.db == nil || v == nil { + return + } + recipients, err := h.db.ListVulnerabilityAlertRecipients(v) + if err != nil { + h.logger.Warn("查询漏洞提醒接收人失败", zap.String("vulnerability_id", v.ID), zap.Error(err)) + return + } + eligible := recipients[:0] + for _, recipient := range recipients { + if robot.SupportsProactive(recipient.Platform) { + eligible = append(eligible, recipient) + } + } + if err := h.db.EnqueueVulnerabilityAlertDeliveries(v.ID, eligible); err != nil { + h.logger.Warn("写入漏洞提醒投递队列失败", zap.String("vulnerability_id", v.ID), zap.Error(err)) + return + } + select { + case h.alertWake <- struct{}{}: + default: + } +} + +// RunVulnerabilityAlertWorker drains the durable outbox. Failed sends use +// exponential backoff and remain retryable across application restarts. +func (h *RobotHandler) RunVulnerabilityAlertWorker(ctx context.Context) { + if h == nil || h.db == nil { + return + } + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + h.processVulnerabilityAlertDeliveries(ctx) + select { + case <-ctx.Done(): + return + case <-ticker.C: + case <-h.alertWake: + } + } +} + +func (h *RobotHandler) processVulnerabilityAlertDeliveries(ctx context.Context) { + deliveries, err := h.db.ListDueVulnerabilityAlertDeliveries(50) + if err != nil { + h.logger.Warn("读取漏洞提醒投递队列失败", zap.Error(err)) + return + } + for _, delivery := range deliveries { + sendCtx, cancel := context.WithTimeout(ctx, 15*time.Second) + err := robot.SendProactive(sendCtx, h.config.Robots, delivery.Platform, delivery.ExternalUserID, formatVulnerabilityRobotAlert(delivery.Vulnerability)) + cancel() + if err != nil { + attempts := delivery.Attempts + 1 + _ = h.db.MarkVulnerabilityAlertDeliveryFailed(delivery.ID, attempts, err) + h.logger.Warn("发送漏洞机器人提醒失败", zap.String("platform", delivery.Platform), zap.String("user_id", delivery.UserID), zap.String("vulnerability_id", delivery.Vulnerability.ID), zap.Int("attempt", attempts), zap.Error(err)) + continue + } + _ = h.db.MarkVulnerabilityAlertDeliverySent(delivery.ID) + h.logger.Info("漏洞机器人提醒已发送", zap.String("platform", delivery.Platform), zap.String("user_id", delivery.UserID), zap.String("vulnerability_id", delivery.Vulnerability.ID)) + } +} + +func (h *RobotHandler) cmdVulnerabilityAlerts(platform, externalUserID, arg string) string { + access, err := h.resolveRobotAccess(platform, externalUserID) + if err != nil { + return h.robotAccessDeniedMessage(platform) + } + arg = strings.ToLower(strings.TrimSpace(arg)) + if arg == "" || arg == "status" || arg == "状态" { + sub, err := h.db.GetVulnerabilityAlertSubscription(access.User.ID) + if err != nil { + return "读取漏洞提醒设置失败,请稍后重试。" + } + state := "已关闭" + if sub.Enabled { + state = "已开启" + } + return fmt.Sprintf("漏洞提醒%s;最低提醒级别:%s(%s)。\n可在 Web「漏洞管理」页面同步修改。", state, vulnerabilityAlertSeverityLabel(sub.MinSeverity), sub.MinSeverity) + } + enabled, severity := true, "" + switch arg { + case "开启", "开", "on", "enable": + severity = "high" + case "关闭", "关", "off", "disable": + enabled, severity = false, "high" + case "仅严重", "严重", "critical": + severity = "critical" + case "高危以上", "高危", "high": + severity = "high" + case "中危以上", "中危", "medium": + severity = "medium" + case "低危以上", "低危", "low": + severity = "low" + case "全部", "all", "info": + severity = "info" + default: + return "用法:漏洞提醒 开启|关闭|仅严重|高危以上|中危以上" + } + current, _ := h.db.GetVulnerabilityAlertSubscription(access.User.ID) + if (arg == "开启" || arg == "开" || arg == "on" || arg == "enable" || !enabled) && current != nil { + severity = current.MinSeverity + } + sub, err := h.db.UpsertVulnerabilityAlertSubscription(access.User.ID, enabled, severity) + if err != nil { + return "更新漏洞提醒失败,请稍后重试。" + } + if !sub.Enabled { + return "已关闭漏洞提醒。Web 端设置已同步。" + } + bindings, _ := h.db.ListRobotUserBindings(access.User.ID) + deliveryReady := false + for _, binding := range bindings { + if binding.Enabled && robot.SupportsProactive(binding.Platform) { + deliveryReady = true + break + } + } + if !deliveryReady { + return fmt.Sprintf("已保存漏洞提醒:%s及以上。当前没有支持主动推送的已绑定账号;请绑定企业微信、飞书、Telegram、Slack 或 Discord。", vulnerabilityAlertSeverityLabel(sub.MinSeverity)) + } + return fmt.Sprintf("已开启漏洞提醒:%s及以上漏洞将通过已绑定机器人推送。Web 端设置已同步。", vulnerabilityAlertSeverityLabel(sub.MinSeverity)) +} diff --git a/internal/handler/vulnerability_alert_test.go b/internal/handler/vulnerability_alert_test.go new file mode 100644 index 00000000..b3596808 --- /dev/null +++ b/internal/handler/vulnerability_alert_test.go @@ -0,0 +1,51 @@ +package handler + +import ( + "path/filepath" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "go.uber.org/zap" +) + +func TestRobotVulnerabilityAlertCommandSharesSubscription(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "alert-command.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) + } + user, err := db.CreateRBACUser("alert-user", "Alert User", "hash", true, []string{database.RBACSystemRoleOperator}) + if err != nil { + t.Fatal(err) + } + if err := db.CreateRobotBindingCode(user.ID, "alert-code", time.Now().Add(time.Minute)); err != nil { + t.Fatal(err) + } + if _, err := db.ConsumeRobotBindingCode("wecom", "external-user", "alert-code"); err != nil { + t.Fatal(err) + } + h := NewRobotHandler(&config.Config{}, db, nil, zap.NewNop()) + + if got := h.HandleMessage("wecom", "external-user", "漏洞提醒 高危以上"); !strings.Contains(got, "已开启") { + t.Fatalf("enable reply: %s", got) + } + sub, err := db.GetVulnerabilityAlertSubscription(user.ID) + if err != nil || !sub.Enabled || sub.MinSeverity != "high" { + t.Fatalf("web subscription not updated: %#v %v", sub, err) + } + if got := h.HandleMessage("wecom", "external-user", "vuln alerts off"); !strings.Contains(got, "已关闭") { + t.Fatalf("disable reply: %s", got) + } + sub, _ = db.GetVulnerabilityAlertSubscription(user.ID) + if sub.Enabled { + t.Fatalf("subscription remained enabled: %#v", sub) + } +}