mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-04-21 18:26:38 +02:00
Delete handler directory
This commit is contained in:
公明
-2579
File diff suppressed because it is too large
Load Diff
@@ -1,173 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/attackchain"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// AttackChainHandler 攻击链处理器
|
||||
type AttackChainHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
openAIConfig *config.OpenAIConfig
|
||||
mu sync.RWMutex // 保护 openAIConfig 的并发访问
|
||||
// 用于防止同一对话的并发生成
|
||||
generatingLocks sync.Map // map[string]*sync.Mutex
|
||||
}
|
||||
|
||||
// NewAttackChainHandler 创建新的攻击链处理器
|
||||
func NewAttackChainHandler(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.Logger) *AttackChainHandler {
|
||||
return &AttackChainHandler{
|
||||
db: db,
|
||||
logger: logger,
|
||||
openAIConfig: openAIConfig,
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateConfig 更新OpenAI配置
|
||||
func (h *AttackChainHandler) UpdateConfig(cfg *config.OpenAIConfig) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.openAIConfig = cfg
|
||||
h.logger.Info("AttackChainHandler配置已更新",
|
||||
zap.String("base_url", cfg.BaseURL),
|
||||
zap.String("model", cfg.Model),
|
||||
)
|
||||
}
|
||||
|
||||
// getOpenAIConfig 获取OpenAI配置（线程安全）
|
||||
func (h *AttackChainHandler) getOpenAIConfig() *config.OpenAIConfig {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return h.openAIConfig
|
||||
}
|
||||
|
||||
// GetAttackChain 获取攻击链（按需生成）
|
||||
// GET /api/attack-chain/:conversationId
|
||||
func (h *AttackChainHandler) GetAttackChain(c *gin.Context) {
|
||||
conversationID := c.Param("conversationId")
|
||||
if conversationID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"})
|
||||
return
|
||||
}
|
||||
|
||||
// 检查对话是否存在
|
||||
_, err := h.db.GetConversation(conversationID)
|
||||
if err != nil {
|
||||
h.logger.Warn("对话不存在", zap.String("conversationId", conversationID), zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
// 先尝试从数据库加载（如果已生成过）
|
||||
openAIConfig := h.getOpenAIConfig()
|
||||
builder := attackchain.NewBuilder(h.db, openAIConfig, h.logger)
|
||||
chain, err := builder.LoadChainFromDatabase(conversationID)
|
||||
if err == nil && len(chain.Nodes) > 0 {
|
||||
// 如果已存在，直接返回
|
||||
h.logger.Info("返回已存在的攻击链", zap.String("conversationId", conversationID))
|
||||
c.JSON(http.StatusOK, chain)
|
||||
return
|
||||
}
|
||||
|
||||
// 如果不存在，则生成新的攻击链（按需生成）
|
||||
// 使用锁机制防止同一对话的并发生成
|
||||
lockInterface, _ := h.generatingLocks.LoadOrStore(conversationID, &sync.Mutex{})
|
||||
lock := lockInterface.(*sync.Mutex)
|
||||
|
||||
// 尝试获取锁，如果正在生成则返回错误
|
||||
acquired := lock.TryLock()
|
||||
if !acquired {
|
||||
h.logger.Info("攻击链正在生成中，请稍后再试", zap.String("conversationId", conversationID))
|
||||
c.JSON(http.StatusConflict, gin.H{"error": "攻击链正在生成中，请稍后再试"})
|
||||
return
|
||||
}
|
||||
defer lock.Unlock()
|
||||
|
||||
// 再次检查是否已生成（可能在等待锁的过程中已经生成完成）
|
||||
chain, err = builder.LoadChainFromDatabase(conversationID)
|
||||
if err == nil && len(chain.Nodes) > 0 {
|
||||
h.logger.Info("返回已存在的攻击链（在锁等待期间已生成）", zap.String("conversationId", conversationID))
|
||||
c.JSON(http.StatusOK, chain)
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("开始生成攻击链", zap.String("conversationId", conversationID))
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
chain, err = builder.BuildChainFromConversation(ctx, conversationID)
|
||||
if err != nil {
|
||||
h.logger.Error("生成攻击链失败", zap.String("conversationId", conversationID), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "生成攻击链失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 生成完成后，从锁映射中删除（可选，保留也可以用于防止短时间内重复生成）
|
||||
// h.generatingLocks.Delete(conversationID)
|
||||
|
||||
c.JSON(http.StatusOK, chain)
|
||||
}
|
||||
|
||||
// RegenerateAttackChain 重新生成攻击链
|
||||
// POST /api/attack-chain/:conversationId/regenerate
|
||||
func (h *AttackChainHandler) RegenerateAttackChain(c *gin.Context) {
|
||||
conversationID := c.Param("conversationId")
|
||||
if conversationID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"})
|
||||
return
|
||||
}
|
||||
|
||||
// 检查对话是否存在
|
||||
_, err := h.db.GetConversation(conversationID)
|
||||
if err != nil {
|
||||
h.logger.Warn("对话不存在", zap.String("conversationId", conversationID), zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
// 删除旧的攻击链
|
||||
if err := h.db.DeleteAttackChain(conversationID); err != nil {
|
||||
h.logger.Warn("删除旧攻击链失败", zap.Error(err))
|
||||
}
|
||||
|
||||
// 使用锁机制防止并发生成
|
||||
lockInterface, _ := h.generatingLocks.LoadOrStore(conversationID, &sync.Mutex{})
|
||||
lock := lockInterface.(*sync.Mutex)
|
||||
|
||||
acquired := lock.TryLock()
|
||||
if !acquired {
|
||||
h.logger.Info("攻击链正在生成中，请稍后再试", zap.String("conversationId", conversationID))
|
||||
c.JSON(http.StatusConflict, gin.H{"error": "攻击链正在生成中，请稍后再试"})
|
||||
return
|
||||
}
|
||||
defer lock.Unlock()
|
||||
|
||||
// 生成新的攻击链
|
||||
h.logger.Info("重新生成攻击链", zap.String("conversationId", conversationID))
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
openAIConfig := h.getOpenAIConfig()
|
||||
builder := attackchain.NewBuilder(h.db, openAIConfig, h.logger)
|
||||
chain, err := builder.BuildChainFromConversation(ctx, conversationID)
|
||||
if err != nil {
|
||||
h.logger.Error("生成攻击链失败", zap.String("conversationId", conversationID), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "生成攻击链失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, chain)
|
||||
}
|
||||
|
||||
-156
@@ -1,156 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// AuthHandler handles authentication-related endpoints.
|
||||
type AuthHandler struct {
|
||||
manager *security.AuthManager
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewAuthHandler creates a new AuthHandler.
|
||||
func NewAuthHandler(manager *security.AuthManager, cfg *config.Config, configPath string, logger *zap.Logger) *AuthHandler {
|
||||
return &AuthHandler{
|
||||
manager: manager,
|
||||
config: cfg,
|
||||
configPath: configPath,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
type loginRequest struct {
|
||||
Password string `json:"password" binding:"required"`
|
||||
}
|
||||
|
||||
type changePasswordRequest struct {
|
||||
OldPassword string `json:"oldPassword"`
|
||||
NewPassword string `json:"newPassword"`
|
||||
}
|
||||
|
||||
// Login verifies password and returns a session token.
|
||||
func (h *AuthHandler) Login(c *gin.Context) {
|
||||
var req loginRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "密码不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
token, expiresAt, err := h.manager.Authenticate(req.Password)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "密码错误"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"token": token,
|
||||
"expires_at": expiresAt.UTC().Format(time.RFC3339),
|
||||
"session_duration_hr": h.manager.SessionDurationHours(),
|
||||
})
|
||||
}
|
||||
|
||||
// Logout revokes the current session token.
|
||||
func (h *AuthHandler) Logout(c *gin.Context) {
|
||||
token := c.GetString(security.ContextAuthTokenKey)
|
||||
if token == "" {
|
||||
authHeader := c.GetHeader("Authorization")
|
||||
if len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "Bearer ") {
|
||||
token = strings.TrimSpace(authHeader[7:])
|
||||
} else {
|
||||
token = strings.TrimSpace(authHeader)
|
||||
}
|
||||
}
|
||||
|
||||
h.manager.RevokeToken(token)
|
||||
c.JSON(http.StatusOK, gin.H{"message": "已退出登录"})
|
||||
}
|
||||
|
||||
// ChangePassword updates the login password.
|
||||
func (h *AuthHandler) ChangePassword(c *gin.Context) {
|
||||
var req changePasswordRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "参数无效"})
|
||||
return
|
||||
}
|
||||
|
||||
oldPassword := strings.TrimSpace(req.OldPassword)
|
||||
newPassword := strings.TrimSpace(req.NewPassword)
|
||||
|
||||
if oldPassword == "" || newPassword == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "当前密码和新密码均不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
if len(newPassword) < 8 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "新密码长度至少需要 8 位"})
|
||||
return
|
||||
}
|
||||
|
||||
if oldPassword == newPassword {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "新密码不能与旧密码相同"})
|
||||
return
|
||||
}
|
||||
|
||||
if !h.manager.CheckPassword(oldPassword) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "当前密码不正确"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := config.PersistAuthPassword(h.configPath, newPassword); err != nil {
|
||||
if h.logger != nil {
|
||||
h.logger.Error("保存新密码失败", zap.Error(err))
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存新密码失败，请重试"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.manager.UpdateConfig(newPassword, h.config.Auth.SessionDurationHours); err != nil {
|
||||
if h.logger != nil {
|
||||
h.logger.Error("更新认证配置失败", zap.Error(err))
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "更新认证配置失败"})
|
||||
return
|
||||
}
|
||||
|
||||
h.config.Auth.Password = newPassword
|
||||
h.config.Auth.GeneratedPassword = ""
|
||||
h.config.Auth.GeneratedPasswordPersisted = false
|
||||
h.config.Auth.GeneratedPasswordPersistErr = ""
|
||||
|
||||
if h.logger != nil {
|
||||
h.logger.Info("登录密码已更新，所有会话已失效")
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "密码已更新，请使用新密码重新登录"})
|
||||
}
|
||||
|
||||
// Validate returns the current session status.
|
||||
func (h *AuthHandler) Validate(c *gin.Context) {
|
||||
token := c.GetString(security.ContextAuthTokenKey)
|
||||
if token == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "会话无效"})
|
||||
return
|
||||
}
|
||||
|
||||
session, ok := h.manager.ValidateToken(token)
|
||||
if !ok {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "会话已过期"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"token": session.Token,
|
||||
"expires_at": session.ExpiresAt.UTC().Format(time.RFC3339),
|
||||
})
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,813 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// RegisterBatchTaskMCPTools 注册批量任务队列相关 MCP 工具（需传入已初始化 DB 的 AgentHandler）
|
||||
func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *zap.Logger) {
|
||||
if mcpServer == nil || h == nil || logger == nil {
|
||||
return
|
||||
}
|
||||
|
||||
reg := func(tool mcp.Tool, fn func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error)) {
|
||||
mcpServer.RegisterTool(tool, fn)
|
||||
}
|
||||
|
||||
// --- list ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskList,
|
||||
Description: "列出批量任务队列（精简摘要，省上下文）。含队列元数据、子任务 id/status/截断后的 message、各状态计数。完整子任务（含 result/error/conversationId/时间等）请用 batch_task_get(queue_id)。",
|
||||
ShortDescription: "列出批量任务队列",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"status": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "筛选状态：all（默认）、pending、running、paused、completed、cancelled",
|
||||
"enum": []string{"all", "pending", "running", "paused", "completed", "cancelled"},
|
||||
},
|
||||
"keyword": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "按队列 ID 或标题模糊搜索",
|
||||
},
|
||||
"page": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "页码，从 1 开始，默认 1",
|
||||
},
|
||||
"page_size": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "每页条数，默认 20，最大 100",
|
||||
},
|
||||
},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
status := mcpArgString(args, "status")
|
||||
if status == "" {
|
||||
status = "all"
|
||||
}
|
||||
keyword := mcpArgString(args, "keyword")
|
||||
page := int(mcpArgFloat(args, "page"))
|
||||
if page <= 0 {
|
||||
page = 1
|
||||
}
|
||||
pageSize := int(mcpArgFloat(args, "page_size"))
|
||||
if pageSize <= 0 {
|
||||
pageSize = 20
|
||||
}
|
||||
if pageSize > 100 {
|
||||
pageSize = 100
|
||||
}
|
||||
offset := (page - 1) * pageSize
|
||||
if offset > 100000 {
|
||||
offset = 100000
|
||||
}
|
||||
queues, total, err := h.batchTaskManager.ListQueues(pageSize, offset, status, keyword)
|
||||
if err != nil {
|
||||
return batchMCPTextResult(fmt.Sprintf("列出队列失败: %v", err), true), nil
|
||||
}
|
||||
totalPages := (total + pageSize - 1) / pageSize
|
||||
if totalPages == 0 {
|
||||
totalPages = 1
|
||||
}
|
||||
slim := make([]batchTaskQueueMCPListItem, 0, len(queues))
|
||||
for _, q := range queues {
|
||||
if q == nil {
|
||||
continue
|
||||
}
|
||||
slim = append(slim, toBatchTaskQueueMCPListItem(q))
|
||||
}
|
||||
payload := map[string]interface{}{
|
||||
"queues": slim,
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": pageSize,
|
||||
"total_pages": totalPages,
|
||||
}
|
||||
logger.Info("MCP batch_task_list", zap.String("status", status), zap.Int("total", total))
|
||||
return batchMCPJSONResult(payload)
|
||||
})
|
||||
|
||||
// --- get ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskGet,
|
||||
Description: "根据 queue_id 获取单个批量任务队列详情（含子任务列表、Cron、调度开关与最近错误信息）。",
|
||||
ShortDescription: "获取批量任务队列详情",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
queue, ok := h.batchTaskManager.GetBatchQueue(qid)
|
||||
if !ok {
|
||||
return batchMCPTextResult("队列不存在: "+qid, true), nil
|
||||
}
|
||||
return batchMCPJSONResult(queue)
|
||||
})
|
||||
|
||||
// --- create ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskCreate,
|
||||
Description: `创建新的批量任务队列。任务列表使用 tasks（字符串数组）或 tasks_text（多行，每行一条）。
|
||||
agent_mode: single（默认）或 multi（需系统启用多代理）。schedule_mode: manual（默认）或 cron；为 cron 时必须提供 cron_expr（如 "0 */6 * * *"）。
|
||||
默认创建后不会立即执行。可通过 execute_now=true 在创建后立即启动；也可后续调用 batch_task_start 手工启动。Cron 队列若需按表达式自动触发下一轮，还需保持调度开关开启（可用 batch_task_schedule_enabled）。`,
|
||||
ShortDescription: "创建批量任务队列（可选立即执行）",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"title": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "可选标题",
|
||||
},
|
||||
"role": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "角色名称，空表示默认",
|
||||
},
|
||||
"tasks": map[string]interface{}{
|
||||
"type": "array",
|
||||
"description": "任务指令列表，每项一条",
|
||||
"items": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
"tasks_text": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "多行文本，每行一条任务（与 tasks 二选一）",
|
||||
},
|
||||
"agent_mode": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "single 或 multi",
|
||||
"enum": []string{"single", "multi"},
|
||||
},
|
||||
"schedule_mode": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "manual 或 cron",
|
||||
"enum": []string{"manual", "cron"},
|
||||
},
|
||||
"cron_expr": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "schedule_mode 为 cron 时必填。标准 5 段格式：分钟 小时 日 月 星期，例如 \"0 */6 * * *\"（每6小时）、\"30 2 * * 1-5\"（工作日凌晨2:30）",
|
||||
},
|
||||
"execute_now": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "是否创建后立即执行，默认 false",
|
||||
},
|
||||
},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
tasks, errMsg := batchMCPTasksFromArgs(args)
|
||||
if errMsg != "" {
|
||||
return batchMCPTextResult(errMsg, true), nil
|
||||
}
|
||||
title := mcpArgString(args, "title")
|
||||
role := mcpArgString(args, "role")
|
||||
agentMode := normalizeBatchQueueAgentMode(mcpArgString(args, "agent_mode"))
|
||||
scheduleMode := normalizeBatchQueueScheduleMode(mcpArgString(args, "schedule_mode"))
|
||||
cronExpr := strings.TrimSpace(mcpArgString(args, "cron_expr"))
|
||||
var nextRunAt *time.Time
|
||||
if scheduleMode == "cron" {
|
||||
if cronExpr == "" {
|
||||
return batchMCPTextResult("Cron 调度模式下 cron_expr 不能为空", true), nil
|
||||
}
|
||||
sch, err := h.batchCronParser.Parse(cronExpr)
|
||||
if err != nil {
|
||||
return batchMCPTextResult("无效的 Cron 表达式: "+err.Error(), true), nil
|
||||
}
|
||||
n := sch.Next(time.Now())
|
||||
nextRunAt = &n
|
||||
}
|
||||
executeNow, ok := mcpArgBool(args, "execute_now")
|
||||
if !ok {
|
||||
executeNow = false
|
||||
}
|
||||
queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, nextRunAt, tasks)
|
||||
if createErr != nil {
|
||||
return batchMCPTextResult("创建队列失败: "+createErr.Error(), true), nil
|
||||
}
|
||||
started := false
|
||||
if executeNow {
|
||||
ok, err := h.startBatchQueueExecution(queue.ID, false)
|
||||
if !ok {
|
||||
return batchMCPTextResult("队列不存在: "+queue.ID, true), nil
|
||||
}
|
||||
if err != nil {
|
||||
return batchMCPTextResult("创建成功但启动失败: "+err.Error(), true), nil
|
||||
}
|
||||
started = true
|
||||
if refreshed, exists := h.batchTaskManager.GetBatchQueue(queue.ID); exists {
|
||||
queue = refreshed
|
||||
}
|
||||
}
|
||||
logger.Info("MCP batch_task_create", zap.String("queueId", queue.ID), zap.Int("taskCount", len(tasks)))
|
||||
return batchMCPJSONResult(map[string]interface{}{
|
||||
"queue_id": queue.ID,
|
||||
"queue": queue,
|
||||
"started": started,
|
||||
"execute_now": executeNow,
|
||||
"reminder": func() string {
|
||||
if started {
|
||||
return "队列已创建并立即启动。"
|
||||
}
|
||||
return "队列已创建，当前为 pending。需要开始执行时请调用 MCP 工具 batch_task_start（queue_id 同上）。Cron 自动调度需 schedule_enabled 为 true，可用 batch_task_schedule_enabled。"
|
||||
}(),
|
||||
})
|
||||
})
|
||||
|
||||
// --- start ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskStart,
|
||||
Description: `启动或继续执行批量任务队列（pending / paused）。
|
||||
与 batch_task_create 配合使用：仅创建队列不会自动执行，需调用本工具才会开始跑子任务。`,
|
||||
ShortDescription: "启动/继续批量任务队列（创建后需调用才会执行）",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
ok, err := h.startBatchQueueExecution(qid, false)
|
||||
if !ok {
|
||||
return batchMCPTextResult("队列不存在: "+qid, true), nil
|
||||
}
|
||||
if err != nil {
|
||||
return batchMCPTextResult("启动失败: "+err.Error(), true), nil
|
||||
}
|
||||
logger.Info("MCP batch_task_start", zap.String("queueId", qid))
|
||||
return batchMCPTextResult("已提交启动，队列将开始执行。", false), nil
|
||||
})
|
||||
|
||||
// --- rerun (reset + start for completed/cancelled queues) ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskRerun,
|
||||
Description: "重跑已完成或已取消的批量任务队列。会重置所有子任务状态后重新执行一轮。",
|
||||
ShortDescription: "重跑批量任务队列",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
queue, exists := h.batchTaskManager.GetBatchQueue(qid)
|
||||
if !exists {
|
||||
return batchMCPTextResult("队列不存在: "+qid, true), nil
|
||||
}
|
||||
if queue.Status != "completed" && queue.Status != "cancelled" {
|
||||
return batchMCPTextResult("仅已完成或已取消的队列可以重跑，当前状态: "+queue.Status, true), nil
|
||||
}
|
||||
if !h.batchTaskManager.ResetQueueForRerun(qid) {
|
||||
return batchMCPTextResult("重置队列失败", true), nil
|
||||
}
|
||||
ok, err := h.startBatchQueueExecution(qid, false)
|
||||
if !ok {
|
||||
return batchMCPTextResult("启动失败", true), nil
|
||||
}
|
||||
if err != nil {
|
||||
return batchMCPTextResult("启动失败: "+err.Error(), true), nil
|
||||
}
|
||||
logger.Info("MCP batch_task_rerun", zap.String("queueId", qid))
|
||||
return batchMCPTextResult("已重置并重新启动队列。", false), nil
|
||||
})
|
||||
|
||||
// --- pause ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskPause,
|
||||
Description: "暂停正在运行的批量任务队列（当前子任务会被取消）。",
|
||||
ShortDescription: "暂停批量任务队列",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
if !h.batchTaskManager.PauseQueue(qid) {
|
||||
return batchMCPTextResult("无法暂停：队列不存在或当前非 running 状态", true), nil
|
||||
}
|
||||
logger.Info("MCP batch_task_pause", zap.String("queueId", qid))
|
||||
return batchMCPTextResult("队列已暂停。", false), nil
|
||||
})
|
||||
|
||||
// --- delete queue ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskDelete,
|
||||
Description: "删除批量任务队列及其子任务记录。",
|
||||
ShortDescription: "删除批量任务队列",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
if !h.batchTaskManager.DeleteQueue(qid) {
|
||||
return batchMCPTextResult("删除失败：队列不存在", true), nil
|
||||
}
|
||||
logger.Info("MCP batch_task_delete", zap.String("queueId", qid))
|
||||
return batchMCPTextResult("队列已删除。", false), nil
|
||||
})
|
||||
|
||||
// --- update metadata (title/role/agentMode) ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskUpdateMetadata,
|
||||
Description: "修改批量任务队列的标题、角色和代理模式。仅在队列非 running 状态下可修改。",
|
||||
ShortDescription: "修改批量任务队列标题/角色/代理模式",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
"title": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "新标题（空字符串清除标题）",
|
||||
},
|
||||
"role": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "新角色名（空字符串使用默认角色）",
|
||||
},
|
||||
"agent_mode": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "代理模式：single（单代理 ReAct）或 multi（多代理）",
|
||||
"enum": []string{"single", "multi"},
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
title := mcpArgString(args, "title")
|
||||
role := mcpArgString(args, "role")
|
||||
agentMode := mcpArgString(args, "agent_mode")
|
||||
if err := h.batchTaskManager.UpdateQueueMetadata(qid, title, role, agentMode); err != nil {
|
||||
return batchMCPTextResult(err.Error(), true), nil
|
||||
}
|
||||
updated, _ := h.batchTaskManager.GetBatchQueue(qid)
|
||||
logger.Info("MCP batch_task_update_metadata", zap.String("queueId", qid))
|
||||
return batchMCPJSONResult(updated)
|
||||
})
|
||||
|
||||
// --- update schedule ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskUpdateSchedule,
|
||||
Description: `修改批量任务队列的调度方式和 Cron 表达式。仅在队列非 running 状态下可修改。
|
||||
schedule_mode 为 cron 时必须提供有效 cron_expr；为 manual 时会清除 Cron 配置。`,
|
||||
ShortDescription: "修改批量任务调度配置（Cron 表达式）",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
"schedule_mode": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "manual 或 cron",
|
||||
"enum": []string{"manual", "cron"},
|
||||
},
|
||||
"cron_expr": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Cron 表达式（schedule_mode 为 cron 时必填）。标准 5 段格式：分钟 小时 日 月 星期，如 \"0 */6 * * *\"（每6小时）、\"30 2 * * 1-5\"（工作日凌晨2:30）",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id", "schedule_mode"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
queue, exists := h.batchTaskManager.GetBatchQueue(qid)
|
||||
if !exists {
|
||||
return batchMCPTextResult("队列不存在: "+qid, true), nil
|
||||
}
|
||||
if queue.Status == "running" {
|
||||
return batchMCPTextResult("队列正在运行中，无法修改调度配置", true), nil
|
||||
}
|
||||
scheduleMode := normalizeBatchQueueScheduleMode(mcpArgString(args, "schedule_mode"))
|
||||
cronExpr := strings.TrimSpace(mcpArgString(args, "cron_expr"))
|
||||
var nextRunAt *time.Time
|
||||
if scheduleMode == "cron" {
|
||||
if cronExpr == "" {
|
||||
return batchMCPTextResult("Cron 调度模式下 cron_expr 不能为空", true), nil
|
||||
}
|
||||
sch, err := h.batchCronParser.Parse(cronExpr)
|
||||
if err != nil {
|
||||
return batchMCPTextResult("无效的 Cron 表达式: "+err.Error(), true), nil
|
||||
}
|
||||
n := sch.Next(time.Now())
|
||||
nextRunAt = &n
|
||||
}
|
||||
h.batchTaskManager.UpdateQueueSchedule(qid, scheduleMode, cronExpr, nextRunAt)
|
||||
updated, _ := h.batchTaskManager.GetBatchQueue(qid)
|
||||
logger.Info("MCP batch_task_update_schedule", zap.String("queueId", qid), zap.String("scheduleMode", scheduleMode), zap.String("cronExpr", cronExpr))
|
||||
return batchMCPJSONResult(updated)
|
||||
})
|
||||
|
||||
// --- schedule enabled ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskScheduleEnabled,
|
||||
Description: `设置是否允许 Cron 自动触发该队列。关闭后仍保留 Cron 表达式，仅停止定时自动跑；可用手工「启动」执行。
|
||||
仅对 schedule_mode 为 cron 的队列有意义。`,
|
||||
ShortDescription: "开关批量任务 Cron 自动调度",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
"schedule_enabled": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "true 允许定时触发，false 仅手工执行",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id", "schedule_enabled"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
en, ok := mcpArgBool(args, "schedule_enabled")
|
||||
if !ok {
|
||||
return batchMCPTextResult("schedule_enabled 必须为布尔值", true), nil
|
||||
}
|
||||
if _, exists := h.batchTaskManager.GetBatchQueue(qid); !exists {
|
||||
return batchMCPTextResult("队列不存在", true), nil
|
||||
}
|
||||
if !h.batchTaskManager.SetScheduleEnabled(qid, en) {
|
||||
return batchMCPTextResult("更新失败", true), nil
|
||||
}
|
||||
queue, _ := h.batchTaskManager.GetBatchQueue(qid)
|
||||
logger.Info("MCP batch_task_schedule_enabled", zap.String("queueId", qid), zap.Bool("enabled", en))
|
||||
return batchMCPJSONResult(queue)
|
||||
})
|
||||
|
||||
// --- add task ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskAdd,
|
||||
Description: "向处于 pending 状态的队列追加一条子任务。",
|
||||
ShortDescription: "批量队列添加子任务",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
"message": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "任务指令内容",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id", "message"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
msg := strings.TrimSpace(mcpArgString(args, "message"))
|
||||
if qid == "" || msg == "" {
|
||||
return batchMCPTextResult("queue_id 与 message 均不能为空", true), nil
|
||||
}
|
||||
task, err := h.batchTaskManager.AddTaskToQueue(qid, msg)
|
||||
if err != nil {
|
||||
return batchMCPTextResult(err.Error(), true), nil
|
||||
}
|
||||
queue, _ := h.batchTaskManager.GetBatchQueue(qid)
|
||||
logger.Info("MCP batch_task_add_task", zap.String("queueId", qid), zap.String("taskId", task.ID))
|
||||
return batchMCPJSONResult(map[string]interface{}{"task": task, "queue": queue})
|
||||
})
|
||||
|
||||
// --- update task ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskUpdate,
|
||||
Description: "修改 pending 队列中仍为 pending 的子任务文案。",
|
||||
ShortDescription: "更新批量子任务内容",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
"task_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "子任务 ID",
|
||||
},
|
||||
"message": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "新的任务指令",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id", "task_id", "message"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
tid := mcpArgString(args, "task_id")
|
||||
msg := strings.TrimSpace(mcpArgString(args, "message"))
|
||||
if qid == "" || tid == "" || msg == "" {
|
||||
return batchMCPTextResult("queue_id、task_id、message 均不能为空", true), nil
|
||||
}
|
||||
if err := h.batchTaskManager.UpdateTaskMessage(qid, tid, msg); err != nil {
|
||||
return batchMCPTextResult(err.Error(), true), nil
|
||||
}
|
||||
queue, _ := h.batchTaskManager.GetBatchQueue(qid)
|
||||
logger.Info("MCP batch_task_update_task", zap.String("queueId", qid), zap.String("taskId", tid))
|
||||
return batchMCPJSONResult(queue)
|
||||
})
|
||||
|
||||
// --- remove task ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskRemove,
|
||||
Description: "从 pending 队列中删除仍为 pending 的子任务。",
|
||||
ShortDescription: "删除批量子任务",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
"task_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "子任务 ID",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id", "task_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
tid := mcpArgString(args, "task_id")
|
||||
if qid == "" || tid == "" {
|
||||
return batchMCPTextResult("queue_id 与 task_id 均不能为空", true), nil
|
||||
}
|
||||
if err := h.batchTaskManager.DeleteTask(qid, tid); err != nil {
|
||||
return batchMCPTextResult(err.Error(), true), nil
|
||||
}
|
||||
queue, _ := h.batchTaskManager.GetBatchQueue(qid)
|
||||
logger.Info("MCP batch_task_remove_task", zap.String("queueId", qid), zap.String("taskId", tid))
|
||||
return batchMCPJSONResult(queue)
|
||||
})
|
||||
|
||||
logger.Info("批量任务 MCP 工具已注册", zap.Int("count", 12))
|
||||
}
|
||||
|
||||
// --- batch_task_list 精简结构（避免把每条子任务的 result 等大段文本塞进列表上下文） ---
|
||||
|
||||
const mcpBatchListTaskMessageMaxRunes = 160
|
||||
|
||||
// batchTaskMCPListSummary 列表中的子任务摘要（完整字段用 batch_task_get）
|
||||
type batchTaskMCPListSummary struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Message string `json:"message,omitempty"`
|
||||
}
|
||||
|
||||
// batchTaskQueueMCPListItem 列表中的队列摘要
|
||||
type batchTaskQueueMCPListItem struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Role string `json:"role,omitempty"`
|
||||
AgentMode string `json:"agentMode"`
|
||||
ScheduleMode string `json:"scheduleMode"`
|
||||
CronExpr string `json:"cronExpr,omitempty"`
|
||||
NextRunAt *time.Time `json:"nextRunAt,omitempty"`
|
||||
ScheduleEnabled bool `json:"scheduleEnabled"`
|
||||
LastScheduleTriggerAt *time.Time `json:"lastScheduleTriggerAt,omitempty"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
StartedAt *time.Time `json:"startedAt,omitempty"`
|
||||
CompletedAt *time.Time `json:"completedAt,omitempty"`
|
||||
CurrentIndex int `json:"currentIndex"`
|
||||
TaskTotal int `json:"task_total"`
|
||||
TaskCounts map[string]int `json:"task_counts"`
|
||||
Tasks []batchTaskMCPListSummary `json:"tasks"`
|
||||
}
|
||||
|
||||
func truncateStringRunes(s string, maxRunes int) string {
|
||||
if maxRunes <= 0 {
|
||||
return ""
|
||||
}
|
||||
n := 0
|
||||
for i := range s {
|
||||
if n == maxRunes {
|
||||
out := strings.TrimSpace(s[:i])
|
||||
if out == "" {
|
||||
return "…"
|
||||
}
|
||||
return out + "…"
|
||||
}
|
||||
n++
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
const mcpBatchListMaxTasksPerQueue = 200 // 列表中每个队列最多返回的子任务摘要数
|
||||
|
||||
func toBatchTaskQueueMCPListItem(q *BatchTaskQueue) batchTaskQueueMCPListItem {
|
||||
counts := map[string]int{
|
||||
"pending": 0,
|
||||
"running": 0,
|
||||
"completed": 0,
|
||||
"failed": 0,
|
||||
"cancelled": 0,
|
||||
}
|
||||
tasks := make([]batchTaskMCPListSummary, 0, len(q.Tasks))
|
||||
for _, t := range q.Tasks {
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
counts[t.Status]++
|
||||
// 列表视图限制子任务摘要数量，完整列表通过 batch_task_get 查看
|
||||
if len(tasks) < mcpBatchListMaxTasksPerQueue {
|
||||
tasks = append(tasks, batchTaskMCPListSummary{
|
||||
ID: t.ID,
|
||||
Status: t.Status,
|
||||
Message: truncateStringRunes(t.Message, mcpBatchListTaskMessageMaxRunes),
|
||||
})
|
||||
}
|
||||
}
|
||||
return batchTaskQueueMCPListItem{
|
||||
ID: q.ID,
|
||||
Title: q.Title,
|
||||
Role: q.Role,
|
||||
AgentMode: q.AgentMode,
|
||||
ScheduleMode: q.ScheduleMode,
|
||||
CronExpr: q.CronExpr,
|
||||
NextRunAt: q.NextRunAt,
|
||||
ScheduleEnabled: q.ScheduleEnabled,
|
||||
LastScheduleTriggerAt: q.LastScheduleTriggerAt,
|
||||
Status: q.Status,
|
||||
CreatedAt: q.CreatedAt,
|
||||
StartedAt: q.StartedAt,
|
||||
CompletedAt: q.CompletedAt,
|
||||
CurrentIndex: q.CurrentIndex,
|
||||
TaskTotal: len(tasks),
|
||||
TaskCounts: counts,
|
||||
Tasks: tasks,
|
||||
}
|
||||
}
|
||||
|
||||
func batchMCPTextResult(text string, isErr bool) *mcp.ToolResult {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: text}},
|
||||
IsError: isErr,
|
||||
}
|
||||
}
|
||||
|
||||
func batchMCPJSONResult(v interface{}) (*mcp.ToolResult, error) {
|
||||
b, err := json.MarshalIndent(v, "", " ")
|
||||
if err != nil {
|
||||
return batchMCPTextResult(fmt.Sprintf("JSON 编码失败: %v", err), true), nil
|
||||
}
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: string(b)}}}, nil
|
||||
}
|
||||
|
||||
func batchMCPTasksFromArgs(args map[string]interface{}) ([]string, string) {
|
||||
if raw, ok := args["tasks"]; ok && raw != nil {
|
||||
switch t := raw.(type) {
|
||||
case []interface{}:
|
||||
out := make([]string, 0, len(t))
|
||||
for _, x := range t {
|
||||
if s, ok := x.(string); ok {
|
||||
if tr := strings.TrimSpace(s); tr != "" {
|
||||
out = append(out, tr)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(out) > 0 {
|
||||
return out, ""
|
||||
}
|
||||
}
|
||||
}
|
||||
if txt := mcpArgString(args, "tasks_text"); txt != "" {
|
||||
lines := strings.Split(txt, "\n")
|
||||
out := make([]string, 0, len(lines))
|
||||
for _, line := range lines {
|
||||
if tr := strings.TrimSpace(line); tr != "" {
|
||||
out = append(out, tr)
|
||||
}
|
||||
}
|
||||
if len(out) > 0 {
|
||||
return out, ""
|
||||
}
|
||||
}
|
||||
return nil, "需要提供 tasks（字符串数组）或 tasks_text（多行文本，每行一条任务）"
|
||||
}
|
||||
|
||||
func mcpArgString(args map[string]interface{}, key string) string {
|
||||
v, ok := args[key]
|
||||
if !ok || v == nil {
|
||||
return ""
|
||||
}
|
||||
switch t := v.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(t)
|
||||
case float64:
|
||||
return strings.TrimSpace(strconv.FormatFloat(t, 'f', -1, 64))
|
||||
case json.Number:
|
||||
return strings.TrimSpace(t.String())
|
||||
default:
|
||||
return strings.TrimSpace(fmt.Sprint(t))
|
||||
}
|
||||
}
|
||||
|
||||
func mcpArgFloat(args map[string]interface{}, key string) float64 {
|
||||
v, ok := args[key]
|
||||
if !ok || v == nil {
|
||||
return 0
|
||||
}
|
||||
switch t := v.(type) {
|
||||
case float64:
|
||||
return t
|
||||
case int:
|
||||
return float64(t)
|
||||
case int64:
|
||||
return float64(t)
|
||||
case json.Number:
|
||||
f, _ := t.Float64()
|
||||
return f
|
||||
case string:
|
||||
f, _ := strconv.ParseFloat(strings.TrimSpace(t), 64)
|
||||
return f
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func mcpArgBool(args map[string]interface{}, key string) (val bool, ok bool) {
|
||||
v, exists := args[key]
|
||||
if !exists {
|
||||
return false, false
|
||||
}
|
||||
switch t := v.(type) {
|
||||
case bool:
|
||||
return t, true
|
||||
case string:
|
||||
s := strings.ToLower(strings.TrimSpace(t))
|
||||
if s == "true" || s == "1" || s == "yes" {
|
||||
return true, true
|
||||
}
|
||||
if s == "false" || s == "0" || s == "no" {
|
||||
return false, true
|
||||
}
|
||||
case float64:
|
||||
return t != 0, true
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
@@ -1,512 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
chatUploadsRootDirName = "chat_uploads"
|
||||
maxChatUploadEditBytes = 2 * 1024 * 1024 // 文本编辑上限
|
||||
)
|
||||
|
||||
// ChatUploadsHandler 对话中上传附件（chat_uploads 目录）的管理 API
|
||||
type ChatUploadsHandler struct {
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewChatUploadsHandler 创建处理器
|
||||
func NewChatUploadsHandler(logger *zap.Logger) *ChatUploadsHandler {
|
||||
return &ChatUploadsHandler{logger: logger}
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) absRoot() (string, error) {
|
||||
cwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Abs(filepath.Join(cwd, chatUploadsRootDirName))
|
||||
}
|
||||
|
||||
// resolveUnderChatUploads 校验 relativePath（使用 / 分隔）对应文件必须在 chat_uploads 根下
|
||||
func (h *ChatUploadsHandler) resolveUnderChatUploads(relativePath string) (abs string, err error) {
|
||||
root, err := h.absRoot()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
rel := strings.TrimSpace(relativePath)
|
||||
if rel == "" {
|
||||
return "", fmt.Errorf("empty path")
|
||||
}
|
||||
rel = filepath.Clean(filepath.FromSlash(rel))
|
||||
if rel == "." || strings.HasPrefix(rel, "..") {
|
||||
return "", fmt.Errorf("invalid path")
|
||||
}
|
||||
full := filepath.Join(root, rel)
|
||||
full, err = filepath.Abs(full)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
rootAbs, _ := filepath.Abs(root)
|
||||
if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) {
|
||||
return "", fmt.Errorf("path escapes chat_uploads root")
|
||||
}
|
||||
return full, nil
|
||||
}
|
||||
|
||||
// ChatUploadFileItem 列表项
|
||||
type ChatUploadFileItem struct {
|
||||
RelativePath string `json:"relativePath"`
|
||||
AbsolutePath string `json:"absolutePath"` // 服务器上的绝对路径，便于在对话中引用（与附件落盘路径一致）
|
||||
Name string `json:"name"`
|
||||
Size int64 `json:"size"`
|
||||
ModifiedUnix int64 `json:"modifiedUnix"`
|
||||
Date string `json:"date"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
// SubPath 为日期、会话目录之下的子路径（不含文件名），如 date/conv/a/b/file 则为 "a/b"；无嵌套则为 ""。
|
||||
SubPath string `json:"subPath"`
|
||||
}
|
||||
|
||||
// List GET /api/chat-uploads
|
||||
func (h *ChatUploadsHandler) List(c *gin.Context) {
|
||||
conversationFilter := strings.TrimSpace(c.Query("conversation"))
|
||||
root, err := h.absRoot()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
// 保证根目录存在，否则「按文件夹」浏览时无法 mkdir，且首次列表为空时界面无路径工具栏
|
||||
if err := os.MkdirAll(root, 0755); err != nil {
|
||||
h.logger.Warn("创建 chat_uploads 根目录失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
var files []ChatUploadFileItem
|
||||
var folders []string
|
||||
err = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil {
|
||||
return walkErr
|
||||
}
|
||||
rel, err := filepath.Rel(root, path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rel == "." {
|
||||
return nil
|
||||
}
|
||||
relSlash := filepath.ToSlash(rel)
|
||||
if d.IsDir() {
|
||||
folders = append(folders, relSlash)
|
||||
return nil
|
||||
}
|
||||
info, err := d.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
parts := strings.Split(relSlash, "/")
|
||||
var dateStr, convID string
|
||||
if len(parts) >= 2 {
|
||||
dateStr = parts[0]
|
||||
}
|
||||
if len(parts) >= 3 {
|
||||
convID = parts[1]
|
||||
}
|
||||
var subPath string
|
||||
if len(parts) >= 4 {
|
||||
subPath = strings.Join(parts[2:len(parts)-1], "/")
|
||||
}
|
||||
if conversationFilter != "" && convID != conversationFilter {
|
||||
return nil
|
||||
}
|
||||
absPath, _ := filepath.Abs(path)
|
||||
files = append(files, ChatUploadFileItem{
|
||||
RelativePath: relSlash,
|
||||
AbsolutePath: absPath,
|
||||
Name: d.Name(),
|
||||
Size: info.Size(),
|
||||
ModifiedUnix: info.ModTime().Unix(),
|
||||
Date: dateStr,
|
||||
ConversationID: convID,
|
||||
SubPath: subPath,
|
||||
})
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
h.logger.Warn("列举对话附件失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if conversationFilter != "" {
|
||||
filteredFolders := make([]string, 0, len(folders))
|
||||
for _, rel := range folders {
|
||||
parts := strings.Split(rel, "/")
|
||||
if len(parts) >= 2 && parts[1] == conversationFilter {
|
||||
filteredFolders = append(filteredFolders, rel)
|
||||
continue
|
||||
}
|
||||
if len(parts) == 1 {
|
||||
prefix := rel + "/"
|
||||
for _, f := range files {
|
||||
if strings.HasPrefix(f.RelativePath, prefix) {
|
||||
filteredFolders = append(filteredFolders, rel)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
folders = filteredFolders
|
||||
}
|
||||
sort.Strings(folders)
|
||||
sort.Slice(files, func(i, j int) bool {
|
||||
return files[i].ModifiedUnix > files[j].ModifiedUnix
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{"files": files, "folders": folders})
|
||||
}
|
||||
|
||||
// Download GET /api/chat-uploads/download?path=...
|
||||
func (h *ChatUploadsHandler) Download(c *gin.Context) {
|
||||
p := c.Query("path")
|
||||
abs, err := h.resolveUnderChatUploads(p)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
st, err := os.Stat(abs)
|
||||
if err != nil || st.IsDir() {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
|
||||
return
|
||||
}
|
||||
c.FileAttachment(abs, filepath.Base(abs))
|
||||
}
|
||||
|
||||
type chatUploadPathBody struct {
|
||||
Path string `json:"path"`
|
||||
}
|
||||
|
||||
// Delete DELETE /api/chat-uploads
|
||||
func (h *ChatUploadsHandler) Delete(c *gin.Context) {
|
||||
var body chatUploadPathBody
|
||||
if err := c.ShouldBindJSON(&body); err != nil || strings.TrimSpace(body.Path) == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
abs, err := h.resolveUnderChatUploads(body.Path)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
st, err := os.Stat(abs)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if st.IsDir() {
|
||||
if err := os.RemoveAll(abs); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
} else {
|
||||
if err := os.Remove(abs); err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
type chatUploadMkdirBody struct {
|
||||
Parent string `json:"parent"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// Mkdir POST /api/chat-uploads/mkdir — 在 parent 目录下新建子目录（parent 为 chat_uploads 下相对路径，空表示根目录；name 为单段目录名）
|
||||
func (h *ChatUploadsHandler) Mkdir(c *gin.Context) {
|
||||
var body chatUploadMkdirBody
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
name := strings.TrimSpace(body.Name)
|
||||
if name == "" || strings.ContainsAny(name, `/\`) || name == "." || name == ".." {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid name"})
|
||||
return
|
||||
}
|
||||
if utf8.RuneCountInString(name) > 200 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name too long"})
|
||||
return
|
||||
}
|
||||
|
||||
parent := strings.TrimSpace(body.Parent)
|
||||
parent = filepath.ToSlash(filepath.Clean(filepath.FromSlash(parent)))
|
||||
parent = strings.Trim(parent, "/")
|
||||
if parent == "." {
|
||||
parent = ""
|
||||
}
|
||||
|
||||
root, err := h.absRoot()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if parent != "" {
|
||||
absParent, err := h.resolveUnderChatUploads(parent)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
st, err := os.Stat(absParent)
|
||||
if err != nil || !st.IsDir() {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "parent not found"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
var rel string
|
||||
if parent == "" {
|
||||
rel = name
|
||||
} else {
|
||||
rel = parent + "/" + name
|
||||
}
|
||||
absNew, err := h.resolveUnderChatUploads(rel)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if _, err := os.Stat(absNew); err == nil {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": "already exists"})
|
||||
return
|
||||
}
|
||||
if err := os.Mkdir(absNew, 0755); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
relOut, _ := filepath.Rel(root, absNew)
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "relativePath": filepath.ToSlash(relOut)})
|
||||
}
|
||||
|
||||
type chatUploadRenameBody struct {
|
||||
Path string `json:"path"`
|
||||
NewName string `json:"newName"`
|
||||
}
|
||||
|
||||
// Rename PUT /api/chat-uploads/rename
|
||||
func (h *ChatUploadsHandler) Rename(c *gin.Context) {
|
||||
var body chatUploadRenameBody
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
newName := strings.TrimSpace(body.NewName)
|
||||
if newName == "" || strings.ContainsAny(newName, `/\`) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid newName"})
|
||||
return
|
||||
}
|
||||
abs, err := h.resolveUnderChatUploads(body.Path)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
dir := filepath.Dir(abs)
|
||||
newAbs := filepath.Join(dir, filepath.Base(newName))
|
||||
root, _ := h.absRoot()
|
||||
newAbs, _ = filepath.Abs(newAbs)
|
||||
if newAbs != root && !strings.HasPrefix(newAbs, root+string(filepath.Separator)) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid target path"})
|
||||
return
|
||||
}
|
||||
if err := os.Rename(abs, newAbs); err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
newRel, _ := filepath.Rel(root, newAbs)
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "relativePath": filepath.ToSlash(newRel)})
|
||||
}
|
||||
|
||||
type chatUploadContentBody struct {
|
||||
Path string `json:"path"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// GetContent GET /api/chat-uploads/content?path=...
|
||||
func (h *ChatUploadsHandler) GetContent(c *gin.Context) {
|
||||
p := c.Query("path")
|
||||
abs, err := h.resolveUnderChatUploads(p)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
st, err := os.Stat(abs)
|
||||
if err != nil || st.IsDir() {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
|
||||
return
|
||||
}
|
||||
if st.Size() > maxChatUploadEditBytes {
|
||||
c.JSON(http.StatusRequestEntityTooLarge, gin.H{"error": "file too large for editor"})
|
||||
return
|
||||
}
|
||||
b, err := os.ReadFile(abs)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if !utf8.Valid(b) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "binary file not editable in UI"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"content": string(b)})
|
||||
}
|
||||
|
||||
// PutContent PUT /api/chat-uploads/content
|
||||
func (h *ChatUploadsHandler) PutContent(c *gin.Context) {
|
||||
var body chatUploadContentBody
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
if !utf8.ValidString(body.Content) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "content must be valid UTF-8"})
|
||||
return
|
||||
}
|
||||
if len(body.Content) > maxChatUploadEditBytes {
|
||||
c.JSON(http.StatusRequestEntityTooLarge, gin.H{"error": "content too large"})
|
||||
return
|
||||
}
|
||||
abs, err := h.resolveUnderChatUploads(body.Path)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := os.WriteFile(abs, []byte(body.Content), 0644); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func chatUploadShortRand(n int) string {
|
||||
const letters = "0123456789abcdef"
|
||||
b := make([]byte, n)
|
||||
_, _ = rand.Read(b)
|
||||
for i := range b {
|
||||
b[i] = letters[int(b[i])%len(letters)]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// Upload POST /api/chat-uploads multipart: file；conversationId 可选；relativeDir 可选（chat_uploads 下目录的相对路径，将文件直接上传至该目录）
|
||||
func (h *ChatUploadsHandler) Upload(c *gin.Context) {
|
||||
fh, err := c.FormFile("file")
|
||||
if err != nil || fh == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "missing file"})
|
||||
return
|
||||
}
|
||||
root, err := h.absRoot()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
var targetDir string
|
||||
targetRel := strings.TrimSpace(c.PostForm("relativeDir"))
|
||||
if targetRel != "" {
|
||||
absDir, err := h.resolveUnderChatUploads(targetRel)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
st, err := os.Stat(absDir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
if err := os.MkdirAll(absDir, 0755); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
} else {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
} else if !st.IsDir() {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "relativeDir is not a directory"})
|
||||
return
|
||||
}
|
||||
targetDir = absDir
|
||||
} else {
|
||||
convID := strings.TrimSpace(c.PostForm("conversationId"))
|
||||
convDir := convID
|
||||
if convDir == "" {
|
||||
convDir = "_manual"
|
||||
} else {
|
||||
convDir = strings.ReplaceAll(convDir, string(filepath.Separator), "_")
|
||||
}
|
||||
dateStr := time.Now().Format("2006-01-02")
|
||||
targetDir = filepath.Join(root, dateStr, convDir)
|
||||
if err := os.MkdirAll(targetDir, 0755); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
}
|
||||
baseName := filepath.Base(fh.Filename)
|
||||
if baseName == "" || baseName == "." {
|
||||
baseName = "file"
|
||||
}
|
||||
baseName = strings.ReplaceAll(baseName, string(filepath.Separator), "_")
|
||||
ext := filepath.Ext(baseName)
|
||||
nameNoExt := strings.TrimSuffix(baseName, ext)
|
||||
suffix := fmt.Sprintf("_%s_%s", time.Now().Format("150405"), chatUploadShortRand(6))
|
||||
var unique string
|
||||
if ext != "" {
|
||||
unique = nameNoExt + suffix + ext
|
||||
} else {
|
||||
unique = baseName + suffix
|
||||
}
|
||||
fullPath := filepath.Join(targetDir, unique)
|
||||
src, err := fh.Open()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
defer src.Close()
|
||||
dst, err := os.Create(fullPath)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
defer dst.Close()
|
||||
if _, err := io.Copy(dst, src); err != nil {
|
||||
_ = os.Remove(fullPath)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
rel, _ := filepath.Rel(root, fullPath)
|
||||
absSaved, _ := filepath.Abs(fullPath)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"relativePath": filepath.ToSlash(rel),
|
||||
"absolutePath": absSaved,
|
||||
"name": unique,
|
||||
})
|
||||
}
|
||||
-1594
File diff suppressed because it is too large
Load Diff
@@ -1,233 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// ConversationHandler 对话处理器
|
||||
type ConversationHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewConversationHandler 创建新的对话处理器
|
||||
func NewConversationHandler(db *database.DB, logger *zap.Logger) *ConversationHandler {
|
||||
return &ConversationHandler{
|
||||
db: db,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateConversationRequest 创建对话请求
|
||||
type CreateConversationRequest struct {
|
||||
Title string `json:"title"`
|
||||
}
|
||||
|
||||
// CreateConversation 创建新对话
|
||||
func (h *ConversationHandler) CreateConversation(c *gin.Context) {
|
||||
var req CreateConversationRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
title := req.Title
|
||||
if title == "" {
|
||||
title = "新对话"
|
||||
}
|
||||
|
||||
conv, err := h.db.CreateConversation(title)
|
||||
if err != nil {
|
||||
h.logger.Error("创建对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, conv)
|
||||
}
|
||||
|
||||
// ListConversations 列出对话
|
||||
func (h *ConversationHandler) ListConversations(c *gin.Context) {
|
||||
limitStr := c.DefaultQuery("limit", "50")
|
||||
offsetStr := c.DefaultQuery("offset", "0")
|
||||
search := c.Query("search") // 获取搜索参数
|
||||
|
||||
limit, _ := strconv.Atoi(limitStr)
|
||||
offset, _ := strconv.Atoi(offsetStr)
|
||||
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 50
|
||||
}
|
||||
|
||||
conversations, err := h.db.ListConversations(limit, offset, search)
|
||||
if err != nil {
|
||||
h.logger.Error("获取对话列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, conversations)
|
||||
}
|
||||
|
||||
// GetConversation 获取对话
|
||||
func (h *ConversationHandler) GetConversation(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
// 默认轻量加载，只有用户需要展开详情时再按需拉取
|
||||
// include_process_details=1/true 时返回全量 processDetails（兼容旧行为）
|
||||
includeStr := c.DefaultQuery("include_process_details", "0")
|
||||
include := includeStr == "1" || includeStr == "true" || includeStr == "yes"
|
||||
|
||||
var (
|
||||
conv *database.Conversation
|
||||
err error
|
||||
)
|
||||
if include {
|
||||
conv, err = h.db.GetConversation(id)
|
||||
} else {
|
||||
conv, err = h.db.GetConversationLite(id)
|
||||
}
|
||||
if err != nil {
|
||||
h.logger.Error("获取对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, conv)
|
||||
}
|
||||
|
||||
// GetMessageProcessDetails 获取指定消息的过程详情（按需加载）
|
||||
func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) {
|
||||
messageID := c.Param("id")
|
||||
if messageID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "message id required"})
|
||||
return
|
||||
}
|
||||
|
||||
details, err := h.db.GetProcessDetails(messageID)
|
||||
if err != nil {
|
||||
h.logger.Error("获取过程详情失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 转换为前端期望的 JSON 结构（与 GetConversation 中 processDetails 结构一致）
|
||||
out := make([]map[string]interface{}, 0, len(details))
|
||||
for _, d := range details {
|
||||
var data interface{}
|
||||
if d.Data != "" {
|
||||
if err := json.Unmarshal([]byte(d.Data), &data); err != nil {
|
||||
h.logger.Warn("解析过程详情数据失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
out = append(out, map[string]interface{}{
|
||||
"id": d.ID,
|
||||
"messageId": d.MessageID,
|
||||
"conversationId": d.ConversationID,
|
||||
"eventType": d.EventType,
|
||||
"message": d.Message,
|
||||
"data": data,
|
||||
"createdAt": d.CreatedAt,
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"processDetails": out})
|
||||
}
|
||||
|
||||
// UpdateConversationRequest 更新对话请求
|
||||
type UpdateConversationRequest struct {
|
||||
Title string `json:"title"`
|
||||
}
|
||||
|
||||
// UpdateConversation 更新对话
|
||||
func (h *ConversationHandler) UpdateConversation(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
var req UpdateConversationRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if req.Title == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "标题不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateConversationTitle(id, req.Title); err != nil {
|
||||
h.logger.Error("更新对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 返回更新后的对话
|
||||
conv, err := h.db.GetConversation(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取更新后的对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, conv)
|
||||
}
|
||||
|
||||
// DeleteConversation 删除对话
|
||||
func (h *ConversationHandler) DeleteConversation(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
if err := h.db.DeleteConversation(id); err != nil {
|
||||
h.logger.Error("删除对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
// DeleteTurnRequest 删除一轮对话（POST /api/conversations/:id/delete-turn）
|
||||
type DeleteTurnRequest struct {
|
||||
MessageID string `json:"messageId"`
|
||||
}
|
||||
|
||||
// DeleteConversationTurn 删除锚点消息所在轮次（从该轮 user 到下一轮 user 之前），并清空 last_react_*。
|
||||
func (h *ConversationHandler) DeleteConversationTurn(c *gin.Context) {
|
||||
conversationID := c.Param("id")
|
||||
if conversationID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "conversation id required"})
|
||||
return
|
||||
}
|
||||
|
||||
var req DeleteTurnRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.MessageID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "messageId required"})
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := h.db.GetConversation(conversationID); err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
deletedIDs, err := h.db.DeleteConversationTurn(conversationID, req.MessageID)
|
||||
if err != nil {
|
||||
h.logger.Warn("删除对话轮次失败",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("messageId", req.MessageID),
|
||||
zap.Error(err),
|
||||
)
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"deletedMessageIds": deletedIDs,
|
||||
"message": "ok",
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,542 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// ExternalMCPHandler 外部MCP处理器
|
||||
type ExternalMCPHandler struct {
|
||||
manager *mcp.ExternalMCPManager
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewExternalMCPHandler 创建外部MCP处理器
|
||||
func NewExternalMCPHandler(manager *mcp.ExternalMCPManager, cfg *config.Config, configPath string, logger *zap.Logger) *ExternalMCPHandler {
|
||||
return &ExternalMCPHandler{
|
||||
manager: manager,
|
||||
config: cfg,
|
||||
configPath: configPath,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// GetExternalMCPs 获取所有外部MCP配置
|
||||
func (h *ExternalMCPHandler) GetExternalMCPs(c *gin.Context) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
|
||||
configs := h.manager.GetConfigs()
|
||||
|
||||
// 获取所有外部MCP的工具数量
|
||||
toolCounts := h.manager.GetToolCounts()
|
||||
|
||||
// 转换为响应格式
|
||||
result := make(map[string]ExternalMCPResponse)
|
||||
for name, cfg := range configs {
|
||||
client, exists := h.manager.GetClient(name)
|
||||
status := "disconnected"
|
||||
if exists {
|
||||
status = client.GetStatus()
|
||||
} else if h.isEnabled(cfg) {
|
||||
status = "disconnected"
|
||||
} else {
|
||||
status = "disabled"
|
||||
}
|
||||
|
||||
toolCount := toolCounts[name]
|
||||
errorMsg := ""
|
||||
if status == "error" {
|
||||
errorMsg = h.manager.GetError(name)
|
||||
}
|
||||
|
||||
result[name] = ExternalMCPResponse{
|
||||
Config: cfg,
|
||||
Status: status,
|
||||
ToolCount: toolCount,
|
||||
Error: errorMsg,
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"servers": result,
|
||||
"stats": h.manager.GetStats(),
|
||||
})
|
||||
}
|
||||
|
||||
// GetExternalMCP 获取单个外部MCP配置
|
||||
func (h *ExternalMCPHandler) GetExternalMCP(c *gin.Context) {
|
||||
name := c.Param("name")
|
||||
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
|
||||
configs := h.manager.GetConfigs()
|
||||
cfg, exists := configs[name]
|
||||
if !exists {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "外部MCP配置不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
client, clientExists := h.manager.GetClient(name)
|
||||
status := "disconnected"
|
||||
if clientExists {
|
||||
status = client.GetStatus()
|
||||
} else if h.isEnabled(cfg) {
|
||||
status = "disconnected"
|
||||
} else {
|
||||
status = "disabled"
|
||||
}
|
||||
|
||||
// 获取工具数量
|
||||
toolCount := 0
|
||||
if clientExists && client.IsConnected() {
|
||||
if count, err := h.manager.GetToolCount(name); err == nil {
|
||||
toolCount = count
|
||||
}
|
||||
}
|
||||
|
||||
// 获取错误信息
|
||||
errorMsg := ""
|
||||
if status == "error" {
|
||||
errorMsg = h.manager.GetError(name)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, ExternalMCPResponse{
|
||||
Config: cfg,
|
||||
Status: status,
|
||||
ToolCount: toolCount,
|
||||
Error: errorMsg,
|
||||
})
|
||||
}
|
||||
|
||||
// AddOrUpdateExternalMCP 添加或更新外部MCP配置
|
||||
func (h *ExternalMCPHandler) AddOrUpdateExternalMCP(c *gin.Context) {
|
||||
var req AddOrUpdateExternalMCPRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
name := c.Param("name")
|
||||
if name == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
// 验证配置
|
||||
if err := h.validateConfig(req.Config); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
// 添加或更新配置
|
||||
if err := h.manager.AddOrUpdateConfig(name, req.Config); err != nil {
|
||||
h.logger.Error("添加或更新外部MCP配置失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "添加或更新配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 更新内存中的配置
|
||||
if h.config.ExternalMCP.Servers == nil {
|
||||
h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig)
|
||||
}
|
||||
|
||||
// 如果用户提供了 disabled 或 enabled 字段，保留它们以保持向后兼容
|
||||
// 同时将值迁移到 external_mcp_enable
|
||||
cfg := req.Config
|
||||
|
||||
if req.Config.Disabled {
|
||||
// 用户设置了 disabled: true
|
||||
cfg.ExternalMCPEnable = false
|
||||
cfg.Disabled = true
|
||||
cfg.Enabled = false
|
||||
} else if req.Config.Enabled {
|
||||
// 用户设置了 enabled: true
|
||||
cfg.ExternalMCPEnable = true
|
||||
cfg.Enabled = true
|
||||
cfg.Disabled = false
|
||||
} else if !req.Config.ExternalMCPEnable {
|
||||
// 用户没有设置任何字段，且 external_mcp_enable 为 false
|
||||
// 检查现有配置是否有旧字段
|
||||
if existingCfg, exists := h.config.ExternalMCP.Servers[name]; exists {
|
||||
// 保留现有的旧字段
|
||||
cfg.Enabled = existingCfg.Enabled
|
||||
cfg.Disabled = existingCfg.Disabled
|
||||
}
|
||||
} else {
|
||||
// 用户通过新字段启用了（external_mcp_enable: true），但没有设置旧字段
|
||||
// 为了向后兼容，我们设置 enabled: true
|
||||
// 这样即使原始配置中有 disabled: false，也会被转换为 enabled: true
|
||||
cfg.Enabled = true
|
||||
cfg.Disabled = false
|
||||
}
|
||||
|
||||
h.config.ExternalMCP.Servers[name] = cfg
|
||||
|
||||
// 保存到配置文件
|
||||
if err := h.saveConfig(); err != nil {
|
||||
h.logger.Error("保存配置失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("外部MCP配置已更新", zap.String("name", name))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "配置已更新"})
|
||||
}
|
||||
|
||||
// DeleteExternalMCP 删除外部MCP配置
|
||||
func (h *ExternalMCPHandler) DeleteExternalMCP(c *gin.Context) {
|
||||
name := c.Param("name")
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
// 移除配置
|
||||
if err := h.manager.RemoveConfig(name); err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "配置不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
// 从内存配置中删除
|
||||
if h.config.ExternalMCP.Servers != nil {
|
||||
delete(h.config.ExternalMCP.Servers, name)
|
||||
}
|
||||
|
||||
// 保存到配置文件
|
||||
if err := h.saveConfig(); err != nil {
|
||||
h.logger.Error("保存配置失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("外部MCP配置已删除", zap.String("name", name))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "配置已删除"})
|
||||
}
|
||||
|
||||
// StartExternalMCP 启动外部MCP
|
||||
func (h *ExternalMCPHandler) StartExternalMCP(c *gin.Context) {
|
||||
name := c.Param("name")
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
// 更新配置为启用
|
||||
if h.config.ExternalMCP.Servers == nil {
|
||||
h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig)
|
||||
}
|
||||
cfg := h.config.ExternalMCP.Servers[name]
|
||||
cfg.ExternalMCPEnable = true
|
||||
h.config.ExternalMCP.Servers[name] = cfg
|
||||
|
||||
// 保存到配置文件
|
||||
if err := h.saveConfig(); err != nil {
|
||||
h.logger.Error("保存配置失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 启动客户端（立即创建客户端并设置状态为connecting，实际连接在后台进行）
|
||||
h.logger.Info("开始启动外部MCP", zap.String("name", name))
|
||||
if err := h.manager.StartClient(name); err != nil {
|
||||
h.logger.Error("启动外部MCP失败", zap.String("name", name), zap.Error(err))
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"error": err.Error(),
|
||||
"status": "error",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取客户端状态（应该是connecting）
|
||||
client, exists := h.manager.GetClient(name)
|
||||
status := "connecting"
|
||||
if exists {
|
||||
status = client.GetStatus()
|
||||
}
|
||||
|
||||
// 立即返回，不等待连接完成
|
||||
// 客户端会在后台异步连接，用户可以通过状态查询接口查看连接状态
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "外部MCP启动请求已提交，正在后台连接中",
|
||||
"status": status,
|
||||
})
|
||||
}
|
||||
|
||||
// StopExternalMCP 停止外部MCP
|
||||
func (h *ExternalMCPHandler) StopExternalMCP(c *gin.Context) {
|
||||
name := c.Param("name")
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
// 停止客户端
|
||||
if err := h.manager.StopClient(name); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 更新配置
|
||||
if h.config.ExternalMCP.Servers == nil {
|
||||
h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig)
|
||||
}
|
||||
cfg := h.config.ExternalMCP.Servers[name]
|
||||
cfg.ExternalMCPEnable = false
|
||||
h.config.ExternalMCP.Servers[name] = cfg
|
||||
|
||||
// 保存到配置文件
|
||||
if err := h.saveConfig(); err != nil {
|
||||
h.logger.Error("保存配置失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("外部MCP已停止", zap.String("name", name))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "外部MCP已停止"})
|
||||
}
|
||||
|
||||
// GetExternalMCPStats 获取统计信息
|
||||
func (h *ExternalMCPHandler) GetExternalMCPStats(c *gin.Context) {
|
||||
stats := h.manager.GetStats()
|
||||
c.JSON(http.StatusOK, stats)
|
||||
}
|
||||
|
||||
// validateConfig 验证配置
|
||||
func (h *ExternalMCPHandler) validateConfig(cfg config.ExternalMCPServerConfig) error {
|
||||
transport := cfg.Transport
|
||||
if transport == "" {
|
||||
// 如果没有指定transport，根据是否有command或url判断
|
||||
if cfg.Command != "" {
|
||||
transport = "stdio"
|
||||
} else if cfg.URL != "" {
|
||||
transport = "http"
|
||||
} else {
|
||||
return fmt.Errorf("需要指定command（stdio模式）或url（http/sse模式）")
|
||||
}
|
||||
}
|
||||
|
||||
switch transport {
|
||||
case "http":
|
||||
if cfg.URL == "" {
|
||||
return fmt.Errorf("HTTP模式需要URL")
|
||||
}
|
||||
case "stdio":
|
||||
if cfg.Command == "" {
|
||||
return fmt.Errorf("stdio模式需要command")
|
||||
}
|
||||
case "sse":
|
||||
if cfg.URL == "" {
|
||||
return fmt.Errorf("SSE模式需要URL")
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("不支持的传输模式: %s，支持的模式: http, stdio, sse", transport)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// isEnabled 检查是否启用
|
||||
func (h *ExternalMCPHandler) isEnabled(cfg config.ExternalMCPServerConfig) bool {
|
||||
// 优先使用 ExternalMCPEnable 字段
|
||||
// 如果没有设置，检查旧的 enabled/disabled 字段（向后兼容）
|
||||
if cfg.ExternalMCPEnable {
|
||||
return true
|
||||
}
|
||||
// 向后兼容：检查旧字段
|
||||
if cfg.Disabled {
|
||||
return false
|
||||
}
|
||||
if cfg.Enabled {
|
||||
return true
|
||||
}
|
||||
// 都没有设置，默认为启用
|
||||
return true
|
||||
}
|
||||
|
||||
// saveConfig 保存配置到文件
|
||||
func (h *ExternalMCPHandler) saveConfig() error {
|
||||
// 读取现有配置文件并创建备份
|
||||
data, err := os.ReadFile(h.configPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("读取配置文件失败: %w", err)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(h.configPath+".backup", data, 0644); err != nil {
|
||||
h.logger.Warn("创建配置备份失败", zap.Error(err))
|
||||
}
|
||||
|
||||
root, err := loadYAMLDocument(h.configPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("解析配置文件失败: %w", err)
|
||||
}
|
||||
|
||||
// 在更新前，读取原始配置中的 enabled/disabled 字段，以便保持向后兼容
|
||||
originalConfigs := make(map[string]map[string]bool)
|
||||
externalMCPNode := findMapValue(root.Content[0], "external_mcp")
|
||||
if externalMCPNode != nil && externalMCPNode.Kind == yaml.MappingNode {
|
||||
serversNode := findMapValue(externalMCPNode, "servers")
|
||||
if serversNode != nil && serversNode.Kind == yaml.MappingNode {
|
||||
// 遍历现有的服务器配置，保存 enabled/disabled 字段
|
||||
for i := 0; i < len(serversNode.Content); i += 2 {
|
||||
if i+1 >= len(serversNode.Content) {
|
||||
break
|
||||
}
|
||||
nameNode := serversNode.Content[i]
|
||||
serverNode := serversNode.Content[i+1]
|
||||
if nameNode.Kind == yaml.ScalarNode && serverNode.Kind == yaml.MappingNode {
|
||||
serverName := nameNode.Value
|
||||
originalConfigs[serverName] = make(map[string]bool)
|
||||
// 检查是否有 enabled 字段
|
||||
if enabledVal := findBoolInMap(serverNode, "enabled"); enabledVal != nil {
|
||||
originalConfigs[serverName]["enabled"] = *enabledVal
|
||||
}
|
||||
// 检查是否有 disabled 字段
|
||||
if disabledVal := findBoolInMap(serverNode, "disabled"); disabledVal != nil {
|
||||
originalConfigs[serverName]["disabled"] = *disabledVal
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 更新外部MCP配置
|
||||
updateExternalMCPConfig(root, h.config.ExternalMCP, originalConfigs)
|
||||
|
||||
if err := writeYAMLDocument(h.configPath, root); err != nil {
|
||||
return fmt.Errorf("保存配置文件失败: %w", err)
|
||||
}
|
||||
|
||||
h.logger.Info("配置已保存", zap.String("path", h.configPath))
|
||||
return nil
|
||||
}
|
||||
|
||||
// updateExternalMCPConfig 更新外部MCP配置
|
||||
func updateExternalMCPConfig(doc *yaml.Node, cfg config.ExternalMCPConfig, originalConfigs map[string]map[string]bool) {
|
||||
root := doc.Content[0]
|
||||
externalMCPNode := ensureMap(root, "external_mcp")
|
||||
serversNode := ensureMap(externalMCPNode, "servers")
|
||||
|
||||
// 清空现有服务器配置
|
||||
serversNode.Content = nil
|
||||
|
||||
// 添加新的服务器配置
|
||||
for name, serverCfg := range cfg.Servers {
|
||||
// 添加服务器名称键
|
||||
nameNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: name}
|
||||
serverNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
|
||||
serversNode.Content = append(serversNode.Content, nameNode, serverNode)
|
||||
|
||||
// 设置服务器配置字段
|
||||
if serverCfg.Command != "" {
|
||||
setStringInMap(serverNode, "command", serverCfg.Command)
|
||||
}
|
||||
if len(serverCfg.Args) > 0 {
|
||||
setStringArrayInMap(serverNode, "args", serverCfg.Args)
|
||||
}
|
||||
// 保存 env 字段（环境变量）
|
||||
if serverCfg.Env != nil && len(serverCfg.Env) > 0 {
|
||||
envNode := ensureMap(serverNode, "env")
|
||||
for envKey, envValue := range serverCfg.Env {
|
||||
setStringInMap(envNode, envKey, envValue)
|
||||
}
|
||||
}
|
||||
if serverCfg.Transport != "" {
|
||||
setStringInMap(serverNode, "transport", serverCfg.Transport)
|
||||
}
|
||||
if serverCfg.URL != "" {
|
||||
setStringInMap(serverNode, "url", serverCfg.URL)
|
||||
}
|
||||
// 保存 headers 字段（HTTP/SSE 请求头）
|
||||
if serverCfg.Headers != nil && len(serverCfg.Headers) > 0 {
|
||||
headersNode := ensureMap(serverNode, "headers")
|
||||
for k, v := range serverCfg.Headers {
|
||||
setStringInMap(headersNode, k, v)
|
||||
}
|
||||
}
|
||||
if serverCfg.Description != "" {
|
||||
setStringInMap(serverNode, "description", serverCfg.Description)
|
||||
}
|
||||
if serverCfg.Timeout > 0 {
|
||||
setIntInMap(serverNode, "timeout", serverCfg.Timeout)
|
||||
}
|
||||
// 保存 external_mcp_enable 字段（新字段）
|
||||
setBoolInMap(serverNode, "external_mcp_enable", serverCfg.ExternalMCPEnable)
|
||||
// 保存 tool_enabled 字段（每个工具的启用状态）
|
||||
if serverCfg.ToolEnabled != nil && len(serverCfg.ToolEnabled) > 0 {
|
||||
toolEnabledNode := ensureMap(serverNode, "tool_enabled")
|
||||
for toolName, enabled := range serverCfg.ToolEnabled {
|
||||
setBoolInMap(toolEnabledNode, toolName, enabled)
|
||||
}
|
||||
}
|
||||
// 保留旧的 enabled/disabled 字段以保持向后兼容
|
||||
originalFields, hasOriginal := originalConfigs[name]
|
||||
|
||||
// 如果原始配置中有 enabled 字段，保留它
|
||||
if hasOriginal {
|
||||
if enabledVal, hasEnabled := originalFields["enabled"]; hasEnabled {
|
||||
setBoolInMap(serverNode, "enabled", enabledVal)
|
||||
}
|
||||
// 如果原始配置中有 disabled 字段，保留它
|
||||
// 注意：由于 omitempty，disabled: false 不会被保存，但 disabled: true 会被保存
|
||||
if disabledVal, hasDisabled := originalFields["disabled"]; hasDisabled {
|
||||
if disabledVal {
|
||||
setBoolInMap(serverNode, "disabled", disabledVal)
|
||||
} else {
|
||||
// 如果原始配置中有 disabled: false，我们保存 enabled: true 来等效表示
|
||||
// 因为 disabled: false 等价于 enabled: true
|
||||
setBoolInMap(serverNode, "enabled", true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 如果用户在当前请求中明确设置了这些字段，也保存它们
|
||||
if serverCfg.Enabled {
|
||||
setBoolInMap(serverNode, "enabled", serverCfg.Enabled)
|
||||
}
|
||||
if serverCfg.Disabled {
|
||||
setBoolInMap(serverNode, "disabled", serverCfg.Disabled)
|
||||
} else if !hasOriginal && serverCfg.ExternalMCPEnable {
|
||||
// 如果用户通过新字段启用了，且原始配置中没有旧字段，保存 enabled: true 以保持向后兼容
|
||||
setBoolInMap(serverNode, "enabled", true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// setStringArrayInMap 设置字符串数组
|
||||
func setStringArrayInMap(mapNode *yaml.Node, key string, values []string) {
|
||||
_, valueNode := ensureKeyValue(mapNode, key)
|
||||
valueNode.Kind = yaml.SequenceNode
|
||||
valueNode.Tag = "!!seq"
|
||||
valueNode.Content = nil
|
||||
for _, v := range values {
|
||||
itemNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: v}
|
||||
valueNode.Content = append(valueNode.Content, itemNode)
|
||||
}
|
||||
}
|
||||
|
||||
// AddOrUpdateExternalMCPRequest 添加或更新外部MCP请求
|
||||
type AddOrUpdateExternalMCPRequest struct {
|
||||
Config config.ExternalMCPServerConfig `json:"config"`
|
||||
}
|
||||
|
||||
// ExternalMCPResponse 外部MCP响应
|
||||
type ExternalMCPResponse struct {
|
||||
Config config.ExternalMCPServerConfig `json:"config"`
|
||||
Status string `json:"status"` // "connected", "disconnected", "disabled", "error", "connecting"
|
||||
ToolCount int `json:"tool_count"` // 工具数量
|
||||
Error string `json:"error,omitempty"` // 错误信息（仅在status为error时存在）
|
||||
}
|
||||
@@ -1,518 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func setupTestRouter() (*gin.Engine, *ExternalMCPHandler, string) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
|
||||
// 创建临时配置文件
|
||||
tmpFile, err := os.CreateTemp("", "test-config-*.yaml")
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
tmpFile.WriteString("server:\n host: 0.0.0.0\n port: 8080\n")
|
||||
tmpFile.Close()
|
||||
configPath := tmpFile.Name()
|
||||
|
||||
logger := zap.NewNop()
|
||||
manager := mcp.NewExternalMCPManager(logger)
|
||||
cfg := &config.Config{
|
||||
ExternalMCP: config.ExternalMCPConfig{
|
||||
Servers: make(map[string]config.ExternalMCPServerConfig),
|
||||
},
|
||||
}
|
||||
|
||||
handler := NewExternalMCPHandler(manager, cfg, configPath, logger)
|
||||
|
||||
api := router.Group("/api")
|
||||
api.GET("/external-mcp", handler.GetExternalMCPs)
|
||||
api.GET("/external-mcp/stats", handler.GetExternalMCPStats)
|
||||
api.GET("/external-mcp/:name", handler.GetExternalMCP)
|
||||
api.PUT("/external-mcp/:name", handler.AddOrUpdateExternalMCP)
|
||||
api.DELETE("/external-mcp/:name", handler.DeleteExternalMCP)
|
||||
api.POST("/external-mcp/:name/start", handler.StartExternalMCP)
|
||||
api.POST("/external-mcp/:name/stop", handler.StopExternalMCP)
|
||||
|
||||
return router, handler, configPath
|
||||
}
|
||||
|
||||
func cleanupTestConfig(configPath string) {
|
||||
os.Remove(configPath)
|
||||
os.Remove(configPath + ".backup")
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_AddOrUpdateExternalMCP_Stdio(t *testing.T) {
|
||||
router, _, configPath := setupTestRouter()
|
||||
defer cleanupTestConfig(configPath)
|
||||
|
||||
// 测试添加stdio模式的配置
|
||||
configJSON := `{
|
||||
"command": "python3",
|
||||
"args": ["/path/to/script.py", "--server", "http://example.com"],
|
||||
"description": "Test stdio MCP",
|
||||
"timeout": 300,
|
||||
"enabled": true
|
||||
}`
|
||||
|
||||
var configObj config.ExternalMCPServerConfig
|
||||
if err := json.Unmarshal([]byte(configJSON), &configObj); err != nil {
|
||||
t.Fatalf("解析配置JSON失败: %v", err)
|
||||
}
|
||||
|
||||
reqBody := AddOrUpdateExternalMCPRequest{
|
||||
Config: configObj,
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(reqBody)
|
||||
req := httptest.NewRequest("PUT", "/api/external-mcp/test-stdio", bytes.NewBuffer(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// 验证配置已添加
|
||||
req2 := httptest.NewRequest("GET", "/api/external-mcp/test-stdio", nil)
|
||||
w2 := httptest.NewRecorder()
|
||||
router.ServeHTTP(w2, req2)
|
||||
|
||||
if w2.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String())
|
||||
}
|
||||
|
||||
var response ExternalMCPResponse
|
||||
if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if response.Config.Command != "python3" {
|
||||
t.Errorf("期望command为python3，实际%s", response.Config.Command)
|
||||
}
|
||||
if len(response.Config.Args) != 3 {
|
||||
t.Errorf("期望args长度为3，实际%d", len(response.Config.Args))
|
||||
}
|
||||
if response.Config.Description != "Test stdio MCP" {
|
||||
t.Errorf("期望description为'Test stdio MCP'，实际%s", response.Config.Description)
|
||||
}
|
||||
if response.Config.Timeout != 300 {
|
||||
t.Errorf("期望timeout为300，实际%d", response.Config.Timeout)
|
||||
}
|
||||
if !response.Config.Enabled {
|
||||
t.Error("期望enabled为true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_AddOrUpdateExternalMCP_HTTP(t *testing.T) {
|
||||
router, _, configPath := setupTestRouter()
|
||||
defer cleanupTestConfig(configPath)
|
||||
|
||||
// 测试添加HTTP模式的配置
|
||||
configJSON := `{
|
||||
"transport": "http",
|
||||
"url": "http://127.0.0.1:8081/mcp",
|
||||
"enabled": true
|
||||
}`
|
||||
|
||||
var configObj config.ExternalMCPServerConfig
|
||||
if err := json.Unmarshal([]byte(configJSON), &configObj); err != nil {
|
||||
t.Fatalf("解析配置JSON失败: %v", err)
|
||||
}
|
||||
|
||||
reqBody := AddOrUpdateExternalMCPRequest{
|
||||
Config: configObj,
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(reqBody)
|
||||
req := httptest.NewRequest("PUT", "/api/external-mcp/test-http", bytes.NewBuffer(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// 验证配置已添加
|
||||
req2 := httptest.NewRequest("GET", "/api/external-mcp/test-http", nil)
|
||||
w2 := httptest.NewRecorder()
|
||||
router.ServeHTTP(w2, req2)
|
||||
|
||||
if w2.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String())
|
||||
}
|
||||
|
||||
var response ExternalMCPResponse
|
||||
if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if response.Config.Transport != "http" {
|
||||
t.Errorf("期望transport为http，实际%s", response.Config.Transport)
|
||||
}
|
||||
if response.Config.URL != "http://127.0.0.1:8081/mcp" {
|
||||
t.Errorf("期望url为'http://127.0.0.1:8081/mcp'，实际%s", response.Config.URL)
|
||||
}
|
||||
if !response.Config.Enabled {
|
||||
t.Error("期望enabled为true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_AddOrUpdateExternalMCP_InvalidConfig(t *testing.T) {
|
||||
router, _, configPath := setupTestRouter()
|
||||
defer cleanupTestConfig(configPath)
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
configJSON string
|
||||
expectedErr string
|
||||
}{
|
||||
{
|
||||
name: "缺少command和url",
|
||||
configJSON: `{"enabled": true}`,
|
||||
expectedErr: "需要指定command（stdio模式）或url（http/sse模式）",
|
||||
},
|
||||
{
|
||||
name: "stdio模式缺少command",
|
||||
configJSON: `{"args": ["test"], "enabled": true}`,
|
||||
expectedErr: "stdio模式需要command",
|
||||
},
|
||||
{
|
||||
name: "http模式缺少url",
|
||||
configJSON: `{"transport": "http", "enabled": true}`,
|
||||
expectedErr: "HTTP模式需要URL",
|
||||
},
|
||||
{
|
||||
name: "无效的transport",
|
||||
configJSON: `{"transport": "invalid", "enabled": true}`,
|
||||
expectedErr: "不支持的传输模式",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var configObj config.ExternalMCPServerConfig
|
||||
if err := json.Unmarshal([]byte(tc.configJSON), &configObj); err != nil {
|
||||
t.Fatalf("解析配置JSON失败: %v", err)
|
||||
}
|
||||
|
||||
reqBody := AddOrUpdateExternalMCPRequest{
|
||||
Config: configObj,
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(reqBody)
|
||||
req := httptest.NewRequest("PUT", "/api/external-mcp/test-invalid", bytes.NewBuffer(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("期望状态码400，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var response map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
errorMsg := response["error"].(string)
|
||||
// 对于stdio模式缺少command的情况，错误信息可能略有不同
|
||||
if tc.name == "stdio模式缺少command" {
|
||||
if !strings.Contains(errorMsg, "stdio") && !strings.Contains(errorMsg, "command") {
|
||||
t.Errorf("期望错误信息包含'stdio'或'command'，实际'%s'", errorMsg)
|
||||
}
|
||||
} else if !strings.Contains(errorMsg, tc.expectedErr) {
|
||||
t.Errorf("期望错误信息包含'%s'，实际'%s'", tc.expectedErr, errorMsg)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_DeleteExternalMCP(t *testing.T) {
|
||||
router, handler, configPath := setupTestRouter()
|
||||
defer cleanupTestConfig(configPath)
|
||||
|
||||
// 先添加一个配置
|
||||
configObj := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Enabled: true,
|
||||
}
|
||||
handler.manager.AddOrUpdateConfig("test-delete", configObj)
|
||||
|
||||
// 删除配置
|
||||
req := httptest.NewRequest("DELETE", "/api/external-mcp/test-delete", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// 验证配置已删除
|
||||
req2 := httptest.NewRequest("GET", "/api/external-mcp/test-delete", nil)
|
||||
w2 := httptest.NewRecorder()
|
||||
router.ServeHTTP(w2, req2)
|
||||
|
||||
if w2.Code != http.StatusNotFound {
|
||||
t.Errorf("期望状态码404，实际%d: %s", w2.Code, w2.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_GetExternalMCPs(t *testing.T) {
|
||||
router, handler, _ := setupTestRouter()
|
||||
|
||||
// 添加多个配置
|
||||
handler.manager.AddOrUpdateConfig("test1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Enabled: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("test2", config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
Enabled: false,
|
||||
})
|
||||
|
||||
req := httptest.NewRequest("GET", "/api/external-mcp", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var response map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
servers := response["servers"].(map[string]interface{})
|
||||
if len(servers) != 2 {
|
||||
t.Errorf("期望2个服务器，实际%d", len(servers))
|
||||
}
|
||||
if _, ok := servers["test1"]; !ok {
|
||||
t.Error("期望包含test1")
|
||||
}
|
||||
if _, ok := servers["test2"]; !ok {
|
||||
t.Error("期望包含test2")
|
||||
}
|
||||
|
||||
stats := response["stats"].(map[string]interface{})
|
||||
if int(stats["total"].(float64)) != 2 {
|
||||
t.Errorf("期望总数为2，实际%d", int(stats["total"].(float64)))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_GetExternalMCPStats(t *testing.T) {
|
||||
router, handler, _ := setupTestRouter()
|
||||
|
||||
// 添加配置
|
||||
handler.manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Enabled: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
Enabled: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Enabled: false,
|
||||
Disabled: true,
|
||||
})
|
||||
|
||||
req := httptest.NewRequest("GET", "/api/external-mcp/stats", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var stats map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &stats); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if int(stats["total"].(float64)) != 3 {
|
||||
t.Errorf("期望总数为3，实际%d", int(stats["total"].(float64)))
|
||||
}
|
||||
if int(stats["enabled"].(float64)) != 2 {
|
||||
t.Errorf("期望启用数为2，实际%d", int(stats["enabled"].(float64)))
|
||||
}
|
||||
if int(stats["disabled"].(float64)) != 1 {
|
||||
t.Errorf("期望停用数为1，实际%d", int(stats["disabled"].(float64)))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_StartStopExternalMCP(t *testing.T) {
|
||||
router, handler, configPath := setupTestRouter()
|
||||
defer cleanupTestConfig(configPath)
|
||||
|
||||
// 添加一个禁用的配置
|
||||
handler.manager.AddOrUpdateConfig("test-start-stop", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Enabled: false,
|
||||
Disabled: true,
|
||||
})
|
||||
|
||||
// 测试启动（可能会失败，因为没有真实的服务器）
|
||||
req := httptest.NewRequest("POST", "/api/external-mcp/test-start-stop/start", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
// 启动可能会失败，但应该返回合理的状态码
|
||||
if w.Code != http.StatusOK {
|
||||
// 如果启动失败，应该是400或500
|
||||
if w.Code != http.StatusBadRequest && w.Code != http.StatusInternalServerError {
|
||||
t.Errorf("期望状态码200/400/500，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// 测试停止
|
||||
req2 := httptest.NewRequest("POST", "/api/external-mcp/test-start-stop/stop", nil)
|
||||
w2 := httptest.NewRecorder()
|
||||
router.ServeHTTP(w2, req2)
|
||||
|
||||
if w2.Code != http.StatusOK {
|
||||
t.Errorf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_GetExternalMCP_NotFound(t *testing.T) {
|
||||
router, _, _ := setupTestRouter()
|
||||
|
||||
req := httptest.NewRequest("GET", "/api/external-mcp/nonexistent", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Errorf("期望状态码404，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_DeleteExternalMCP_NotFound(t *testing.T) {
|
||||
router, _, configPath := setupTestRouter()
|
||||
defer cleanupTestConfig(configPath)
|
||||
|
||||
req := httptest.NewRequest("DELETE", "/api/external-mcp/nonexistent", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
// 删除不存在的配置可能返回200（幂等操作）或404，都是合理的
|
||||
if w.Code != http.StatusNotFound && w.Code != http.StatusOK {
|
||||
t.Errorf("期望状态码404或200，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_AddOrUpdateExternalMCP_EmptyName(t *testing.T) {
|
||||
router, _, _ := setupTestRouter()
|
||||
|
||||
configObj := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Enabled: true,
|
||||
}
|
||||
|
||||
reqBody := AddOrUpdateExternalMCPRequest{
|
||||
Config: configObj,
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(reqBody)
|
||||
req := httptest.NewRequest("PUT", "/api/external-mcp/", bytes.NewBuffer(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
// 空名称应该返回404或400
|
||||
if w.Code != http.StatusNotFound && w.Code != http.StatusBadRequest {
|
||||
t.Errorf("期望状态码404或400，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_AddOrUpdateExternalMCP_InvalidJSON(t *testing.T) {
|
||||
router, _, _ := setupTestRouter()
|
||||
|
||||
// 发送无效的JSON
|
||||
body := []byte(`{"config": invalid json}`)
|
||||
req := httptest.NewRequest("PUT", "/api/external-mcp/test", bytes.NewBuffer(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("期望状态码400，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPHandler_UpdateExistingConfig(t *testing.T) {
|
||||
router, handler, configPath := setupTestRouter()
|
||||
defer cleanupTestConfig(configPath)
|
||||
|
||||
// 先添加配置
|
||||
config1 := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Enabled: true,
|
||||
}
|
||||
handler.manager.AddOrUpdateConfig("test-update", config1)
|
||||
|
||||
// 更新配置
|
||||
config2 := config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
Enabled: true,
|
||||
}
|
||||
|
||||
reqBody := AddOrUpdateExternalMCPRequest{
|
||||
Config: config2,
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(reqBody)
|
||||
req := httptest.NewRequest("PUT", "/api/external-mcp/test-update", bytes.NewBuffer(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// 验证配置已更新
|
||||
req2 := httptest.NewRequest("GET", "/api/external-mcp/test-update", nil)
|
||||
w2 := httptest.NewRecorder()
|
||||
router.ServeHTTP(w2, req2)
|
||||
|
||||
if w2.Code != http.StatusOK {
|
||||
t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String())
|
||||
}
|
||||
|
||||
var response ExternalMCPResponse
|
||||
if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if response.Config.URL != "http://127.0.0.1:8081/mcp" {
|
||||
t.Errorf("期望url为'http://127.0.0.1:8081/mcp'，实际%s", response.Config.URL)
|
||||
}
|
||||
if response.Config.Command != "" {
|
||||
t.Errorf("期望command为空，实际%s", response.Config.Command)
|
||||
}
|
||||
}
|
||||
-467
@@ -1,467 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
openaiClient "cyberstrike-ai/internal/openai"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type FofaHandler struct {
|
||||
cfg *config.Config
|
||||
logger *zap.Logger
|
||||
client *http.Client
|
||||
openAIClient *openaiClient.Client
|
||||
}
|
||||
|
||||
func NewFofaHandler(cfg *config.Config, logger *zap.Logger) *FofaHandler {
|
||||
// LLM 请求通常比 FOFA 查询更慢一点，单独给一个更宽松的超时。
|
||||
llmHTTPClient := &http.Client{Timeout: 2 * time.Minute}
|
||||
var llmCfg *config.OpenAIConfig
|
||||
if cfg != nil {
|
||||
llmCfg = &cfg.OpenAI
|
||||
}
|
||||
return &FofaHandler{
|
||||
cfg: cfg,
|
||||
logger: logger,
|
||||
client: &http.Client{Timeout: 30 * time.Second},
|
||||
openAIClient: openaiClient.NewClient(llmCfg, llmHTTPClient, logger),
|
||||
}
|
||||
}
|
||||
|
||||
type fofaSearchRequest struct {
|
||||
Query string `json:"query" binding:"required"`
|
||||
Size int `json:"size,omitempty"`
|
||||
Page int `json:"page,omitempty"`
|
||||
Fields string `json:"fields,omitempty"`
|
||||
Full bool `json:"full,omitempty"`
|
||||
}
|
||||
|
||||
type fofaParseRequest struct {
|
||||
Text string `json:"text" binding:"required"`
|
||||
}
|
||||
|
||||
type fofaParseResponse struct {
|
||||
Query string `json:"query"`
|
||||
Explanation string `json:"explanation,omitempty"`
|
||||
Warnings []string `json:"warnings,omitempty"`
|
||||
}
|
||||
|
||||
type fofaAPIResponse struct {
|
||||
Error bool `json:"error"`
|
||||
ErrMsg string `json:"errmsg"`
|
||||
Size int `json:"size"`
|
||||
Page int `json:"page"`
|
||||
Total int `json:"total"`
|
||||
Mode string `json:"mode"`
|
||||
Query string `json:"query"`
|
||||
Results [][]interface{} `json:"results"`
|
||||
}
|
||||
|
||||
type fofaSearchResponse struct {
|
||||
Query string `json:"query"`
|
||||
Size int `json:"size"`
|
||||
Page int `json:"page"`
|
||||
Total int `json:"total"`
|
||||
Fields []string `json:"fields"`
|
||||
ResultsCount int `json:"results_count"`
|
||||
Results []map[string]interface{} `json:"results"`
|
||||
}
|
||||
|
||||
func (h *FofaHandler) resolveCredentials() (email, apiKey string) {
|
||||
// 优先环境变量（便于容器部署），其次配置文件
|
||||
email = strings.TrimSpace(os.Getenv("FOFA_EMAIL"))
|
||||
apiKey = strings.TrimSpace(os.Getenv("FOFA_API_KEY"))
|
||||
if email != "" && apiKey != "" {
|
||||
return email, apiKey
|
||||
}
|
||||
if h.cfg != nil {
|
||||
if email == "" {
|
||||
email = strings.TrimSpace(h.cfg.FOFA.Email)
|
||||
}
|
||||
if apiKey == "" {
|
||||
apiKey = strings.TrimSpace(h.cfg.FOFA.APIKey)
|
||||
}
|
||||
}
|
||||
return email, apiKey
|
||||
}
|
||||
|
||||
func (h *FofaHandler) resolveBaseURL() string {
|
||||
if h.cfg != nil {
|
||||
if v := strings.TrimSpace(h.cfg.FOFA.BaseURL); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return "https://fofa.info/api/v1/search/all"
|
||||
}
|
||||
|
||||
// ParseNaturalLanguage 将自然语言解析为 FOFA 查询语法（仅生成，不执行查询）
|
||||
func (h *FofaHandler) ParseNaturalLanguage(c *gin.Context) {
|
||||
var req fofaParseRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
req.Text = strings.TrimSpace(req.Text)
|
||||
if req.Text == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "text 不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
if h.cfg == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "系统配置未初始化"})
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(h.cfg.OpenAI.APIKey) == "" || strings.TrimSpace(h.cfg.OpenAI.Model) == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"error": "未配置 AI 模型：请在系统设置中填写 openai.api_key 与 openai.model（支持 OpenAI 兼容 API，如 DeepSeek）",
|
||||
"need": []string{"openai.api_key", "openai.model"},
|
||||
})
|
||||
return
|
||||
}
|
||||
if h.openAIClient == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "AI 客户端未初始化"})
|
||||
return
|
||||
}
|
||||
|
||||
systemPrompt := strings.TrimSpace(`
|
||||
你是“FOFA 查询语法生成器”。任务：把用户输入的自然语言搜索意图，转换成 FOFA 查询语法。
|
||||
|
||||
输出要求（非常重要）：
|
||||
1) 只输出 JSON（不要 markdown、不要代码块、不要额外解释文本）
|
||||
2) JSON 结构必须是：
|
||||
{
|
||||
"query": "string，FOFA查询语法（可直接粘贴到 FOFA 或本系统查询框）",
|
||||
"explanation": "string，可选，解释你如何映射字段/逻辑",
|
||||
"warnings": ["string"...] 可选，列出歧义/风险/需要人工确认的点
|
||||
}
|
||||
3) 如果用户输入本身已经是 FOFA 查询语法（或非常接近 FOFA 语法的表达式），应当“原样返回”为 query：
|
||||
- 不要擅自改写字段名、操作符、括号结构
|
||||
- 不要改写任何字符串值（尤其是地理位置类值），不要做缩写/同义词替换/翻译/音译
|
||||
|
||||
查询语法要点（来自 FOFA 语法参考）：
|
||||
- 逻辑连接符：&&（与）、||（或），必要时用 () 包住子表达式以确认优先级（括号优先级最高）
|
||||
- 当同一层级同时出现 && 与 ||（混用）时，用 () 明确优先级（避免歧义）
|
||||
- 比较/匹配：
|
||||
- = 匹配；当字段="" 时，可查询“不存在该字段”或“值为空”的情况
|
||||
- == 完全匹配；当字段=="" 时，可查询“字段存在且值为空”的情况
|
||||
- != 不匹配；当字段!="" 时，可查询“值不为空”的情况
|
||||
- *= 模糊匹配；可使用 * 或 ? 进行搜索
|
||||
- 直接输入关键词（不带字段）会在标题、HTML内容、HTTP头、URL字段中搜索；但当意图明确时优先用字段表达（更可控、更准确）
|
||||
|
||||
字段示例速查（来自用户提供的案例，可直接套用/拼接）：
|
||||
- 高级搜索操作符示例：
|
||||
- title="beijing" （= 匹配）
|
||||
- title=="" （== 完全匹配，字段存在且值为空）
|
||||
- title="" （= 匹配，可能表示字段不存在或值为空）
|
||||
- title!="" （!= 不匹配，可用于值不为空）
|
||||
- title*="*Home*" （*= 模糊匹配，用 * 或 ?）
|
||||
- (app="Apache" || app="Nginx") && country="CN" （混用 && / || 时用括号）
|
||||
- 基础类（General）：
|
||||
- ip="1.1.1.1"
|
||||
- ip="220.181.111.1/24"
|
||||
- ip="2600:9000:202a:2600:18:4ab7:f600:93a1"
|
||||
- port="6379"
|
||||
- domain="qq.com"
|
||||
- host=".fofa.info"
|
||||
- os="centos"
|
||||
- server="Microsoft-IIS/10"
|
||||
- asn="19551"
|
||||
- org="LLC Baxet"
|
||||
- is_domain=true / is_domain=false
|
||||
- is_ipv6=true / is_ipv6=false
|
||||
- 标记类（Special Label）：
|
||||
- app="Microsoft-Exchange"
|
||||
- fid="sSXXGNUO2FefBTcCLIT/2Q=="
|
||||
- product="NGINX"
|
||||
- product="Roundcube-Webmail" && product.version="1.6.10"
|
||||
- category="服务"
|
||||
- type="service" / type="subdomain"
|
||||
- cloud_name="Aliyundun"
|
||||
- is_cloud=true / is_cloud=false
|
||||
- is_fraud=true / is_fraud=false
|
||||
- is_honeypot=true / is_honeypot=false
|
||||
- 协议类（type=service）：
|
||||
- protocol="quic"
|
||||
- banner="users"
|
||||
- banner_hash="7330105010150477363"
|
||||
- banner_fid="zRpqmn0FXQRjZpH8MjMX55zpMy9SgsW8"
|
||||
- base_protocol="udp" / base_protocol="tcp"
|
||||
- 网站类（type=subdomain）：
|
||||
- title="beijing"
|
||||
- header="elastic"
|
||||
- header_hash="1258854265"
|
||||
- body="网络空间测绘"
|
||||
- body_hash="-2090962452"
|
||||
- js_name="js/jquery.js"
|
||||
- js_md5="82ac3f14327a8b7ba49baa208d4eaa15"
|
||||
- cname="customers.spektrix.com"
|
||||
- cname_domain="siteforce.com"
|
||||
- icon_hash="-247388890"
|
||||
- status_code="402"
|
||||
- icp="京ICP证030173号"
|
||||
- sdk_hash="Are3qNnP2Eqn7q5kAoUO3l+w3mgVIytO"
|
||||
- 地理位置（Location）：
|
||||
- country="CN" 或 country="中国"
|
||||
- region="Zhejiang" 或 region="浙江"（仅支持中国地区中文）
|
||||
- city="Hangzhou"
|
||||
- 证书类（Certificate）：
|
||||
- cert="baidu"
|
||||
- cert.subject="Oracle Corporation"
|
||||
- cert.issuer="DigiCert"
|
||||
- cert.subject.org="Oracle Corporation"
|
||||
- cert.subject.cn="baidu.com"
|
||||
- cert.issuer.org="cPanel, Inc."
|
||||
- cert.issuer.cn="Synology Inc. CA"
|
||||
- cert.domain="huawei.com"
|
||||
- cert.is_equal=true / cert.is_equal=false
|
||||
- cert.is_valid=true / cert.is_valid=false
|
||||
- cert.is_match=true / cert.is_match=false
|
||||
- cert.is_expired=true / cert.is_expired=false
|
||||
- jarm="2ad2ad0002ad2ad22c2ad2ad2ad2ad2eac92ec34bcc0cf7520e97547f83e81"
|
||||
- tls.version="TLS 1.3"
|
||||
- tls.ja3s="15af977ce25de452b96affa2addb1036"
|
||||
- cert.sn="356078156165546797850343536942784588840297"
|
||||
- cert.not_after.after="2025-03-01" / cert.not_after.before="2025-03-01"
|
||||
- cert.not_before.after="2025-03-01" / cert.not_before.before="2025-03-01"
|
||||
- 时间类（Last update time）：
|
||||
- after="2023-01-01"
|
||||
- before="2023-12-01"
|
||||
- after="2023-01-01" && before="2023-12-01"
|
||||
- 独立IP语法（需配合 ip_filter / ip_exclude）：
|
||||
- ip_filter(banner="SSH-2.0-OpenSSH_6.7p2") && ip_filter(icon_hash="-1057022626")
|
||||
- ip_filter(banner="SSH-2.0-OpenSSH_6.7p2" && asn="3462") && ip_exclude(title="EdgeOS")
|
||||
- port_size="6" / port_size_gt="6" / port_size_lt="12"
|
||||
- ip_ports="80,161"
|
||||
- ip_country="CN"
|
||||
- ip_region="Zhejiang"
|
||||
- ip_city="Hangzhou"
|
||||
- ip_after="2021-03-18"
|
||||
- ip_before="2019-09-09"
|
||||
|
||||
生成约束与注意事项：
|
||||
- 字符串值一律用英文双引号包裹，例如 title="登录"、country="CN"
|
||||
- 字符串值保持字面一致：不要缩写（例如 city="beijing" 不要变成 city="BJ"），不要用别名（例如 Beijing/Peking），不要擅自翻译/音译/改写大小写
|
||||
- 地理位置字段（country/region/city）更倾向于“按用户给定值输出”；不确定合法取值时，不要猜测，把备选写进 warnings
|
||||
- 不要捏造不存在的 FOFA 字段；不确定时把不确定点写进 warnings，并输出一个保守的 query
|
||||
- 当用户描述里有“多个与/或条件”，优先加 () 明确优先级，例如：(app="Apache" || app="Nginx") && country="CN"
|
||||
- 当用户缺少关键条件导致范围过大或歧义（如地点/协议/端口/服务类型未说明），允许 query 为空字符串，并在 warnings 里明确需要补充的信息
|
||||
`)
|
||||
|
||||
userPrompt := fmt.Sprintf("自然语言意图：%s", req.Text)
|
||||
|
||||
requestBody := map[string]interface{}{
|
||||
"model": h.cfg.OpenAI.Model,
|
||||
"messages": []map[string]interface{}{
|
||||
{"role": "system", "content": systemPrompt},
|
||||
{"role": "user", "content": userPrompt},
|
||||
},
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 1200,
|
||||
}
|
||||
|
||||
// OpenAI 返回结构：只需要 choices[0].message.content
|
||||
var apiResponse struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), 90*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := h.openAIClient.ChatCompletion(ctx, requestBody, &apiResponse); err != nil {
|
||||
var apiErr *openaiClient.APIError
|
||||
if errors.As(err, &apiErr) {
|
||||
h.logger.Warn("FOFA自然语言解析：LLM返回错误", zap.Int("status", apiErr.StatusCode))
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "AI 解析失败（上游返回非 200），请检查模型配置或稍后重试"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "AI 解析失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if len(apiResponse.Choices) == 0 {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "AI 未返回有效结果"})
|
||||
return
|
||||
}
|
||||
|
||||
content := strings.TrimSpace(apiResponse.Choices[0].Message.Content)
|
||||
// 兼容模型偶尔返回 ```json ... ``` 的情况
|
||||
content = strings.TrimPrefix(content, "```json")
|
||||
content = strings.TrimPrefix(content, "```")
|
||||
content = strings.TrimSuffix(content, "```")
|
||||
content = strings.TrimSpace(content)
|
||||
|
||||
var parsed fofaParseResponse
|
||||
if err := json.Unmarshal([]byte(content), &parsed); err != nil {
|
||||
// 直接回传一部分原文，方便排查，但避免太大
|
||||
snippet := content
|
||||
if len(snippet) > 1200 {
|
||||
snippet = snippet[:1200]
|
||||
}
|
||||
c.JSON(http.StatusBadGateway, gin.H{
|
||||
"error": "AI 返回内容无法解析为 JSON，请稍后重试或换个描述方式",
|
||||
"snippet": snippet,
|
||||
})
|
||||
return
|
||||
}
|
||||
parsed.Query = strings.TrimSpace(parsed.Query)
|
||||
if parsed.Query == "" {
|
||||
// query 允许为空（表示需求不明确），但前端需要明确提示
|
||||
if len(parsed.Warnings) == 0 {
|
||||
parsed.Warnings = []string{"需求信息不足，未能生成可用的 FOFA 查询语法，请补充关键条件（如国家/端口/产品/域名等）。"}
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, parsed)
|
||||
}
|
||||
|
||||
// Search FOFA 查询（后端代理，避免前端暴露 key）
|
||||
func (h *FofaHandler) Search(c *gin.Context) {
|
||||
var req fofaSearchRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
req.Query = strings.TrimSpace(req.Query)
|
||||
if req.Query == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "query 不能为空"})
|
||||
return
|
||||
}
|
||||
if req.Size <= 0 {
|
||||
req.Size = 100
|
||||
}
|
||||
if req.Page <= 0 {
|
||||
req.Page = 1
|
||||
}
|
||||
// FOFA 接口 size 上限和账户权限相关，这里只做一个合理的保护
|
||||
if req.Size > 10000 {
|
||||
req.Size = 10000
|
||||
}
|
||||
if req.Fields == "" {
|
||||
req.Fields = "host,ip,port,domain,title,protocol,country,province,city,server"
|
||||
}
|
||||
|
||||
email, apiKey := h.resolveCredentials()
|
||||
if email == "" || apiKey == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"error": "FOFA 未配置：请在系统设置中填写 FOFA Email/API Key，或设置环境变量 FOFA_EMAIL/FOFA_API_KEY",
|
||||
"need": []string{"fofa.email", "fofa.api_key"},
|
||||
"env_key": []string{"FOFA_EMAIL", "FOFA_API_KEY"},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
baseURL := h.resolveBaseURL()
|
||||
qb64 := base64.StdEncoding.EncodeToString([]byte(req.Query))
|
||||
|
||||
u, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "FOFA base_url 无效: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
params := u.Query()
|
||||
params.Set("email", email)
|
||||
params.Set("key", apiKey)
|
||||
params.Set("qbase64", qb64)
|
||||
params.Set("size", fmt.Sprintf("%d", req.Size))
|
||||
params.Set("page", fmt.Sprintf("%d", req.Page))
|
||||
params.Set("fields", strings.TrimSpace(req.Fields))
|
||||
if req.Full {
|
||||
params.Set("full", "true")
|
||||
} else {
|
||||
// 明确传 false，便于排查
|
||||
params.Set("full", "false")
|
||||
}
|
||||
u.RawQuery = params.Encode()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(c.Request.Context(), http.MethodGet, u.String(), nil)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "创建请求失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := h.client.Do(httpReq)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "请求 FOFA 失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": fmt.Sprintf("FOFA 返回非 2xx: %d", resp.StatusCode)})
|
||||
return
|
||||
}
|
||||
|
||||
var apiResp fofaAPIResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "解析 FOFA 响应失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if apiResp.Error {
|
||||
msg := strings.TrimSpace(apiResp.ErrMsg)
|
||||
if msg == "" {
|
||||
msg = "FOFA 返回错误"
|
||||
}
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": msg})
|
||||
return
|
||||
}
|
||||
|
||||
fields := splitAndCleanCSV(req.Fields)
|
||||
results := make([]map[string]interface{}, 0, len(apiResp.Results))
|
||||
for _, row := range apiResp.Results {
|
||||
item := make(map[string]interface{}, len(fields))
|
||||
for i, f := range fields {
|
||||
if i < len(row) {
|
||||
item[f] = row[i]
|
||||
} else {
|
||||
item[f] = nil
|
||||
}
|
||||
}
|
||||
results = append(results, item)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, fofaSearchResponse{
|
||||
Query: req.Query,
|
||||
Size: apiResp.Size,
|
||||
Page: apiResp.Page,
|
||||
Total: apiResp.Total,
|
||||
Fields: fields,
|
||||
ResultsCount: len(results),
|
||||
Results: results,
|
||||
})
|
||||
}
|
||||
|
||||
func splitAndCleanCSV(s string) []string {
|
||||
parts := strings.Split(s, ",")
|
||||
out := make([]string, 0, len(parts))
|
||||
seen := make(map[string]struct{}, len(parts))
|
||||
for _, p := range parts {
|
||||
v := strings.TrimSpace(p)
|
||||
if v == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[v]; ok {
|
||||
continue
|
||||
}
|
||||
seen[v] = struct{}{}
|
||||
out = append(out, v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -1,320 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// GroupHandler 分组处理器
|
||||
type GroupHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewGroupHandler 创建新的分组处理器
|
||||
func NewGroupHandler(db *database.DB, logger *zap.Logger) *GroupHandler {
|
||||
return &GroupHandler{
|
||||
db: db,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateGroupRequest 创建分组请求
|
||||
type CreateGroupRequest struct {
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon"`
|
||||
}
|
||||
|
||||
// CreateGroup 创建分组
|
||||
func (h *GroupHandler) CreateGroup(c *gin.Context) {
|
||||
var req CreateGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if req.Name == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "分组名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
group, err := h.db.CreateGroup(req.Name, req.Icon)
|
||||
if err != nil {
|
||||
h.logger.Error("创建分组失败", zap.Error(err))
|
||||
// 如果是名称重复错误，返回400状态码
|
||||
if err.Error() == "分组名称已存在" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "分组名称已存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, group)
|
||||
}
|
||||
|
||||
// ListGroups 列出所有分组
|
||||
func (h *GroupHandler) ListGroups(c *gin.Context) {
|
||||
groups, err := h.db.ListGroups()
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, groups)
|
||||
}
|
||||
|
||||
// GetGroup 获取分组
|
||||
func (h *GroupHandler) GetGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
group, err := h.db.GetGroup(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组失败", zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "分组不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, group)
|
||||
}
|
||||
|
||||
// UpdateGroupRequest 更新分组请求
|
||||
type UpdateGroupRequest struct {
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon"`
|
||||
}
|
||||
|
||||
// UpdateGroup 更新分组
|
||||
func (h *GroupHandler) UpdateGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
var req UpdateGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if req.Name == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "分组名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateGroup(id, req.Name, req.Icon); err != nil {
|
||||
h.logger.Error("更新分组失败", zap.Error(err))
|
||||
// 如果是名称重复错误，返回400状态码
|
||||
if err.Error() == "分组名称已存在" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "分组名称已存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
group, err := h.db.GetGroup(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取更新后的分组失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, group)
|
||||
}
|
||||
|
||||
// DeleteGroup 删除分组
|
||||
func (h *GroupHandler) DeleteGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
if err := h.db.DeleteGroup(id); err != nil {
|
||||
h.logger.Error("删除分组失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
// AddConversationToGroupRequest 添加对话到分组请求
|
||||
type AddConversationToGroupRequest struct {
|
||||
ConversationID string `json:"conversationId"`
|
||||
GroupID string `json:"groupId"`
|
||||
}
|
||||
|
||||
// AddConversationToGroup 将对话添加到分组
|
||||
func (h *GroupHandler) AddConversationToGroup(c *gin.Context) {
|
||||
var req AddConversationToGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.AddConversationToGroup(req.ConversationID, req.GroupID); err != nil {
|
||||
h.logger.Error("添加对话到分组失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "添加成功"})
|
||||
}
|
||||
|
||||
// RemoveConversationFromGroup 从分组中移除对话
|
||||
func (h *GroupHandler) RemoveConversationFromGroup(c *gin.Context) {
|
||||
conversationID := c.Param("conversationId")
|
||||
groupID := c.Param("id")
|
||||
|
||||
if err := h.db.RemoveConversationFromGroup(conversationID, groupID); err != nil {
|
||||
h.logger.Error("从分组中移除对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "移除成功"})
|
||||
}
|
||||
|
||||
// GroupConversation 分组对话响应结构
|
||||
type GroupConversation struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
Pinned bool `json:"pinned"`
|
||||
GroupPinned bool `json:"groupPinned"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
// GetGroupConversations 获取分组中的所有对话
|
||||
func (h *GroupHandler) GetGroupConversations(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
searchQuery := c.Query("search") // 获取搜索参数
|
||||
|
||||
var conversations []*database.Conversation
|
||||
var err error
|
||||
|
||||
// 如果有搜索关键词，使用搜索方法；否则使用普通方法
|
||||
if searchQuery != "" {
|
||||
conversations, err = h.db.SearchConversationsByGroup(groupID, searchQuery)
|
||||
} else {
|
||||
conversations, err = h.db.GetConversationsByGroup(groupID)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取每个对话在分组中的置顶状态
|
||||
groupConvs := make([]GroupConversation, 0, len(conversations))
|
||||
for _, conv := range conversations {
|
||||
// 查询分组内置顶状态
|
||||
var groupPinned int
|
||||
err := h.db.QueryRow(
|
||||
"SELECT COALESCE(pinned, 0) FROM conversation_group_mappings WHERE conversation_id = ? AND group_id = ?",
|
||||
conv.ID, groupID,
|
||||
).Scan(&groupPinned)
|
||||
if err != nil {
|
||||
h.logger.Warn("查询分组内置顶状态失败", zap.String("conversationId", conv.ID), zap.Error(err))
|
||||
groupPinned = 0
|
||||
}
|
||||
|
||||
groupConvs = append(groupConvs, GroupConversation{
|
||||
ID: conv.ID,
|
||||
Title: conv.Title,
|
||||
Pinned: conv.Pinned,
|
||||
GroupPinned: groupPinned != 0,
|
||||
CreatedAt: conv.CreatedAt,
|
||||
UpdatedAt: conv.UpdatedAt,
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, groupConvs)
|
||||
}
|
||||
|
||||
// GetAllMappings 批量获取所有分组映射（消除前端 N+1 请求）
|
||||
func (h *GroupHandler) GetAllMappings(c *gin.Context) {
|
||||
mappings, err := h.db.GetAllGroupMappings()
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组映射失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, mappings)
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedRequest 更新对话置顶状态请求
|
||||
type UpdateConversationPinnedRequest struct {
|
||||
Pinned bool `json:"pinned"`
|
||||
}
|
||||
|
||||
// UpdateConversationPinned 更新对话置顶状态
|
||||
func (h *GroupHandler) UpdateConversationPinned(c *gin.Context) {
|
||||
conversationID := c.Param("id")
|
||||
|
||||
var req UpdateConversationPinnedRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateConversationPinned(conversationID, req.Pinned); err != nil {
|
||||
h.logger.Error("更新对话置顶状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "更新成功"})
|
||||
}
|
||||
|
||||
// UpdateGroupPinnedRequest 更新分组置顶状态请求
|
||||
type UpdateGroupPinnedRequest struct {
|
||||
Pinned bool `json:"pinned"`
|
||||
}
|
||||
|
||||
// UpdateGroupPinned 更新分组置顶状态
|
||||
func (h *GroupHandler) UpdateGroupPinned(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
|
||||
var req UpdateGroupPinnedRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateGroupPinned(groupID, req.Pinned); err != nil {
|
||||
h.logger.Error("更新分组置顶状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "更新成功"})
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedInGroupRequest 更新分组对话置顶状态请求
|
||||
type UpdateConversationPinnedInGroupRequest struct {
|
||||
Pinned bool `json:"pinned"`
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedInGroup 更新对话在分组中的置顶状态
|
||||
func (h *GroupHandler) UpdateConversationPinnedInGroup(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
conversationID := c.Param("conversationId")
|
||||
|
||||
var req UpdateConversationPinnedInGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateConversationPinnedInGroup(conversationID, groupID, req.Pinned); err != nil {
|
||||
h.logger.Error("更新分组对话置顶状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "更新成功"})
|
||||
}
|
||||
@@ -1,517 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/knowledge"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// KnowledgeHandler 知识库处理器
|
||||
type KnowledgeHandler struct {
|
||||
manager *knowledge.Manager
|
||||
retriever *knowledge.Retriever
|
||||
indexer *knowledge.Indexer
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewKnowledgeHandler 创建新的知识库处理器
|
||||
func NewKnowledgeHandler(
|
||||
manager *knowledge.Manager,
|
||||
retriever *knowledge.Retriever,
|
||||
indexer *knowledge.Indexer,
|
||||
db *database.DB,
|
||||
logger *zap.Logger,
|
||||
) *KnowledgeHandler {
|
||||
return &KnowledgeHandler{
|
||||
manager: manager,
|
||||
retriever: retriever,
|
||||
indexer: indexer,
|
||||
db: db,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// GetCategories 获取所有分类
|
||||
func (h *KnowledgeHandler) GetCategories(c *gin.Context) {
|
||||
categories, err := h.manager.GetCategories()
|
||||
if err != nil {
|
||||
h.logger.Error("获取分类失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"categories": categories})
|
||||
}
|
||||
|
||||
// GetItems 获取知识项列表（支持按分类分页和关键字搜索，默认不返回完整内容）
|
||||
func (h *KnowledgeHandler) GetItems(c *gin.Context) {
|
||||
category := c.Query("category")
|
||||
searchKeyword := c.Query("search") // 搜索关键字
|
||||
|
||||
// 如果提供了搜索关键字，执行关键字搜索（在所有数据中搜索）
|
||||
if searchKeyword != "" {
|
||||
items, err := h.manager.SearchItemsByKeyword(searchKeyword, category)
|
||||
if err != nil {
|
||||
h.logger.Error("搜索知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 按分类分组结果
|
||||
groupedByCategory := make(map[string][]*knowledge.KnowledgeItemSummary)
|
||||
for _, item := range items {
|
||||
cat := item.Category
|
||||
if cat == "" {
|
||||
cat = "未分类"
|
||||
}
|
||||
groupedByCategory[cat] = append(groupedByCategory[cat], item)
|
||||
}
|
||||
|
||||
// 转换为 CategoryWithItems 格式
|
||||
categoriesWithItems := make([]*knowledge.CategoryWithItems, 0, len(groupedByCategory))
|
||||
for cat, catItems := range groupedByCategory {
|
||||
categoriesWithItems = append(categoriesWithItems, &knowledge.CategoryWithItems{
|
||||
Category: cat,
|
||||
ItemCount: len(catItems),
|
||||
Items: catItems,
|
||||
})
|
||||
}
|
||||
|
||||
// 按分类名称排序
|
||||
for i := 0; i < len(categoriesWithItems)-1; i++ {
|
||||
for j := i + 1; j < len(categoriesWithItems); j++ {
|
||||
if categoriesWithItems[i].Category > categoriesWithItems[j].Category {
|
||||
categoriesWithItems[i], categoriesWithItems[j] = categoriesWithItems[j], categoriesWithItems[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"categories": categoriesWithItems,
|
||||
"total": len(categoriesWithItems),
|
||||
"search": searchKeyword,
|
||||
"is_search": true,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// 分页模式：categoryPage=true 表示按分类分页，否则按项分页（向后兼容）
|
||||
categoryPageMode := c.Query("categoryPage") != "false" // 默认使用分类分页
|
||||
|
||||
// 分页参数
|
||||
limit := 50 // 默认每页 50 条（分类分页时为分类数，项分页时为项数）
|
||||
offset := 0
|
||||
if limitStr := c.Query("limit"); limitStr != "" {
|
||||
if parsed, err := parseInt(limitStr); err == nil && parsed > 0 && parsed <= 500 {
|
||||
limit = parsed
|
||||
}
|
||||
}
|
||||
if offsetStr := c.Query("offset"); offsetStr != "" {
|
||||
if parsed, err := parseInt(offsetStr); err == nil && parsed >= 0 {
|
||||
offset = parsed
|
||||
}
|
||||
}
|
||||
|
||||
// 如果指定了 category 参数，且使用分类分页模式，则只返回该分类
|
||||
if category != "" && categoryPageMode {
|
||||
// 单分类模式：返回该分类的所有知识项（不分页）
|
||||
items, total, err := h.manager.GetItemsSummary(category, 0, 0)
|
||||
if err != nil {
|
||||
h.logger.Error("获取知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 包装成分类结构
|
||||
categoriesWithItems := []*knowledge.CategoryWithItems{
|
||||
{
|
||||
Category: category,
|
||||
ItemCount: total,
|
||||
Items: items,
|
||||
},
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"categories": categoriesWithItems,
|
||||
"total": 1, // 只有一个分类
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if categoryPageMode {
|
||||
// 按分类分页模式（默认）
|
||||
// limit 表示每页分类数，推荐 5-10 个分类
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 10 // 默认每页 10 个分类
|
||||
}
|
||||
|
||||
categoriesWithItems, totalCategories, err := h.manager.GetCategoriesWithItems(limit, offset)
|
||||
if err != nil {
|
||||
h.logger.Error("获取分类知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"categories": categoriesWithItems,
|
||||
"total": totalCategories,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// 按项分页模式（向后兼容）
|
||||
// 是否包含完整内容（默认 false，只返回摘要）
|
||||
includeContent := c.Query("includeContent") == "true"
|
||||
|
||||
if includeContent {
|
||||
// 返回完整内容（向后兼容）
|
||||
items, err := h.manager.GetItemsWithOptions(category, limit, offset, true)
|
||||
if err != nil {
|
||||
h.logger.Error("获取知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取总数
|
||||
total, err := h.manager.GetItemsCount(category)
|
||||
if err != nil {
|
||||
h.logger.Warn("获取知识项总数失败", zap.Error(err))
|
||||
total = len(items)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"items": items,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
})
|
||||
} else {
|
||||
// 返回摘要（不包含完整内容，推荐方式）
|
||||
items, total, err := h.manager.GetItemsSummary(category, limit, offset)
|
||||
if err != nil {
|
||||
h.logger.Error("获取知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"items": items,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// GetItem 获取单个知识项
|
||||
func (h *KnowledgeHandler) GetItem(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
item, err := h.manager.GetItem(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, item)
|
||||
}
|
||||
|
||||
// CreateItem 创建知识项
|
||||
func (h *KnowledgeHandler) CreateItem(c *gin.Context) {
|
||||
var req struct {
|
||||
Category string `json:"category" binding:"required"`
|
||||
Title string `json:"title" binding:"required"`
|
||||
Content string `json:"content" binding:"required"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
item, err := h.manager.CreateItem(req.Category, req.Title, req.Content)
|
||||
if err != nil {
|
||||
h.logger.Error("创建知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 异步索引
|
||||
go func() {
|
||||
ctx := context.Background()
|
||||
if err := h.indexer.IndexItem(ctx, item.ID); err != nil {
|
||||
h.logger.Warn("索引知识项失败", zap.String("itemId", item.ID), zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
c.JSON(http.StatusOK, item)
|
||||
}
|
||||
|
||||
// UpdateItem 更新知识项
|
||||
func (h *KnowledgeHandler) UpdateItem(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
var req struct {
|
||||
Category string `json:"category" binding:"required"`
|
||||
Title string `json:"title" binding:"required"`
|
||||
Content string `json:"content" binding:"required"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
item, err := h.manager.UpdateItem(id, req.Category, req.Title, req.Content)
|
||||
if err != nil {
|
||||
h.logger.Error("更新知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 异步重新索引
|
||||
go func() {
|
||||
ctx := context.Background()
|
||||
if err := h.indexer.IndexItem(ctx, item.ID); err != nil {
|
||||
h.logger.Warn("重新索引知识项失败", zap.String("itemId", item.ID), zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
c.JSON(http.StatusOK, item)
|
||||
}
|
||||
|
||||
// DeleteItem 删除知识项
|
||||
func (h *KnowledgeHandler) DeleteItem(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
if err := h.manager.DeleteItem(id); err != nil {
|
||||
h.logger.Error("删除知识项失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
// RebuildIndex 重建索引
|
||||
func (h *KnowledgeHandler) RebuildIndex(c *gin.Context) {
|
||||
// 异步重建索引
|
||||
go func() {
|
||||
ctx := context.Background()
|
||||
if err := h.indexer.RebuildIndex(ctx); err != nil {
|
||||
h.logger.Error("重建索引失败", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "索引重建已开始，将在后台进行"})
|
||||
}
|
||||
|
||||
// ScanKnowledgeBase 扫描知识库
|
||||
func (h *KnowledgeHandler) ScanKnowledgeBase(c *gin.Context) {
|
||||
itemsToIndex, err := h.manager.ScanKnowledgeBase()
|
||||
if err != nil {
|
||||
h.logger.Error("扫描知识库失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if len(itemsToIndex) == 0 {
|
||||
c.JSON(http.StatusOK, gin.H{"message": "扫描完成，没有需要索引的新项或更新项"})
|
||||
return
|
||||
}
|
||||
|
||||
// 异步索引新添加或更新的项（增量索引）
|
||||
go func() {
|
||||
ctx := context.Background()
|
||||
h.logger.Info("开始增量索引", zap.Int("count", len(itemsToIndex)))
|
||||
failedCount := 0
|
||||
consecutiveFailures := 0
|
||||
var firstFailureItemID string
|
||||
var firstFailureError error
|
||||
|
||||
for i, itemID := range itemsToIndex {
|
||||
if err := h.indexer.IndexItem(ctx, itemID); err != nil {
|
||||
failedCount++
|
||||
consecutiveFailures++
|
||||
|
||||
// 只在第一个失败时记录详细日志
|
||||
if consecutiveFailures == 1 {
|
||||
firstFailureItemID = itemID
|
||||
firstFailureError = err
|
||||
h.logger.Warn("索引知识项失败",
|
||||
zap.String("itemId", itemID),
|
||||
zap.Int("totalItems", len(itemsToIndex)),
|
||||
zap.Error(err),
|
||||
)
|
||||
}
|
||||
|
||||
// 如果连续失败 2 次，立即停止增量索引
|
||||
if consecutiveFailures >= 2 {
|
||||
h.logger.Error("连续索引失败次数过多，立即停止增量索引",
|
||||
zap.Int("consecutiveFailures", consecutiveFailures),
|
||||
zap.Int("totalItems", len(itemsToIndex)),
|
||||
zap.Int("processedItems", i+1),
|
||||
zap.String("firstFailureItemId", firstFailureItemID),
|
||||
zap.Error(firstFailureError),
|
||||
)
|
||||
break
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// 成功时重置连续失败计数
|
||||
if consecutiveFailures > 0 {
|
||||
consecutiveFailures = 0
|
||||
firstFailureItemID = ""
|
||||
firstFailureError = nil
|
||||
}
|
||||
|
||||
// 减少进度日志频率
|
||||
if (i+1)%10 == 0 || i+1 == len(itemsToIndex) {
|
||||
h.logger.Info("索引进度", zap.Int("current", i+1), zap.Int("total", len(itemsToIndex)), zap.Int("failed", failedCount))
|
||||
}
|
||||
}
|
||||
h.logger.Info("增量索引完成", zap.Int("totalItems", len(itemsToIndex)), zap.Int("failedCount", failedCount))
|
||||
}()
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": fmt.Sprintf("扫描完成，开始索引 %d 个新添加或更新的知识项", len(itemsToIndex)),
|
||||
"items_to_index": len(itemsToIndex),
|
||||
})
|
||||
}
|
||||
|
||||
// GetRetrievalLogs 获取检索日志
|
||||
func (h *KnowledgeHandler) GetRetrievalLogs(c *gin.Context) {
|
||||
conversationID := c.Query("conversationId")
|
||||
messageID := c.Query("messageId")
|
||||
limit := 50 // 默认 50 条
|
||||
|
||||
if limitStr := c.Query("limit"); limitStr != "" {
|
||||
if parsed, err := parseInt(limitStr); err == nil && parsed > 0 {
|
||||
limit = parsed
|
||||
}
|
||||
}
|
||||
|
||||
logs, err := h.manager.GetRetrievalLogs(conversationID, messageID, limit)
|
||||
if err != nil {
|
||||
h.logger.Error("获取检索日志失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"logs": logs})
|
||||
}
|
||||
|
||||
// DeleteRetrievalLog 删除检索日志
|
||||
func (h *KnowledgeHandler) DeleteRetrievalLog(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
if err := h.manager.DeleteRetrievalLog(id); err != nil {
|
||||
h.logger.Error("删除检索日志失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
// GetIndexStatus 获取索引状态
|
||||
func (h *KnowledgeHandler) GetIndexStatus(c *gin.Context) {
|
||||
status, err := h.manager.GetIndexStatus()
|
||||
if err != nil {
|
||||
h.logger.Error("获取索引状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取索引器的错误信息
|
||||
if h.indexer != nil {
|
||||
lastError, lastErrorTime := h.indexer.GetLastError()
|
||||
if lastError != "" {
|
||||
// 如果错误是最近发生的（5 分钟内），则返回错误信息
|
||||
if time.Since(lastErrorTime) < 5*time.Minute {
|
||||
status["last_error"] = lastError
|
||||
status["last_error_time"] = lastErrorTime.Format(time.RFC3339)
|
||||
}
|
||||
}
|
||||
|
||||
// 获取重建索引状态
|
||||
isRebuilding, totalItems, current, failed, lastItemID, lastChunks, startTime := h.indexer.GetRebuildStatus()
|
||||
if isRebuilding {
|
||||
status["is_rebuilding"] = true
|
||||
status["rebuild_total"] = totalItems
|
||||
status["rebuild_current"] = current
|
||||
status["rebuild_failed"] = failed
|
||||
status["rebuild_start_time"] = startTime.Format(time.RFC3339)
|
||||
if lastItemID != "" {
|
||||
status["rebuild_last_item_id"] = lastItemID
|
||||
}
|
||||
if lastChunks > 0 {
|
||||
status["rebuild_last_chunks"] = lastChunks
|
||||
}
|
||||
// 重建中时，is_complete 为 false
|
||||
status["is_complete"] = false
|
||||
// 计算重建进度百分比
|
||||
if totalItems > 0 {
|
||||
status["progress_percent"] = float64(current) / float64(totalItems) * 100
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, status)
|
||||
}
|
||||
|
||||
// Search 搜索知识库（用于 API 调用，Agent 内部使用 Retriever）
|
||||
func (h *KnowledgeHandler) Search(c *gin.Context) {
|
||||
var req knowledge.SearchRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Retriever.Search 经 Eino VectorEinoRetriever，与 MCP 工具链一致。
|
||||
results, err := h.retriever.Search(c.Request.Context(), &req)
|
||||
if err != nil {
|
||||
h.logger.Error("搜索知识库失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"results": results})
|
||||
}
|
||||
|
||||
// GetStats 获取知识库统计信息
|
||||
func (h *KnowledgeHandler) GetStats(c *gin.Context) {
|
||||
totalCategories, totalItems, err := h.manager.GetStats()
|
||||
if err != nil {
|
||||
h.logger.Error("获取知识库统计信息失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"enabled": true,
|
||||
"total_categories": totalCategories,
|
||||
"total_items": totalItems,
|
||||
})
|
||||
}
|
||||
|
||||
// 辅助函数：解析整数
|
||||
func parseInt(s string) (int, error) {
|
||||
var result int
|
||||
_, err := fmt.Sscanf(s, "%d", &result)
|
||||
return result, err
|
||||
}
|
||||
@@ -1,299 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agents"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
var markdownAgentFilenameRe = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_.-]*\.md$`)
|
||||
|
||||
// MarkdownAgentsHandler 管理 agents 目录下子代理 Markdown（增删改查）。
|
||||
type MarkdownAgentsHandler struct {
|
||||
dir string
|
||||
}
|
||||
|
||||
// NewMarkdownAgentsHandler dir 须为已解析的绝对路径。
|
||||
func NewMarkdownAgentsHandler(dir string) *MarkdownAgentsHandler {
|
||||
return &MarkdownAgentsHandler{dir: strings.TrimSpace(dir)}
|
||||
}
|
||||
|
||||
func (h *MarkdownAgentsHandler) safeJoin(filename string) (string, error) {
|
||||
filename = strings.TrimSpace(filename)
|
||||
if filename == "" || !markdownAgentFilenameRe.MatchString(filename) {
|
||||
return "", fmt.Errorf("非法文件名")
|
||||
}
|
||||
clean := filepath.Clean(filename)
|
||||
if clean != filename || strings.Contains(clean, "..") {
|
||||
return "", fmt.Errorf("非法文件名")
|
||||
}
|
||||
return filepath.Join(h.dir, clean), nil
|
||||
}
|
||||
|
||||
// existingOtherOrchestrator 若目录中已有别的主代理文件，返回其文件名；writingBasename 为当前正在写入的文件名时视为同一文件不冲突。
|
||||
func existingOtherOrchestrator(dir, writingBasename string) (other string, err error) {
|
||||
load, err := agents.LoadMarkdownAgentsDir(dir)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if load.Orchestrator == nil {
|
||||
return "", nil
|
||||
}
|
||||
if strings.EqualFold(load.Orchestrator.Filename, writingBasename) {
|
||||
return "", nil
|
||||
}
|
||||
return load.Orchestrator.Filename, nil
|
||||
}
|
||||
|
||||
// ListMarkdownAgents GET /api/multi-agent/markdown-agents
|
||||
func (h *MarkdownAgentsHandler) ListMarkdownAgents(c *gin.Context) {
|
||||
if h.dir == "" {
|
||||
c.JSON(http.StatusOK, gin.H{"agents": []any{}, "dir": "", "error": "未配置 agents 目录"})
|
||||
return
|
||||
}
|
||||
files, err := agents.LoadMarkdownAgentFiles(h.dir)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
out := make([]gin.H, 0, len(files))
|
||||
for _, fa := range files {
|
||||
sub := fa.Config
|
||||
out = append(out, gin.H{
|
||||
"filename": fa.Filename,
|
||||
"id": sub.ID,
|
||||
"name": sub.Name,
|
||||
"description": sub.Description,
|
||||
"is_orchestrator": fa.IsOrchestrator,
|
||||
"kind": sub.Kind,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"agents": out, "dir": h.dir})
|
||||
}
|
||||
|
||||
// GetMarkdownAgent GET /api/multi-agent/markdown-agents/:filename
|
||||
func (h *MarkdownAgentsHandler) GetMarkdownAgent(c *gin.Context) {
|
||||
filename := c.Param("filename")
|
||||
path, err := h.safeJoin(filename)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "文件不存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
sub, err := agents.ParseMarkdownSubAgent(filename, string(b))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
isOrch := agents.IsOrchestratorMarkdown(filename, agents.FrontMatter{Kind: sub.Kind})
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"filename": filename,
|
||||
"raw": string(b),
|
||||
"id": sub.ID,
|
||||
"name": sub.Name,
|
||||
"description": sub.Description,
|
||||
"tools": sub.RoleTools,
|
||||
"instruction": sub.Instruction,
|
||||
"bind_role": sub.BindRole,
|
||||
"max_iterations": sub.MaxIterations,
|
||||
"kind": sub.Kind,
|
||||
"is_orchestrator": isOrch,
|
||||
})
|
||||
}
|
||||
|
||||
type markdownAgentBody struct {
|
||||
Filename string `json:"filename"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Tools []string `json:"tools"`
|
||||
Instruction string `json:"instruction"`
|
||||
BindRole string `json:"bind_role"`
|
||||
MaxIterations int `json:"max_iterations"`
|
||||
Kind string `json:"kind"`
|
||||
Raw string `json:"raw"`
|
||||
}
|
||||
|
||||
// CreateMarkdownAgent POST /api/multi-agent/markdown-agents
|
||||
func (h *MarkdownAgentsHandler) CreateMarkdownAgent(c *gin.Context) {
|
||||
if h.dir == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "未配置 agents 目录"})
|
||||
return
|
||||
}
|
||||
var body markdownAgentBody
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
filename := strings.TrimSpace(body.Filename)
|
||||
if filename == "" {
|
||||
if strings.EqualFold(strings.TrimSpace(body.Kind), "orchestrator") {
|
||||
filename = agents.OrchestratorMarkdownFilename
|
||||
} else {
|
||||
base := agents.SlugID(body.Name)
|
||||
if base == "" {
|
||||
base = "agent"
|
||||
}
|
||||
filename = base + ".md"
|
||||
}
|
||||
}
|
||||
path, err := h.safeJoin(filename)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if _, err := os.Stat(path); err == nil {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": "文件已存在"})
|
||||
return
|
||||
}
|
||||
sub := config.MultiAgentSubConfig{
|
||||
ID: strings.TrimSpace(body.ID),
|
||||
Name: strings.TrimSpace(body.Name),
|
||||
Description: strings.TrimSpace(body.Description),
|
||||
Instruction: strings.TrimSpace(body.Instruction),
|
||||
RoleTools: body.Tools,
|
||||
BindRole: strings.TrimSpace(body.BindRole),
|
||||
MaxIterations: body.MaxIterations,
|
||||
Kind: strings.TrimSpace(body.Kind),
|
||||
}
|
||||
if strings.EqualFold(filepath.Base(path), agents.OrchestratorMarkdownFilename) && sub.Kind == "" {
|
||||
sub.Kind = "orchestrator"
|
||||
}
|
||||
if sub.ID == "" {
|
||||
sub.ID = agents.SlugID(sub.Name)
|
||||
}
|
||||
if sub.Name == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name 必填"})
|
||||
return
|
||||
}
|
||||
var out []byte
|
||||
if strings.TrimSpace(body.Raw) != "" {
|
||||
out = []byte(body.Raw)
|
||||
} else {
|
||||
out, err = agents.BuildMarkdownFile(sub)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
}
|
||||
if want := agents.WantsMarkdownOrchestrator(filepath.Base(path), body.Kind, string(out)); want {
|
||||
other, oerr := existingOtherOrchestrator(h.dir, filepath.Base(path))
|
||||
if oerr != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": oerr.Error()})
|
||||
return
|
||||
}
|
||||
if other != "" {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": fmt.Sprintf("已存在主代理定义：%s，请先删除或取消其主代理标记", other)})
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := os.MkdirAll(h.dir, 0755); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := os.WriteFile(path, out, 0644); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"filename": filepath.Base(path), "message": "已创建"})
|
||||
}
|
||||
|
||||
// UpdateMarkdownAgent PUT /api/multi-agent/markdown-agents/:filename
|
||||
func (h *MarkdownAgentsHandler) UpdateMarkdownAgent(c *gin.Context) {
|
||||
filename := c.Param("filename")
|
||||
path, err := h.safeJoin(filename)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
var body markdownAgentBody
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
sub := config.MultiAgentSubConfig{
|
||||
ID: strings.TrimSpace(body.ID),
|
||||
Name: strings.TrimSpace(body.Name),
|
||||
Description: strings.TrimSpace(body.Description),
|
||||
Instruction: strings.TrimSpace(body.Instruction),
|
||||
RoleTools: body.Tools,
|
||||
BindRole: strings.TrimSpace(body.BindRole),
|
||||
MaxIterations: body.MaxIterations,
|
||||
Kind: strings.TrimSpace(body.Kind),
|
||||
}
|
||||
if strings.EqualFold(filename, agents.OrchestratorMarkdownFilename) && sub.Kind == "" {
|
||||
sub.Kind = "orchestrator"
|
||||
}
|
||||
if sub.Name == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name 必填"})
|
||||
return
|
||||
}
|
||||
if sub.ID == "" {
|
||||
sub.ID = agents.SlugID(sub.Name)
|
||||
}
|
||||
var out []byte
|
||||
if strings.TrimSpace(body.Raw) != "" {
|
||||
out = []byte(body.Raw)
|
||||
} else {
|
||||
out, err = agents.BuildMarkdownFile(sub)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
}
|
||||
if want := agents.WantsMarkdownOrchestrator(filename, body.Kind, string(out)); want {
|
||||
other, oerr := existingOtherOrchestrator(h.dir, filename)
|
||||
if oerr != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": oerr.Error()})
|
||||
return
|
||||
}
|
||||
if other != "" {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": fmt.Sprintf("已存在主代理定义：%s，请先删除或取消其主代理标记", other)})
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := os.WriteFile(path, out, 0644); err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "文件不存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "已保存"})
|
||||
}
|
||||
|
||||
// DeleteMarkdownAgent DELETE /api/multi-agent/markdown-agents/:filename
|
||||
func (h *MarkdownAgentsHandler) DeleteMarkdownAgent(c *gin.Context) {
|
||||
filename := c.Param("filename")
|
||||
path, err := h.safeJoin(filename)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := os.Remove(path); err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "文件不存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "已删除"})
|
||||
}
|
||||
@@ -1,420 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// MonitorHandler 监控处理器
|
||||
type MonitorHandler struct {
|
||||
mcpServer *mcp.Server
|
||||
externalMCPMgr *mcp.ExternalMCPManager
|
||||
executor *security.Executor
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewMonitorHandler 创建新的监控处理器
|
||||
func NewMonitorHandler(mcpServer *mcp.Server, executor *security.Executor, db *database.DB, logger *zap.Logger) *MonitorHandler {
|
||||
return &MonitorHandler{
|
||||
mcpServer: mcpServer,
|
||||
externalMCPMgr: nil, // 将在创建后设置
|
||||
executor: executor,
|
||||
db: db,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// SetExternalMCPManager 设置外部MCP管理器
|
||||
func (h *MonitorHandler) SetExternalMCPManager(mgr *mcp.ExternalMCPManager) {
|
||||
h.externalMCPMgr = mgr
|
||||
}
|
||||
|
||||
// MonitorResponse 监控响应
|
||||
type MonitorResponse struct {
|
||||
Executions []*mcp.ToolExecution `json:"executions"`
|
||||
Stats map[string]*mcp.ToolStats `json:"stats"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Total int `json:"total,omitempty"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PageSize int `json:"page_size,omitempty"`
|
||||
TotalPages int `json:"total_pages,omitempty"`
|
||||
}
|
||||
|
||||
// Monitor 获取监控信息
|
||||
func (h *MonitorHandler) Monitor(c *gin.Context) {
|
||||
// 解析分页参数
|
||||
page := 1
|
||||
pageSize := 20
|
||||
if pageStr := c.Query("page"); pageStr != "" {
|
||||
if p, err := strconv.Atoi(pageStr); err == nil && p > 0 {
|
||||
page = p
|
||||
}
|
||||
}
|
||||
if pageSizeStr := c.Query("page_size"); pageSizeStr != "" {
|
||||
if ps, err := strconv.Atoi(pageSizeStr); err == nil && ps > 0 && ps <= 100 {
|
||||
pageSize = ps
|
||||
}
|
||||
}
|
||||
|
||||
// 解析状态筛选参数
|
||||
status := c.Query("status")
|
||||
// 解析工具筛选参数
|
||||
toolName := c.Query("tool")
|
||||
|
||||
executions, total := h.loadExecutionsWithPagination(page, pageSize, status, toolName)
|
||||
stats := h.loadStats()
|
||||
|
||||
totalPages := (total + pageSize - 1) / pageSize
|
||||
if totalPages == 0 {
|
||||
totalPages = 1
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, MonitorResponse{
|
||||
Executions: executions,
|
||||
Stats: stats,
|
||||
Timestamp: time.Now(),
|
||||
Total: total,
|
||||
Page: page,
|
||||
PageSize: pageSize,
|
||||
TotalPages: totalPages,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) loadExecutions() []*mcp.ToolExecution {
|
||||
executions, _ := h.loadExecutionsWithPagination(1, 1000, "", "")
|
||||
return executions
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) {
|
||||
if h.db == nil {
|
||||
allExecutions := h.mcpServer.GetAllExecutions()
|
||||
// 如果指定了状态筛选或工具筛选，先进行筛选
|
||||
if status != "" || toolName != "" {
|
||||
filtered := make([]*mcp.ToolExecution, 0)
|
||||
for _, exec := range allExecutions {
|
||||
matchStatus := status == "" || exec.Status == status
|
||||
// 支持部分匹配（模糊搜索）
|
||||
matchTool := toolName == "" || strings.Contains(strings.ToLower(exec.ToolName), strings.ToLower(toolName))
|
||||
if matchStatus && matchTool {
|
||||
filtered = append(filtered, exec)
|
||||
}
|
||||
}
|
||||
allExecutions = filtered
|
||||
}
|
||||
total := len(allExecutions)
|
||||
offset := (page - 1) * pageSize
|
||||
end := offset + pageSize
|
||||
if end > total {
|
||||
end = total
|
||||
}
|
||||
if offset >= total {
|
||||
return []*mcp.ToolExecution{}, total
|
||||
}
|
||||
return allExecutions[offset:end], total
|
||||
}
|
||||
|
||||
offset := (page - 1) * pageSize
|
||||
executions, err := h.db.LoadToolExecutionsWithPagination(offset, pageSize, status, toolName)
|
||||
if err != nil {
|
||||
h.logger.Warn("从数据库加载执行记录失败，回退到内存数据", zap.Error(err))
|
||||
allExecutions := h.mcpServer.GetAllExecutions()
|
||||
// 如果指定了状态筛选或工具筛选，先进行筛选
|
||||
if status != "" || toolName != "" {
|
||||
filtered := make([]*mcp.ToolExecution, 0)
|
||||
for _, exec := range allExecutions {
|
||||
matchStatus := status == "" || exec.Status == status
|
||||
// 支持部分匹配（模糊搜索）
|
||||
matchTool := toolName == "" || strings.Contains(strings.ToLower(exec.ToolName), strings.ToLower(toolName))
|
||||
if matchStatus && matchTool {
|
||||
filtered = append(filtered, exec)
|
||||
}
|
||||
}
|
||||
allExecutions = filtered
|
||||
}
|
||||
total := len(allExecutions)
|
||||
offset := (page - 1) * pageSize
|
||||
end := offset + pageSize
|
||||
if end > total {
|
||||
end = total
|
||||
}
|
||||
if offset >= total {
|
||||
return []*mcp.ToolExecution{}, total
|
||||
}
|
||||
return allExecutions[offset:end], total
|
||||
}
|
||||
|
||||
// 获取总数（考虑状态筛选和工具筛选）
|
||||
total, err := h.db.CountToolExecutions(status, toolName)
|
||||
if err != nil {
|
||||
h.logger.Warn("获取执行记录总数失败", zap.Error(err))
|
||||
// 回退：使用已加载的记录数估算
|
||||
total = offset + len(executions)
|
||||
if len(executions) == pageSize {
|
||||
total = offset + len(executions) + 1
|
||||
}
|
||||
}
|
||||
|
||||
return executions, total
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) loadStats() map[string]*mcp.ToolStats {
|
||||
// 合并内部MCP服务器和外部MCP管理器的统计信息
|
||||
stats := make(map[string]*mcp.ToolStats)
|
||||
|
||||
// 加载内部MCP服务器的统计信息
|
||||
if h.db == nil {
|
||||
internalStats := h.mcpServer.GetStats()
|
||||
for k, v := range internalStats {
|
||||
stats[k] = v
|
||||
}
|
||||
} else {
|
||||
dbStats, err := h.db.LoadToolStats()
|
||||
if err != nil {
|
||||
h.logger.Warn("从数据库加载统计信息失败，回退到内存数据", zap.Error(err))
|
||||
internalStats := h.mcpServer.GetStats()
|
||||
for k, v := range internalStats {
|
||||
stats[k] = v
|
||||
}
|
||||
} else {
|
||||
for k, v := range dbStats {
|
||||
stats[k] = v
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 合并外部MCP管理器的统计信息
|
||||
if h.externalMCPMgr != nil {
|
||||
externalStats := h.externalMCPMgr.GetToolStats()
|
||||
for k, v := range externalStats {
|
||||
// 如果已存在，合并统计信息
|
||||
if existing, exists := stats[k]; exists {
|
||||
existing.TotalCalls += v.TotalCalls
|
||||
existing.SuccessCalls += v.SuccessCalls
|
||||
existing.FailedCalls += v.FailedCalls
|
||||
// 使用最新的调用时间
|
||||
if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) {
|
||||
existing.LastCallTime = v.LastCallTime
|
||||
}
|
||||
} else {
|
||||
stats[k] = v
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return stats
|
||||
}
|
||||
|
||||
|
||||
// GetExecution 获取特定执行记录
|
||||
func (h *MonitorHandler) GetExecution(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
// 先从内部MCP服务器查找
|
||||
exec, exists := h.mcpServer.GetExecution(id)
|
||||
if exists {
|
||||
c.JSON(http.StatusOK, exec)
|
||||
return
|
||||
}
|
||||
|
||||
// 如果找不到，尝试从外部MCP管理器查找
|
||||
if h.externalMCPMgr != nil {
|
||||
exec, exists = h.externalMCPMgr.GetExecution(id)
|
||||
if exists {
|
||||
c.JSON(http.StatusOK, exec)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 如果都找不到，尝试从数据库查找（如果使用数据库存储）
|
||||
if h.db != nil {
|
||||
exec, err := h.db.GetToolExecution(id)
|
||||
if err == nil && exec != nil {
|
||||
c.JSON(http.StatusOK, exec)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "执行记录未找到"})
|
||||
}
|
||||
|
||||
// BatchGetToolNames 批量获取工具执行的工具名称（消除前端 N+1 请求）
|
||||
func (h *MonitorHandler) BatchGetToolNames(c *gin.Context) {
|
||||
var req struct {
|
||||
IDs []string `json:"ids"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
result := make(map[string]string, len(req.IDs))
|
||||
for _, id := range req.IDs {
|
||||
// 先从内部MCP服务器查找
|
||||
if exec, exists := h.mcpServer.GetExecution(id); exists {
|
||||
result[id] = exec.ToolName
|
||||
continue
|
||||
}
|
||||
// 再从外部MCP管理器查找
|
||||
if h.externalMCPMgr != nil {
|
||||
if exec, exists := h.externalMCPMgr.GetExecution(id); exists {
|
||||
result[id] = exec.ToolName
|
||||
continue
|
||||
}
|
||||
}
|
||||
// 最后从数据库查找
|
||||
if h.db != nil {
|
||||
if exec, err := h.db.GetToolExecution(id); err == nil && exec != nil {
|
||||
result[id] = exec.ToolName
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, result)
|
||||
}
|
||||
|
||||
// GetStats 获取统计信息
|
||||
func (h *MonitorHandler) GetStats(c *gin.Context) {
|
||||
stats := h.loadStats()
|
||||
c.JSON(http.StatusOK, stats)
|
||||
}
|
||||
|
||||
// DeleteExecution 删除执行记录
|
||||
func (h *MonitorHandler) DeleteExecution(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if id == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
// 如果使用数据库，先获取执行记录信息，然后删除并更新统计
|
||||
if h.db != nil {
|
||||
// 先获取执行记录信息（用于更新统计）
|
||||
exec, err := h.db.GetToolExecution(id)
|
||||
if err != nil {
|
||||
// 如果找不到记录，可能已经被删除，直接返回成功
|
||||
h.logger.Warn("执行记录不存在，可能已被删除", zap.String("executionId", id), zap.Error(err))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "执行记录不存在或已被删除"})
|
||||
return
|
||||
}
|
||||
|
||||
// 删除执行记录
|
||||
err = h.db.DeleteToolExecution(id)
|
||||
if err != nil {
|
||||
h.logger.Error("删除执行记录失败", zap.Error(err), zap.String("executionId", id))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "删除执行记录失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 更新统计信息（减少相应的计数）
|
||||
totalCalls := 1
|
||||
successCalls := 0
|
||||
failedCalls := 0
|
||||
if exec.Status == "failed" {
|
||||
failedCalls = 1
|
||||
} else if exec.Status == "completed" {
|
||||
successCalls = 1
|
||||
}
|
||||
|
||||
if exec.ToolName != "" {
|
||||
if err := h.db.DecreaseToolStats(exec.ToolName, totalCalls, successCalls, failedCalls); err != nil {
|
||||
h.logger.Warn("更新统计信息失败", zap.Error(err), zap.String("toolName", exec.ToolName))
|
||||
// 不返回错误，因为记录已经删除成功
|
||||
}
|
||||
}
|
||||
|
||||
h.logger.Info("执行记录已从数据库删除", zap.String("executionId", id), zap.String("toolName", exec.ToolName))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除"})
|
||||
return
|
||||
}
|
||||
|
||||
// 如果不使用数据库，尝试从内存中删除（内部MCP服务器）
|
||||
// 注意：内存中的记录可能已经被清理，所以这里只记录日志
|
||||
h.logger.Info("尝试删除内存中的执行记录", zap.String("executionId", id))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除（如果存在）"})
|
||||
}
|
||||
|
||||
// DeleteExecutions 批量删除执行记录
|
||||
func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
|
||||
var request struct {
|
||||
IDs []string `json:"ids"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "请求参数无效: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if len(request.IDs) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID列表不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
// 如果使用数据库，先获取执行记录信息，然后删除并更新统计
|
||||
if h.db != nil {
|
||||
// 先获取执行记录信息（用于更新统计）
|
||||
executions, err := h.db.GetToolExecutionsByIds(request.IDs)
|
||||
if err != nil {
|
||||
h.logger.Error("获取执行记录失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "获取执行记录失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 按工具名称分组统计需要减少的数量
|
||||
toolStats := make(map[string]struct {
|
||||
totalCalls int
|
||||
successCalls int
|
||||
failedCalls int
|
||||
})
|
||||
|
||||
for _, exec := range executions {
|
||||
if exec.ToolName == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
stats := toolStats[exec.ToolName]
|
||||
stats.totalCalls++
|
||||
if exec.Status == "failed" {
|
||||
stats.failedCalls++
|
||||
} else if exec.Status == "completed" {
|
||||
stats.successCalls++
|
||||
}
|
||||
toolStats[exec.ToolName] = stats
|
||||
}
|
||||
|
||||
// 批量删除执行记录
|
||||
err = h.db.DeleteToolExecutions(request.IDs)
|
||||
if err != nil {
|
||||
h.logger.Error("批量删除执行记录失败", zap.Error(err), zap.Int("count", len(request.IDs)))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "批量删除执行记录失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 更新统计信息（减少相应的计数）
|
||||
for toolName, stats := range toolStats {
|
||||
if err := h.db.DecreaseToolStats(toolName, stats.totalCalls, stats.successCalls, stats.failedCalls); err != nil {
|
||||
h.logger.Warn("更新统计信息失败", zap.Error(err), zap.String("toolName", toolName))
|
||||
// 不返回错误，因为记录已经删除成功
|
||||
}
|
||||
}
|
||||
|
||||
h.logger.Info("批量删除执行记录成功", zap.Int("count", len(request.IDs)))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "成功删除执行记录", "deleted": len(executions)})
|
||||
return
|
||||
}
|
||||
|
||||
// 如果不使用数据库，尝试从内存中删除（内部MCP服务器）
|
||||
// 注意：内存中的记录可能已经被清理，所以这里只记录日志
|
||||
h.logger.Info("尝试批量删除内存中的执行记录", zap.Int("count", len(request.IDs)))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除（如果存在）"})
|
||||
}
|
||||
|
||||
|
||||
@@ -1,316 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/multiagent"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// MultiAgentLoopStream Eino DeepAgent 流式对话（需 config.multi_agent.enabled）。
|
||||
func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
if h.config == nil || !h.config.MultiAgent.Enabled {
|
||||
ev := StreamEvent{Type: "error", Message: "多代理未启用，请在设置或 config.yaml 中开启 multi_agent.enabled"}
|
||||
b, _ := json.Marshal(ev)
|
||||
fmt.Fprintf(c.Writer, "data: %s\n\n", b)
|
||||
done := StreamEvent{Type: "done", Message: ""}
|
||||
db, _ := json.Marshal(done)
|
||||
fmt.Fprintf(c.Writer, "data: %s\n\n", db)
|
||||
if flusher, ok := c.Writer.(http.Flusher); ok {
|
||||
flusher.Flush()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
var req ChatRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
event := StreamEvent{Type: "error", Message: "请求参数错误: " + err.Error()}
|
||||
b, _ := json.Marshal(event)
|
||||
fmt.Fprintf(c.Writer, "data: %s\n\n", b)
|
||||
c.Writer.Flush()
|
||||
return
|
||||
}
|
||||
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
|
||||
// 用于在 sendEvent 中判断是否为用户主动停止导致的取消。
|
||||
// 注意：baseCtx 会在后面创建；该变量用于闭包提前捕获引用。
|
||||
var baseCtx context.Context
|
||||
|
||||
clientDisconnected := false
|
||||
// 与 sseKeepalive 共用：禁止并发写 ResponseWriter，否则会破坏 chunked 编码（ERR_INVALID_CHUNKED_ENCODING）。
|
||||
var sseWriteMu sync.Mutex
|
||||
sendEvent := func(eventType, message string, data interface{}) {
|
||||
if clientDisconnected {
|
||||
return
|
||||
}
|
||||
// 用户主动停止时，Eino 可能仍会并发上报 eventType=="error"。
|
||||
// 为避免 UI 看到“取消错误 + cancelled 文案”两条回复，这里直接丢弃取消对应的 error。
|
||||
if eventType == "error" && baseCtx != nil && errors.Is(context.Cause(baseCtx), ErrTaskCancelled) {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
clientDisconnected = true
|
||||
return
|
||||
default:
|
||||
}
|
||||
ev := StreamEvent{Type: eventType, Message: message, Data: data}
|
||||
b, _ := json.Marshal(ev)
|
||||
sseWriteMu.Lock()
|
||||
_, err := fmt.Fprintf(c.Writer, "data: %s\n\n", b)
|
||||
if err != nil {
|
||||
sseWriteMu.Unlock()
|
||||
clientDisconnected = true
|
||||
return
|
||||
}
|
||||
if flusher, ok := c.Writer.(http.Flusher); ok {
|
||||
flusher.Flush()
|
||||
} else {
|
||||
c.Writer.Flush()
|
||||
}
|
||||
sseWriteMu.Unlock()
|
||||
}
|
||||
|
||||
h.logger.Info("收到 Eino DeepAgent 流式请求",
|
||||
zap.String("conversationId", req.ConversationID),
|
||||
)
|
||||
|
||||
prep, err := h.prepareMultiAgentSession(&req)
|
||||
if err != nil {
|
||||
sendEvent("error", err.Error(), nil)
|
||||
sendEvent("done", "", nil)
|
||||
return
|
||||
}
|
||||
if prep.CreatedNew {
|
||||
sendEvent("conversation", "会话已创建", map[string]interface{}{
|
||||
"conversationId": prep.ConversationID,
|
||||
})
|
||||
}
|
||||
|
||||
conversationID := prep.ConversationID
|
||||
assistantMessageID := prep.AssistantMessageID
|
||||
|
||||
if prep.UserMessageID != "" {
|
||||
sendEvent("message_saved", "", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"userMessageId": prep.UserMessageID,
|
||||
})
|
||||
}
|
||||
|
||||
progressCallback := h.createProgressCallback(conversationID, assistantMessageID, sendEvent)
|
||||
|
||||
baseCtx, cancelWithCause := context.WithCancelCause(context.Background())
|
||||
taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
defer timeoutCancel()
|
||||
defer cancelWithCause(nil)
|
||||
|
||||
if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil {
|
||||
var errorMsg string
|
||||
if errors.Is(err, ErrTaskAlreadyRunning) {
|
||||
errorMsg = "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。"
|
||||
sendEvent("error", errorMsg, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"errorType": "task_already_running",
|
||||
})
|
||||
} else {
|
||||
errorMsg = "❌ 无法启动任务: " + err.Error()
|
||||
sendEvent("error", errorMsg, nil)
|
||||
}
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ? WHERE id = ?", errorMsg, assistantMessageID)
|
||||
}
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
return
|
||||
}
|
||||
|
||||
taskStatus := "completed"
|
||||
defer h.tasks.FinishTask(conversationID, taskStatus)
|
||||
|
||||
sendEvent("progress", "正在启动 Eino DeepAgent...", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
|
||||
stopKeepalive := make(chan struct{})
|
||||
go sseKeepalive(c, stopKeepalive, &sseWriteMu)
|
||||
defer close(stopKeepalive)
|
||||
|
||||
result, runErr := multiagent.RunDeepAgent(
|
||||
taskCtx,
|
||||
h.config,
|
||||
&h.config.MultiAgent,
|
||||
h.agent,
|
||||
h.logger,
|
||||
conversationID,
|
||||
prep.FinalMessage,
|
||||
prep.History,
|
||||
prep.RoleTools,
|
||||
progressCallback,
|
||||
h.agentsMarkdownDir,
|
||||
)
|
||||
|
||||
if runErr != nil {
|
||||
cause := context.Cause(baseCtx)
|
||||
if errors.Is(cause, ErrTaskCancelled) {
|
||||
taskStatus = "cancelled"
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
cancelMsg := "任务已被用户取消，后续操作已停止。"
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ? WHERE id = ?", cancelMsg, assistantMessageID)
|
||||
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil)
|
||||
}
|
||||
sendEvent("cancelled", cancelMsg, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr))
|
||||
taskStatus = "failed"
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
errMsg := "执行失败: " + runErr.Error()
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ? WHERE id = ?", errMsg, assistantMessageID)
|
||||
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil)
|
||||
}
|
||||
sendEvent("error", errMsg, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
return
|
||||
}
|
||||
|
||||
if assistantMessageID != "" {
|
||||
mcpIDsJSON := ""
|
||||
if len(result.MCPExecutionIDs) > 0 {
|
||||
jsonData, _ := json.Marshal(result.MCPExecutionIDs)
|
||||
mcpIDsJSON = string(jsonData)
|
||||
}
|
||||
_, _ = h.db.Exec(
|
||||
"UPDATE messages SET content = ?, mcp_execution_ids = ? WHERE id = ?",
|
||||
result.Response,
|
||||
mcpIDsJSON,
|
||||
assistantMessageID,
|
||||
)
|
||||
}
|
||||
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存 ReAct 数据失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
sendEvent("response", result.Response, map[string]interface{}{
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": "eino_deep",
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
}
|
||||
|
||||
// MultiAgentLoop Eino DeepAgent 非流式对话（与 POST /api/agent-loop 对齐，需 multi_agent.enabled）。
|
||||
func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
if h.config == nil || !h.config.MultiAgent.Enabled {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "多代理未启用，请在 config.yaml 中设置 multi_agent.enabled: true"})
|
||||
return
|
||||
}
|
||||
|
||||
var req ChatRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("收到 Eino DeepAgent 非流式请求", zap.String("conversationId", req.ConversationID))
|
||||
|
||||
prep, err := h.prepareMultiAgentSession(&req)
|
||||
if err != nil {
|
||||
status, msg := multiAgentHTTPErrorStatus(err)
|
||||
c.JSON(status, gin.H{"error": msg})
|
||||
return
|
||||
}
|
||||
|
||||
result, runErr := multiagent.RunDeepAgent(
|
||||
c.Request.Context(),
|
||||
h.config,
|
||||
&h.config.MultiAgent,
|
||||
h.agent,
|
||||
h.logger,
|
||||
prep.ConversationID,
|
||||
prep.FinalMessage,
|
||||
prep.History,
|
||||
prep.RoleTools,
|
||||
nil,
|
||||
h.agentsMarkdownDir,
|
||||
)
|
||||
if runErr != nil {
|
||||
h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr))
|
||||
errMsg := "执行失败: " + runErr.Error()
|
||||
if prep.AssistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ? WHERE id = ?", errMsg, prep.AssistantMessageID)
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg})
|
||||
return
|
||||
}
|
||||
|
||||
if prep.AssistantMessageID != "" {
|
||||
mcpIDsJSON := ""
|
||||
if len(result.MCPExecutionIDs) > 0 {
|
||||
jsonData, _ := json.Marshal(result.MCPExecutionIDs)
|
||||
mcpIDsJSON = string(jsonData)
|
||||
}
|
||||
_, _ = h.db.Exec(
|
||||
"UPDATE messages SET content = ?, mcp_execution_ids = ? WHERE id = ?",
|
||||
result.Response,
|
||||
mcpIDsJSON,
|
||||
prep.AssistantMessageID,
|
||||
)
|
||||
}
|
||||
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(prep.ConversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存 ReAct 数据失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, ChatResponse{
|
||||
Response: result.Response,
|
||||
MCPExecutionIDs: result.MCPExecutionIDs,
|
||||
ConversationID: prep.ConversationID,
|
||||
Time: time.Now(),
|
||||
})
|
||||
}
|
||||
|
||||
func multiAgentHTTPErrorStatus(err error) (int, string) {
|
||||
msg := err.Error()
|
||||
switch {
|
||||
case strings.Contains(msg, "对话不存在"):
|
||||
return http.StatusNotFound, msg
|
||||
case strings.Contains(msg, "未找到该 WebShell"):
|
||||
return http.StatusBadRequest, msg
|
||||
case strings.Contains(msg, "附件最多"):
|
||||
return http.StatusBadRequest, msg
|
||||
case strings.Contains(msg, "保存用户消息失败"), strings.Contains(msg, "创建对话失败"):
|
||||
return http.StatusInternalServerError, msg
|
||||
case strings.Contains(msg, "保存上传文件失败"):
|
||||
return http.StatusInternalServerError, msg
|
||||
default:
|
||||
return http.StatusBadRequest, msg
|
||||
}
|
||||
}
|
||||
@@ -1,138 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// multiAgentPrepared 多代理请求在调用 Eino 前的会话与消息准备结果。
|
||||
type multiAgentPrepared struct {
|
||||
ConversationID string
|
||||
CreatedNew bool
|
||||
History []agent.ChatMessage
|
||||
FinalMessage string
|
||||
RoleTools []string
|
||||
AssistantMessageID string
|
||||
UserMessageID string
|
||||
}
|
||||
|
||||
func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest) (*multiAgentPrepared, error) {
|
||||
if len(req.Attachments) > maxAttachments {
|
||||
return nil, fmt.Errorf("附件最多 %d 个", maxAttachments)
|
||||
}
|
||||
|
||||
conversationID := strings.TrimSpace(req.ConversationID)
|
||||
createdNew := false
|
||||
if conversationID == "" {
|
||||
title := safeTruncateString(req.Message, 50)
|
||||
var conv *database.Conversation
|
||||
var err error
|
||||
if strings.TrimSpace(req.WebShellConnectionID) != "" {
|
||||
conv, err = h.db.CreateConversationWithWebshell(strings.TrimSpace(req.WebShellConnectionID), title)
|
||||
} else {
|
||||
conv, err = h.db.CreateConversation(title)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建对话失败: %w", err)
|
||||
}
|
||||
conversationID = conv.ID
|
||||
createdNew = true
|
||||
} else {
|
||||
if _, err := h.db.GetConversation(conversationID); err != nil {
|
||||
return nil, 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 := req.Message
|
||||
var roleTools []string
|
||||
if req.WebShellConnectionID != "" {
|
||||
conn, errConn := h.db.GetWebshellConnection(strings.TrimSpace(req.WebShellConnectionID))
|
||||
if errConn != nil || conn == nil {
|
||||
h.logger.Warn("WebShell AI 助手：未找到连接", zap.String("id", req.WebShellConnectionID), zap.Error(errConn))
|
||||
return nil, fmt.Errorf("未找到该 WebShell 连接")
|
||||
}
|
||||
remark := conn.Remark
|
||||
if remark == "" {
|
||||
remark = conn.URL
|
||||
}
|
||||
finalMessage = fmt.Sprintf("[WebShell 助手上下文] 当前连接 ID：%s，备注：%s。可用工具（仅在该连接上操作时使用，connection_id 填 \"%s\"）：webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write、record_vulnerability、list_knowledge_risk_types、search_knowledge_base。Skills 包请使用 Eino 多代理内置 `skill` 工具。\n\n用户请求：%s",
|
||||
conn.ID, remark, conn.ID, req.Message)
|
||||
roleTools = []string{
|
||||
builtin.ToolWebshellExec,
|
||||
builtin.ToolWebshellFileList,
|
||||
builtin.ToolWebshellFileRead,
|
||||
builtin.ToolWebshellFileWrite,
|
||||
builtin.ToolRecordVulnerability,
|
||||
builtin.ToolListKnowledgeRiskTypes,
|
||||
builtin.ToolSearchKnowledgeBase,
|
||||
}
|
||||
} else if req.Role != "" && req.Role != "默认" && h.config != nil && h.config.Roles != nil {
|
||||
if role, exists := h.config.Roles[req.Role]; exists && role.Enabled {
|
||||
if role.UserPrompt != "" {
|
||||
finalMessage = role.UserPrompt + "\n\n" + req.Message
|
||||
}
|
||||
roleTools = role.Tools
|
||||
}
|
||||
}
|
||||
|
||||
var savedPaths []string
|
||||
if len(req.Attachments) > 0 {
|
||||
var aerr error
|
||||
savedPaths, aerr = saveAttachmentsToDateAndConversationDir(req.Attachments, conversationID, h.logger)
|
||||
if aerr != nil {
|
||||
return nil, fmt.Errorf("保存上传文件失败: %w", aerr)
|
||||
}
|
||||
}
|
||||
finalMessage = appendAttachmentsToMessage(finalMessage, req.Attachments, savedPaths)
|
||||
|
||||
userContent := userMessageContentForStorage(req.Message, req.Attachments, savedPaths)
|
||||
userMsgRow, uerr := h.db.AddMessage(conversationID, "user", userContent, nil)
|
||||
if uerr != nil {
|
||||
h.logger.Error("保存用户消息失败", zap.Error(uerr))
|
||||
return nil, fmt.Errorf("保存用户消息失败: %w", uerr)
|
||||
}
|
||||
userMessageID := ""
|
||||
if userMsgRow != nil {
|
||||
userMessageID = userMsgRow.ID
|
||||
}
|
||||
|
||||
assistantMsg, aerr := h.db.AddMessage(conversationID, "assistant", "处理中...", nil)
|
||||
var assistantMessageID string
|
||||
if aerr != nil {
|
||||
h.logger.Warn("创建助手消息占位失败", zap.Error(aerr))
|
||||
} else if assistantMsg != nil {
|
||||
assistantMessageID = assistantMsg.ID
|
||||
}
|
||||
|
||||
return &multiAgentPrepared{
|
||||
ConversationID: conversationID,
|
||||
CreatedNew: createdNew,
|
||||
History: agentHistoryMessages,
|
||||
FinalMessage: finalMessage,
|
||||
RoleTools: roleTools,
|
||||
AssistantMessageID: assistantMessageID,
|
||||
UserMessageID: userMessageID,
|
||||
}, nil
|
||||
}
|
||||
-4596
File diff suppressed because it is too large
Load Diff
@@ -1,139 +0,0 @@
|
||||
package handler
|
||||
|
||||
// apiDocI18n 为 OpenAPI 文档提供 x-i18n-* 扩展键，供前端 apiDocs 国际化使用。
|
||||
// 前端通过 apiDocs.tags.* / apiDocs.summary.* / apiDocs.response.* 翻译。
|
||||
|
||||
var apiDocI18nTagToKey = map[string]string{
|
||||
"认证": "auth", "对话管理": "conversationManagement", "对话交互": "conversationInteraction",
|
||||
"批量任务": "batchTasks", "对话分组": "conversationGroups", "漏洞管理": "vulnerabilityManagement",
|
||||
"角色管理": "roleManagement", "Skills管理": "skillsManagement", "监控": "monitoring",
|
||||
"配置管理": "configManagement", "外部MCP管理": "externalMCPManagement", "攻击链": "attackChain",
|
||||
"知识库": "knowledgeBase", "MCP": "mcp",
|
||||
}
|
||||
|
||||
var apiDocI18nSummaryToKey = map[string]string{
|
||||
"用户登录": "login", "用户登出": "logout", "修改密码": "changePassword", "验证Token": "validateToken",
|
||||
"创建对话": "createConversation", "列出对话": "listConversations", "查看对话详情": "getConversationDetail",
|
||||
"更新对话": "updateConversation", "删除对话": "deleteConversation", "获取对话结果": "getConversationResult",
|
||||
"发送消息并获取AI回复（非流式）": "sendMessageNonStream", "发送消息并获取AI回复（流式）": "sendMessageStream",
|
||||
"取消任务": "cancelTask", "列出运行中的任务": "listRunningTasks", "列出已完成的任务": "listCompletedTasks",
|
||||
"创建批量任务队列": "createBatchQueue", "列出批量任务队列": "listBatchQueues", "获取批量任务队列": "getBatchQueue",
|
||||
"删除批量任务队列": "deleteBatchQueue", "启动批量任务队列": "startBatchQueue", "暂停批量任务队列": "pauseBatchQueue",
|
||||
"添加任务到队列": "addTaskToQueue", "SQL注入扫描": "sqlInjectionScan", "端口扫描": "portScan",
|
||||
"更新批量任务": "updateBatchTask", "删除批量任务": "deleteBatchTask",
|
||||
"创建分组": "createGroup", "列出分组": "listGroups", "获取分组": "getGroup", "更新分组": "updateGroup",
|
||||
"删除分组": "deleteGroup", "获取分组中的对话": "getGroupConversations", "添加对话到分组": "addConversationToGroup",
|
||||
"从分组移除对话": "removeConversationFromGroup",
|
||||
"列出漏洞": "listVulnerabilities", "创建漏洞": "createVulnerability", "获取漏洞统计": "getVulnerabilityStats",
|
||||
"获取漏洞": "getVulnerability", "更新漏洞": "updateVulnerability", "删除漏洞": "deleteVulnerability",
|
||||
"列出角色": "listRoles", "创建角色": "createRole", "获取角色": "getRole", "更新角色": "updateRole", "删除角色": "deleteRole",
|
||||
"获取可用Skills列表": "getAvailableSkills", "列出Skills": "listSkills", "创建Skill": "createSkill",
|
||||
"获取Skill统计": "getSkillStats", "清空Skill统计": "clearSkillStats", "获取Skill": "getSkill",
|
||||
"更新Skill": "updateSkill", "删除Skill": "deleteSkill", "获取绑定角色": "getBoundRoles",
|
||||
"获取监控信息": "getMonitorInfo", "获取执行记录": "getExecutionRecords", "删除执行记录": "deleteExecutionRecord",
|
||||
"批量删除执行记录": "batchDeleteExecutionRecords", "获取统计信息": "getStats",
|
||||
"获取配置": "getConfig", "更新配置": "updateConfig", "获取工具配置": "getToolConfig", "应用配置": "applyConfig",
|
||||
"列出外部MCP": "listExternalMCP", "获取外部MCP统计": "getExternalMCPStats", "获取外部MCP": "getExternalMCP",
|
||||
"添加或更新外部MCP": "addOrUpdateExternalMCP", "stdio模式配置": "stdioModeConfig", "SSE模式配置": "sseModeConfig",
|
||||
"删除外部MCP": "deleteExternalMCP", "启动外部MCP": "startExternalMCP", "停止外部MCP": "stopExternalMCP",
|
||||
"获取攻击链": "getAttackChain", "重新生成攻击链": "regenerateAttackChain",
|
||||
"设置对话置顶": "pinConversation", "设置分组置顶": "pinGroup", "设置分组中对话的置顶": "pinGroupConversation",
|
||||
"获取分类": "getCategories", "列出知识项": "listKnowledgeItems", "创建知识项": "createKnowledgeItem",
|
||||
"获取知识项": "getKnowledgeItem", "更新知识项": "updateKnowledgeItem", "删除知识项": "deleteKnowledgeItem",
|
||||
"获取索引状态": "getIndexStatus", "重建索引": "rebuildIndex", "扫描知识库": "scanKnowledgeBase",
|
||||
"搜索知识库": "searchKnowledgeBase", "基础搜索": "basicSearch", "按风险类型搜索": "searchByRiskType",
|
||||
"获取检索日志": "getRetrievalLogs", "删除检索日志": "deleteRetrievalLog",
|
||||
"MCP端点": "mcpEndpoint", "列出所有工具": "listAllTools", "调用工具": "invokeTool", "初始化连接": "initConnection",
|
||||
"成功响应": "successResponse", "错误响应": "errorResponse",
|
||||
}
|
||||
|
||||
var apiDocI18nResponseDescToKey = map[string]string{
|
||||
"获取成功": "getSuccess", "未授权": "unauthorized", "未授权，需要有效的Token": "unauthorizedToken",
|
||||
"创建成功": "createSuccess", "请求参数错误": "badRequest", "对话不存在": "conversationNotFound",
|
||||
"对话不存在或结果不存在": "conversationOrResultNotFound", "请求参数错误（如task为空）": "badRequestTaskEmpty",
|
||||
"请求参数错误或分组名称已存在": "badRequestGroupNameExists", "分组不存在": "groupNotFound",
|
||||
"请求参数错误（如配置格式不正确、缺少必需字段等）": "badRequestConfig",
|
||||
"请求参数错误（如query为空）": "badRequestQueryEmpty", "方法不允许（仅支持POST请求）": "methodNotAllowed",
|
||||
"登录成功": "loginSuccess", "密码错误": "invalidPassword", "登出成功": "logoutSuccess",
|
||||
"密码修改成功": "passwordChanged", "Token有效": "tokenValid", "Token无效或已过期": "tokenInvalid",
|
||||
"对话创建成功": "conversationCreated", "服务器内部错误": "internalError", "更新成功": "updateSuccess",
|
||||
"删除成功": "deleteSuccess", "队列不存在": "queueNotFound", "启动成功": "startSuccess",
|
||||
"暂停成功": "pauseSuccess", "添加成功": "addSuccess",
|
||||
"任务不存在": "taskNotFound", "对话或分组不存在": "conversationOrGroupNotFound",
|
||||
"取消请求已提交": "cancelSubmitted", "未找到正在执行的任务": "noRunningTask",
|
||||
"消息发送成功，返回AI回复": "messageSent", "流式响应（Server-Sent Events）": "streamResponse",
|
||||
}
|
||||
|
||||
// enrichSpecWithI18nKeys 在 spec 的每个 operation 上写入 x-i18n-tags、x-i18n-summary，
|
||||
// 在每个 response 上写入 x-i18n-description，供前端按 key 做国际化。
|
||||
func enrichSpecWithI18nKeys(spec map[string]interface{}) {
|
||||
paths, _ := spec["paths"].(map[string]interface{})
|
||||
if paths == nil {
|
||||
return
|
||||
}
|
||||
for _, pathItem := range paths {
|
||||
pm, _ := pathItem.(map[string]interface{})
|
||||
if pm == nil {
|
||||
continue
|
||||
}
|
||||
for _, method := range []string{"get", "post", "put", "delete", "patch"} {
|
||||
opVal, ok := pm[method]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
op, _ := opVal.(map[string]interface{})
|
||||
if op == nil {
|
||||
continue
|
||||
}
|
||||
// x-i18n-tags: 与 tags 一一对应的 i18n 键数组（spec 中 tags 为 []string）
|
||||
switch tags := op["tags"].(type) {
|
||||
case []string:
|
||||
if len(tags) > 0 {
|
||||
keys := make([]string, 0, len(tags))
|
||||
for _, s := range tags {
|
||||
if k := apiDocI18nTagToKey[s]; k != "" {
|
||||
keys = append(keys, k)
|
||||
} else {
|
||||
keys = append(keys, s)
|
||||
}
|
||||
}
|
||||
op["x-i18n-tags"] = keys
|
||||
}
|
||||
case []interface{}:
|
||||
if len(tags) > 0 {
|
||||
keys := make([]interface{}, 0, len(tags))
|
||||
for _, t := range tags {
|
||||
if s, ok := t.(string); ok {
|
||||
if k := apiDocI18nTagToKey[s]; k != "" {
|
||||
keys = append(keys, k)
|
||||
} else {
|
||||
keys = append(keys, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(keys) > 0 {
|
||||
op["x-i18n-tags"] = keys
|
||||
}
|
||||
}
|
||||
}
|
||||
// x-i18n-summary
|
||||
if summary, _ := op["summary"].(string); summary != "" {
|
||||
if k := apiDocI18nSummaryToKey[summary]; k != "" {
|
||||
op["x-i18n-summary"] = k
|
||||
}
|
||||
}
|
||||
// responses -> 每个 status -> x-i18n-description
|
||||
if respMap, _ := op["responses"].(map[string]interface{}); respMap != nil {
|
||||
for _, rv := range respMap {
|
||||
if r, _ := rv.(map[string]interface{}); r != nil {
|
||||
if desc, _ := r["description"].(string); desc != "" {
|
||||
if k := apiDocI18nResponseDescToKey[desc]; k != "" {
|
||||
r["x-i18n-description"] = k
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,907 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"crypto/sha1"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"encoding/xml"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"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 = "当前"
|
||||
robotCmdStop = "停止"
|
||||
robotCmdRoles = "角色"
|
||||
robotCmdRolesList = "角色列表"
|
||||
robotCmdSwitchRole = "切换角色"
|
||||
robotCmdDelete = "删除"
|
||||
robotCmdVersion = "版本"
|
||||
)
|
||||
|
||||
// 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
|
||||
sessionRoles map[string]string // key: "platform_userID", value: roleName（默认"默认"）
|
||||
cancelMu sync.Mutex // 保护 runningCancels
|
||||
runningCancels map[string]context.CancelFunc // key: "platform_userID", 用于停止命令中断任务
|
||||
}
|
||||
|
||||
// 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),
|
||||
sessionRoles: make(map[string]string),
|
||||
runningCancels: make(map[string]context.CancelFunc),
|
||||
}
|
||||
}
|
||||
|
||||
// sessionKey 生成会话 key
|
||||
func (h *RobotHandler) sessionKey(platform, userID string) string {
|
||||
return platform + "_" + userID
|
||||
}
|
||||
|
||||
// getOrCreateConversation 获取或创建当前会话，title 用于新对话的标题（取用户首条消息前50字）
|
||||
func (h *RobotHandler) getOrCreateConversation(platform, userID, title string) (convID string, isNew bool) {
|
||||
h.mu.RLock()
|
||||
convID = h.sessions[h.sessionKey(platform, userID)]
|
||||
h.mu.RUnlock()
|
||||
if convID != "" {
|
||||
return convID, false
|
||||
}
|
||||
t := strings.TrimSpace(title)
|
||||
if t == "" {
|
||||
t = "新对话 " + time.Now().Format("01-02 15:04")
|
||||
} else {
|
||||
t = safeTruncateString(t, 50)
|
||||
}
|
||||
conv, err := h.db.CreateConversation(t)
|
||||
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()
|
||||
}
|
||||
|
||||
// getRole 获取当前用户使用的角色，未设置时返回"默认"
|
||||
func (h *RobotHandler) getRole(platform, userID string) string {
|
||||
h.mu.RLock()
|
||||
role := h.sessionRoles[h.sessionKey(platform, userID)]
|
||||
h.mu.RUnlock()
|
||||
if role == "" {
|
||||
return "默认"
|
||||
}
|
||||
return role
|
||||
}
|
||||
|
||||
// setRole 设置当前用户使用的角色
|
||||
func (h *RobotHandler) setRole(platform, userID, roleName string) {
|
||||
h.mu.Lock()
|
||||
h.sessionRoles[h.sessionKey(platform, userID)] = roleName
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// clearConversation 清空当前会话（切换到新对话）
|
||||
func (h *RobotHandler) clearConversation(platform, userID string) (newConvID string) {
|
||||
title := "新对话 " + time.Now().Format("01-02 15:04")
|
||||
conv, err := h.db.CreateConversation(title)
|
||||
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 "请输入内容或发送「帮助」/ help 查看命令。"
|
||||
}
|
||||
|
||||
// 先尝试作为命令处理（支持中英文）
|
||||
if cmdReply, ok := h.handleRobotCommand(platform, userID, text); ok {
|
||||
return cmdReply
|
||||
}
|
||||
|
||||
// 普通消息：走 Agent
|
||||
convID, _ := h.getOrCreateConversation(platform, userID, text)
|
||||
if convID == "" {
|
||||
return "无法创建或获取对话，请稍后再试。"
|
||||
}
|
||||
// 若对话标题为「新对话 xx:xx」格式（由「新对话」命令创建），将标题更新为首条消息内容，与 Web 端体验一致
|
||||
if conv, err := h.db.GetConversation(convID); err == nil && strings.HasPrefix(conv.Title, "新对话 ") {
|
||||
newTitle := safeTruncateString(text, 50)
|
||||
if newTitle != "" {
|
||||
_ = h.db.UpdateConversationTitle(convID, newTitle)
|
||||
}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
sk := h.sessionKey(platform, userID)
|
||||
h.cancelMu.Lock()
|
||||
h.runningCancels[sk] = cancel
|
||||
h.cancelMu.Unlock()
|
||||
defer func() {
|
||||
cancel()
|
||||
h.cancelMu.Lock()
|
||||
delete(h.runningCancels, sk)
|
||||
h.cancelMu.Unlock()
|
||||
}()
|
||||
role := h.getRole(platform, userID)
|
||||
resp, newConvID, err := h.agentHandler.ProcessMessageForRobot(ctx, convID, text, role)
|
||||
if err != nil {
|
||||
h.logger.Warn("机器人 Agent 执行失败", zap.String("platform", platform), zap.String("userID", userID), zap.Error(err))
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return "任务已取消。"
|
||||
}
|
||||
return "处理失败: " + err.Error()
|
||||
}
|
||||
if newConvID != convID {
|
||||
h.setConversation(platform, userID, newConvID)
|
||||
}
|
||||
return resp
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdHelp() string {
|
||||
return "**【CyberStrikeAI 机器人命令】**\n\n" +
|
||||
"- `帮助` `help` — 显示本帮助 | Show this help\n" +
|
||||
"- `列表` `list` — 列出所有对话标题与 ID | List conversations\n" +
|
||||
"- `切换 <ID>` `switch <ID>` — 指定对话继续 | Switch to conversation\n" +
|
||||
"- `新对话` `new` — 开启新对话 | Start new conversation\n" +
|
||||
"- `清空` `clear` — 清空当前上下文 | Clear context\n" +
|
||||
"- `当前` `current` — 显示当前对话 ID 与标题 | Show current conversation\n" +
|
||||
"- `停止` `stop` — 中断当前任务 | Stop running task\n" +
|
||||
"- `角色` `roles` — 列出所有可用角色 | List roles\n" +
|
||||
"- `角色 <名>` `role <name>` — 切换当前角色 | Switch role\n" +
|
||||
"- `删除 <ID>` `delete <ID>` — 删除指定对话 | Delete conversation\n" +
|
||||
"- `版本` `version` — 显示当前版本号 | Show version\n\n" +
|
||||
"---\n" +
|
||||
"除以上命令外，直接输入内容将发送给 AI 进行渗透测试/安全分析。\n" +
|
||||
"Otherwise, send any text for AI penetration testing / security analysis."
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdList() 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) cmdStop(platform, userID string) string {
|
||||
sk := h.sessionKey(platform, userID)
|
||||
h.cancelMu.Lock()
|
||||
cancel, ok := h.runningCancels[sk]
|
||||
if ok {
|
||||
delete(h.runningCancels, sk)
|
||||
cancel()
|
||||
}
|
||||
h.cancelMu.Unlock()
|
||||
if !ok {
|
||||
return "当前没有正在执行的任务。"
|
||||
}
|
||||
return "已停止当前任务。"
|
||||
}
|
||||
|
||||
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 + "（获取标题失败）"
|
||||
}
|
||||
role := h.getRole(platform, userID)
|
||||
return fmt.Sprintf("当前对话：「%s」\nID: %s\n当前角色: %s", conv.Title, conv.ID, role)
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdRoles() string {
|
||||
if h.config.Roles == nil || len(h.config.Roles) == 0 {
|
||||
return "暂无可用角色。"
|
||||
}
|
||||
names := make([]string, 0, len(h.config.Roles))
|
||||
for name, role := range h.config.Roles {
|
||||
if role.Enabled {
|
||||
names = append(names, name)
|
||||
}
|
||||
}
|
||||
if len(names) == 0 {
|
||||
return "暂无可用角色。"
|
||||
}
|
||||
sort.Slice(names, func(i, j int) bool {
|
||||
if names[i] == "默认" {
|
||||
return true
|
||||
}
|
||||
if names[j] == "默认" {
|
||||
return false
|
||||
}
|
||||
return names[i] < names[j]
|
||||
})
|
||||
var b strings.Builder
|
||||
b.WriteString("【角色列表】\n")
|
||||
for _, name := range names {
|
||||
role := h.config.Roles[name]
|
||||
desc := role.Description
|
||||
if desc == "" {
|
||||
desc = "无描述"
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("· %s — %s\n", name, desc))
|
||||
}
|
||||
return strings.TrimSuffix(b.String(), "\n")
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdSwitchRole(platform, userID, roleName string) string {
|
||||
if roleName == "" {
|
||||
return "请指定角色名称，例如：角色 渗透测试"
|
||||
}
|
||||
if h.config.Roles == nil {
|
||||
return "暂无可用角色。"
|
||||
}
|
||||
role, exists := h.config.Roles[roleName]
|
||||
if !exists {
|
||||
return fmt.Sprintf("角色「%s」不存在。发送「角色」查看可用角色。", roleName)
|
||||
}
|
||||
if !role.Enabled {
|
||||
return fmt.Sprintf("角色「%s」已禁用。", roleName)
|
||||
}
|
||||
h.setRole(platform, userID, roleName)
|
||||
return fmt.Sprintf("已切换到角色：「%s」\n%s", roleName, role.Description)
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdDelete(platform, userID, convID string) string {
|
||||
if convID == "" {
|
||||
return "请指定对话 ID，例如：删除 xxx-xxx-xxx"
|
||||
}
|
||||
sk := h.sessionKey(platform, userID)
|
||||
h.mu.RLock()
|
||||
currentConvID := h.sessions[sk]
|
||||
h.mu.RUnlock()
|
||||
if convID == currentConvID {
|
||||
// 删除当前对话时，先清空会话绑定
|
||||
h.mu.Lock()
|
||||
delete(h.sessions, sk)
|
||||
h.mu.Unlock()
|
||||
}
|
||||
if err := h.db.DeleteConversation(convID); err != nil {
|
||||
return "删除失败: " + err.Error()
|
||||
}
|
||||
return fmt.Sprintf("已删除对话 ID: %s", convID)
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdVersion() string {
|
||||
v := h.config.Version
|
||||
if v == "" {
|
||||
v = "未知"
|
||||
}
|
||||
return "CyberStrikeAI " + v
|
||||
}
|
||||
|
||||
// handleRobotCommand 处理机器人内置命令；若匹配到命令返回 (回复内容, true)，否则返回 ("", false)
|
||||
func (h *RobotHandler) handleRobotCommand(platform, userID, text string) (string, bool) {
|
||||
switch {
|
||||
case text == robotCmdHelp || text == "help" || text == "？" || text == "?":
|
||||
return h.cmdHelp(), true
|
||||
case text == robotCmdList || text == robotCmdListAlt || text == "list":
|
||||
return h.cmdList(), true
|
||||
case strings.HasPrefix(text, robotCmdSwitch+" ") || strings.HasPrefix(text, robotCmdContinue+" ") || strings.HasPrefix(text, "switch ") || strings.HasPrefix(text, "continue "):
|
||||
var id string
|
||||
switch {
|
||||
case strings.HasPrefix(text, robotCmdSwitch+" "):
|
||||
id = strings.TrimSpace(text[len(robotCmdSwitch)+1:])
|
||||
case strings.HasPrefix(text, robotCmdContinue+" "):
|
||||
id = strings.TrimSpace(text[len(robotCmdContinue)+1:])
|
||||
case strings.HasPrefix(text, "switch "):
|
||||
id = strings.TrimSpace(text[7:])
|
||||
default:
|
||||
id = strings.TrimSpace(text[9:])
|
||||
}
|
||||
return h.cmdSwitch(platform, userID, id), true
|
||||
case text == robotCmdNew || text == "new":
|
||||
return h.cmdNew(platform, userID), true
|
||||
case text == robotCmdClear || text == "clear":
|
||||
return h.cmdClear(platform, userID), true
|
||||
case text == robotCmdCurrent || text == "current":
|
||||
return h.cmdCurrent(platform, userID), true
|
||||
case text == robotCmdStop || text == "stop":
|
||||
return h.cmdStop(platform, userID), true
|
||||
case text == robotCmdRoles || text == robotCmdRolesList || text == "roles":
|
||||
return h.cmdRoles(), true
|
||||
case strings.HasPrefix(text, robotCmdRoles+" ") || strings.HasPrefix(text, robotCmdSwitchRole+" ") || strings.HasPrefix(text, "role "):
|
||||
var roleName string
|
||||
switch {
|
||||
case strings.HasPrefix(text, robotCmdRoles+" "):
|
||||
roleName = strings.TrimSpace(text[len(robotCmdRoles)+1:])
|
||||
case strings.HasPrefix(text, robotCmdSwitchRole+" "):
|
||||
roleName = strings.TrimSpace(text[len(robotCmdSwitchRole)+1:])
|
||||
default:
|
||||
roleName = strings.TrimSpace(text[5:])
|
||||
}
|
||||
return h.cmdSwitchRole(platform, userID, roleName), true
|
||||
case strings.HasPrefix(text, robotCmdDelete+" ") || strings.HasPrefix(text, "delete "):
|
||||
var convID string
|
||||
if strings.HasPrefix(text, robotCmdDelete+" ") {
|
||||
convID = strings.TrimSpace(text[len(robotCmdDelete)+1:])
|
||||
} else {
|
||||
convID = strings.TrimSpace(text[7:])
|
||||
}
|
||||
return h.cmdDelete(platform, userID, convID), true
|
||||
case text == robotCmdVersion || text == "version":
|
||||
return h.cmdVersion(), true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
// —————— 企业微信 ——————
|
||||
|
||||
// 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（仅用于兼容，当前使用手动构造 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
|
||||
}
|
||||
// Gin 的 Query() 会自动 URL 解码，拿到的就是正确的 base64 字符串
|
||||
echostr := c.Query("echostr")
|
||||
msgSignature := c.Query("msg_signature")
|
||||
timestamp := c.Query("timestamp")
|
||||
nonce := c.Query("nonce")
|
||||
|
||||
// 验证签名：将 token、timestamp、nonce、echostr 四个参数排序后拼接计算 SHA1
|
||||
signature := h.signWecomRequest(h.config.Robots.Wecom.Token, timestamp, nonce, echostr)
|
||||
if signature != msgSignature {
|
||||
h.logger.Warn("企业微信 URL 验证签名失败", zap.String("expected", msgSignature), zap.String("got", signature))
|
||||
c.String(http.StatusBadRequest, "invalid signature")
|
||||
return
|
||||
}
|
||||
|
||||
if echostr == "" {
|
||||
c.String(http.StatusBadRequest, "missing echostr")
|
||||
return
|
||||
}
|
||||
|
||||
// 如果配置了 EncodingAESKey，说明是加密模式，需要解密 echostr
|
||||
if h.config.Robots.Wecom.EncodingAESKey != "" {
|
||||
decrypted, err := wecomDecrypt(h.config.Robots.Wecom.EncodingAESKey, echostr)
|
||||
if err != nil {
|
||||
h.logger.Warn("企业微信 echostr 解密失败", zap.Error(err))
|
||||
c.String(http.StatusBadRequest, "decrypt failed")
|
||||
return
|
||||
}
|
||||
c.String(http.StatusOK, string(decrypted))
|
||||
return
|
||||
}
|
||||
|
||||
// 明文模式直接返回 echostr
|
||||
c.String(http.StatusOK, echostr)
|
||||
}
|
||||
|
||||
// signWecomRequest 生成企业微信请求签名
|
||||
// 企业微信签名算法：将 token、timestamp、nonce、echostr 四个值排序后拼接成字符串，再计算 SHA1
|
||||
func (h *RobotHandler) signWecomRequest(token, timestamp, nonce, echostr string) string {
|
||||
strs := []string{token, timestamp, nonce, echostr}
|
||||
sort.Strings(strs)
|
||||
s := strings.Join(strs, "")
|
||||
hash := sha1.Sum([]byte(s))
|
||||
return fmt.Sprintf("%x", hash)
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// wecomEncrypt 企业微信消息加密（AES-256-CBC，PKCS7，明文格式：16字节随机+4字节长度+消息+corpID）
|
||||
func wecomEncrypt(encodingAESKey, message, corpID string) (string, error) {
|
||||
key, err := base64.StdEncoding.DecodeString(encodingAESKey + "=")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(key) != 32 {
|
||||
return "", fmt.Errorf("encoding_aes_key 解码后应为 32 字节")
|
||||
}
|
||||
// 构造明文：16 字节随机 + 4 字节长度 (大端) + 消息 + corpID
|
||||
random := make([]byte, 16)
|
||||
if _, err := rand.Read(random); err != nil {
|
||||
// 降级方案：使用时间戳生成随机数
|
||||
for i := range random {
|
||||
random[i] = byte(time.Now().UnixNano() % 256)
|
||||
}
|
||||
}
|
||||
msgLen := len(message)
|
||||
msgBytes := []byte(message)
|
||||
corpBytes := []byte(corpID)
|
||||
plain := make([]byte, 16+4+msgLen+len(corpBytes))
|
||||
copy(plain[:16], random)
|
||||
binary.BigEndian.PutUint32(plain[16:20], uint32(msgLen))
|
||||
copy(plain[20:20+msgLen], msgBytes)
|
||||
copy(plain[20+msgLen:], corpBytes)
|
||||
// PKCS7 填充
|
||||
padding := aes.BlockSize - len(plain)%aes.BlockSize
|
||||
pad := bytes.Repeat([]byte{byte(padding)}, padding)
|
||||
plain = append(plain, pad...)
|
||||
// AES-256-CBC 加密
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
iv := key[:16]
|
||||
ciphertext := make([]byte, len(plain))
|
||||
mode := cipher.NewCBCEncrypter(block, iv)
|
||||
mode.CryptBlocks(ciphertext, plain)
|
||||
return base64.StdEncoding.EncodeToString(ciphertext), nil
|
||||
}
|
||||
|
||||
// HandleWecomPOST 企业微信消息回调（POST），支持明文与加密模式
|
||||
func (h *RobotHandler) HandleWecomPOST(c *gin.Context) {
|
||||
if !h.config.Robots.Wecom.Enabled {
|
||||
h.logger.Debug("企业微信机器人未启用，跳过请求")
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
// 从 URL 获取签名参数（加密模式回复时需要用到）
|
||||
timestamp := c.Query("timestamp")
|
||||
nonce := c.Query("nonce")
|
||||
msgSignature := c.Query("msg_signature")
|
||||
|
||||
// 先读取请求体，后续解析/签名验证都会用到
|
||||
bodyRaw, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
h.logger.Warn("企业微信 POST 读取请求体失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
h.logger.Debug("企业微信 POST 收到请求", zap.String("body", string(bodyRaw)))
|
||||
|
||||
// 验证请求签名防止伪造。企业微信签名算法同 URL 验证，使用 token、timestamp、nonce、 Encrypt 四个字段
|
||||
// 若配置了 Token 则必须校验签名，避免未授权请求触发 Agent（防止平台被接管）
|
||||
token := h.config.Robots.Wecom.Token
|
||||
if token != "" {
|
||||
if msgSignature == "" {
|
||||
h.logger.Warn("企业微信 POST 缺少签名，已拒绝（需配置 token 并确保回调携带 msg_signature）")
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
var tmp wecomXML
|
||||
if err := xml.Unmarshal(bodyRaw, &tmp); err != nil {
|
||||
h.logger.Warn("企业微信 POST 签名验证前解析 XML 失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
expected := h.signWecomRequest(token, timestamp, nonce, tmp.Encrypt)
|
||||
if expected != msgSignature {
|
||||
h.logger.Warn("企业微信 POST 签名验证失败", zap.String("expected", expected), zap.String("got", msgSignature))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
var body wecomXML
|
||||
if err := xml.Unmarshal(bodyRaw, &body); err != nil {
|
||||
h.logger.Warn("企业微信 POST 解析 XML 失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
h.logger.Debug("企业微信 XML 解析成功", zap.String("ToUserName", body.ToUserName), zap.String("FromUserName", body.FromUserName), zap.String("MsgType", body.MsgType), zap.String("Content", body.Content), zap.String("Encrypt", body.Encrypt))
|
||||
|
||||
// 保存企业 ID（用于明文模式回复）
|
||||
enterpriseID := body.ToUserName
|
||||
|
||||
// 加密模式：先解密再解析内层 XML
|
||||
if body.Encrypt != "" && h.config.Robots.Wecom.EncodingAESKey != "" {
|
||||
h.logger.Debug("企业微信进入加密模式解密流程")
|
||||
decrypted, err := wecomDecrypt(h.config.Robots.Wecom.EncodingAESKey, body.Encrypt)
|
||||
if err != nil {
|
||||
h.logger.Warn("企业微信消息解密失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
h.logger.Debug("企业微信解密成功", zap.String("decrypted", string(decrypted)))
|
||||
if err := xml.Unmarshal(decrypted, &body); err != nil {
|
||||
h.logger.Warn("企业微信解密后 XML 解析失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
h.logger.Debug("企业微信内层 XML 解析成功", zap.String("FromUserName", body.FromUserName), zap.String("Content", body.Content))
|
||||
}
|
||||
|
||||
userID := body.FromUserName
|
||||
text := strings.TrimSpace(body.Content)
|
||||
|
||||
// 限制回复内容长度（企业微信限制 2048 字节）
|
||||
maxReplyLen := 2000
|
||||
limitReply := func(s string) string {
|
||||
if len(s) > maxReplyLen {
|
||||
return s[:maxReplyLen] + "\n\n（内容过长，已截断）"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
if body.MsgType != "text" {
|
||||
h.logger.Debug("企业微信收到非文本消息", zap.String("MsgType", body.MsgType))
|
||||
h.sendWecomReply(c, userID, enterpriseID, limitReply("暂仅支持文本消息，请发送文字。"), timestamp, nonce)
|
||||
return
|
||||
}
|
||||
|
||||
// 文本消息：先判断是否为内置命令（如 帮助/列表/新对话 等），这类命令处理很快，可以直接走被动回复，避免依赖主动发送 API。
|
||||
if cmdReply, ok := h.handleRobotCommand("wecom", userID, text); ok {
|
||||
h.logger.Debug("企业微信收到命令消息，走被动回复", zap.String("userID", userID), zap.String("text", text))
|
||||
h.sendWecomReply(c, userID, enterpriseID, limitReply(cmdReply), timestamp, nonce)
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Debug("企业微信开始处理消息（异步 AI）", zap.String("userID", userID), zap.String("text", text))
|
||||
|
||||
// 企业微信被动回复有 5 秒超时限制，而 AI 调用通常超过该时长。
|
||||
// 这里采用推荐做法：立即返回 success（或空串），然后通过主动发送接口推送完整回复。
|
||||
c.String(http.StatusOK, "success")
|
||||
|
||||
// 异步处理消息并通过企业微信主动消息接口发送结果
|
||||
go func() {
|
||||
reply := h.HandleMessage("wecom", userID, text)
|
||||
reply = limitReply(reply)
|
||||
h.logger.Debug("企业微信消息处理完成", zap.String("userID", userID), zap.String("reply", reply))
|
||||
// 调用企业微信 API 主动发送消息
|
||||
h.sendWecomMessageViaAPI(userID, enterpriseID, reply)
|
||||
}()
|
||||
}
|
||||
|
||||
// sendWecomReply 发送企业微信回复（加密模式自动加密）
|
||||
// 参数：toUser=用户 ID, fromUser=企业 ID（明文模式）/CorpID（加密模式）, content=回复内容，timestamp/nonce=请求参数
|
||||
func (h *RobotHandler) sendWecomReply(c *gin.Context, toUser, fromUser, content, timestamp, nonce string) {
|
||||
// 加密模式：判断 EncodingAESKey 是否配置
|
||||
if h.config.Robots.Wecom.EncodingAESKey != "" {
|
||||
// 加密模式使用 CorpID 进行加密
|
||||
corpID := h.config.Robots.Wecom.CorpID
|
||||
if corpID == "" {
|
||||
h.logger.Warn("企业微信加密模式缺少 CorpID 配置")
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
|
||||
// 构造完整的明文 XML 回复（格式严格按企业微信文档要求）
|
||||
plainResp := fmt.Sprintf(`<xml>
|
||||
<ToUserName><![CDATA[%s]]></ToUserName>
|
||||
<FromUserName><![CDATA[%s]]></FromUserName>
|
||||
<CreateTime>%d</CreateTime>
|
||||
<MsgType><![CDATA[text]]></MsgType>
|
||||
<Content><![CDATA[%s]]></Content>
|
||||
</xml>`, toUser, fromUser, time.Now().Unix(), content)
|
||||
|
||||
encrypted, err := wecomEncrypt(h.config.Robots.Wecom.EncodingAESKey, plainResp, corpID)
|
||||
if err != nil {
|
||||
h.logger.Warn("企业微信回复加密失败", zap.Error(err))
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
// 使用请求中的 timestamp/nonce 生成签名（企业微信要求回复时使用与请求相同的 timestamp 和 nonce）
|
||||
msgSignature := h.signWecomRequest(h.config.Robots.Wecom.Token, timestamp, nonce, encrypted)
|
||||
|
||||
h.logger.Debug("企业微信发送加密回复",
|
||||
zap.String("Encrypt", encrypted[:50]+"..."),
|
||||
zap.String("MsgSignature", msgSignature),
|
||||
zap.String("TimeStamp", timestamp),
|
||||
zap.String("Nonce", nonce))
|
||||
|
||||
// 加密模式仅返回 4 个核心字段（企业微信官方要求）
|
||||
xmlResp := fmt.Sprintf(`<xml><Encrypt><![CDATA[%s]]></Encrypt><MsgSignature><![CDATA[%s]]></MsgSignature><TimeStamp><![CDATA[%s]]></TimeStamp><Nonce><![CDATA[%s]]></Nonce></xml>`, encrypted, msgSignature, timestamp, nonce)
|
||||
// also log the final response body so we can cross-check with the
|
||||
// network traffic or developer console
|
||||
h.logger.Debug("企业微信加密回复包", zap.String("xml", xmlResp))
|
||||
// for additional confidence, decrypt the payload ourselves and log it
|
||||
if dec, err2 := wecomDecrypt(h.config.Robots.Wecom.EncodingAESKey, encrypted); err2 == nil {
|
||||
h.logger.Debug("企业微信加密回复解密检查", zap.String("plain", string(dec)))
|
||||
} else {
|
||||
h.logger.Warn("企业微信加密回复解密检查失败", zap.Error(err2))
|
||||
}
|
||||
|
||||
// 使用 c.Writer.Write 直接写入响应，避免 c.String 的转义问题
|
||||
c.Writer.WriteHeader(http.StatusOK)
|
||||
// use text/xml as that's what WeCom examples show
|
||||
c.Writer.Header().Set("Content-Type", "text/xml; charset=utf-8")
|
||||
_, _ = c.Writer.Write([]byte(xmlResp))
|
||||
h.logger.Debug("企业微信加密回复已发送")
|
||||
return
|
||||
}
|
||||
|
||||
// 明文模式
|
||||
h.logger.Debug("企业微信发送明文回复", zap.String("ToUserName", toUser), zap.String("FromUserName", fromUser), zap.String("Content", content[:50]+"..."))
|
||||
|
||||
// 手动构造 XML 响应（使用 CDATA 包裹所有字段，并包含 AgentID）
|
||||
xmlResp := fmt.Sprintf(`<xml>
|
||||
<ToUserName><![CDATA[%s]]></ToUserName>
|
||||
<FromUserName><![CDATA[%s]]></FromUserName>
|
||||
<CreateTime>%d</CreateTime>
|
||||
<MsgType><![CDATA[text]]></MsgType>
|
||||
<Content><![CDATA[%s]]></Content>
|
||||
</xml>`, toUser, fromUser, time.Now().Unix(), content)
|
||||
|
||||
// log the exact plaintext response for debugging
|
||||
h.logger.Debug("企业微信明文回复包", zap.String("xml", xmlResp))
|
||||
|
||||
// use text/xml as recommended by WeCom docs
|
||||
c.Header("Content-Type", "text/xml; charset=utf-8")
|
||||
c.String(http.StatusOK, xmlResp)
|
||||
h.logger.Debug("企业微信明文回复已发送")
|
||||
}
|
||||
|
||||
// —————— 测试接口（需登录，用于验证机器人逻辑，无需钉钉/飞书客户端） ——————
|
||||
|
||||
// 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})
|
||||
}
|
||||
|
||||
// sendWecomMessageViaAPI 通过企业微信 API 主动发送消息（用于异步处理后的结果发送）
|
||||
func (h *RobotHandler) sendWecomMessageViaAPI(toUser, toParty, content string) {
|
||||
if !h.config.Robots.Wecom.Enabled {
|
||||
return
|
||||
}
|
||||
|
||||
secret := h.config.Robots.Wecom.Secret
|
||||
corpID := h.config.Robots.Wecom.CorpID
|
||||
agentID := h.config.Robots.Wecom.AgentID
|
||||
|
||||
if secret == "" || corpID == "" {
|
||||
h.logger.Warn("企业微信主动 API 缺少 secret 或 corpID 配置")
|
||||
return
|
||||
}
|
||||
|
||||
// 第 1 步：获取 access_token
|
||||
tokenURL := fmt.Sprintf("https://qyapi.weixin.qq.com/cgi-bin/gettoken?corpid=%s&corpsecret=%s", corpID, secret)
|
||||
resp, err := http.Get(tokenURL)
|
||||
if err != nil {
|
||||
h.logger.Warn("企业微信获取 token 失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var tokenResp struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
ErrCode int `json:"errcode"`
|
||||
ErrMsg string `json:"errmsg"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
|
||||
h.logger.Warn("企业微信 token 响应解析失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if tokenResp.ErrCode != 0 {
|
||||
h.logger.Warn("企业微信 token 获取错误", zap.String("errmsg", tokenResp.ErrMsg), zap.Int("errcode", tokenResp.ErrCode))
|
||||
return
|
||||
}
|
||||
|
||||
// 第 2 步：构造发送消息请求
|
||||
msgReq := map[string]interface{}{
|
||||
"touser": toUser,
|
||||
"msgtype": "text",
|
||||
"agentid": agentID,
|
||||
"text": map[string]interface{}{
|
||||
"content": content,
|
||||
},
|
||||
}
|
||||
|
||||
msgBody, err := json.Marshal(msgReq)
|
||||
if err != nil {
|
||||
h.logger.Warn("企业微信消息序列化失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
// 第 3 步：发送消息
|
||||
sendURL := fmt.Sprintf("https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token=%s", tokenResp.AccessToken)
|
||||
msgResp, err := http.Post(sendURL, "application/json", bytes.NewReader(msgBody))
|
||||
if err != nil {
|
||||
h.logger.Warn("企业微信主动发送消息失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
defer msgResp.Body.Close()
|
||||
|
||||
var sendResp struct {
|
||||
ErrCode int `json:"errcode"`
|
||||
ErrMsg string `json:"errmsg"`
|
||||
InvalidUser string `json:"invaliduser"`
|
||||
MsgID string `json:"msgid"`
|
||||
}
|
||||
if err := json.NewDecoder(msgResp.Body).Decode(&sendResp); err != nil {
|
||||
h.logger.Warn("企业微信发送响应解析失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
if sendResp.ErrCode == 0 {
|
||||
h.logger.Debug("企业微信主动发送消息成功", zap.String("msgid", sendResp.MsgID))
|
||||
} else {
|
||||
h.logger.Warn("企业微信主动发送消息失败", zap.String("errmsg", sendResp.ErrMsg), zap.Int("errcode", sendResp.ErrCode), zap.String("invaliduser", sendResp.InvalidUser))
|
||||
}
|
||||
}
|
||||
|
||||
// —————— 钉钉 ——————
|
||||
|
||||
// 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{})
|
||||
}
|
||||
-487
@@ -1,487 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// RoleHandler 角色处理器
|
||||
type RoleHandler struct {
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
skillsManager SkillsManager // Skills管理器接口（可选）
|
||||
}
|
||||
|
||||
// SkillsManager Skills管理器接口
|
||||
type SkillsManager interface {
|
||||
ListSkills() ([]string, error)
|
||||
}
|
||||
|
||||
// NewRoleHandler 创建新的角色处理器
|
||||
func NewRoleHandler(cfg *config.Config, configPath string, logger *zap.Logger) *RoleHandler {
|
||||
return &RoleHandler{
|
||||
config: cfg,
|
||||
configPath: configPath,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// SetSkillsManager 设置Skills管理器
|
||||
func (h *RoleHandler) SetSkillsManager(manager SkillsManager) {
|
||||
h.skillsManager = manager
|
||||
}
|
||||
|
||||
// GetSkills 获取所有可用的skills列表
|
||||
func (h *RoleHandler) GetSkills(c *gin.Context) {
|
||||
if h.skillsManager == nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skills": []string{},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
skills, err := h.skillsManager.ListSkills()
|
||||
if err != nil {
|
||||
h.logger.Warn("获取skills列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skills": []string{},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skills": skills,
|
||||
})
|
||||
}
|
||||
|
||||
// GetRoles 获取所有角色
|
||||
func (h *RoleHandler) GetRoles(c *gin.Context) {
|
||||
if h.config.Roles == nil {
|
||||
h.config.Roles = make(map[string]config.RoleConfig)
|
||||
}
|
||||
|
||||
roles := make([]config.RoleConfig, 0, len(h.config.Roles))
|
||||
for key, role := range h.config.Roles {
|
||||
// 确保角色的key与name一致
|
||||
if role.Name == "" {
|
||||
role.Name = key
|
||||
}
|
||||
roles = append(roles, role)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"roles": roles,
|
||||
})
|
||||
}
|
||||
|
||||
// GetRole 获取单个角色
|
||||
func (h *RoleHandler) GetRole(c *gin.Context) {
|
||||
roleName := c.Param("name")
|
||||
if roleName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "角色名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
if h.config.Roles == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
role, exists := h.config.Roles[roleName]
|
||||
if !exists {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
// 确保角色的name与key一致
|
||||
if role.Name == "" {
|
||||
role.Name = roleName
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"role": role,
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateRole 更新角色
|
||||
func (h *RoleHandler) UpdateRole(c *gin.Context) {
|
||||
roleName := c.Param("name")
|
||||
if roleName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "角色名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
var req config.RoleConfig
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 确保角色名称与请求中的name一致
|
||||
if req.Name == "" {
|
||||
req.Name = roleName
|
||||
}
|
||||
|
||||
// 初始化Roles map
|
||||
if h.config.Roles == nil {
|
||||
h.config.Roles = make(map[string]config.RoleConfig)
|
||||
}
|
||||
|
||||
// 删除所有与角色name相同但key不同的旧角色（避免重复）
|
||||
// 使用角色name作为key，确保唯一性
|
||||
finalKey := req.Name
|
||||
keysToDelete := make([]string, 0)
|
||||
for key := range h.config.Roles {
|
||||
// 如果key与最终的key不同，但name相同，则标记为删除
|
||||
if key != finalKey {
|
||||
role := h.config.Roles[key]
|
||||
// 确保角色的name字段正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = key
|
||||
}
|
||||
if role.Name == req.Name {
|
||||
keysToDelete = append(keysToDelete, key)
|
||||
}
|
||||
}
|
||||
}
|
||||
// 删除旧的角色
|
||||
for _, key := range keysToDelete {
|
||||
delete(h.config.Roles, key)
|
||||
h.logger.Info("删除重复的角色", zap.String("oldKey", key), zap.String("name", req.Name))
|
||||
}
|
||||
|
||||
// 如果当前更新的key与最终key不同，也需要删除旧的
|
||||
if roleName != finalKey {
|
||||
delete(h.config.Roles, roleName)
|
||||
}
|
||||
|
||||
// 如果角色名称改变，需要删除旧文件
|
||||
if roleName != finalKey {
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
rolesDir := h.config.RolesDir
|
||||
if rolesDir == "" {
|
||||
rolesDir = "roles" // 默认目录
|
||||
}
|
||||
|
||||
// 如果是相对路径，相对于配置文件所在目录
|
||||
if !filepath.IsAbs(rolesDir) {
|
||||
rolesDir = filepath.Join(configDir, rolesDir)
|
||||
}
|
||||
|
||||
// 删除旧的角色文件
|
||||
oldSafeFileName := sanitizeFileName(roleName)
|
||||
oldRoleFileYaml := filepath.Join(rolesDir, oldSafeFileName+".yaml")
|
||||
oldRoleFileYml := filepath.Join(rolesDir, oldSafeFileName+".yml")
|
||||
|
||||
if _, err := os.Stat(oldRoleFileYaml); err == nil {
|
||||
if err := os.Remove(oldRoleFileYaml); err != nil {
|
||||
h.logger.Warn("删除旧角色配置文件失败", zap.String("file", oldRoleFileYaml), zap.Error(err))
|
||||
}
|
||||
}
|
||||
if _, err := os.Stat(oldRoleFileYml); err == nil {
|
||||
if err := os.Remove(oldRoleFileYml); err != nil {
|
||||
h.logger.Warn("删除旧角色配置文件失败", zap.String("file", oldRoleFileYml), zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 使用角色name作为key来保存（确保唯一性）
|
||||
h.config.Roles[finalKey] = req
|
||||
|
||||
// 保存配置到文件
|
||||
if err := h.saveConfig(); err != nil {
|
||||
h.logger.Error("保存配置失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("更新角色", zap.String("oldKey", roleName), zap.String("newKey", finalKey), zap.String("name", req.Name))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "角色已更新",
|
||||
"role": req,
|
||||
})
|
||||
}
|
||||
|
||||
// CreateRole 创建新角色
|
||||
func (h *RoleHandler) CreateRole(c *gin.Context) {
|
||||
var req config.RoleConfig
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if req.Name == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "角色名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
// 初始化Roles map
|
||||
if h.config.Roles == nil {
|
||||
h.config.Roles = make(map[string]config.RoleConfig)
|
||||
}
|
||||
|
||||
// 检查角色是否已存在
|
||||
if _, exists := h.config.Roles[req.Name]; exists {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "角色已存在"})
|
||||
return
|
||||
}
|
||||
|
||||
// 创建角色（默认启用）
|
||||
if !req.Enabled {
|
||||
req.Enabled = true
|
||||
}
|
||||
|
||||
h.config.Roles[req.Name] = req
|
||||
|
||||
// 保存配置到文件
|
||||
if err := h.saveConfig(); err != nil {
|
||||
h.logger.Error("保存配置失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("创建角色", zap.String("roleName", req.Name))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "角色已创建",
|
||||
"role": req,
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteRole 删除角色
|
||||
func (h *RoleHandler) DeleteRole(c *gin.Context) {
|
||||
roleName := c.Param("name")
|
||||
if roleName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "角色名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
if h.config.Roles == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
if _, exists := h.config.Roles[roleName]; !exists {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
// 不允许删除"默认"角色
|
||||
if roleName == "默认" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "不能删除默认角色"})
|
||||
return
|
||||
}
|
||||
|
||||
delete(h.config.Roles, roleName)
|
||||
|
||||
// 删除对应的角色文件
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
rolesDir := h.config.RolesDir
|
||||
if rolesDir == "" {
|
||||
rolesDir = "roles" // 默认目录
|
||||
}
|
||||
|
||||
// 如果是相对路径，相对于配置文件所在目录
|
||||
if !filepath.IsAbs(rolesDir) {
|
||||
rolesDir = filepath.Join(configDir, rolesDir)
|
||||
}
|
||||
|
||||
// 尝试删除角色文件（.yaml 和 .yml）
|
||||
safeFileName := sanitizeFileName(roleName)
|
||||
roleFileYaml := filepath.Join(rolesDir, safeFileName+".yaml")
|
||||
roleFileYml := filepath.Join(rolesDir, safeFileName+".yml")
|
||||
|
||||
// 删除 .yaml 文件（如果存在）
|
||||
if _, err := os.Stat(roleFileYaml); err == nil {
|
||||
if err := os.Remove(roleFileYaml); err != nil {
|
||||
h.logger.Warn("删除角色配置文件失败", zap.String("file", roleFileYaml), zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已删除角色配置文件", zap.String("file", roleFileYaml))
|
||||
}
|
||||
}
|
||||
|
||||
// 删除 .yml 文件（如果存在）
|
||||
if _, err := os.Stat(roleFileYml); err == nil {
|
||||
if err := os.Remove(roleFileYml); err != nil {
|
||||
h.logger.Warn("删除角色配置文件失败", zap.String("file", roleFileYml), zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已删除角色配置文件", zap.String("file", roleFileYml))
|
||||
}
|
||||
}
|
||||
|
||||
h.logger.Info("删除角色", zap.String("roleName", roleName))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "角色已删除",
|
||||
})
|
||||
}
|
||||
|
||||
// saveConfig 保存配置到目录中的文件
|
||||
func (h *RoleHandler) saveConfig() error {
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
rolesDir := h.config.RolesDir
|
||||
if rolesDir == "" {
|
||||
rolesDir = "roles" // 默认目录
|
||||
}
|
||||
|
||||
// 如果是相对路径，相对于配置文件所在目录
|
||||
if !filepath.IsAbs(rolesDir) {
|
||||
rolesDir = filepath.Join(configDir, rolesDir)
|
||||
}
|
||||
|
||||
// 确保目录存在
|
||||
if err := os.MkdirAll(rolesDir, 0755); err != nil {
|
||||
return fmt.Errorf("创建角色目录失败: %w", err)
|
||||
}
|
||||
|
||||
// 保存每个角色到独立的文件
|
||||
if h.config.Roles != nil {
|
||||
for roleName, role := range h.config.Roles {
|
||||
// 确保角色名称正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = roleName
|
||||
}
|
||||
|
||||
// 使用角色名称作为文件名（安全化文件名，避免特殊字符）
|
||||
safeFileName := sanitizeFileName(role.Name)
|
||||
roleFile := filepath.Join(rolesDir, safeFileName+".yaml")
|
||||
|
||||
// 将角色配置序列化为YAML
|
||||
roleData, err := yaml.Marshal(&role)
|
||||
if err != nil {
|
||||
h.logger.Error("序列化角色配置失败", zap.String("role", roleName), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
// 处理icon字段：确保包含\U的icon值被引号包围（YAML需要引号才能正确解析Unicode转义）
|
||||
roleDataStr := string(roleData)
|
||||
if role.Icon != "" && strings.HasPrefix(role.Icon, "\\U") {
|
||||
// 匹配 icon: \UXXXXXXXX 格式（没有引号），排除已经有引号的情况
|
||||
// 使用负向前瞻确保后面没有引号，或者直接匹配没有引号的情况
|
||||
re := regexp.MustCompile(`(?m)^(icon:\s+)(\\U[0-9A-F]{8})(\s*)$`)
|
||||
roleDataStr = re.ReplaceAllString(roleDataStr, `${1}"${2}"${3}`)
|
||||
roleData = []byte(roleDataStr)
|
||||
}
|
||||
|
||||
// 写入文件
|
||||
if err := os.WriteFile(roleFile, roleData, 0644); err != nil {
|
||||
h.logger.Error("保存角色配置文件失败", zap.String("role", roleName), zap.String("file", roleFile), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
h.logger.Info("角色配置已保存到文件", zap.String("role", roleName), zap.String("file", roleFile))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// sanitizeFileName 将角色名称转换为安全的文件名
|
||||
func sanitizeFileName(name string) string {
|
||||
// 替换可能不安全的字符
|
||||
replacer := map[rune]string{
|
||||
'/': "_",
|
||||
'\\': "_",
|
||||
':': "_",
|
||||
'*': "_",
|
||||
'?': "_",
|
||||
'"': "_",
|
||||
'<': "_",
|
||||
'>': "_",
|
||||
'|': "_",
|
||||
' ': "_",
|
||||
}
|
||||
|
||||
var result []rune
|
||||
for _, r := range name {
|
||||
if replacement, ok := replacer[r]; ok {
|
||||
result = append(result, []rune(replacement)...)
|
||||
} else {
|
||||
result = append(result, r)
|
||||
}
|
||||
}
|
||||
|
||||
fileName := string(result)
|
||||
// 如果文件名为空，使用默认名称
|
||||
if fileName == "" {
|
||||
fileName = "role"
|
||||
}
|
||||
|
||||
return fileName
|
||||
}
|
||||
|
||||
// updateRolesConfig 更新角色配置
|
||||
func updateRolesConfig(doc *yaml.Node, cfg config.RolesConfig) {
|
||||
root := doc.Content[0]
|
||||
rolesNode := ensureMap(root, "roles")
|
||||
|
||||
// 清空现有角色
|
||||
if rolesNode.Kind == yaml.MappingNode {
|
||||
rolesNode.Content = nil
|
||||
}
|
||||
|
||||
// 添加新角色（使用name作为key，确保唯一性）
|
||||
if cfg.Roles != nil {
|
||||
// 先建立一个以name为key的map，去重（保留最后一个）
|
||||
rolesByName := make(map[string]config.RoleConfig)
|
||||
for roleKey, role := range cfg.Roles {
|
||||
// 确保角色的name字段正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = roleKey
|
||||
}
|
||||
// 使用name作为最终key，如果有多个key对应相同的name，只保留最后一个
|
||||
rolesByName[role.Name] = role
|
||||
}
|
||||
|
||||
// 将去重后的角色写入YAML
|
||||
for roleName, role := range rolesByName {
|
||||
roleNode := ensureMap(rolesNode, roleName)
|
||||
setStringInMap(roleNode, "name", role.Name)
|
||||
setStringInMap(roleNode, "description", role.Description)
|
||||
setStringInMap(roleNode, "user_prompt", role.UserPrompt)
|
||||
if role.Icon != "" {
|
||||
setStringInMap(roleNode, "icon", role.Icon)
|
||||
}
|
||||
setBoolInMap(roleNode, "enabled", role.Enabled)
|
||||
|
||||
// 添加工具列表（优先使用tools字段）
|
||||
if len(role.Tools) > 0 {
|
||||
toolsNode := ensureArray(roleNode, "tools")
|
||||
toolsNode.Content = nil
|
||||
for _, toolKey := range role.Tools {
|
||||
toolNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: toolKey}
|
||||
toolsNode.Content = append(toolsNode.Content, toolNode)
|
||||
}
|
||||
} else if len(role.MCPs) > 0 {
|
||||
// 向后兼容：如果没有tools但有mcps，保存mcps
|
||||
mcpsNode := ensureArray(roleNode, "mcps")
|
||||
mcpsNode.Content = nil
|
||||
for _, mcpName := range role.MCPs {
|
||||
mcpNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: mcpName}
|
||||
mcpsNode.Content = append(mcpsNode.Content, mcpNode)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ensureArray 确保数组中存在指定key的数组节点
|
||||
func ensureArray(parent *yaml.Node, key string) *yaml.Node {
|
||||
_, valueNode := ensureKeyValue(parent, key)
|
||||
if valueNode.Kind != yaml.SequenceNode {
|
||||
valueNode.Kind = yaml.SequenceNode
|
||||
valueNode.Tag = "!!seq"
|
||||
valueNode.Content = nil
|
||||
}
|
||||
return valueNode
|
||||
}
|
||||
@@ -1,758 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/skillpackage"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// SkillsHandler Skills处理器（磁盘 + Eino 规范；运行时由 Eino ADK skill 中间件加载）
|
||||
type SkillsHandler struct {
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
db *database.DB // 数据库连接（遗留统计；MCP list/read 已移除）
|
||||
}
|
||||
|
||||
// NewSkillsHandler 创建新的Skills处理器
|
||||
func NewSkillsHandler(cfg *config.Config, configPath string, logger *zap.Logger) *SkillsHandler {
|
||||
return &SkillsHandler{
|
||||
config: cfg,
|
||||
configPath: configPath,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *SkillsHandler) skillsRootAbs() string {
|
||||
skillsDir := h.config.SkillsDir
|
||||
if skillsDir == "" {
|
||||
skillsDir = "skills"
|
||||
}
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
if !filepath.IsAbs(skillsDir) {
|
||||
skillsDir = filepath.Join(configDir, skillsDir)
|
||||
}
|
||||
return skillsDir
|
||||
}
|
||||
|
||||
// SetDB 设置数据库连接（用于获取调用统计）
|
||||
func (h *SkillsHandler) SetDB(db *database.DB) {
|
||||
h.db = db
|
||||
}
|
||||
|
||||
// GetSkills 获取所有skills列表（支持分页和搜索）
|
||||
func (h *SkillsHandler) GetSkills(c *gin.Context) {
|
||||
allSummaries, err := skillpackage.ListSkillSummaries(h.skillsRootAbs())
|
||||
if err != nil {
|
||||
h.logger.Error("获取skills列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
searchKeyword := strings.TrimSpace(c.Query("search"))
|
||||
|
||||
allSkillsInfo := make([]map[string]interface{}, 0, len(allSummaries))
|
||||
for _, s := range allSummaries {
|
||||
skillInfo := map[string]interface{}{
|
||||
"id": s.ID,
|
||||
"name": s.Name,
|
||||
"dir_name": s.DirName,
|
||||
"description": s.Description,
|
||||
"version": s.Version,
|
||||
"path": s.Path,
|
||||
"tags": s.Tags,
|
||||
"triggers": s.Triggers,
|
||||
"script_count": s.ScriptCount,
|
||||
"file_count": s.FileCount,
|
||||
"progressive": s.Progressive,
|
||||
"file_size": s.FileSize,
|
||||
"mod_time": s.ModTime,
|
||||
}
|
||||
allSkillsInfo = append(allSkillsInfo, skillInfo)
|
||||
}
|
||||
|
||||
filteredSkillsInfo := allSkillsInfo
|
||||
if searchKeyword != "" {
|
||||
keywordLower := strings.ToLower(searchKeyword)
|
||||
filteredSkillsInfo = make([]map[string]interface{}, 0)
|
||||
for _, skillInfo := range allSkillsInfo {
|
||||
id := strings.ToLower(fmt.Sprintf("%v", skillInfo["id"]))
|
||||
name := strings.ToLower(fmt.Sprintf("%v", skillInfo["name"]))
|
||||
description := strings.ToLower(fmt.Sprintf("%v", skillInfo["description"]))
|
||||
path := strings.ToLower(fmt.Sprintf("%v", skillInfo["path"]))
|
||||
version := strings.ToLower(fmt.Sprintf("%v", skillInfo["version"]))
|
||||
tagsJoined := ""
|
||||
if tags, ok := skillInfo["tags"].([]string); ok {
|
||||
tagsJoined = strings.ToLower(strings.Join(tags, " "))
|
||||
}
|
||||
trigJoined := ""
|
||||
if tr, ok := skillInfo["triggers"].([]string); ok {
|
||||
trigJoined = strings.ToLower(strings.Join(tr, " "))
|
||||
}
|
||||
if strings.Contains(id, keywordLower) ||
|
||||
strings.Contains(name, keywordLower) ||
|
||||
strings.Contains(description, keywordLower) ||
|
||||
strings.Contains(path, keywordLower) ||
|
||||
strings.Contains(version, keywordLower) ||
|
||||
strings.Contains(tagsJoined, keywordLower) ||
|
||||
strings.Contains(trigJoined, keywordLower) {
|
||||
filteredSkillsInfo = append(filteredSkillsInfo, skillInfo)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 分页参数
|
||||
limit := 20 // 默认每页20条
|
||||
offset := 0
|
||||
if limitStr := c.Query("limit"); limitStr != "" {
|
||||
if parsed, err := parseInt(limitStr); err == nil && parsed > 0 {
|
||||
// 允许更大的limit用于搜索场景，但设置一个合理的上限（10000）
|
||||
if parsed <= 10000 {
|
||||
limit = parsed
|
||||
} else {
|
||||
limit = 10000
|
||||
}
|
||||
}
|
||||
}
|
||||
if offsetStr := c.Query("offset"); offsetStr != "" {
|
||||
if parsed, err := parseInt(offsetStr); err == nil && parsed >= 0 {
|
||||
offset = parsed
|
||||
}
|
||||
}
|
||||
|
||||
// 计算分页范围
|
||||
total := len(filteredSkillsInfo)
|
||||
start := offset
|
||||
end := offset + limit
|
||||
if start > total {
|
||||
start = total
|
||||
}
|
||||
if end > total {
|
||||
end = total
|
||||
}
|
||||
|
||||
// 获取当前页的skill列表
|
||||
var paginatedSkillsInfo []map[string]interface{}
|
||||
if start < end {
|
||||
paginatedSkillsInfo = filteredSkillsInfo[start:end]
|
||||
} else {
|
||||
paginatedSkillsInfo = []map[string]interface{}{}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skills": paginatedSkillsInfo,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
})
|
||||
}
|
||||
|
||||
// GetSkill 获取单个skill的详细信息
|
||||
func (h *SkillsHandler) GetSkill(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
resPath := strings.TrimSpace(c.Query("resource_path"))
|
||||
if resPath == "" {
|
||||
resPath = strings.TrimSpace(c.Query("skill_script_path"))
|
||||
}
|
||||
if resPath != "" {
|
||||
content, err := skillpackage.ReadScriptText(h.skillsRootAbs(), skillName, resPath, 0)
|
||||
if err != nil {
|
||||
h.logger.Warn("读取skill资源失败", zap.String("skill", skillName), zap.String("path", resPath), zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skill": map[string]interface{}{
|
||||
"id": skillName,
|
||||
},
|
||||
"resource": map[string]interface{}{
|
||||
"path": resPath,
|
||||
"content": content,
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
depthStr := strings.ToLower(strings.TrimSpace(c.DefaultQuery("depth", "full")))
|
||||
section := strings.TrimSpace(c.Query("section"))
|
||||
opt := skillpackage.LoadOptions{Section: section}
|
||||
switch depthStr {
|
||||
case "summary":
|
||||
opt.Depth = "summary"
|
||||
case "full", "":
|
||||
opt.Depth = "full"
|
||||
default:
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "depth 仅支持 summary 或 full"})
|
||||
return
|
||||
}
|
||||
|
||||
skill, err := skillpackage.LoadSkill(h.skillsRootAbs(), skillName, opt)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载skill失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "skill不存在: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
skillPath := skill.Path
|
||||
skillFile := filepath.Join(skillPath, "SKILL.md")
|
||||
|
||||
fileInfo, _ := os.Stat(skillFile)
|
||||
var fileSize int64
|
||||
var modTime string
|
||||
if fileInfo != nil {
|
||||
fileSize = fileInfo.Size()
|
||||
modTime = fileInfo.ModTime().Format("2006-01-02 15:04:05")
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skill": map[string]interface{}{
|
||||
"id": skill.DirName,
|
||||
"name": skill.Name,
|
||||
"description": skill.Description,
|
||||
"content": skill.Content,
|
||||
"path": skill.Path,
|
||||
"version": skill.Version,
|
||||
"tags": skill.Tags,
|
||||
"scripts": skill.Scripts,
|
||||
"sections": skill.Sections,
|
||||
"package_files": skill.PackageFiles,
|
||||
"file_size": fileSize,
|
||||
"mod_time": modTime,
|
||||
"depth": depthStr,
|
||||
"section": section,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// ListSkillPackageFiles lists all files in a skill directory (Agent Skills layout).
|
||||
func (h *SkillsHandler) ListSkillPackageFiles(c *gin.Context) {
|
||||
skillID := c.Param("name")
|
||||
files, err := skillpackage.ListPackageFiles(h.skillsRootAbs(), skillID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"files": files})
|
||||
}
|
||||
|
||||
// GetSkillPackageFile returns one file by relative path (?path=).
|
||||
func (h *SkillsHandler) GetSkillPackageFile(c *gin.Context) {
|
||||
skillID := c.Param("name")
|
||||
rel := strings.TrimSpace(c.Query("path"))
|
||||
if rel == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "query path is required"})
|
||||
return
|
||||
}
|
||||
b, err := skillpackage.ReadPackageFile(h.skillsRootAbs(), skillID, rel, 0)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"path": rel, "content": string(b)})
|
||||
}
|
||||
|
||||
// PutSkillPackageFile writes a file inside the skill package.
|
||||
func (h *SkillsHandler) PutSkillPackageFile(c *gin.Context) {
|
||||
skillID := c.Param("name")
|
||||
var req struct {
|
||||
Path string `json:"path" binding:"required"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if req.Path == "SKILL.md" {
|
||||
if err := skillpackage.ValidateSkillMDPackage([]byte(req.Content), skillID); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := skillpackage.WritePackageFile(h.skillsRootAbs(), skillID, req.Path, []byte(req.Content)); err != nil {
|
||||
h.logger.Error("写入 skill 文件失败", zap.String("skill", skillID), zap.String("path", req.Path), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "saved", "path": req.Path})
|
||||
}
|
||||
|
||||
// GetSkillBoundRoles 获取绑定指定skill的角色列表
|
||||
func (h *SkillsHandler) GetSkillBoundRoles(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
boundRoles := h.getRolesBoundToSkill(skillName)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skill": skillName,
|
||||
"bound_roles": boundRoles,
|
||||
"bound_count": len(boundRoles),
|
||||
})
|
||||
}
|
||||
|
||||
// getRolesBoundToSkill 获取绑定指定skill的角色列表（不修改配置）
|
||||
func (h *SkillsHandler) getRolesBoundToSkill(skillName string) []string {
|
||||
if h.config.Roles == nil {
|
||||
return []string{}
|
||||
}
|
||||
|
||||
boundRoles := make([]string, 0)
|
||||
for roleName, role := range h.config.Roles {
|
||||
// 确保角色名称正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = roleName
|
||||
}
|
||||
|
||||
// 检查角色的Skills列表中是否包含该skill
|
||||
if len(role.Skills) > 0 {
|
||||
for _, skill := range role.Skills {
|
||||
if skill == skillName {
|
||||
boundRoles = append(boundRoles, roleName)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return boundRoles
|
||||
}
|
||||
|
||||
// CreateSkill 创建新 skill（标准 Agent Skills：生成 SKILL.md + YAML front matter）
|
||||
func (h *SkillsHandler) CreateSkill(c *gin.Context) {
|
||||
var req struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description" binding:"required"`
|
||||
Content string `json:"content" binding:"required"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if !isValidSkillName(req.Name) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill 目录名须为小写字母、数字、连字符（与 Agent Skills name 一致）"})
|
||||
return
|
||||
}
|
||||
|
||||
manifest := &skillpackage.SkillManifest{
|
||||
Name: req.Name,
|
||||
Description: strings.TrimSpace(req.Description),
|
||||
}
|
||||
skillMD, err := skillpackage.BuildSkillMD(manifest, req.Content)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := skillpackage.ValidateSkillMDPackage(skillMD, req.Name); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
skillDir := filepath.Join(h.skillsRootAbs(), req.Name)
|
||||
if err := os.MkdirAll(skillDir, 0755); err != nil {
|
||||
h.logger.Error("创建skill目录失败", zap.String("skill", req.Name), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "创建skill目录失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := os.Stat(filepath.Join(skillDir, "SKILL.md")); err == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill已存在"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), skillMD, 0644); err != nil {
|
||||
h.logger.Error("创建 SKILL.md 失败", zap.String("skill", req.Name), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "创建 SKILL.md 失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("创建skill成功", zap.String("skill", req.Name))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "skill已创建",
|
||||
"skill": map[string]interface{}{
|
||||
"name": req.Name,
|
||||
"path": skillDir,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateSkill 更新 SKILL.md（保留 front matter 中除 description 外的字段；可选覆盖 description）
|
||||
func (h *SkillsHandler) UpdateSkill(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Description string `json:"description"`
|
||||
Content string `json:"content" binding:"required"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
mdPath := filepath.Join(h.skillsRootAbs(), skillName, "SKILL.md")
|
||||
raw, err := os.ReadFile(mdPath)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "skill不存在: " + err.Error()})
|
||||
return
|
||||
}
|
||||
m, _, err := skillpackage.ParseSkillMD(raw)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if req.Description != "" {
|
||||
m.Description = strings.TrimSpace(req.Description)
|
||||
}
|
||||
skillMD, err := skillpackage.BuildSkillMD(m, req.Content)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := skillpackage.ValidateSkillMDPackage(skillMD, skillName); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
skillDir := filepath.Join(h.skillsRootAbs(), skillName)
|
||||
|
||||
if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), skillMD, 0644); err != nil {
|
||||
h.logger.Error("更新 SKILL.md 失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "更新 SKILL.md 失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("更新skill成功", zap.String("skill", skillName))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "skill已更新",
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteSkill 删除skill
|
||||
func (h *SkillsHandler) DeleteSkill(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
// 检查是否有角色绑定了该skill，如果有则自动移除绑定
|
||||
affectedRoles := h.removeSkillFromRoles(skillName)
|
||||
if len(affectedRoles) > 0 {
|
||||
h.logger.Info("从角色中移除skill绑定",
|
||||
zap.String("skill", skillName),
|
||||
zap.Strings("roles", affectedRoles))
|
||||
}
|
||||
|
||||
skillDir := filepath.Join(h.skillsRootAbs(), skillName)
|
||||
if err := os.RemoveAll(skillDir); err != nil {
|
||||
h.logger.Error("删除skill失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "删除skill失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
responseMsg := "skill已删除"
|
||||
if len(affectedRoles) > 0 {
|
||||
responseMsg = fmt.Sprintf("skill已删除，已自动从 %d 个角色中移除绑定: %s",
|
||||
len(affectedRoles), strings.Join(affectedRoles, ", "))
|
||||
}
|
||||
|
||||
h.logger.Info("删除skill成功", zap.String("skill", skillName))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": responseMsg,
|
||||
"affected_roles": affectedRoles,
|
||||
})
|
||||
}
|
||||
|
||||
// GetSkillStats 获取skills调用统计信息
|
||||
func (h *SkillsHandler) GetSkillStats(c *gin.Context) {
|
||||
skillList, err := skillpackage.ListSkillDirNames(h.skillsRootAbs())
|
||||
if err != nil {
|
||||
h.logger.Error("获取skills列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
skillsDir := h.skillsRootAbs()
|
||||
|
||||
// 从数据库加载调用统计
|
||||
var skillStatsMap map[string]*database.SkillStats
|
||||
if h.db != nil {
|
||||
dbStats, err := h.db.LoadSkillStats()
|
||||
if err != nil {
|
||||
h.logger.Warn("从数据库加载Skills统计信息失败", zap.Error(err))
|
||||
skillStatsMap = make(map[string]*database.SkillStats)
|
||||
} else {
|
||||
skillStatsMap = dbStats
|
||||
}
|
||||
} else {
|
||||
skillStatsMap = make(map[string]*database.SkillStats)
|
||||
}
|
||||
|
||||
// 构建统计信息（包含所有skills，即使没有调用记录）
|
||||
statsList := make([]map[string]interface{}, 0, len(skillList))
|
||||
totalCalls := 0
|
||||
totalSuccess := 0
|
||||
totalFailed := 0
|
||||
|
||||
for _, skillName := range skillList {
|
||||
stat, exists := skillStatsMap[skillName]
|
||||
if !exists {
|
||||
stat = &database.SkillStats{
|
||||
SkillName: skillName,
|
||||
TotalCalls: 0,
|
||||
SuccessCalls: 0,
|
||||
FailedCalls: 0,
|
||||
}
|
||||
}
|
||||
|
||||
totalCalls += stat.TotalCalls
|
||||
totalSuccess += stat.SuccessCalls
|
||||
totalFailed += stat.FailedCalls
|
||||
|
||||
lastCallTimeStr := ""
|
||||
if stat.LastCallTime != nil {
|
||||
lastCallTimeStr = stat.LastCallTime.Format("2006-01-02 15:04:05")
|
||||
}
|
||||
|
||||
statsList = append(statsList, map[string]interface{}{
|
||||
"skill_name": stat.SkillName,
|
||||
"total_calls": stat.TotalCalls,
|
||||
"success_calls": stat.SuccessCalls,
|
||||
"failed_calls": stat.FailedCalls,
|
||||
"last_call_time": lastCallTimeStr,
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"total_skills": len(skillList),
|
||||
"total_calls": totalCalls,
|
||||
"total_success": totalSuccess,
|
||||
"total_failed": totalFailed,
|
||||
"skills_dir": skillsDir,
|
||||
"stats": statsList,
|
||||
})
|
||||
}
|
||||
|
||||
// ClearSkillStats 清空所有Skills统计信息
|
||||
func (h *SkillsHandler) ClearSkillStats(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "数据库连接未配置"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.ClearSkillStats(); err != nil {
|
||||
h.logger.Error("清空Skills统计信息失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "清空统计信息失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("已清空所有Skills统计信息")
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "已清空所有Skills统计信息",
|
||||
})
|
||||
}
|
||||
|
||||
// ClearSkillStatsByName 清空指定skill的统计信息
|
||||
func (h *SkillsHandler) ClearSkillStatsByName(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "数据库连接未配置"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.ClearSkillStatsByName(skillName); err != nil {
|
||||
h.logger.Error("清空指定skill统计信息失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "清空统计信息失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("已清空指定skill统计信息", zap.String("skill", skillName))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": fmt.Sprintf("已清空skill '%s' 的统计信息", skillName),
|
||||
})
|
||||
}
|
||||
|
||||
// removeSkillFromRoles 从所有角色中移除指定的skill绑定
|
||||
// 返回受影响角色名称列表
|
||||
func (h *SkillsHandler) removeSkillFromRoles(skillName string) []string {
|
||||
if h.config.Roles == nil {
|
||||
return []string{}
|
||||
}
|
||||
|
||||
affectedRoles := make([]string, 0)
|
||||
rolesToUpdate := make(map[string]config.RoleConfig)
|
||||
|
||||
// 遍历所有角色，查找并移除skill绑定
|
||||
for roleName, role := range h.config.Roles {
|
||||
// 确保角色名称正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = roleName
|
||||
}
|
||||
|
||||
// 检查角色的Skills列表中是否包含要删除的skill
|
||||
if len(role.Skills) > 0 {
|
||||
updated := false
|
||||
newSkills := make([]string, 0, len(role.Skills))
|
||||
for _, skill := range role.Skills {
|
||||
if skill != skillName {
|
||||
newSkills = append(newSkills, skill)
|
||||
} else {
|
||||
updated = true
|
||||
}
|
||||
}
|
||||
if updated {
|
||||
role.Skills = newSkills
|
||||
rolesToUpdate[roleName] = role
|
||||
affectedRoles = append(affectedRoles, roleName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 如果有角色需要更新，保存到文件
|
||||
if len(rolesToUpdate) > 0 {
|
||||
// 更新内存中的配置
|
||||
for roleName, role := range rolesToUpdate {
|
||||
h.config.Roles[roleName] = role
|
||||
}
|
||||
// 保存更新后的角色配置到文件
|
||||
if err := h.saveRolesConfig(); err != nil {
|
||||
h.logger.Error("保存角色配置失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
return affectedRoles
|
||||
}
|
||||
|
||||
// saveRolesConfig 保存角色配置到文件（从SkillsHandler调用）
|
||||
func (h *SkillsHandler) saveRolesConfig() error {
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
rolesDir := h.config.RolesDir
|
||||
if rolesDir == "" {
|
||||
rolesDir = "roles" // 默认目录
|
||||
}
|
||||
|
||||
// 如果是相对路径，相对于配置文件所在目录
|
||||
if !filepath.IsAbs(rolesDir) {
|
||||
rolesDir = filepath.Join(configDir, rolesDir)
|
||||
}
|
||||
|
||||
// 确保目录存在
|
||||
if err := os.MkdirAll(rolesDir, 0755); err != nil {
|
||||
return fmt.Errorf("创建角色目录失败: %w", err)
|
||||
}
|
||||
|
||||
// 保存每个角色到独立的文件
|
||||
if h.config.Roles != nil {
|
||||
for roleName, role := range h.config.Roles {
|
||||
// 确保角色名称正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = roleName
|
||||
}
|
||||
|
||||
// 使用角色名称作为文件名（安全化文件名，避免特殊字符）
|
||||
safeFileName := sanitizeRoleFileName(role.Name)
|
||||
roleFile := filepath.Join(rolesDir, safeFileName+".yaml")
|
||||
|
||||
// 将角色配置序列化为YAML
|
||||
roleData, err := yaml.Marshal(&role)
|
||||
if err != nil {
|
||||
h.logger.Error("序列化角色配置失败", zap.String("role", roleName), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
// 处理icon字段：确保包含\U的icon值被引号包围（YAML需要引号才能正确解析Unicode转义）
|
||||
roleDataStr := string(roleData)
|
||||
if role.Icon != "" && strings.HasPrefix(role.Icon, "\\U") {
|
||||
// 匹配 icon: \UXXXXXXXX 格式（没有引号），排除已经有引号的情况
|
||||
re := regexp.MustCompile(`(?m)^(icon:\s+)(\\U[0-9A-F]{8})(\s*)$`)
|
||||
roleDataStr = re.ReplaceAllString(roleDataStr, `${1}"${2}"${3}`)
|
||||
roleData = []byte(roleDataStr)
|
||||
}
|
||||
|
||||
// 写入文件
|
||||
if err := os.WriteFile(roleFile, roleData, 0644); err != nil {
|
||||
h.logger.Error("保存角色配置文件失败", zap.String("role", roleName), zap.String("file", roleFile), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
h.logger.Info("角色配置已保存到文件", zap.String("role", roleName), zap.String("file", roleFile))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// sanitizeRoleFileName 将角色名称转换为安全的文件名
|
||||
func sanitizeRoleFileName(name string) string {
|
||||
// 替换可能不安全的字符
|
||||
replacer := map[rune]string{
|
||||
'/': "_",
|
||||
'\\': "_",
|
||||
':': "_",
|
||||
'*': "_",
|
||||
'?': "_",
|
||||
'"': "_",
|
||||
'<': "_",
|
||||
'>': "_",
|
||||
'|': "_",
|
||||
' ': "_",
|
||||
}
|
||||
|
||||
var result []rune
|
||||
for _, r := range name {
|
||||
if replacement, ok := replacer[r]; ok {
|
||||
result = append(result, []rune(replacement)...)
|
||||
} else {
|
||||
result = append(result, r)
|
||||
}
|
||||
}
|
||||
|
||||
fileName := string(result)
|
||||
// 如果文件名为空，使用默认名称
|
||||
if fileName == "" {
|
||||
fileName = "role"
|
||||
}
|
||||
|
||||
return fileName
|
||||
}
|
||||
|
||||
// isValidSkillName 验证 skill 目录名（与 Agent Skills 的 name 字段一致：小写、数字、连字符）
|
||||
func isValidSkillName(name string) bool {
|
||||
if name == "" || len(name) > 100 {
|
||||
return false
|
||||
}
|
||||
for _, r := range name {
|
||||
if !((r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-') {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -1,58 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// sseInterval is how often we write on long SSE streams. Shorter intervals help NATs and
|
||||
// some proxies that treat connections as idle; 10s is a reasonable balance with traffic.
|
||||
const sseKeepaliveInterval = 10 * time.Second
|
||||
|
||||
// sseKeepalive sends periodic SSE traffic so proxies (e.g. nginx proxy_read_timeout), NATs,
|
||||
// and load balancers do not close long-running streams. Some intermediaries ignore comment-only
|
||||
// lines, so we send both a comment and a minimal data frame (type heartbeat) per tick.
|
||||
//
|
||||
// writeMu must be the same mutex used by sendEvent for this request: concurrent writes to
|
||||
// http.ResponseWriter break chunked transfer encoding (browser: net::ERR_INVALID_CHUNKED_ENCODING).
|
||||
func sseKeepalive(c *gin.Context, stop <-chan struct{}, writeMu *sync.Mutex) {
|
||||
if writeMu == nil {
|
||||
return
|
||||
}
|
||||
ticker := time.NewTicker(sseKeepaliveInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-c.Request.Context().Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-c.Request.Context().Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
writeMu.Lock()
|
||||
if _, err := fmt.Fprintf(c.Writer, ": keepalive\n\n"); err != nil {
|
||||
writeMu.Unlock()
|
||||
return
|
||||
}
|
||||
// data: frame so strict proxies still see downstream bytes (comments alone may not reset timers)
|
||||
if _, err := fmt.Fprintf(c.Writer, `data: {"type":"heartbeat"}`+"\n\n"); err != nil {
|
||||
writeMu.Unlock()
|
||||
return
|
||||
}
|
||||
if flusher, ok := c.Writer.(http.Flusher); ok {
|
||||
flusher.Flush()
|
||||
}
|
||||
writeMu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,276 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrTaskCancelled 用户取消任务的错误
|
||||
var ErrTaskCancelled = errors.New("agent task cancelled by user")
|
||||
|
||||
// ErrTaskAlreadyRunning 会话已有任务正在执行
|
||||
var ErrTaskAlreadyRunning = errors.New("agent task already running for conversation")
|
||||
|
||||
// AgentTask 描述正在运行的Agent任务
|
||||
type AgentTask struct {
|
||||
ConversationID string `json:"conversationId"`
|
||||
Message string `json:"message,omitempty"`
|
||||
StartedAt time.Time `json:"startedAt"`
|
||||
Status string `json:"status"`
|
||||
CancellingAt time.Time `json:"-"` // 进入 cancelling 状态的时间，用于清理长时间卡住的任务
|
||||
|
||||
cancel func(error)
|
||||
}
|
||||
|
||||
// CompletedTask 已完成的任务（用于历史记录）
|
||||
type CompletedTask struct {
|
||||
ConversationID string `json:"conversationId"`
|
||||
Message string `json:"message,omitempty"`
|
||||
StartedAt time.Time `json:"startedAt"`
|
||||
CompletedAt time.Time `json:"completedAt"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
// AgentTaskManager 管理正在运行的Agent任务
|
||||
type AgentTaskManager struct {
|
||||
mu sync.RWMutex
|
||||
tasks map[string]*AgentTask
|
||||
completedTasks []*CompletedTask // 最近完成的任务历史
|
||||
maxHistorySize int // 最大历史记录数
|
||||
historyRetention time.Duration // 历史记录保留时间
|
||||
}
|
||||
|
||||
const (
|
||||
// cancellingStuckThreshold 处于「取消中」超过此时长则强制从运行列表移除。正常取消会在当前步骤内返回，
|
||||
// 超过则视为卡住，尽快释放会话。常见做法多为 30–60s 内释放。
|
||||
cancellingStuckThreshold = 45 * time.Second
|
||||
// cancellingStuckThresholdLegacy 未记录 CancellingAt 时用 StartedAt 判断的兜底时长
|
||||
cancellingStuckThresholdLegacy = 2 * time.Minute
|
||||
cleanupInterval = 15 * time.Second // 与上面阈值配合，最长约 60s 内移除
|
||||
)
|
||||
|
||||
// NewAgentTaskManager 创建任务管理器
|
||||
func NewAgentTaskManager() *AgentTaskManager {
|
||||
m := &AgentTaskManager{
|
||||
tasks: make(map[string]*AgentTask),
|
||||
completedTasks: make([]*CompletedTask, 0),
|
||||
maxHistorySize: 50, // 最多保留50条历史记录
|
||||
historyRetention: 24 * time.Hour, // 保留24小时
|
||||
}
|
||||
go m.runStuckCancellingCleanup()
|
||||
return m
|
||||
}
|
||||
|
||||
// runStuckCancellingCleanup 定期将长时间处于「取消中」的任务强制结束，避免卡住无法发新消息
|
||||
func (m *AgentTaskManager) runStuckCancellingCleanup() {
|
||||
ticker := time.NewTicker(cleanupInterval)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
m.cleanupStuckCancelling()
|
||||
}
|
||||
}
|
||||
|
||||
func (m *AgentTaskManager) cleanupStuckCancelling() {
|
||||
m.mu.Lock()
|
||||
var toFinish []string
|
||||
now := time.Now()
|
||||
for id, task := range m.tasks {
|
||||
if task.Status != "cancelling" {
|
||||
continue
|
||||
}
|
||||
var elapsed time.Duration
|
||||
if !task.CancellingAt.IsZero() {
|
||||
elapsed = now.Sub(task.CancellingAt)
|
||||
if elapsed < cancellingStuckThreshold {
|
||||
continue
|
||||
}
|
||||
} else {
|
||||
elapsed = now.Sub(task.StartedAt)
|
||||
if elapsed < cancellingStuckThresholdLegacy {
|
||||
continue
|
||||
}
|
||||
}
|
||||
toFinish = append(toFinish, id)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
for _, id := range toFinish {
|
||||
m.FinishTask(id, "cancelled")
|
||||
}
|
||||
}
|
||||
|
||||
// StartTask 注册并开始一个新的任务
|
||||
func (m *AgentTaskManager) StartTask(conversationID, message string, cancel context.CancelCauseFunc) (*AgentTask, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if _, exists := m.tasks[conversationID]; exists {
|
||||
return nil, ErrTaskAlreadyRunning
|
||||
}
|
||||
|
||||
task := &AgentTask{
|
||||
ConversationID: conversationID,
|
||||
Message: message,
|
||||
StartedAt: time.Now(),
|
||||
Status: "running",
|
||||
cancel: func(err error) {
|
||||
if cancel != nil {
|
||||
cancel(err)
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
m.tasks[conversationID] = task
|
||||
return task, nil
|
||||
}
|
||||
|
||||
// CancelTask 取消指定会话的任务。若任务已在取消中，仍返回 (true, nil) 以便接口幂等、前端不报错。
|
||||
func (m *AgentTaskManager) CancelTask(conversationID string, cause error) (bool, error) {
|
||||
m.mu.Lock()
|
||||
task, exists := m.tasks[conversationID]
|
||||
if !exists {
|
||||
m.mu.Unlock()
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// 如果已经处于取消流程，视为成功（幂等），避免前端重复点击报「未找到任务」
|
||||
if task.Status == "cancelling" {
|
||||
m.mu.Unlock()
|
||||
return true, nil
|
||||
}
|
||||
|
||||
task.Status = "cancelling"
|
||||
task.CancellingAt = time.Now()
|
||||
cancel := task.cancel
|
||||
m.mu.Unlock()
|
||||
|
||||
if cause == nil {
|
||||
cause = ErrTaskCancelled
|
||||
}
|
||||
if cancel != nil {
|
||||
cancel(cause)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// UpdateTaskStatus 更新任务状态但不删除任务（用于在发送事件前更新状态）
|
||||
func (m *AgentTaskManager) UpdateTaskStatus(conversationID string, status string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
task, exists := m.tasks[conversationID]
|
||||
if !exists {
|
||||
return
|
||||
}
|
||||
|
||||
if status != "" {
|
||||
task.Status = status
|
||||
}
|
||||
}
|
||||
|
||||
// FinishTask 完成任务并从管理器中移除
|
||||
func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
task, exists := m.tasks[conversationID]
|
||||
if !exists {
|
||||
return
|
||||
}
|
||||
|
||||
if finalStatus != "" {
|
||||
task.Status = finalStatus
|
||||
}
|
||||
|
||||
// 保存到历史记录
|
||||
completedTask := &CompletedTask{
|
||||
ConversationID: task.ConversationID,
|
||||
Message: task.Message,
|
||||
StartedAt: task.StartedAt,
|
||||
CompletedAt: time.Now(),
|
||||
Status: finalStatus,
|
||||
}
|
||||
|
||||
// 添加到历史记录
|
||||
m.completedTasks = append(m.completedTasks, completedTask)
|
||||
|
||||
// 清理过期和过多的历史记录
|
||||
m.cleanupHistory()
|
||||
|
||||
// 从运行任务中移除
|
||||
delete(m.tasks, conversationID)
|
||||
}
|
||||
|
||||
// cleanupHistory 清理过期的历史记录
|
||||
func (m *AgentTaskManager) cleanupHistory() {
|
||||
now := time.Now()
|
||||
cutoffTime := now.Add(-m.historyRetention)
|
||||
|
||||
// 过滤掉过期的记录
|
||||
validTasks := make([]*CompletedTask, 0, len(m.completedTasks))
|
||||
for _, task := range m.completedTasks {
|
||||
if task.CompletedAt.After(cutoffTime) {
|
||||
validTasks = append(validTasks, task)
|
||||
}
|
||||
}
|
||||
|
||||
// 如果仍然超过最大数量，只保留最新的
|
||||
if len(validTasks) > m.maxHistorySize {
|
||||
// 按完成时间排序，保留最新的
|
||||
// 由于是追加的，最新的在最后，所以直接取最后N个
|
||||
start := len(validTasks) - m.maxHistorySize
|
||||
validTasks = validTasks[start:]
|
||||
}
|
||||
|
||||
m.completedTasks = validTasks
|
||||
}
|
||||
|
||||
// GetActiveTasks 返回所有正在运行的任务
|
||||
func (m *AgentTaskManager) GetActiveTasks() []*AgentTask {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
result := make([]*AgentTask, 0, len(m.tasks))
|
||||
for _, task := range m.tasks {
|
||||
result = append(result, &AgentTask{
|
||||
ConversationID: task.ConversationID,
|
||||
Message: task.Message,
|
||||
StartedAt: task.StartedAt,
|
||||
Status: task.Status,
|
||||
})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// GetCompletedTasks 返回最近完成的任务历史
|
||||
func (m *AgentTaskManager) GetCompletedTasks() []*CompletedTask {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
// 清理过期记录（只读锁，不影响其他操作）
|
||||
// 注意：这里不能直接调用cleanupHistory，因为需要写锁
|
||||
// 所以返回时过滤过期记录
|
||||
now := time.Now()
|
||||
cutoffTime := now.Add(-m.historyRetention)
|
||||
|
||||
result := make([]*CompletedTask, 0, len(m.completedTasks))
|
||||
for _, task := range m.completedTasks {
|
||||
if task.CompletedAt.After(cutoffTime) {
|
||||
result = append(result, task)
|
||||
}
|
||||
}
|
||||
|
||||
// 按完成时间倒序排序（最新的在前）
|
||||
// 由于是追加的，最新的在最后，需要反转
|
||||
for i, j := 0, len(result)-1; i < j; i, j = i+1, j-1 {
|
||||
result[i], result[j] = result[j], result[i]
|
||||
}
|
||||
|
||||
// 限制返回数量
|
||||
if len(result) > m.maxHistorySize {
|
||||
result = result[:m.maxHistorySize]
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
@@ -1,257 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
terminalMaxCommandLen = 4096
|
||||
terminalMaxOutputLen = 256 * 1024 // 256KB
|
||||
terminalTimeout = 30 * time.Minute
|
||||
)
|
||||
|
||||
// TerminalHandler 处理系统设置中的终端命令执行
|
||||
type TerminalHandler struct {
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// maskTerminalCommand 对可能包含敏感信息的终端命令做脱敏，避免在日志中直接记录密码等内容
|
||||
func maskTerminalCommand(cmd string) string {
|
||||
trimmed := strings.TrimSpace(cmd)
|
||||
lower := strings.ToLower(trimmed)
|
||||
if strings.Contains(lower, "sudo") || strings.Contains(lower, "password") {
|
||||
return "[masked sensitive terminal command]"
|
||||
}
|
||||
if len(trimmed) > 256 {
|
||||
return trimmed[:256] + "..."
|
||||
}
|
||||
return trimmed
|
||||
}
|
||||
|
||||
// NewTerminalHandler 创建终端处理器
|
||||
func NewTerminalHandler(logger *zap.Logger) *TerminalHandler {
|
||||
return &TerminalHandler{logger: logger}
|
||||
}
|
||||
|
||||
// RunCommandRequest 执行命令请求
|
||||
type RunCommandRequest struct {
|
||||
Command string `json:"command"`
|
||||
Shell string `json:"shell,omitempty"`
|
||||
Cwd string `json:"cwd,omitempty"`
|
||||
}
|
||||
|
||||
// RunCommandResponse 执行命令响应
|
||||
type RunCommandResponse struct {
|
||||
Stdout string `json:"stdout"`
|
||||
Stderr string `json:"stderr"`
|
||||
ExitCode int `json:"exit_code"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// RunCommand 执行终端命令（需登录）
|
||||
func (h *TerminalHandler) RunCommand(c *gin.Context) {
|
||||
var req RunCommandRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "请求体无效，需要 command 字段"})
|
||||
return
|
||||
}
|
||||
|
||||
cmdStr := strings.TrimSpace(req.Command)
|
||||
if cmdStr == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "command 不能为空"})
|
||||
return
|
||||
}
|
||||
if len(cmdStr) > terminalMaxCommandLen {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "命令过长"})
|
||||
return
|
||||
}
|
||||
|
||||
shell := req.Shell
|
||||
if shell == "" {
|
||||
if runtime.GOOS == "windows" {
|
||||
shell = "cmd"
|
||||
} else {
|
||||
shell = "sh"
|
||||
}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), terminalTimeout)
|
||||
defer cancel()
|
||||
|
||||
var cmd *exec.Cmd
|
||||
if runtime.GOOS == "windows" {
|
||||
cmd = exec.CommandContext(ctx, "cmd", "/c", cmdStr)
|
||||
} else {
|
||||
cmd = exec.CommandContext(ctx, shell, "-c", cmdStr)
|
||||
// 无 TTY 时设置 COLUMNS/TERM，使 ping 等工具的 usage 排版与真实终端一致
|
||||
cmd.Env = append(os.Environ(), "COLUMNS=256", "LINES=40", "TERM=xterm-256color")
|
||||
}
|
||||
|
||||
if req.Cwd != "" {
|
||||
absCwd, err := filepath.Abs(req.Cwd)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "工作目录无效"})
|
||||
return
|
||||
}
|
||||
cur, _ := os.Getwd()
|
||||
curAbs, _ := filepath.Abs(cur)
|
||||
rel, err := filepath.Rel(curAbs, absCwd)
|
||||
if err != nil || strings.HasPrefix(rel, "..") || rel == ".." {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "工作目录必须在当前进程目录下"})
|
||||
return
|
||||
}
|
||||
cmd.Dir = absCwd
|
||||
}
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
err := cmd.Run()
|
||||
stdoutBytes := stdout.Bytes()
|
||||
stderrBytes := stderr.Bytes()
|
||||
|
||||
// 限制输出长度，防止内存占用过大（复制后截断，避免修改原 buffer）
|
||||
truncSuffix := []byte("\n...(输出已截断)\n")
|
||||
if len(stdoutBytes) > terminalMaxOutputLen {
|
||||
tmp := make([]byte, terminalMaxOutputLen+len(truncSuffix))
|
||||
n := copy(tmp, stdoutBytes[:terminalMaxOutputLen])
|
||||
copy(tmp[n:], truncSuffix)
|
||||
stdoutBytes = tmp
|
||||
}
|
||||
if len(stderrBytes) > terminalMaxOutputLen {
|
||||
tmp := make([]byte, terminalMaxOutputLen+len(truncSuffix))
|
||||
n := copy(tmp, stderrBytes[:terminalMaxOutputLen])
|
||||
copy(tmp[n:], truncSuffix)
|
||||
stderrBytes = tmp
|
||||
}
|
||||
|
||||
exitCode := 0
|
||||
if err != nil {
|
||||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||||
exitCode = exitErr.ExitCode()
|
||||
} else {
|
||||
exitCode = -1
|
||||
}
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
so := strings.ReplaceAll(string(stdoutBytes), "\r\n", "\n")
|
||||
so = strings.ReplaceAll(so, "\r", "\n")
|
||||
se := strings.ReplaceAll(string(stderrBytes), "\r\n", "\n")
|
||||
se = strings.ReplaceAll(se, "\r", "\n")
|
||||
resp := RunCommandResponse{
|
||||
Stdout: so,
|
||||
Stderr: se,
|
||||
ExitCode: -1,
|
||||
Error: "命令执行超时（" + terminalTimeout.String() + "）",
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
return
|
||||
}
|
||||
h.logger.Debug("终端命令执行异常", zap.String("command", maskTerminalCommand(cmdStr)), zap.Error(err))
|
||||
}
|
||||
|
||||
// 统一为 \n，避免前端因 \r 出现错位/对角线排版
|
||||
stdoutStr := strings.ReplaceAll(string(stdoutBytes), "\r\n", "\n")
|
||||
stdoutStr = strings.ReplaceAll(stdoutStr, "\r", "\n")
|
||||
stderrStr := strings.ReplaceAll(string(stderrBytes), "\r\n", "\n")
|
||||
stderrStr = strings.ReplaceAll(stderrStr, "\r", "\n")
|
||||
|
||||
resp := RunCommandResponse{
|
||||
Stdout: stdoutStr,
|
||||
Stderr: stderrStr,
|
||||
ExitCode: exitCode,
|
||||
}
|
||||
if err != nil && exitCode != 0 {
|
||||
resp.Error = err.Error()
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// streamEvent SSE 事件
|
||||
type streamEvent struct {
|
||||
T string `json:"t"` // "out" | "err" | "exit"
|
||||
D string `json:"d,omitempty"`
|
||||
C int `json:"c"` // exit code（不用 omitempty，否则 0 不序列化导致前端显示 [exit undefined]）
|
||||
}
|
||||
|
||||
// RunCommandStream 流式执行命令，输出实时推送到前端（SSE）
|
||||
func (h *TerminalHandler) RunCommandStream(c *gin.Context) {
|
||||
var req RunCommandRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "请求体无效，需要 command 字段"})
|
||||
return
|
||||
}
|
||||
cmdStr := strings.TrimSpace(req.Command)
|
||||
if cmdStr == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "command 不能为空"})
|
||||
return
|
||||
}
|
||||
if len(cmdStr) > terminalMaxCommandLen {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "命令过长"})
|
||||
return
|
||||
}
|
||||
shell := req.Shell
|
||||
if shell == "" {
|
||||
if runtime.GOOS == "windows" {
|
||||
shell = "cmd"
|
||||
} else {
|
||||
shell = "sh"
|
||||
}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), terminalTimeout)
|
||||
defer cancel()
|
||||
|
||||
var cmd *exec.Cmd
|
||||
if runtime.GOOS == "windows" {
|
||||
cmd = exec.CommandContext(ctx, "cmd", "/c", cmdStr)
|
||||
} else {
|
||||
cmd = exec.CommandContext(ctx, shell, "-c", cmdStr)
|
||||
cmd.Env = append(os.Environ(), "COLUMNS=256", "LINES=40", "TERM=xterm-256color")
|
||||
}
|
||||
if req.Cwd != "" {
|
||||
absCwd, err := filepath.Abs(req.Cwd)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "工作目录无效"})
|
||||
return
|
||||
}
|
||||
cur, _ := os.Getwd()
|
||||
curAbs, _ := filepath.Abs(cur)
|
||||
rel, err := filepath.Rel(curAbs, absCwd)
|
||||
if err != nil || strings.HasPrefix(rel, "..") || rel == ".." {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "工作目录必须在当前进程目录下"})
|
||||
return
|
||||
}
|
||||
cmd.Dir = absCwd
|
||||
}
|
||||
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
c.Writer.WriteHeader(http.StatusOK)
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
|
||||
sendEvent := func(ev streamEvent) {
|
||||
body, _ := json.Marshal(ev)
|
||||
c.SSEvent("", string(body))
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
runCommandStreamImpl(cmd, sendEvent, ctx)
|
||||
}
|
||||
@@ -1,46 +0,0 @@
|
||||
//go:build !windows
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"os/exec"
|
||||
"strings"
|
||||
|
||||
"github.com/creack/pty"
|
||||
)
|
||||
|
||||
const ptyCols = 256
|
||||
const ptyRows = 40
|
||||
|
||||
// runCommandStreamImpl 在 Unix 下用 PTY 执行，使 ping 等命令按终端宽度排版（isatty 为真）
|
||||
func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx context.Context) {
|
||||
ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: ptyCols, Rows: ptyRows})
|
||||
if err != nil {
|
||||
sendEvent(streamEvent{T: "exit", C: -1})
|
||||
return
|
||||
}
|
||||
defer ptmx.Close()
|
||||
|
||||
normalize := func(s string) string {
|
||||
s = strings.ReplaceAll(s, "\r\n", "\n")
|
||||
return strings.ReplaceAll(s, "\r", "\n")
|
||||
}
|
||||
sc := bufio.NewScanner(ptmx)
|
||||
for sc.Scan() {
|
||||
sendEvent(streamEvent{T: "out", D: normalize(sc.Text())})
|
||||
}
|
||||
exitCode := 0
|
||||
if err := cmd.Wait(); err != nil {
|
||||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||||
exitCode = exitErr.ExitCode()
|
||||
} else {
|
||||
exitCode = -1
|
||||
}
|
||||
}
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
exitCode = -1
|
||||
}
|
||||
sendEvent(streamEvent{T: "exit", C: exitCode})
|
||||
}
|
||||
@@ -1,65 +0,0 @@
|
||||
//go:build windows
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// runCommandStreamImpl 在 Windows 下用 stdout/stderr 管道执行
|
||||
func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx context.Context) {
|
||||
stdoutPipe, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
sendEvent(streamEvent{T: "exit", C: -1})
|
||||
return
|
||||
}
|
||||
stderrPipe, err := cmd.StderrPipe()
|
||||
if err != nil {
|
||||
sendEvent(streamEvent{T: "exit", C: -1})
|
||||
return
|
||||
}
|
||||
if err := cmd.Start(); err != nil {
|
||||
sendEvent(streamEvent{T: "exit", C: -1})
|
||||
return
|
||||
}
|
||||
|
||||
normalize := func(s string) string {
|
||||
s = strings.ReplaceAll(s, "\r\n", "\n")
|
||||
return strings.ReplaceAll(s, "\r", "\n")
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
sc := bufio.NewScanner(stdoutPipe)
|
||||
for sc.Scan() {
|
||||
sendEvent(streamEvent{T: "out", D: normalize(sc.Text())})
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
sc := bufio.NewScanner(stderrPipe)
|
||||
for sc.Scan() {
|
||||
sendEvent(streamEvent{T: "err", D: normalize(sc.Text())})
|
||||
}
|
||||
}()
|
||||
|
||||
wg.Wait()
|
||||
exitCode := 0
|
||||
if err := cmd.Wait(); err != nil {
|
||||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||||
exitCode = exitErr.ExitCode()
|
||||
} else {
|
||||
exitCode = -1
|
||||
}
|
||||
}
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
exitCode = -1
|
||||
}
|
||||
sendEvent(streamEvent{T: "exit", C: exitCode})
|
||||
}
|
||||
@@ -1,112 +0,0 @@
|
||||
//go:build !windows
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"time"
|
||||
|
||||
"github.com/creack/pty"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// terminalResize is sent by the frontend when the xterm.js terminal is resized.
|
||||
type terminalResize struct {
|
||||
Type string `json:"type"`
|
||||
Cols uint16 `json:"cols"`
|
||||
Rows uint16 `json:"rows"`
|
||||
}
|
||||
|
||||
// wsUpgrader 仅用于系统设置中的终端 WebSocket，会复用已有的登录保护（JWT 中间件在上层路由组）
|
||||
var wsUpgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
// 由于已在 Gin 路由层做了认证，这里放宽 Origin，方便在同一域名下通过 HTTPS/WSS 访问
|
||||
return true
|
||||
},
|
||||
}
|
||||
|
||||
// RunCommandWS 提供真正交互式 Shell：基于 WebSocket + PTY 的长会话
|
||||
// 前端建立 WebSocket 连接后，所有键盘输入都会透传到 Shell，Shell 的输出也会实时写回前端。
|
||||
func (h *TerminalHandler) RunCommandWS(c *gin.Context) {
|
||||
conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
// 启动交互式 Shell，这里优先使用 bash，找不到则退回 sh
|
||||
shell := "bash"
|
||||
if _, err := exec.LookPath(shell); err != nil {
|
||||
shell = "sh"
|
||||
}
|
||||
cmd := exec.Command(shell)
|
||||
cmd.Env = append(os.Environ(),
|
||||
"COLUMNS=80",
|
||||
"LINES=24",
|
||||
"TERM=xterm-256color",
|
||||
)
|
||||
|
||||
// Use 80x24 as a safe default; the frontend will send the actual size immediately after connecting.
|
||||
ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: 80, Rows: 24})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer ptmx.Close()
|
||||
|
||||
// Shell -> WebSocket：将 PTY 输出实时发给前端
|
||||
doneChan := make(chan struct{})
|
||||
go func() {
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, err := ptmx.Read(buf)
|
||||
if n > 0 {
|
||||
_ = conn.WriteMessage(websocket.BinaryMessage, buf[:n])
|
||||
}
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
close(doneChan)
|
||||
}()
|
||||
|
||||
// WebSocket -> Shell：将前端输入写入 PTY（包括 sudo 密码、Ctrl+C 等）
|
||||
conn.SetReadLimit(64 * 1024)
|
||||
_ = conn.SetReadDeadline(time.Now().Add(terminalTimeout))
|
||||
conn.SetPongHandler(func(string) error {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(terminalTimeout))
|
||||
return nil
|
||||
})
|
||||
|
||||
for {
|
||||
msgType, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
_ = cmd.Process.Kill()
|
||||
break
|
||||
}
|
||||
if msgType != websocket.TextMessage && msgType != websocket.BinaryMessage {
|
||||
continue
|
||||
}
|
||||
if len(data) == 0 {
|
||||
continue
|
||||
}
|
||||
// Check if this is a resize message (JSON with type:"resize")
|
||||
if msgType == websocket.TextMessage && len(data) > 0 && data[0] == '{' {
|
||||
var resize terminalResize
|
||||
if json.Unmarshal(data, &resize) == nil && resize.Type == "resize" && resize.Cols > 0 && resize.Rows > 0 {
|
||||
_ = pty.Setsize(ptmx, &pty.Winsize{Cols: resize.Cols, Rows: resize.Rows})
|
||||
continue
|
||||
}
|
||||
}
|
||||
if _, err := ptmx.Write(data); err != nil {
|
||||
_ = cmd.Process.Kill()
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
<-doneChan
|
||||
}
|
||||
|
||||
@@ -1,263 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// VulnerabilityHandler 漏洞处理器
|
||||
type VulnerabilityHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewVulnerabilityHandler 创建新的漏洞处理器
|
||||
func NewVulnerabilityHandler(db *database.DB, logger *zap.Logger) *VulnerabilityHandler {
|
||||
return &VulnerabilityHandler{
|
||||
db: db,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateVulnerabilityRequest 创建漏洞请求
|
||||
type CreateVulnerabilityRequest struct {
|
||||
ConversationID string `json:"conversation_id" binding:"required"`
|
||||
Title string `json:"title" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity" binding:"required"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
Target string `json:"target"`
|
||||
Proof string `json:"proof"`
|
||||
Impact string `json:"impact"`
|
||||
Recommendation string `json:"recommendation"`
|
||||
}
|
||||
|
||||
// CreateVulnerability 创建漏洞
|
||||
func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) {
|
||||
var req CreateVulnerabilityRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: req.ConversationID,
|
||||
Title: req.Title,
|
||||
Description: req.Description,
|
||||
Severity: req.Severity,
|
||||
Status: req.Status,
|
||||
Type: req.Type,
|
||||
Target: req.Target,
|
||||
Proof: req.Proof,
|
||||
Impact: req.Impact,
|
||||
Recommendation: req.Recommendation,
|
||||
}
|
||||
|
||||
created, err := h.db.CreateVulnerability(vuln)
|
||||
if err != nil {
|
||||
h.logger.Error("创建漏洞失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, created)
|
||||
}
|
||||
|
||||
// GetVulnerability 获取漏洞
|
||||
func (h *VulnerabilityHandler) GetVulnerability(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
vuln, err := h.db.GetVulnerability(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞失败", zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "漏洞不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, vuln)
|
||||
}
|
||||
|
||||
// ListVulnerabilitiesResponse 漏洞列表响应
|
||||
type ListVulnerabilitiesResponse struct {
|
||||
Vulnerabilities []*database.Vulnerability `json:"vulnerabilities"`
|
||||
Total int `json:"total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"page_size"`
|
||||
TotalPages int `json:"total_pages"`
|
||||
}
|
||||
|
||||
// ListVulnerabilities 列出漏洞
|
||||
func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
limitStr := c.DefaultQuery("limit", "20")
|
||||
offsetStr := c.DefaultQuery("offset", "0")
|
||||
pageStr := c.Query("page")
|
||||
id := c.Query("id")
|
||||
conversationID := c.Query("conversation_id")
|
||||
severity := c.Query("severity")
|
||||
status := c.Query("status")
|
||||
|
||||
limit, _ := strconv.Atoi(limitStr)
|
||||
offset, _ := strconv.Atoi(offsetStr)
|
||||
page := 1
|
||||
|
||||
// 如果提供了page参数，优先使用page计算offset
|
||||
if pageStr != "" {
|
||||
if p, err := strconv.Atoi(pageStr); err == nil && p > 0 {
|
||||
page = p
|
||||
offset = (page - 1) * limit
|
||||
}
|
||||
}
|
||||
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 20
|
||||
}
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
|
||||
// 获取总数
|
||||
total, err := h.db.CountVulnerabilities(id, conversationID, severity, status)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞总数失败", zap.Error(err))
|
||||
// 继续执行，使用0作为总数
|
||||
total = 0
|
||||
}
|
||||
|
||||
// 获取漏洞列表
|
||||
vulnerabilities, err := h.db.ListVulnerabilities(limit, offset, id, conversationID, severity, status)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 计算总页数
|
||||
totalPages := (total + limit - 1) / limit
|
||||
if totalPages == 0 {
|
||||
totalPages = 1
|
||||
}
|
||||
|
||||
// 如果使用offset计算page，需要重新计算
|
||||
if pageStr == "" {
|
||||
page = (offset / limit) + 1
|
||||
}
|
||||
|
||||
response := ListVulnerabilitiesResponse{
|
||||
Vulnerabilities: vulnerabilities,
|
||||
Total: total,
|
||||
Page: page,
|
||||
PageSize: limit,
|
||||
TotalPages: totalPages,
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, response)
|
||||
}
|
||||
|
||||
// UpdateVulnerabilityRequest 更新漏洞请求
|
||||
type UpdateVulnerabilityRequest struct {
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
Target string `json:"target"`
|
||||
Proof string `json:"proof"`
|
||||
Impact string `json:"impact"`
|
||||
Recommendation string `json:"recommendation"`
|
||||
}
|
||||
|
||||
// UpdateVulnerability 更新漏洞
|
||||
func (h *VulnerabilityHandler) UpdateVulnerability(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
var req UpdateVulnerabilityRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取现有漏洞
|
||||
existing, err := h.db.GetVulnerability(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "漏洞不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
// 更新字段
|
||||
if req.Title != "" {
|
||||
existing.Title = req.Title
|
||||
}
|
||||
if req.Description != "" {
|
||||
existing.Description = req.Description
|
||||
}
|
||||
if req.Severity != "" {
|
||||
existing.Severity = req.Severity
|
||||
}
|
||||
if req.Status != "" {
|
||||
existing.Status = req.Status
|
||||
}
|
||||
if req.Type != "" {
|
||||
existing.Type = req.Type
|
||||
}
|
||||
if req.Target != "" {
|
||||
existing.Target = req.Target
|
||||
}
|
||||
if req.Proof != "" {
|
||||
existing.Proof = req.Proof
|
||||
}
|
||||
if req.Impact != "" {
|
||||
existing.Impact = req.Impact
|
||||
}
|
||||
if req.Recommendation != "" {
|
||||
existing.Recommendation = req.Recommendation
|
||||
}
|
||||
|
||||
if err := h.db.UpdateVulnerability(id, existing); err != nil {
|
||||
h.logger.Error("更新漏洞失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 返回更新后的漏洞
|
||||
updated, err := h.db.GetVulnerability(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取更新后的漏洞失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, updated)
|
||||
}
|
||||
|
||||
// DeleteVulnerability 删除漏洞
|
||||
func (h *VulnerabilityHandler) DeleteVulnerability(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
if err := h.db.DeleteVulnerability(id); err != nil {
|
||||
h.logger.Error("删除漏洞失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
// GetVulnerabilityStats 获取漏洞统计
|
||||
func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) {
|
||||
conversationID := c.Query("conversation_id")
|
||||
|
||||
stats, err := h.db.GetVulnerabilityStats(conversationID)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞统计失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, stats)
|
||||
}
|
||||
|
||||
@@ -1,706 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// WebShellHandler 代理执行 WebShell 命令（类似冰蝎/蚁剑），避免前端跨域并统一构建请求
|
||||
type WebShellHandler struct {
|
||||
logger *zap.Logger
|
||||
client *http.Client
|
||||
db *database.DB
|
||||
}
|
||||
|
||||
// NewWebShellHandler 创建 WebShell 处理器，db 可为 nil（连接配置接口将不可用）
|
||||
func NewWebShellHandler(logger *zap.Logger, db *database.DB) *WebShellHandler {
|
||||
return &WebShellHandler{
|
||||
logger: logger,
|
||||
client: &http.Client{
|
||||
Timeout: 30 * time.Second,
|
||||
Transport: &http.Transport{DisableKeepAlives: false},
|
||||
},
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateConnectionRequest 创建连接请求
|
||||
type CreateConnectionRequest struct {
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmd_param"`
|
||||
Remark string `json:"remark"`
|
||||
}
|
||||
|
||||
// UpdateConnectionRequest 更新连接请求
|
||||
type UpdateConnectionRequest struct {
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmd_param"`
|
||||
Remark string `json:"remark"`
|
||||
}
|
||||
|
||||
// ListConnections 列出所有 WebShell 连接（GET /api/webshell/connections）
|
||||
func (h *WebShellHandler) ListConnections(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
list, err := h.db.ListWebshellConnections()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if list == nil {
|
||||
list = []database.WebShellConnection{}
|
||||
}
|
||||
c.JSON(http.StatusOK, list)
|
||||
}
|
||||
|
||||
// CreateConnection 创建 WebShell 连接（POST /api/webshell/connections）
|
||||
func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
var req CreateConnectionRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
req.URL = strings.TrimSpace(req.URL)
|
||||
if req.URL == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url is required"})
|
||||
return
|
||||
}
|
||||
if _, err := url.Parse(req.URL); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||
return
|
||||
}
|
||||
method := strings.ToLower(strings.TrimSpace(req.Method))
|
||||
if method != "get" && method != "post" {
|
||||
method = "post"
|
||||
}
|
||||
shellType := strings.ToLower(strings.TrimSpace(req.Type))
|
||||
if shellType == "" {
|
||||
shellType = "php"
|
||||
}
|
||||
conn := &database.WebShellConnection{
|
||||
ID: "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12],
|
||||
URL: req.URL,
|
||||
Password: strings.TrimSpace(req.Password),
|
||||
Type: shellType,
|
||||
Method: method,
|
||||
CmdParam: strings.TrimSpace(req.CmdParam),
|
||||
Remark: strings.TrimSpace(req.Remark),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := h.db.CreateWebshellConnection(conn); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, conn)
|
||||
}
|
||||
|
||||
// UpdateConnection 更新 WebShell 连接（PUT /api/webshell/connections/:id）
|
||||
func (h *WebShellHandler) UpdateConnection(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if id == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"})
|
||||
return
|
||||
}
|
||||
var req UpdateConnectionRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
req.URL = strings.TrimSpace(req.URL)
|
||||
if req.URL == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url is required"})
|
||||
return
|
||||
}
|
||||
if _, err := url.Parse(req.URL); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||
return
|
||||
}
|
||||
method := strings.ToLower(strings.TrimSpace(req.Method))
|
||||
if method != "get" && method != "post" {
|
||||
method = "post"
|
||||
}
|
||||
shellType := strings.ToLower(strings.TrimSpace(req.Type))
|
||||
if shellType == "" {
|
||||
shellType = "php"
|
||||
}
|
||||
conn := &database.WebShellConnection{
|
||||
ID: id,
|
||||
URL: req.URL,
|
||||
Password: strings.TrimSpace(req.Password),
|
||||
Type: shellType,
|
||||
Method: method,
|
||||
CmdParam: strings.TrimSpace(req.CmdParam),
|
||||
Remark: strings.TrimSpace(req.Remark),
|
||||
}
|
||||
if err := h.db.UpdateWebshellConnection(conn); err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "connection not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
updated, _ := h.db.GetWebshellConnection(id)
|
||||
if updated != nil {
|
||||
c.JSON(http.StatusOK, updated)
|
||||
} else {
|
||||
c.JSON(http.StatusOK, conn)
|
||||
}
|
||||
}
|
||||
|
||||
// DeleteConnection 删除 WebShell 连接（DELETE /api/webshell/connections/:id）
|
||||
func (h *WebShellHandler) DeleteConnection(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if id == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"})
|
||||
return
|
||||
}
|
||||
if err := h.db.DeleteWebshellConnection(id); err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "connection not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
// GetConnectionState 获取 WebShell 连接关联的前端持久化状态（GET /api/webshell/connections/:id/state）
|
||||
func (h *WebShellHandler) GetConnectionState(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if id == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"})
|
||||
return
|
||||
}
|
||||
conn, err := h.db.GetWebshellConnection(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if conn == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "connection not found"})
|
||||
return
|
||||
}
|
||||
stateJSON, err := h.db.GetWebshellConnectionState(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
var state interface{}
|
||||
if err := json.Unmarshal([]byte(stateJSON), &state); err != nil {
|
||||
state = map[string]interface{}{}
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"state": state})
|
||||
}
|
||||
|
||||
// SaveConnectionState 保存 WebShell 连接关联的前端持久化状态（PUT /api/webshell/connections/:id/state）
|
||||
func (h *WebShellHandler) SaveConnectionState(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if id == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"})
|
||||
return
|
||||
}
|
||||
conn, err := h.db.GetWebshellConnection(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if conn == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "connection not found"})
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
State json.RawMessage `json:"state"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
raw := req.State
|
||||
if len(raw) == 0 {
|
||||
raw = json.RawMessage(`{}`)
|
||||
}
|
||||
if len(raw) > 2*1024*1024 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "state payload too large (max 2MB)"})
|
||||
return
|
||||
}
|
||||
var anyJSON interface{}
|
||||
if err := json.Unmarshal(raw, &anyJSON); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "state must be valid json"})
|
||||
return
|
||||
}
|
||||
if err := h.db.UpsertWebshellConnectionState(id, string(raw)); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
// GetAIHistory 获取指定 WebShell 连接的 AI 助手对话历史（GET /api/webshell/connections/:id/ai-history）
|
||||
func (h *WebShellHandler) GetAIHistory(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if id == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"})
|
||||
return
|
||||
}
|
||||
conv, err := h.db.GetConversationByWebshellConnectionID(id)
|
||||
if err != nil {
|
||||
h.logger.Warn("获取 WebShell AI 对话失败", zap.String("connectionId", id), zap.Error(err))
|
||||
c.JSON(http.StatusOK, gin.H{"conversationId": nil, "messages": []database.Message{}})
|
||||
return
|
||||
}
|
||||
if conv == nil {
|
||||
c.JSON(http.StatusOK, gin.H{"conversationId": nil, "messages": []database.Message{}})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"conversationId": conv.ID, "messages": conv.Messages})
|
||||
}
|
||||
|
||||
// ListAIConversations 列出该 WebShell 连接下的所有 AI 对话（供侧边栏）
|
||||
func (h *WebShellHandler) ListAIConversations(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if id == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"})
|
||||
return
|
||||
}
|
||||
list, err := h.db.ListConversationsByWebshellConnectionID(id)
|
||||
if err != nil {
|
||||
h.logger.Warn("列出 WebShell AI 对话失败", zap.String("connectionId", id), zap.Error(err))
|
||||
c.JSON(http.StatusOK, []database.WebShellConversationItem{})
|
||||
return
|
||||
}
|
||||
if list == nil {
|
||||
list = []database.WebShellConversationItem{}
|
||||
}
|
||||
c.JSON(http.StatusOK, list)
|
||||
}
|
||||
|
||||
// ExecRequest 执行命令请求（前端传入连接信息 + 命令）
|
||||
type ExecRequest struct {
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"` // php, asp, aspx, jsp, custom
|
||||
Method string `json:"method"` // GET 或 POST，空则默认 POST
|
||||
CmdParam string `json:"cmd_param"` // 命令参数名，如 cmd/xxx，空则默认 cmd
|
||||
Command string `json:"command" binding:"required"`
|
||||
}
|
||||
|
||||
// ExecResponse 执行命令响应
|
||||
type ExecResponse struct {
|
||||
OK bool `json:"ok"`
|
||||
Output string `json:"output"`
|
||||
Error string `json:"error,omitempty"`
|
||||
HTTPCode int `json:"http_code,omitempty"`
|
||||
}
|
||||
|
||||
// FileOpRequest 文件操作请求
|
||||
type FileOpRequest struct {
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"` // GET 或 POST，空则默认 POST
|
||||
CmdParam string `json:"cmd_param"` // 命令参数名，如 cmd/xxx，空则默认 cmd
|
||||
Action string `json:"action" binding:"required"` // list, read, delete, write, mkdir, rename, upload, upload_chunk
|
||||
Path string `json:"path"`
|
||||
TargetPath string `json:"target_path"` // rename 时目标路径
|
||||
Content string `json:"content"` // write/upload 时使用
|
||||
ChunkIndex int `json:"chunk_index"` // upload_chunk 时，0 表示首块
|
||||
}
|
||||
|
||||
// FileOpResponse 文件操作响应
|
||||
type FileOpResponse struct {
|
||||
OK bool `json:"ok"`
|
||||
Output string `json:"output"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) Exec(c *gin.Context) {
|
||||
var req ExecRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
req.URL = strings.TrimSpace(req.URL)
|
||||
req.Command = strings.TrimSpace(req.Command)
|
||||
if req.URL == "" || req.Command == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"})
|
||||
return
|
||||
}
|
||||
|
||||
parsed, err := url.Parse(req.URL)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url: only http(s) allowed"})
|
||||
return
|
||||
}
|
||||
|
||||
useGET := strings.ToUpper(strings.TrimSpace(req.Method)) == "GET"
|
||||
cmdParam := strings.TrimSpace(req.CmdParam)
|
||||
if cmdParam == "" {
|
||||
cmdParam = "cmd"
|
||||
}
|
||||
var httpReq *http.Request
|
||||
if useGET {
|
||||
targetURL := h.buildExecURL(req.URL, req.Type, req.Password, cmdParam, req.Command)
|
||||
httpReq, err = http.NewRequest(http.MethodGet, targetURL, nil)
|
||||
} else {
|
||||
body := h.buildExecBody(req.Type, req.Password, cmdParam, req.Command)
|
||||
httpReq, err = http.NewRequest(http.MethodPost, req.URL, bytes.NewReader(body))
|
||||
httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
}
|
||||
if err != nil {
|
||||
h.logger.Warn("webshell exec NewRequest", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, ExecResponse{OK: false, Error: err.Error()})
|
||||
return
|
||||
}
|
||||
httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)")
|
||||
|
||||
resp, err := h.client.Do(httpReq)
|
||||
if err != nil {
|
||||
h.logger.Warn("webshell exec Do", zap.String("url", req.URL), zap.Error(err))
|
||||
c.JSON(http.StatusOK, ExecResponse{OK: false, Error: err.Error()})
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
out, _ := io.ReadAll(resp.Body)
|
||||
output := string(out)
|
||||
httpCode := resp.StatusCode
|
||||
|
||||
c.JSON(http.StatusOK, ExecResponse{
|
||||
OK: resp.StatusCode == http.StatusOK,
|
||||
Output: output,
|
||||
HTTPCode: httpCode,
|
||||
})
|
||||
}
|
||||
|
||||
// buildExecBody 按常见 WebShell 约定构建 POST 体（多数使用 pass + cmd，可配置命令参数名）
|
||||
func (h *WebShellHandler) buildExecBody(shellType, password, cmdParam, command string) []byte {
|
||||
form := h.execParams(shellType, password, cmdParam, command)
|
||||
return []byte(form.Encode())
|
||||
}
|
||||
|
||||
// buildExecURL 构建 GET 请求的完整 URL（baseURL + ?pass=xxx&cmd=yyy，cmd 可配置）
|
||||
func (h *WebShellHandler) buildExecURL(baseURL, shellType, password, cmdParam, command string) string {
|
||||
form := h.execParams(shellType, password, cmdParam, command)
|
||||
if parsed, err := url.Parse(baseURL); err == nil {
|
||||
parsed.RawQuery = form.Encode()
|
||||
return parsed.String()
|
||||
}
|
||||
return baseURL + "?" + form.Encode()
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) execParams(shellType, password, cmdParam, command string) url.Values {
|
||||
shellType = strings.ToLower(strings.TrimSpace(shellType))
|
||||
if shellType == "" {
|
||||
shellType = "php"
|
||||
}
|
||||
if strings.TrimSpace(cmdParam) == "" {
|
||||
cmdParam = "cmd"
|
||||
}
|
||||
form := url.Values{}
|
||||
form.Set("pass", password)
|
||||
form.Set(cmdParam, command)
|
||||
return form
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) FileOp(c *gin.Context) {
|
||||
var req FileOpRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
req.URL = strings.TrimSpace(req.URL)
|
||||
req.Action = strings.ToLower(strings.TrimSpace(req.Action))
|
||||
if req.URL == "" || req.Action == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"})
|
||||
return
|
||||
}
|
||||
|
||||
parsed, err := url.Parse(req.URL)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url: only http(s) allowed"})
|
||||
return
|
||||
}
|
||||
|
||||
// 通过执行系统命令实现文件操作（与通用一句话兼容）
|
||||
var command string
|
||||
shellType := strings.ToLower(strings.TrimSpace(req.Type))
|
||||
switch req.Action {
|
||||
case "list":
|
||||
path := strings.TrimSpace(req.Path)
|
||||
if path == "" {
|
||||
path = "."
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "dir " + h.escapePath(path)
|
||||
} else {
|
||||
command = "ls -la " + h.escapePath(path)
|
||||
}
|
||||
case "read":
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "type " + h.escapePath(strings.TrimSpace(req.Path))
|
||||
} else {
|
||||
command = "cat " + h.escapePath(strings.TrimSpace(req.Path))
|
||||
}
|
||||
case "delete":
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "del " + h.escapePath(strings.TrimSpace(req.Path))
|
||||
} else {
|
||||
command = "rm -f " + h.escapePath(strings.TrimSpace(req.Path))
|
||||
}
|
||||
case "write":
|
||||
path := h.escapePath(strings.TrimSpace(req.Path))
|
||||
command = "echo " + h.escapeForEcho(req.Content) + " > " + path
|
||||
case "mkdir":
|
||||
path := strings.TrimSpace(req.Path)
|
||||
if path == "" {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "path is required for mkdir"})
|
||||
return
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "md " + h.escapePath(path)
|
||||
} else {
|
||||
command = "mkdir -p " + h.escapePath(path)
|
||||
}
|
||||
case "rename":
|
||||
oldPath := strings.TrimSpace(req.Path)
|
||||
newPath := strings.TrimSpace(req.TargetPath)
|
||||
if oldPath == "" || newPath == "" {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "path and target_path are required for rename"})
|
||||
return
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "move /y " + h.escapePath(oldPath) + " " + h.escapePath(newPath)
|
||||
} else {
|
||||
command = "mv " + h.escapePath(oldPath) + " " + h.escapePath(newPath)
|
||||
}
|
||||
case "upload":
|
||||
path := strings.TrimSpace(req.Path)
|
||||
if path == "" {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "path is required for upload"})
|
||||
return
|
||||
}
|
||||
if len(req.Content) > 512*1024 {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "upload content too large (max 512KB base64)"})
|
||||
return
|
||||
}
|
||||
// base64 仅含 A-Za-z0-9+/=，用单引号包裹安全
|
||||
command = "echo " + "'" + req.Content + "'" + " | base64 -d > " + h.escapePath(path)
|
||||
case "upload_chunk":
|
||||
path := strings.TrimSpace(req.Path)
|
||||
if path == "" {
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "path is required for upload_chunk"})
|
||||
return
|
||||
}
|
||||
redir := ">>"
|
||||
if req.ChunkIndex == 0 {
|
||||
redir = ">"
|
||||
}
|
||||
command = "echo " + "'" + req.Content + "'" + " | base64 -d " + redir + " " + h.escapePath(path)
|
||||
default:
|
||||
c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: "unsupported action: " + req.Action})
|
||||
return
|
||||
}
|
||||
|
||||
useGET := strings.ToUpper(strings.TrimSpace(req.Method)) == "GET"
|
||||
cmdParam := strings.TrimSpace(req.CmdParam)
|
||||
if cmdParam == "" {
|
||||
cmdParam = "cmd"
|
||||
}
|
||||
var httpReq *http.Request
|
||||
if useGET {
|
||||
targetURL := h.buildExecURL(req.URL, req.Type, req.Password, cmdParam, command)
|
||||
httpReq, err = http.NewRequest(http.MethodGet, targetURL, nil)
|
||||
} else {
|
||||
body := h.buildExecBody(req.Type, req.Password, cmdParam, command)
|
||||
httpReq, err = http.NewRequest(http.MethodPost, req.URL, bytes.NewReader(body))
|
||||
httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
}
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, FileOpResponse{OK: false, Error: err.Error()})
|
||||
return
|
||||
}
|
||||
httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)")
|
||||
|
||||
resp, err := h.client.Do(httpReq)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, FileOpResponse{OK: false, Error: err.Error()})
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
out, _ := io.ReadAll(resp.Body)
|
||||
output := string(out)
|
||||
|
||||
c.JSON(http.StatusOK, FileOpResponse{
|
||||
OK: resp.StatusCode == http.StatusOK,
|
||||
Output: output,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) escapePath(p string) string {
|
||||
if p == "" {
|
||||
return "."
|
||||
}
|
||||
// 简单转义空格与敏感字符，避免命令注入
|
||||
return "'" + strings.ReplaceAll(p, "'", "'\\''") + "'"
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) escapeForEcho(s string) string {
|
||||
// 仅用于 write：base64 写入更安全，这里简单用单引号包裹
|
||||
return "'" + strings.ReplaceAll(s, "'", "'\"'\"'") + "'"
|
||||
}
|
||||
|
||||
// ExecWithConnection 在指定 WebShell 连接上执行命令（供 MCP/Agent 等非 HTTP 调用）
|
||||
func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) {
|
||||
if conn == nil {
|
||||
return "", false, "connection is nil"
|
||||
}
|
||||
command = strings.TrimSpace(command)
|
||||
if command == "" {
|
||||
return "", false, "command is required"
|
||||
}
|
||||
useGET := strings.ToUpper(strings.TrimSpace(conn.Method)) == "GET"
|
||||
cmdParam := strings.TrimSpace(conn.CmdParam)
|
||||
if cmdParam == "" {
|
||||
cmdParam = "cmd"
|
||||
}
|
||||
var httpReq *http.Request
|
||||
var err error
|
||||
if useGET {
|
||||
targetURL := h.buildExecURL(conn.URL, conn.Type, conn.Password, cmdParam, command)
|
||||
httpReq, err = http.NewRequest(http.MethodGet, targetURL, nil)
|
||||
} else {
|
||||
body := h.buildExecBody(conn.Type, conn.Password, cmdParam, command)
|
||||
httpReq, err = http.NewRequest(http.MethodPost, conn.URL, bytes.NewReader(body))
|
||||
httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
}
|
||||
if err != nil {
|
||||
return "", false, err.Error()
|
||||
}
|
||||
httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)")
|
||||
resp, err := h.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", false, err.Error()
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
out, _ := io.ReadAll(resp.Body)
|
||||
return string(out), resp.StatusCode == http.StatusOK, ""
|
||||
}
|
||||
|
||||
// FileOpWithConnection 在指定 WebShell 连接上执行文件操作（供 MCP/Agent 调用），支持 list / read / write
|
||||
func (h *WebShellHandler) FileOpWithConnection(conn *database.WebShellConnection, action, path, content, targetPath string) (output string, ok bool, errMsg string) {
|
||||
if conn == nil {
|
||||
return "", false, "connection is nil"
|
||||
}
|
||||
action = strings.ToLower(strings.TrimSpace(action))
|
||||
shellType := strings.ToLower(strings.TrimSpace(conn.Type))
|
||||
if shellType == "" {
|
||||
shellType = "php"
|
||||
}
|
||||
var command string
|
||||
switch action {
|
||||
case "list":
|
||||
if path == "" {
|
||||
path = "."
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "dir " + h.escapePath(strings.TrimSpace(path))
|
||||
} else {
|
||||
command = "ls -la " + h.escapePath(strings.TrimSpace(path))
|
||||
}
|
||||
case "read":
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return "", false, "path is required for read"
|
||||
}
|
||||
if shellType == "asp" || shellType == "aspx" {
|
||||
command = "type " + h.escapePath(path)
|
||||
} else {
|
||||
command = "cat " + h.escapePath(path)
|
||||
}
|
||||
case "write":
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return "", false, "path is required for write"
|
||||
}
|
||||
command = "echo " + h.escapeForEcho(content) + " > " + h.escapePath(path)
|
||||
default:
|
||||
return "", false, "unsupported action: " + action + " (supported: list, read, write)"
|
||||
}
|
||||
useGET := strings.ToUpper(strings.TrimSpace(conn.Method)) == "GET"
|
||||
cmdParam := strings.TrimSpace(conn.CmdParam)
|
||||
if cmdParam == "" {
|
||||
cmdParam = "cmd"
|
||||
}
|
||||
var httpReq *http.Request
|
||||
var err error
|
||||
if useGET {
|
||||
targetURL := h.buildExecURL(conn.URL, conn.Type, conn.Password, cmdParam, command)
|
||||
httpReq, err = http.NewRequest(http.MethodGet, targetURL, nil)
|
||||
} else {
|
||||
body := h.buildExecBody(conn.Type, conn.Password, cmdParam, command)
|
||||
httpReq, err = http.NewRequest(http.MethodPost, conn.URL, bytes.NewReader(body))
|
||||
httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
}
|
||||
if err != nil {
|
||||
return "", false, err.Error()
|
||||
}
|
||||
httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)")
|
||||
resp, err := h.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", false, err.Error()
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
out, _ := io.ReadAll(resp.Body)
|
||||
return string(out), resp.StatusCode == http.StatusOK, ""
|
||||
}
|
||||
Reference in New Issue
Block a user