mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-16 08:57:30 +02:00
Add files via upload
This commit is contained in:
公明
@@ -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))
|
||||
}
|
||||
Reference in New Issue
Block a user