CyberStrikeAI/internal/handler/vulnerability_alert.go

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))
}