mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-07 04:58:03 +02:00
Add files via upload
This commit is contained in:
公明
@@ -14,6 +14,7 @@ import (
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/handler"
|
||||
"cyberstrike-ai/internal/knowledge"
|
||||
"cyberstrike-ai/internal/robot"
|
||||
"cyberstrike-ai/internal/logger"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
@@ -325,6 +326,14 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
|
||||
// 创建OpenAPI处理器
|
||||
conversationHandler := handler.NewConversationHandler(db, log.Logger)
|
||||
robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger)
|
||||
// 飞书/钉钉长连接（无需公网），启用时在后台启动
|
||||
if cfg.Robots.Lark.Enabled && cfg.Robots.Lark.AppID != "" && cfg.Robots.Lark.AppSecret != "" {
|
||||
go robot.StartLark(cfg.Robots.Lark, robotHandler, log.Logger)
|
||||
}
|
||||
if cfg.Robots.Dingtalk.Enabled && cfg.Robots.Dingtalk.ClientID != "" && cfg.Robots.Dingtalk.ClientSecret != "" {
|
||||
go robot.StartDing(cfg.Robots.Dingtalk, robotHandler, log.Logger)
|
||||
}
|
||||
openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, resultStorage, conversationHandler, agentHandler)
|
||||
|
||||
// 创建 App 实例（部分字段稍后填充）
|
||||
@@ -408,6 +417,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
agentHandler,
|
||||
monitorHandler,
|
||||
conversationHandler,
|
||||
robotHandler,
|
||||
groupHandler,
|
||||
configHandler,
|
||||
externalMCPHandler,
|
||||
@@ -472,6 +482,7 @@ func setupRoutes(
|
||||
agentHandler *handler.AgentHandler,
|
||||
monitorHandler *handler.MonitorHandler,
|
||||
conversationHandler *handler.ConversationHandler,
|
||||
robotHandler *handler.RobotHandler,
|
||||
groupHandler *handler.GroupHandler,
|
||||
configHandler *handler.ConfigHandler,
|
||||
externalMCPHandler *handler.ExternalMCPHandler,
|
||||
@@ -497,9 +508,18 @@ func setupRoutes(
|
||||
authRoutes.GET("/validate", security.AuthMiddleware(authManager), authHandler.Validate)
|
||||
}
|
||||
|
||||
// 机器人回调（无需登录，供企业微信/钉钉/飞书服务器调用）
|
||||
api.GET("/robot/wecom", robotHandler.HandleWecomGET)
|
||||
api.POST("/robot/wecom", robotHandler.HandleWecomPOST)
|
||||
api.POST("/robot/dingtalk", robotHandler.HandleDingtalkPOST)
|
||||
api.POST("/robot/lark", robotHandler.HandleLarkPOST)
|
||||
|
||||
protected := api.Group("")
|
||||
protected.Use(security.AuthMiddleware(authManager))
|
||||
{
|
||||
// 机器人测试（需登录）：POST /api/robot/test，body: {"platform":"dingtalk","user_id":"test","text":"帮助"}，用于验证机器人逻辑
|
||||
protected.POST("/robot/test", robotHandler.HandleRobotTest)
|
||||
|
||||
// Agent Loop
|
||||
protected.POST("/agent-loop", agentHandler.AgentLoop)
|
||||
// Agent Loop 流式输出
|
||||
|
||||
@@ -25,11 +25,44 @@ type Config struct {
|
||||
Auth AuthConfig `yaml:"auth"`
|
||||
ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"`
|
||||
Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"`
|
||||
Robots RobotsConfig `yaml:"robots,omitempty" json:"robots,omitempty"` // 企业微信/钉钉/飞书等机器人配置
|
||||
RolesDir string `yaml:"roles_dir,omitempty" json:"roles_dir,omitempty"` // 角色配置文件目录（新方式）
|
||||
Roles map[string]RoleConfig `yaml:"roles,omitempty" json:"roles,omitempty"` // 向后兼容：支持在主配置文件中定义角色
|
||||
SkillsDir string `yaml:"skills_dir,omitempty" json:"skills_dir,omitempty"` // Skills配置文件目录
|
||||
}
|
||||
|
||||
// RobotsConfig 机器人配置（企业微信、钉钉、飞书等）
|
||||
type RobotsConfig struct {
|
||||
Wecom RobotWecomConfig `yaml:"wecom,omitempty" json:"wecom,omitempty"` // 企业微信
|
||||
Dingtalk RobotDingtalkConfig `yaml:"dingtalk,omitempty" json:"dingtalk,omitempty"` // 钉钉
|
||||
Lark RobotLarkConfig `yaml:"lark,omitempty" json:"lark,omitempty"` // 飞书
|
||||
}
|
||||
|
||||
// RobotWecomConfig 企业微信机器人配置
|
||||
type RobotWecomConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
Token string `yaml:"token" json:"token"` // 回调 URL 校验 Token
|
||||
EncodingAESKey string `yaml:"encoding_aes_key" json:"encoding_aes_key"` // EncodingAESKey
|
||||
CorpID string `yaml:"corp_id" json:"corp_id"` // 企业 ID
|
||||
Secret string `yaml:"secret" json:"secret"` // 应用 Secret
|
||||
AgentID int64 `yaml:"agent_id" json:"agent_id"` // 应用 AgentId
|
||||
}
|
||||
|
||||
// RobotDingtalkConfig 钉钉机器人配置
|
||||
type RobotDingtalkConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
ClientID string `yaml:"client_id" json:"client_id"` // 应用 Key (AppKey)
|
||||
ClientSecret string `yaml:"client_secret" json:"client_secret"` // 应用 Secret
|
||||
}
|
||||
|
||||
// RobotLarkConfig 飞书机器人配置
|
||||
type RobotLarkConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
AppID string `yaml:"app_id" json:"app_id"` // 应用 App ID
|
||||
AppSecret string `yaml:"app_secret" json:"app_secret"` // 应用 App Secret
|
||||
VerifyToken string `yaml:"verify_token" json:"verify_token"` // 事件订阅 Verification Token（可选）
|
||||
}
|
||||
|
||||
type ServerConfig struct {
|
||||
Host string `yaml:"host"`
|
||||
Port int `yaml:"port"`
|
||||
|
||||
@@ -259,6 +259,96 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// ProcessMessageForRobot 供机器人（企业微信/钉钉/飞书）调用：与 /api/agent-loop/stream 相同执行路径（含 progressCallback、过程详情），仅不发送 SSE，最后返回完整回复
|
||||
func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationID, message, role string) (response string, convID string, err error) {
|
||||
if conversationID == "" {
|
||||
title := safeTruncateString(message, 50)
|
||||
conv, createErr := h.db.CreateConversation(title)
|
||||
if createErr != nil {
|
||||
return "", "", fmt.Errorf("创建对话失败: %w", createErr)
|
||||
}
|
||||
conversationID = conv.ID
|
||||
} else {
|
||||
if _, getErr := h.db.GetConversation(conversationID); getErr != nil {
|
||||
return "", "", fmt.Errorf("对话不存在")
|
||||
}
|
||||
}
|
||||
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
if err != nil {
|
||||
historyMessages, getErr := h.db.GetMessages(conversationID)
|
||||
if getErr != nil {
|
||||
agentHistoryMessages = []agent.ChatMessage{}
|
||||
} else {
|
||||
agentHistoryMessages = make([]agent.ChatMessage, 0, len(historyMessages))
|
||||
for _, msg := range historyMessages {
|
||||
agentHistoryMessages = append(agentHistoryMessages, agent.ChatMessage{Role: msg.Role, Content: msg.Content})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
finalMessage := message
|
||||
var roleTools, roleSkills []string
|
||||
if role != "" && role != "默认" && h.config.Roles != nil {
|
||||
if r, exists := h.config.Roles[role]; exists && r.Enabled {
|
||||
if r.UserPrompt != "" {
|
||||
finalMessage = r.UserPrompt + "\n\n" + message
|
||||
}
|
||||
roleTools = r.Tools
|
||||
roleSkills = r.Skills
|
||||
}
|
||||
}
|
||||
|
||||
if _, err = h.db.AddMessage(conversationID, "user", message, nil); err != nil {
|
||||
return "", "", fmt.Errorf("保存用户消息失败: %w", err)
|
||||
}
|
||||
|
||||
// 与 agent-loop/stream 一致：先创建助手消息占位，用 progressCallback 写过程详情（不发送 SSE）
|
||||
assistantMsg, err := h.db.AddMessage(conversationID, "assistant", "处理中...", nil)
|
||||
if err != nil {
|
||||
h.logger.Warn("机器人：创建助手消息占位失败", zap.Error(err))
|
||||
}
|
||||
var assistantMessageID string
|
||||
if assistantMsg != nil {
|
||||
assistantMessageID = assistantMsg.ID
|
||||
}
|
||||
progressCallback := h.createProgressCallback(conversationID, assistantMessageID, nil)
|
||||
|
||||
result, err := h.agent.AgentLoopWithProgress(ctx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools, roleSkills)
|
||||
if err != nil {
|
||||
errMsg := "执行失败: " + err.Error()
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ? WHERE id = ?", errMsg, assistantMessageID)
|
||||
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil)
|
||||
}
|
||||
return "", conversationID, err
|
||||
}
|
||||
|
||||
// 更新助手消息内容与 MCP 执行 ID（与 stream 一致）
|
||||
if assistantMessageID != "" {
|
||||
mcpIDsJSON := ""
|
||||
if len(result.MCPExecutionIDs) > 0 {
|
||||
jsonData, _ := json.Marshal(result.MCPExecutionIDs)
|
||||
mcpIDsJSON = string(jsonData)
|
||||
}
|
||||
_, err = h.db.Exec(
|
||||
"UPDATE messages SET content = ?, mcp_execution_ids = ? WHERE id = ?",
|
||||
result.Response, mcpIDsJSON, assistantMessageID,
|
||||
)
|
||||
if err != nil {
|
||||
h.logger.Warn("机器人：更新助手消息失败", zap.Error(err))
|
||||
}
|
||||
} else {
|
||||
if _, err = h.db.AddMessage(conversationID, "assistant", result.Response, result.MCPExecutionIDs); err != nil {
|
||||
h.logger.Warn("机器人：保存助手消息失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
_ = h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput)
|
||||
}
|
||||
return result.Response, conversationID, nil
|
||||
}
|
||||
|
||||
// StreamEvent 流式事件
|
||||
type StreamEvent struct {
|
||||
Type string `json:"type"` // conversation, progress, tool_call, tool_result, response, error, cancelled, done
|
||||
|
||||
@@ -150,6 +150,7 @@ type GetConfigResponse struct {
|
||||
Tools []ToolConfigInfo `json:"tools"`
|
||||
Agent config.AgentConfig `json:"agent"`
|
||||
Knowledge config.KnowledgeConfig `json:"knowledge"`
|
||||
Robots config.RobotsConfig `json:"robots,omitempty"`
|
||||
}
|
||||
|
||||
// ToolConfigInfo 工具配置信息
|
||||
@@ -222,6 +223,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
Tools: tools,
|
||||
Agent: h.config.Agent,
|
||||
Knowledge: h.config.Knowledge,
|
||||
Robots: h.config.Robots,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -479,6 +481,7 @@ type UpdateConfigRequest struct {
|
||||
Tools []ToolEnableStatus `json:"tools,omitempty"`
|
||||
Agent *config.AgentConfig `json:"agent,omitempty"`
|
||||
Knowledge *config.KnowledgeConfig `json:"knowledge,omitempty"`
|
||||
Robots *config.RobotsConfig `json:"robots,omitempty"`
|
||||
}
|
||||
|
||||
// ToolEnableStatus 工具启用状态
|
||||
@@ -555,6 +558,16 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
)
|
||||
}
|
||||
|
||||
// 更新机器人配置
|
||||
if req.Robots != nil {
|
||||
h.config.Robots = *req.Robots
|
||||
h.logger.Info("更新机器人配置",
|
||||
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),
|
||||
)
|
||||
}
|
||||
|
||||
// 更新工具启用状态
|
||||
if req.Tools != nil {
|
||||
// 分离内部工具和外部工具
|
||||
@@ -856,6 +869,7 @@ func (h *ConfigHandler) saveConfig() error {
|
||||
updateOpenAIConfig(root, h.config.OpenAI)
|
||||
updateFOFAConfig(root, h.config.FOFA)
|
||||
updateKnowledgeConfig(root, h.config.Knowledge)
|
||||
updateRobotsConfig(root, h.config.Robots)
|
||||
// 更新外部MCP配置（使用external_mcp.go中的函数，同一包中可直接调用）
|
||||
// 读取原始配置以保持向后兼容
|
||||
originalConfigs := make(map[string]map[string]bool)
|
||||
@@ -1031,6 +1045,30 @@ func updateKnowledgeConfig(doc *yaml.Node, cfg config.KnowledgeConfig) {
|
||||
setFloatInMap(retrievalNode, "hybrid_weight", cfg.Retrieval.HybridWeight)
|
||||
}
|
||||
|
||||
func updateRobotsConfig(doc *yaml.Node, cfg config.RobotsConfig) {
|
||||
root := doc.Content[0]
|
||||
robotsNode := ensureMap(root, "robots")
|
||||
|
||||
wecomNode := ensureMap(robotsNode, "wecom")
|
||||
setBoolInMap(wecomNode, "enabled", cfg.Wecom.Enabled)
|
||||
setStringInMap(wecomNode, "token", cfg.Wecom.Token)
|
||||
setStringInMap(wecomNode, "encoding_aes_key", cfg.Wecom.EncodingAESKey)
|
||||
setStringInMap(wecomNode, "corp_id", cfg.Wecom.CorpID)
|
||||
setStringInMap(wecomNode, "secret", cfg.Wecom.Secret)
|
||||
setIntInMap(wecomNode, "agent_id", int(cfg.Wecom.AgentID))
|
||||
|
||||
dingtalkNode := ensureMap(robotsNode, "dingtalk")
|
||||
setBoolInMap(dingtalkNode, "enabled", cfg.Dingtalk.Enabled)
|
||||
setStringInMap(dingtalkNode, "client_id", cfg.Dingtalk.ClientID)
|
||||
setStringInMap(dingtalkNode, "client_secret", cfg.Dingtalk.ClientSecret)
|
||||
|
||||
larkNode := ensureMap(robotsNode, "lark")
|
||||
setBoolInMap(larkNode, "enabled", cfg.Lark.Enabled)
|
||||
setStringInMap(larkNode, "app_id", cfg.Lark.AppID)
|
||||
setStringInMap(larkNode, "app_secret", cfg.Lark.AppSecret)
|
||||
setStringInMap(larkNode, "verify_token", cfg.Lark.VerifyToken)
|
||||
}
|
||||
|
||||
func ensureMap(parent *yaml.Node, path ...string) *yaml.Node {
|
||||
current := parent
|
||||
for _, key := range path {
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"encoding/xml"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
robotCmdHelp = "帮助"
|
||||
robotCmdList = "列表"
|
||||
robotCmdListAlt = "对话列表"
|
||||
robotCmdSwitch = "切换"
|
||||
robotCmdContinue = "继续"
|
||||
robotCmdNew = "新对话"
|
||||
robotCmdClear = "清空"
|
||||
robotCmdCurrent = "当前"
|
||||
)
|
||||
|
||||
// RobotHandler 企业微信/钉钉/飞书等机器人回调处理
|
||||
type RobotHandler struct {
|
||||
config *config.Config
|
||||
db *database.DB
|
||||
agentHandler *AgentHandler
|
||||
logger *zap.Logger
|
||||
mu sync.RWMutex
|
||||
sessions map[string]string // key: "platform_userID", value: conversationID
|
||||
}
|
||||
|
||||
// NewRobotHandler 创建机器人处理器
|
||||
func NewRobotHandler(cfg *config.Config, db *database.DB, agentHandler *AgentHandler, logger *zap.Logger) *RobotHandler {
|
||||
return &RobotHandler{
|
||||
config: cfg,
|
||||
db: db,
|
||||
agentHandler: agentHandler,
|
||||
logger: logger,
|
||||
sessions: make(map[string]string),
|
||||
}
|
||||
}
|
||||
|
||||
// sessionKey 生成会话 key
|
||||
func (h *RobotHandler) sessionKey(platform, userID string) string {
|
||||
return platform + "_" + userID
|
||||
}
|
||||
|
||||
// getOrCreateConversation 获取或创建当前会话
|
||||
func (h *RobotHandler) getOrCreateConversation(platform, userID string) (convID string, isNew bool) {
|
||||
h.mu.RLock()
|
||||
convID = h.sessions[h.sessionKey(platform, userID)]
|
||||
h.mu.RUnlock()
|
||||
if convID != "" {
|
||||
return convID, false
|
||||
}
|
||||
conv, err := h.db.CreateConversation("机器人对话")
|
||||
if err != nil {
|
||||
h.logger.Warn("创建机器人会话失败", zap.Error(err))
|
||||
return "", false
|
||||
}
|
||||
convID = conv.ID
|
||||
h.mu.Lock()
|
||||
h.sessions[h.sessionKey(platform, userID)] = convID
|
||||
h.mu.Unlock()
|
||||
return convID, true
|
||||
}
|
||||
|
||||
// setConversation 切换当前会话
|
||||
func (h *RobotHandler) setConversation(platform, userID, convID string) {
|
||||
h.mu.Lock()
|
||||
h.sessions[h.sessionKey(platform, userID)] = convID
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// clearConversation 清空当前会话（切换到新对话）
|
||||
func (h *RobotHandler) clearConversation(platform, userID string) (newConvID string) {
|
||||
conv, err := h.db.CreateConversation("新对话")
|
||||
if err != nil {
|
||||
h.logger.Warn("创建新对话失败", zap.Error(err))
|
||||
return ""
|
||||
}
|
||||
h.setConversation(platform, userID, conv.ID)
|
||||
return conv.ID
|
||||
}
|
||||
|
||||
// HandleMessage 处理用户输入，返回回复文本（供各平台 webhook 调用）
|
||||
func (h *RobotHandler) HandleMessage(platform, userID, text string) (reply string) {
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return "请输入内容或发送「帮助」查看命令。"
|
||||
}
|
||||
|
||||
// 命令分发
|
||||
switch {
|
||||
case text == robotCmdHelp || text == "help" || text == "？" || text == "?":
|
||||
return h.cmdHelp()
|
||||
case text == robotCmdList || text == robotCmdListAlt:
|
||||
return h.cmdList(userID)
|
||||
case strings.HasPrefix(text, robotCmdSwitch+" ") || strings.HasPrefix(text, robotCmdContinue+" "):
|
||||
var id string
|
||||
if strings.HasPrefix(text, robotCmdSwitch+" ") {
|
||||
id = strings.TrimSpace(text[len(robotCmdSwitch)+1:])
|
||||
} else {
|
||||
id = strings.TrimSpace(text[len(robotCmdContinue)+1:])
|
||||
}
|
||||
return h.cmdSwitch(platform, userID, id)
|
||||
case text == robotCmdNew:
|
||||
return h.cmdNew(platform, userID)
|
||||
case text == robotCmdClear:
|
||||
return h.cmdClear(platform, userID)
|
||||
case text == robotCmdCurrent:
|
||||
return h.cmdCurrent(platform, userID)
|
||||
}
|
||||
|
||||
// 普通消息：走 Agent
|
||||
convID, _ := h.getOrCreateConversation(platform, userID)
|
||||
if convID == "" {
|
||||
return "无法创建或获取对话，请稍后再试。"
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
defer cancel()
|
||||
resp, newConvID, err := h.agentHandler.ProcessMessageForRobot(ctx, convID, text, "默认")
|
||||
if err != nil {
|
||||
h.logger.Warn("机器人 Agent 执行失败", zap.String("platform", platform), zap.String("userID", userID), zap.Error(err))
|
||||
return "处理失败: " + err.Error()
|
||||
}
|
||||
if newConvID != convID {
|
||||
h.setConversation(platform, userID, newConvID)
|
||||
}
|
||||
return resp
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdHelp() string {
|
||||
return `【CyberStrikeAI 机器人命令】
|
||||
· 帮助 — 显示本帮助
|
||||
· 列表 / 对话列表 — 列出所有对话标题与 ID
|
||||
· 切换 <对话ID> / 继续 <对话ID> — 指定对话继续
|
||||
· 新对话 — 开启新对话
|
||||
· 清空 — 清空当前上下文（等同于新对话）
|
||||
· 当前 — 显示当前对话 ID 与标题
|
||||
除以上命令外，直接输入内容将发送给 AI 进行渗透测试/安全分析。`
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdList(userID string) string {
|
||||
convs, err := h.db.ListConversations(50, 0, "")
|
||||
if err != nil {
|
||||
return "获取对话列表失败: " + err.Error()
|
||||
}
|
||||
if len(convs) == 0 {
|
||||
return "暂无对话。发送任意内容将自动创建新对话。"
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("【对话列表】\n")
|
||||
for i, c := range convs {
|
||||
if i >= 20 {
|
||||
b.WriteString("… 仅显示前 20 条\n")
|
||||
break
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("· %s\n ID: %s\n", c.Title, c.ID))
|
||||
}
|
||||
return strings.TrimSuffix(b.String(), "\n")
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdSwitch(platform, userID, convID string) string {
|
||||
if convID == "" {
|
||||
return "请指定对话 ID，例如：切换 xxx-xxx-xxx"
|
||||
}
|
||||
conv, err := h.db.GetConversation(convID)
|
||||
if err != nil {
|
||||
return "对话不存在或 ID 错误。"
|
||||
}
|
||||
h.setConversation(platform, userID, conv.ID)
|
||||
return fmt.Sprintf("已切换到对话：「%s」\nID: %s", conv.Title, conv.ID)
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdNew(platform, userID string) string {
|
||||
newID := h.clearConversation(platform, userID)
|
||||
if newID == "" {
|
||||
return "创建新对话失败，请重试。"
|
||||
}
|
||||
return "已开启新对话，可直接发送内容。"
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdClear(platform, userID string) string {
|
||||
return h.cmdNew(platform, userID)
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdCurrent(platform, userID string) string {
|
||||
h.mu.RLock()
|
||||
convID := h.sessions[h.sessionKey(platform, userID)]
|
||||
h.mu.RUnlock()
|
||||
if convID == "" {
|
||||
return "当前没有进行中的对话。发送任意内容将创建新对话。"
|
||||
}
|
||||
conv, err := h.db.GetConversation(convID)
|
||||
if err != nil {
|
||||
return "当前对话 ID: " + convID + "（获取标题失败）"
|
||||
}
|
||||
return fmt.Sprintf("当前对话：「%s」\nID: %s", conv.Title, conv.ID)
|
||||
}
|
||||
|
||||
// —————— 企业微信 ——————
|
||||
|
||||
// wecomXML 企业微信回调 XML（明文模式下的简化结构；加密模式需先解密再解析）
|
||||
type wecomXML struct {
|
||||
ToUserName string `xml:"ToUserName"`
|
||||
FromUserName string `xml:"FromUserName"`
|
||||
CreateTime int64 `xml:"CreateTime"`
|
||||
MsgType string `xml:"MsgType"`
|
||||
Content string `xml:"Content"`
|
||||
MsgID string `xml:"MsgId"`
|
||||
AgentID int64 `xml:"AgentID"`
|
||||
Encrypt string `xml:"Encrypt"` // 加密模式下消息在此
|
||||
}
|
||||
|
||||
// wecomReplyXML 被动回复 XML
|
||||
type wecomReplyXML struct {
|
||||
XMLName xml.Name `xml:"xml"`
|
||||
ToUserName string `xml:"ToUserName"`
|
||||
FromUserName string `xml:"FromUserName"`
|
||||
CreateTime int64 `xml:"CreateTime"`
|
||||
MsgType string `xml:"MsgType"`
|
||||
Content string `xml:"Content"`
|
||||
}
|
||||
|
||||
// HandleWecomGET 企业微信 URL 校验（GET）
|
||||
func (h *RobotHandler) HandleWecomGET(c *gin.Context) {
|
||||
if !h.config.Robots.Wecom.Enabled {
|
||||
c.String(http.StatusNotFound, "")
|
||||
return
|
||||
}
|
||||
echostr := c.Query("echostr")
|
||||
if echostr == "" {
|
||||
c.String(http.StatusBadRequest, "missing echostr")
|
||||
return
|
||||
}
|
||||
// 明文模式时企业微信可能直接传 echostr，先直接返回以通过校验
|
||||
c.String(http.StatusOK, echostr)
|
||||
}
|
||||
|
||||
// wecomDecrypt 企业微信消息解密（AES-256-CBC，PKCS7，明文格式：16字节随机+4字节长度+消息+corpID）
|
||||
func wecomDecrypt(encodingAESKey, encryptedB64 string) ([]byte, error) {
|
||||
key, err := base64.StdEncoding.DecodeString(encodingAESKey + "=")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(key) != 32 {
|
||||
return nil, fmt.Errorf("encoding_aes_key 解码后应为 32 字节")
|
||||
}
|
||||
ciphertext, err := base64.StdEncoding.DecodeString(encryptedB64)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
iv := key[:16]
|
||||
mode := cipher.NewCBCDecrypter(block, iv)
|
||||
if len(ciphertext)%aes.BlockSize != 0 {
|
||||
return nil, fmt.Errorf("密文长度不是块大小的倍数")
|
||||
}
|
||||
plain := make([]byte, len(ciphertext))
|
||||
mode.CryptBlocks(plain, ciphertext)
|
||||
// 去除 PKCS7 填充
|
||||
n := int(plain[len(plain)-1])
|
||||
if n < 1 || n > 32 {
|
||||
return nil, fmt.Errorf("无效的 PKCS7 填充")
|
||||
}
|
||||
plain = plain[:len(plain)-n]
|
||||
// 企业微信格式：16 字节随机 + 4 字节长度(大端) + 消息 + corpID
|
||||
if len(plain) < 20 {
|
||||
return nil, fmt.Errorf("明文过短")
|
||||
}
|
||||
msgLen := binary.BigEndian.Uint32(plain[16:20])
|
||||
if int(20+msgLen) > len(plain) {
|
||||
return nil, fmt.Errorf("消息长度越界")
|
||||
}
|
||||
return plain[20 : 20+msgLen], nil
|
||||
}
|
||||
|
||||
// HandleWecomPOST 企业微信消息回调（POST），支持明文与加密模式
|
||||
func (h *RobotHandler) HandleWecomPOST(c *gin.Context) {
|
||||
if !h.config.Robots.Wecom.Enabled {
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
bodyRaw, _ := io.ReadAll(c.Request.Body)
|
||||
var body wecomXML
|
||||
if err := xml.Unmarshal(bodyRaw, &body); err != nil {
|
||||
h.logger.Debug("企业微信 POST 解析 XML 失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
// 加密模式：先解密再解析内层 XML
|
||||
if body.Encrypt != "" && h.config.Robots.Wecom.EncodingAESKey != "" {
|
||||
decrypted, err := wecomDecrypt(h.config.Robots.Wecom.EncodingAESKey, body.Encrypt)
|
||||
if err != nil {
|
||||
h.logger.Warn("企业微信消息解密失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
if err := xml.Unmarshal(decrypted, &body); err != nil {
|
||||
h.logger.Warn("企业微信解密后 XML 解析失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
}
|
||||
if body.MsgType != "text" {
|
||||
c.XML(http.StatusOK, wecomReplyXML{
|
||||
ToUserName: body.FromUserName,
|
||||
FromUserName: body.ToUserName,
|
||||
CreateTime: time.Now().Unix(),
|
||||
MsgType: "text",
|
||||
Content: "暂仅支持文本消息，请发送文字。",
|
||||
})
|
||||
return
|
||||
}
|
||||
userID := body.FromUserName
|
||||
text := strings.TrimSpace(body.Content)
|
||||
reply := h.HandleMessage("wecom", userID, text)
|
||||
// 加密模式需加密回复（此处简化为明文回复；若企业要求加密需再实现加密）
|
||||
c.XML(http.StatusOK, wecomReplyXML{
|
||||
ToUserName: body.FromUserName,
|
||||
FromUserName: body.ToUserName,
|
||||
CreateTime: time.Now().Unix(),
|
||||
MsgType: "text",
|
||||
Content: reply,
|
||||
})
|
||||
}
|
||||
|
||||
// —————— 测试接口（需登录，用于验证机器人逻辑，无需钉钉/飞书客户端） ——————
|
||||
|
||||
// RobotTestRequest 模拟机器人消息请求
|
||||
type RobotTestRequest struct {
|
||||
Platform string `json:"platform"` // 如 "dingtalk"、"lark"、"wecom"
|
||||
UserID string `json:"user_id"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
// HandleRobotTest 供本地验证：POST JSON { "platform", "user_id", "text" }，返回 { "reply": "..." }
|
||||
func (h *RobotHandler) HandleRobotTest(c *gin.Context) {
|
||||
var req RobotTestRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "请求体需为 JSON，包含 platform、user_id、text"})
|
||||
return
|
||||
}
|
||||
platform := strings.TrimSpace(req.Platform)
|
||||
if platform == "" {
|
||||
platform = "test"
|
||||
}
|
||||
userID := strings.TrimSpace(req.UserID)
|
||||
if userID == "" {
|
||||
userID = "test_user"
|
||||
}
|
||||
reply := h.HandleMessage(platform, userID, req.Text)
|
||||
c.JSON(http.StatusOK, gin.H{"reply": reply})
|
||||
}
|
||||
|
||||
// —————— 钉钉 ——————
|
||||
|
||||
// HandleDingtalkPOST 钉钉事件回调（流式接入等）；当前为占位，返回 200
|
||||
func (h *RobotHandler) HandleDingtalkPOST(c *gin.Context) {
|
||||
if !h.config.Robots.Dingtalk.Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{})
|
||||
return
|
||||
}
|
||||
// 钉钉流式/事件回调格式需按官方文档解析并异步回复，此处仅返回 200
|
||||
c.JSON(http.StatusOK, gin.H{"message": "ok"})
|
||||
}
|
||||
|
||||
// —————— 飞书 ——————
|
||||
|
||||
// HandleLarkPOST 飞书事件回调；当前为占位，返回 200；验证时需返回 challenge
|
||||
func (h *RobotHandler) HandleLarkPOST(c *gin.Context) {
|
||||
if !h.config.Robots.Lark.Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{})
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Challenge string `json:"challenge"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err == nil && body.Challenge != "" {
|
||||
c.JSON(http.StatusOK, gin.H{"challenge": body.Challenge})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{})
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
package robot
|
||||
|
||||
// MessageHandler 供飞书/钉钉长连接调用的消息处理接口（由 handler.RobotHandler 实现）
|
||||
type MessageHandler interface {
|
||||
HandleMessage(platform, userID, text string) string
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
package robot
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/open-dingtalk/dingtalk-stream-sdk-go/chatbot"
|
||||
"github.com/open-dingtalk/dingtalk-stream-sdk-go/client"
|
||||
dingutils "github.com/open-dingtalk/dingtalk-stream-sdk-go/utils"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// StartDing 启动钉钉 Stream 长连接（无需公网），收到消息后调用 handler 并通过 SessionWebhook 回复
|
||||
func StartDing(cfg config.RobotDingtalkConfig, h MessageHandler, logger *zap.Logger) {
|
||||
if !cfg.Enabled || cfg.ClientID == "" || cfg.ClientSecret == "" {
|
||||
return
|
||||
}
|
||||
streamClient := client.NewStreamClient(
|
||||
client.WithAppCredential(client.NewAppCredentialConfig(cfg.ClientID, cfg.ClientSecret)),
|
||||
client.WithSubscription(dingutils.SubscriptionTypeKCallback, "/v1.0/im/bot/messages/get",
|
||||
chatbot.NewDefaultChatBotFrameHandler(func(ctx context.Context, msg *chatbot.BotCallbackDataModel) ([]byte, error) {
|
||||
go handleDingMessage(ctx, msg, h, logger)
|
||||
return nil, nil
|
||||
}).OnEventReceived),
|
||||
)
|
||||
logger.Info("钉钉 Stream 正在连接…", zap.String("client_id", cfg.ClientID))
|
||||
go func() {
|
||||
err := streamClient.Start(context.Background())
|
||||
if err != nil {
|
||||
logger.Error("钉钉 Stream 长连接退出", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
logger.Info("钉钉 Stream 已启动（无需公网），等待收消息", zap.String("client_id", cfg.ClientID))
|
||||
}
|
||||
|
||||
func handleDingMessage(ctx context.Context, msg *chatbot.BotCallbackDataModel, h MessageHandler, logger *zap.Logger) {
|
||||
if msg == nil || msg.SessionWebhook == "" {
|
||||
return
|
||||
}
|
||||
content := ""
|
||||
if msg.Text.Content != "" {
|
||||
content = strings.TrimSpace(msg.Text.Content)
|
||||
}
|
||||
if content == "" && msg.Msgtype == "richText" {
|
||||
if cMap, ok := msg.Content.(map[string]interface{}); ok {
|
||||
if rich, ok := cMap["richText"].([]interface{}); ok {
|
||||
for _, c := range rich {
|
||||
if m, ok := c.(map[string]interface{}); ok {
|
||||
if txt, ok := m["text"].(string); ok {
|
||||
content = strings.TrimSpace(txt)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if content == "" {
|
||||
logger.Debug("钉钉消息内容为空，已忽略", zap.String("msgtype", msg.Msgtype))
|
||||
return
|
||||
}
|
||||
logger.Info("钉钉收到消息", zap.String("sender", msg.SenderId), zap.String("content", content))
|
||||
userID := msg.SenderId
|
||||
if userID == "" {
|
||||
userID = msg.ConversationId
|
||||
}
|
||||
reply := h.HandleMessage("dingtalk", userID, content)
|
||||
body := map[string]interface{}{
|
||||
"msgtype": "text",
|
||||
"text": map[string]string{"content": reply},
|
||||
}
|
||||
bodyBytes, _ := json.Marshal(body)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, msg.SessionWebhook, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
logger.Warn("钉钉构造回复请求失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
logger.Warn("钉钉回复请求失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
logger.Warn("钉钉回复非 200", zap.Int("status", resp.StatusCode))
|
||||
return
|
||||
}
|
||||
logger.Debug("钉钉回复成功", zap.String("content_preview", reply))
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package robot
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
|
||||
"github.com/larksuite/oapi-sdk-go/v3/event/dispatcher"
|
||||
larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1"
|
||||
lark "github.com/larksuite/oapi-sdk-go/v3"
|
||||
larkws "github.com/larksuite/oapi-sdk-go/v3/ws"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type larkTextContent struct {
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
// StartLark 启动飞书长连接（无需公网），收到消息后调用 handler 并回复
|
||||
func StartLark(cfg config.RobotLarkConfig, h MessageHandler, logger *zap.Logger) {
|
||||
if !cfg.Enabled || cfg.AppID == "" || cfg.AppSecret == "" {
|
||||
return
|
||||
}
|
||||
larkClient := lark.NewClient(cfg.AppID, cfg.AppSecret)
|
||||
eventHandler := dispatcher.NewEventDispatcher("", "").OnP2MessageReceiveV1(func(ctx context.Context, event *larkim.P2MessageReceiveV1) error {
|
||||
go handleLarkMessage(ctx, event, h, larkClient, logger)
|
||||
return nil
|
||||
})
|
||||
wsClient := larkws.NewClient(cfg.AppID, cfg.AppSecret,
|
||||
larkws.WithEventHandler(eventHandler),
|
||||
larkws.WithLogLevel(larkcore.LogLevelInfo),
|
||||
)
|
||||
go func() {
|
||||
err := wsClient.Start(context.Background())
|
||||
if err != nil {
|
||||
logger.Error("飞书长连接退出", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
logger.Info("飞书长连接已启动（无需公网）", zap.String("app_id", cfg.AppID))
|
||||
}
|
||||
|
||||
func handleLarkMessage(ctx context.Context, event *larkim.P2MessageReceiveV1, h MessageHandler, client *lark.Client, logger *zap.Logger) {
|
||||
if event == nil || event.Event == nil || event.Event.Message == nil || event.Event.Sender == nil || event.Event.Sender.SenderId == nil {
|
||||
return
|
||||
}
|
||||
msg := event.Event.Message
|
||||
msgType := larkcore.StringValue(msg.MessageType)
|
||||
if msgType != larkim.MsgTypeText {
|
||||
logger.Debug("飞书暂仅处理文本消息", zap.String("msg_type", msgType))
|
||||
return
|
||||
}
|
||||
var textBody larkTextContent
|
||||
if err := json.Unmarshal([]byte(larkcore.StringValue(msg.Content)), &textBody); err != nil {
|
||||
logger.Warn("飞书消息 Content 解析失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
text := strings.TrimSpace(textBody.Text)
|
||||
if text == "" {
|
||||
return
|
||||
}
|
||||
userID := ""
|
||||
if event.Event.Sender.SenderId.UserId != nil {
|
||||
userID = *event.Event.Sender.SenderId.UserId
|
||||
}
|
||||
messageID := larkcore.StringValue(msg.MessageId)
|
||||
reply := h.HandleMessage("lark", userID, text)
|
||||
contentBytes, _ := json.Marshal(larkTextContent{Text: reply})
|
||||
_, err := client.Im.Message.Reply(ctx, larkim.NewReplyMessageReqBuilder().
|
||||
MessageId(messageID).
|
||||
Body(larkim.NewReplyMessageReqBodyBuilder().
|
||||
MsgType(larkim.MsgTypeText).
|
||||
Content(string(contentBytes)).
|
||||
Build()).
|
||||
Build())
|
||||
if err != nil {
|
||||
logger.Warn("飞书回复失败", zap.String("message_id", messageID), zap.Error(err))
|
||||
return
|
||||
}
|
||||
logger.Debug("飞书已回复", zap.String("message_id", messageID))
|
||||
}
|
||||
Reference in New Issue
Block a user