From 199392a5d5b3bb0cfa672621141a808f2de3a226 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Tue, 19 May 2026 18:52:22 +0800 Subject: [PATCH] Add files via upload --- internal/app/app.go | 24 ++- internal/handler/config.go | 29 +++ internal/handler/wechat_robot.go | 293 +++++++++++++++++++++++++++++++ 3 files changed, 344 insertions(+), 2 deletions(-) create mode 100644 internal/handler/wechat_robot.go diff --git a/internal/app/app.go b/internal/app/app.go index ca129df4..38ec3d53 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -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 流式输出 diff --git a/internal/handler/config.go b/internal/handler/config.go index b9b1c5d8..f2892aea 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -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) diff --git a/internal/handler/wechat_robot.go b/internal/handler/wechat_robot.go new file mode 100644 index 00000000..93a5ea8f --- /dev/null +++ b/internal/handler/wechat_robot.go @@ -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, + }) +}