Add files via upload

2026-05-20 14:54:54 +02:00 · 2026-05-19 18:52:22 +08:00
parent aafb4cb584
commit 199392a5d5
3 changed files with 344 additions and 2 deletions
@@ -56,6 +56,7 @@ type App struct {
 	robotMu            sync.Mutex                // 保护钉钉/飞书长连接的 cancel
 	dingCancel         context.CancelFunc        // 钉钉 Stream 取消函数，用于配置变更时重启
 	larkCancel         context.CancelFunc        // 飞书长连接取消函数，用于配置变更时重启
+	wechatCancel       context.CancelFunc        // 微信 iLink 长轮询取消函数
 	c2Manager          *c2.Manager               // C2 管理器（未启用 C2 时为 nil）
 	c2Watchdog         *c2.SessionWatchdog       // C2 会话看门狗
 	c2WatchdogCancel   context.CancelFunc        // 看门狗取消函数
@@ -449,9 +450,11 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
 		configHandler.SetRetrieverUpdater(knowledgeRetriever)
 	}

-	// 设置机器人连接重启器，前端应用配置后无需重启服务即可使钉钉/飞书新配置生效
+	// 设置机器人连接重启器，前端应用配置后无需重启服务即可使钉钉/飞书/微信新配置生效
 	configHandler.SetRobotRestarter(app)

+	wechatRobotHandler := handler.NewWechatRobotHandler(cfg, configHandler, log.Logger)
+
 	configHandler.SetC2Runtime(app)
 	configHandler.SetC2ToolRegistrar(func() error {
 		if app.config.C2.EnabledEffective() && app.c2Manager != nil {
@@ -469,6 +472,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
 		notificationHandler,
 		conversationHandler,
 		robotHandler,
+		wechatRobotHandler,
 		groupHandler,
 		configHandler,
 		externalMCPHandler,
@@ -675,9 +679,14 @@ func (a *App) startRobotConnections() {
 		a.dingCancel = cancel
 		go robot.StartDing(ctx, cfg.Robots, a.robotHandler, a.logger.Logger)
 	}
+	if cfg.Robots.Wechat.Enabled && cfg.Robots.Wechat.BotToken != "" {
+		ctx, cancel := context.WithCancel(context.Background())
+		a.wechatCancel = cancel
+		go robot.StartWechat(ctx, cfg.Robots, a.robotHandler, cfg.Version, a.logger.Logger)
+	}
 }

-// RestartRobotConnections 重启钉钉/飞书长连接，使前端应用配置后立即生效（实现 handler.RobotRestarter）
+// RestartRobotConnections 重启钉钉/飞书/微信长连接，使前端应用配置后立即生效（实现 handler.RobotRestarter）
 func (a *App) RestartRobotConnections() {
 	a.robotMu.Lock()
 	if a.dingCancel != nil {
@@ -688,6 +697,10 @@ func (a *App) RestartRobotConnections() {
 		a.larkCancel()
 		a.larkCancel = nil
 	}
+	if a.wechatCancel != nil {
+		a.wechatCancel()
+		a.wechatCancel = nil
+	}
 	a.robotMu.Unlock()
 	// 给旧 goroutine 一点时间退出
 	time.Sleep(200 * time.Millisecond)
@@ -703,6 +716,7 @@ func setupRoutes(
 	notificationHandler *handler.NotificationHandler,
 	conversationHandler *handler.ConversationHandler,
 	robotHandler *handler.RobotHandler,
+	wechatRobotHandler *handler.WechatRobotHandler,
 	groupHandler *handler.GroupHandler,
 	configHandler *handler.ConfigHandler,
 	externalMCPHandler *handler.ExternalMCPHandler,
@@ -751,6 +765,12 @@ func setupRoutes(
 		// 机器人测试（需登录）：POST /api/robot/test，body: {"platform":"dingtalk","user_id":"test","text":"帮助"}，用于验证机器人逻辑
 		protected.POST("/robot/test", robotHandler.HandleRobotTest)

+		// 微信 iLink 扫码绑定（需登录）
+		protected.POST("/robot/wechat/qrcode", wechatRobotHandler.HandleWechatQRCode)
+		protected.GET("/robot/wechat/qrcode/status", wechatRobotHandler.HandleWechatQRCodeStatus)
+		protected.POST("/robot/wechat/qrcode/verify", wechatRobotHandler.HandleWechatVerifyCode)
+		protected.GET("/robot/wechat/status", wechatRobotHandler.HandleWechatStatus)
+
 		// Agent Loop
 		protected.POST("/agent-loop", agentHandler.AgentLoop)
 		// Agent Loop 流式输出
@@ -206,6 +206,25 @@ func (h *ConfigHandler) SetRobotRestarter(restarter RobotRestarter) {
 	h.robotRestarter = restarter
 }

+// ApplyWechatRobotBinding 微信 iLink 扫码绑定成功后写入配置并重启机器人连接
+func (h *ConfigHandler) ApplyWechatRobotBinding(wc config.RobotWechatConfig) error {
+	h.mu.Lock()
+	wc.Enabled = true
+	h.config.Robots.Wechat = wc
+	h.mu.Unlock()
+	if err := h.saveConfig(); err != nil {
+		return err
+	}
+	if h.robotRestarter != nil {
+		h.robotRestarter.RestartRobotConnections()
+	}
+	h.logger.Info("微信机器人绑定已保存",
+		zap.String("ilink_bot_id", wc.ILinkBotID),
+		zap.Bool("enabled", wc.Enabled),
+	)
+	return nil
+}
+
 // GetConfigResponse 获取配置响应
 type GetConfigResponse struct {
 	OpenAI     config.OpenAIConfig     `json:"openai"`
@@ -735,6 +754,7 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
 	if req.Robots != nil {
 		h.config.Robots = *req.Robots
 		h.logger.Info("更新机器人配置",
+			zap.Bool("wechat_enabled", h.config.Robots.Wechat.Enabled),
 			zap.Bool("wecom_enabled", h.config.Robots.Wecom.Enabled),
 			zap.Bool("dingtalk_enabled", h.config.Robots.Dingtalk.Enabled),
 			zap.Bool("lark_enabled", h.config.Robots.Lark.Enabled),
@@ -1481,6 +1501,15 @@ func updateRobotsConfig(doc *yaml.Node, cfg config.RobotsConfig) {
 		setBoolInMap(sessionNode, "strict_user_identity", *cfg.Session.StrictUserIdentity)
 	}

+	wechatNode := ensureMap(robotsNode, "wechat")
+	setBoolInMap(wechatNode, "enabled", cfg.Wechat.Enabled)
+	setStringInMap(wechatNode, "bot_token", cfg.Wechat.BotToken)
+	setStringInMap(wechatNode, "ilink_bot_id", cfg.Wechat.ILinkBotID)
+	setStringInMap(wechatNode, "ilink_user_id", cfg.Wechat.ILinkUserID)
+	setStringInMap(wechatNode, "base_url", cfg.Wechat.BaseURL)
+	setStringInMap(wechatNode, "bot_type", cfg.Wechat.BotType)
+	setStringInMap(wechatNode, "bot_agent", cfg.Wechat.BotAgent)
+
 	wecomNode := ensureMap(robotsNode, "wecom")
 	setBoolInMap(wecomNode, "enabled", cfg.Wecom.Enabled)
 	setStringInMap(wecomNode, "token", cfg.Wecom.Token)
@@ -0,0 +1,293 @@
+package handler
+
+import (
+	"context"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"cyberstrike-ai/internal/config"
+	"cyberstrike-ai/internal/robot/ilink"
+
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	"go.uber.org/zap"
+)
+
+const wechatLoginTTL = 5 * time.Minute
+
+// WechatConfigSaver 绑定成功后写入配置并重启机器人连接
+type WechatConfigSaver interface {
+	ApplyWechatRobotBinding(cfg config.RobotWechatConfig) error
+}
+
+type wechatLoginSession struct {
+	QRCode           string
+	QRCodeImgURL     string
+	PendingVerify    string
+	CurrentBaseURL   string
+	StartedAt        time.Time
+}
+
+// WechatRobotHandler 微信 iLink 机器人（扫码绑定 + 配置）
+type WechatRobotHandler struct {
+	config       *config.Config
+	configSaver  WechatConfigSaver
+	logger       *zap.Logger
+	mu           sync.Mutex
+	logins       map[string]*wechatLoginSession
+}
+
+// NewWechatRobotHandler 创建微信机器人处理器
+func NewWechatRobotHandler(cfg *config.Config, saver WechatConfigSaver, logger *zap.Logger) *WechatRobotHandler {
+	return &WechatRobotHandler{
+		config:      cfg,
+		configSaver: saver,
+		logger:      logger,
+		logins:      make(map[string]*wechatLoginSession),
+	}
+}
+
+func (h *WechatRobotHandler) purgeExpiredLogins() {
+	now := time.Now()
+	for k, v := range h.logins {
+		if now.Sub(v.StartedAt) > wechatLoginTTL {
+			delete(h.logins, k)
+		}
+	}
+}
+
+func (h *WechatRobotHandler) ilinkClient(baseURL string) *ilink.Client {
+	ver := h.config.Version
+	if ver == "" {
+		ver = "1.0.0"
+	}
+	ver = strings.TrimPrefix(strings.TrimSpace(ver), "v")
+	ver = strings.TrimPrefix(ver, "V")
+	wc := h.config.Robots.Wechat
+	return ilink.NewClient(baseURL, wc.BotToken, wc.BotAgent, ilink.BuildClientVersion(ver))
+}
+
+// HandleWechatQRCode POST /api/robot/wechat/qrcode — 生成绑定二维码
+func (h *WechatRobotHandler) HandleWechatQRCode(c *gin.Context) {
+	h.mu.Lock()
+	h.purgeExpiredLogins()
+	h.mu.Unlock()
+
+	var req struct {
+		BotType string `json:"bot_type"`
+	}
+	_ = c.ShouldBindJSON(&req)
+
+	botType := req.BotType
+	if botType == "" {
+		botType = h.config.Robots.Wechat.BotType
+	}
+	if botType == "" {
+		botType = ilink.DefaultBotType
+	}
+	baseURL := h.config.Robots.Wechat.BaseURL
+	if baseURL == "" {
+		baseURL = ilink.DefaultBaseURL
+	}
+
+	var localTokens []string
+	if t := h.config.Robots.Wechat.BotToken; t != "" {
+		localTokens = []string{t}
+	}
+
+	client := h.ilinkClient(baseURL)
+	ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second)
+	defer cancel()
+
+	qr, err := client.GetBotQRCode(ctx, botType, localTokens)
+	if err != nil {
+		h.logger.Warn("获取微信二维码失败", zap.Error(err))
+		c.JSON(http.StatusBadGateway, gin.H{"error": "获取二维码失败: " + err.Error()})
+		return
+	}
+	if qr.QRCode == "" || qr.QRCodeImgContent == "" {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "微信服务器未返回有效二维码"})
+		return
+	}
+
+	sessionKey := uuid.New().String()
+	h.mu.Lock()
+	h.logins[sessionKey] = &wechatLoginSession{
+		QRCode:         qr.QRCode,
+		QRCodeImgURL:   qr.QRCodeImgContent,
+		CurrentBaseURL: baseURL,
+		StartedAt:      time.Now(),
+	}
+	h.mu.Unlock()
+
+	resp := gin.H{
+		"session_key":     sessionKey,
+		"qrcode":          qr.QRCode,
+		"qrcode_open_url": qr.QRCodeImgContent,
+		"message":         "请使用微信扫描二维码并确认绑定",
+	}
+	if dataURL, err := ilink.QRCodeDataURL(qr.QRCodeImgContent, 256); err != nil {
+		h.logger.Warn("生成二维码图片失败", zap.Error(err))
+	} else {
+		resp["qrcode_image_data_url"] = dataURL
+	}
+
+	c.JSON(http.StatusOK, resp)
+}
+
+// HandleWechatQRCodeStatus GET /api/robot/wechat/qrcode/status — 轮询扫码状态
+func (h *WechatRobotHandler) HandleWechatQRCodeStatus(c *gin.Context) {
+	sessionKey := c.Query("session_key")
+	verifyCode := c.Query("verify_code")
+	if sessionKey == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "缺少 session_key"})
+		return
+	}
+
+	h.mu.Lock()
+	sess, ok := h.logins[sessionKey]
+	h.mu.Unlock()
+	if !ok {
+		c.JSON(http.StatusNotFound, gin.H{"error": "登录会话不存在或已过期，请重新生成二维码"})
+		return
+	}
+	if time.Since(sess.StartedAt) > wechatLoginTTL {
+		h.mu.Lock()
+		delete(h.logins, sessionKey)
+		h.mu.Unlock()
+		c.JSON(http.StatusGone, gin.H{"error": "二维码已过期，请重新生成"})
+		return
+	}
+
+	baseURL := sess.CurrentBaseURL
+	if baseURL == "" {
+		baseURL = ilink.DefaultBaseURL
+	}
+	vc := verifyCode
+	if vc == "" {
+		vc = sess.PendingVerify
+	}
+
+	client := h.ilinkClient(baseURL)
+	ctx, cancel := context.WithTimeout(c.Request.Context(), 40*time.Second)
+	defer cancel()
+
+	st, err := client.GetQRCodeStatus(ctx, sess.QRCode, vc)
+	if err != nil {
+		h.logger.Warn("轮询微信二维码状态失败", zap.Error(err))
+		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
+		return
+	}
+
+	switch st.Status {
+	case "wait", "scaned":
+		c.JSON(http.StatusOK, gin.H{"status": st.Status})
+		return
+	case "need_verifycode":
+		c.JSON(http.StatusOK, gin.H{
+			"status":  st.Status,
+			"message": "请在手机微信查看配对数字，并在下方输入",
+		})
+		return
+	case "scaned_but_redirect":
+		if st.RedirectHost != "" {
+			h.mu.Lock()
+			if s, ok := h.logins[sessionKey]; ok {
+				s.CurrentBaseURL = "https://" + st.RedirectHost
+			}
+			h.mu.Unlock()
+		}
+		c.JSON(http.StatusOK, gin.H{"status": st.Status})
+		return
+	case "binded_redirect":
+		h.mu.Lock()
+		delete(h.logins, sessionKey)
+		h.mu.Unlock()
+		c.JSON(http.StatusOK, gin.H{
+			"status":            st.Status,
+			"already_connected": true,
+			"message":           "该微信已绑定过，无需重复绑定",
+		})
+		return
+	case "confirmed":
+		if st.BotToken == "" || st.ILinkBotID == "" {
+			c.JSON(http.StatusBadGateway, gin.H{"error": "绑定确认成功但缺少 bot_token"})
+			return
+		}
+		saveBase := st.BaseURL
+		if saveBase == "" {
+			saveBase = baseURL
+		}
+		wc := h.config.Robots.Wechat
+		wc.Enabled = true
+		wc.BotToken = st.BotToken
+		wc.ILinkBotID = st.ILinkBotID
+		wc.ILinkUserID = st.ILinkUserID
+		wc.BaseURL = saveBase
+		if wc.BotType == "" {
+			wc.BotType = ilink.DefaultBotType
+		}
+		if wc.BotAgent == "" {
+			wc.BotAgent = ilink.DefaultBotAgent
+		}
+		if h.configSaver != nil {
+			if err := h.configSaver.ApplyWechatRobotBinding(wc); err != nil {
+				h.logger.Warn("保存微信机器人配置失败", zap.Error(err))
+				c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
+				return
+			}
+		} else {
+			h.config.Robots.Wechat = wc
+		}
+		h.mu.Lock()
+		delete(h.logins, sessionKey)
+		h.mu.Unlock()
+		c.JSON(http.StatusOK, gin.H{
+			"status":        "confirmed",
+			"message":       "绑定成功，微信机器人已启用",
+			"ilink_bot_id":  st.ILinkBotID,
+			"ilink_user_id": st.ILinkUserID,
+		})
+		return
+	default:
+		c.JSON(http.StatusOK, gin.H{"status": st.Status})
+	}
+}
+
+// HandleWechatVerifyCode POST /api/robot/wechat/qrcode/verify — 提交手机配对数字
+func (h *WechatRobotHandler) HandleWechatVerifyCode(c *gin.Context) {
+	var req struct {
+		SessionKey string `json:"session_key"`
+		VerifyCode string `json:"verify_code"`
+	}
+	if err := c.ShouldBindJSON(&req); err != nil || req.SessionKey == "" || req.VerifyCode == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "需要 session_key 与 verify_code"})
+		return
+	}
+	h.mu.Lock()
+	sess, ok := h.logins[req.SessionKey]
+	if ok {
+		sess.PendingVerify = req.VerifyCode
+	}
+	h.mu.Unlock()
+	if !ok {
+		c.JSON(http.StatusNotFound, gin.H{"error": "登录会话不存在或已过期"})
+		return
+	}
+	c.JSON(http.StatusOK, gin.H{"message": "已提交配对码，请继续等待绑定"})
+}
+
+// HandleWechatStatus GET /api/robot/wechat/status — 当前绑定状态（供前端展示）
+func (h *WechatRobotHandler) HandleWechatStatus(c *gin.Context) {
+	wc := h.config.Robots.Wechat
+	bound := wc.BotToken != "" && wc.ILinkBotID != ""
+	c.JSON(http.StatusOK, gin.H{
+		"enabled":       wc.Enabled,
+		"bound":         bound,
+		"ilink_bot_id":  wc.ILinkBotID,
+		"ilink_user_id": wc.ILinkUserID,
+		"base_url":      wc.BaseURL,
+	})
+}