Files
CyberStrikeAI/internal/handler/config.go
2026-03-23 02:15:46 +08:00

1409 lines
46 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package handler
import (
"bytes"
"context"
"fmt"
"net/http"
"os"
"path/filepath"
"strconv"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/agents"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/knowledge"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
"gopkg.in/yaml.v3"
)
// KnowledgeToolRegistrar 知识库工具注册器接口
type KnowledgeToolRegistrar func() error
// VulnerabilityToolRegistrar 漏洞工具注册器接口
type VulnerabilityToolRegistrar func() error
// WebshellToolRegistrar WebShell 工具注册器接口ApplyConfig 时重新注册)
type WebshellToolRegistrar func() error
// SkillsToolRegistrar Skills工具注册器接口
type SkillsToolRegistrar func() error
// RetrieverUpdater 检索器更新接口
type RetrieverUpdater interface {
UpdateConfig(config *knowledge.RetrievalConfig)
}
// KnowledgeInitializer 知识库初始化器接口
type KnowledgeInitializer func() (*KnowledgeHandler, error)
// AppUpdater App更新接口用于更新App中的知识库组件
type AppUpdater interface {
UpdateKnowledgeComponents(handler *KnowledgeHandler, manager interface{}, retriever interface{}, indexer interface{})
}
// RobotRestarter 机器人连接重启器(用于配置应用后重启钉钉/飞书长连接)
type RobotRestarter interface {
RestartRobotConnections()
}
// ConfigHandler 配置处理器
type ConfigHandler struct {
configPath string
config *config.Config
mcpServer *mcp.Server
executor *security.Executor
agent AgentUpdater // Agent接口用于更新Agent配置
attackChainHandler AttackChainUpdater // 攻击链处理器接口,用于更新配置
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
knowledgeToolRegistrar KnowledgeToolRegistrar // 知识库工具注册器(可选)
vulnerabilityToolRegistrar VulnerabilityToolRegistrar // 漏洞工具注册器(可选)
webshellToolRegistrar WebshellToolRegistrar // WebShell 工具注册器(可选)
skillsToolRegistrar SkillsToolRegistrar // Skills工具注册器可选
retrieverUpdater RetrieverUpdater // 检索器更新器(可选)
knowledgeInitializer KnowledgeInitializer // 知识库初始化器(可选)
appUpdater AppUpdater // App更新器可选
robotRestarter RobotRestarter // 机器人连接重启器可选ApplyConfig 时重启钉钉/飞书
logger *zap.Logger
mu sync.RWMutex
lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置(用于检测变更)
}
// AttackChainUpdater 攻击链处理器更新接口
type AttackChainUpdater interface {
UpdateConfig(cfg *config.OpenAIConfig)
}
// AgentUpdater Agent更新接口
type AgentUpdater interface {
UpdateConfig(cfg *config.OpenAIConfig)
UpdateMaxIterations(maxIterations int)
}
// NewConfigHandler 创建新的配置处理器
func NewConfigHandler(configPath string, cfg *config.Config, mcpServer *mcp.Server, executor *security.Executor, agent AgentUpdater, attackChainHandler AttackChainUpdater, externalMCPMgr *mcp.ExternalMCPManager, logger *zap.Logger) *ConfigHandler {
// 保存初始的嵌入模型配置(如果知识库已启用)
var lastEmbeddingConfig *config.EmbeddingConfig
if cfg.Knowledge.Enabled {
lastEmbeddingConfig = &config.EmbeddingConfig{
Provider: cfg.Knowledge.Embedding.Provider,
Model: cfg.Knowledge.Embedding.Model,
BaseURL: cfg.Knowledge.Embedding.BaseURL,
APIKey: cfg.Knowledge.Embedding.APIKey,
}
}
return &ConfigHandler{
configPath: configPath,
config: cfg,
mcpServer: mcpServer,
executor: executor,
agent: agent,
attackChainHandler: attackChainHandler,
externalMCPMgr: externalMCPMgr,
logger: logger,
lastEmbeddingConfig: lastEmbeddingConfig,
}
}
// SetKnowledgeToolRegistrar 设置知识库工具注册器
func (h *ConfigHandler) SetKnowledgeToolRegistrar(registrar KnowledgeToolRegistrar) {
h.mu.Lock()
defer h.mu.Unlock()
h.knowledgeToolRegistrar = registrar
}
// SetVulnerabilityToolRegistrar 设置漏洞工具注册器
func (h *ConfigHandler) SetVulnerabilityToolRegistrar(registrar VulnerabilityToolRegistrar) {
h.mu.Lock()
defer h.mu.Unlock()
h.vulnerabilityToolRegistrar = registrar
}
// SetWebshellToolRegistrar 设置 WebShell 工具注册器
func (h *ConfigHandler) SetWebshellToolRegistrar(registrar WebshellToolRegistrar) {
h.mu.Lock()
defer h.mu.Unlock()
h.webshellToolRegistrar = registrar
}
// SetSkillsToolRegistrar 设置Skills工具注册器
func (h *ConfigHandler) SetSkillsToolRegistrar(registrar SkillsToolRegistrar) {
h.mu.Lock()
defer h.mu.Unlock()
h.skillsToolRegistrar = registrar
}
// SetRetrieverUpdater 设置检索器更新器
func (h *ConfigHandler) SetRetrieverUpdater(updater RetrieverUpdater) {
h.mu.Lock()
defer h.mu.Unlock()
h.retrieverUpdater = updater
}
// SetKnowledgeInitializer 设置知识库初始化器
func (h *ConfigHandler) SetKnowledgeInitializer(initializer KnowledgeInitializer) {
h.mu.Lock()
defer h.mu.Unlock()
h.knowledgeInitializer = initializer
}
// SetAppUpdater 设置App更新器
func (h *ConfigHandler) SetAppUpdater(updater AppUpdater) {
h.mu.Lock()
defer h.mu.Unlock()
h.appUpdater = updater
}
// SetRobotRestarter 设置机器人连接重启器ApplyConfig 时用于重启钉钉/飞书长连接)
func (h *ConfigHandler) SetRobotRestarter(restarter RobotRestarter) {
h.mu.Lock()
defer h.mu.Unlock()
h.robotRestarter = restarter
}
// GetConfigResponse 获取配置响应
type GetConfigResponse struct {
OpenAI config.OpenAIConfig `json:"openai"`
FOFA config.FofaConfig `json:"fofa"`
MCP config.MCPConfig `json:"mcp"`
Tools []ToolConfigInfo `json:"tools"`
Agent config.AgentConfig `json:"agent"`
Knowledge config.KnowledgeConfig `json:"knowledge"`
Robots config.RobotsConfig `json:"robots,omitempty"`
MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"`
}
// ToolConfigInfo 工具配置信息
type ToolConfigInfo struct {
Name string `json:"name"`
Description string `json:"description"`
Enabled bool `json:"enabled"`
IsExternal bool `json:"is_external,omitempty"` // 是否为外部MCP工具
ExternalMCP string `json:"external_mcp,omitempty"` // 外部MCP名称如果是外部工具
RoleEnabled *bool `json:"role_enabled,omitempty"` // 该工具在当前角色中是否启用nil表示未指定角色或使用所有工具
}
// GetConfig 获取当前配置
func (h *ConfigHandler) GetConfig(c *gin.Context) {
h.mu.RLock()
defer h.mu.RUnlock()
// 获取工具列表(包含内部和外部工具)
// 首先从配置文件获取工具
configToolMap := make(map[string]bool)
tools := make([]ToolConfigInfo, 0, len(h.config.Security.Tools))
for _, tool := range h.config.Security.Tools {
configToolMap[tool.Name] = true
tools = append(tools, ToolConfigInfo{
Name: tool.Name,
Description: h.pickToolDescription(tool.ShortDescription, tool.Description),
Enabled: tool.Enabled,
IsExternal: false,
})
}
// 从MCP服务器获取所有已注册的工具包括直接注册的工具如知识检索工具
if h.mcpServer != nil {
mcpTools := h.mcpServer.GetAllTools()
for _, mcpTool := range mcpTools {
// 跳过已经在配置文件中的工具(避免重复)
if configToolMap[mcpTool.Name] {
continue
}
// 添加直接注册到MCP服务器的工具如知识检索工具
description := mcpTool.ShortDescription
if description == "" {
description = mcpTool.Description
}
if len(description) > 10000 {
description = description[:10000] + "..."
}
tools = append(tools, ToolConfigInfo{
Name: mcpTool.Name,
Description: description,
Enabled: true, // 直接注册的工具默认启用
IsExternal: false,
})
}
}
// 获取外部MCP工具
if h.externalMCPMgr != nil {
ctx := context.Background()
externalTools := h.getExternalMCPTools(ctx)
for _, toolInfo := range externalTools {
tools = append(tools, toolInfo)
}
}
subAgentCount := len(h.config.MultiAgent.SubAgents)
agentsDir := strings.TrimSpace(h.config.AgentsDir)
if agentsDir == "" {
agentsDir = "agents"
}
if !filepath.IsAbs(agentsDir) {
agentsDir = filepath.Join(filepath.Dir(h.configPath), agentsDir)
}
if load, err := agents.LoadMarkdownAgentsDir(agentsDir); err == nil {
subAgentCount = len(agents.MergeYAMLAndMarkdown(h.config.MultiAgent.SubAgents, load.SubAgents))
}
multiPub := config.MultiAgentPublic{
Enabled: h.config.MultiAgent.Enabled,
DefaultMode: h.config.MultiAgent.DefaultMode,
RobotUseMultiAgent: h.config.MultiAgent.RobotUseMultiAgent,
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
SubAgentCount: subAgentCount,
}
if strings.TrimSpace(multiPub.DefaultMode) == "" {
multiPub.DefaultMode = "single"
}
c.JSON(http.StatusOK, GetConfigResponse{
OpenAI: h.config.OpenAI,
FOFA: h.config.FOFA,
MCP: h.config.MCP,
Tools: tools,
Agent: h.config.Agent,
Knowledge: h.config.Knowledge,
Robots: h.config.Robots,
MultiAgent: multiPub,
})
}
// GetToolsResponse 获取工具列表响应(分页)
type GetToolsResponse struct {
Tools []ToolConfigInfo `json:"tools"`
Total int `json:"total"`
TotalEnabled int `json:"total_enabled"` // 已启用的工具总数
Page int `json:"page"`
PageSize int `json:"page_size"`
TotalPages int `json:"total_pages"`
}
// GetTools 获取工具列表(支持分页和搜索)
func (h *ConfigHandler) GetTools(c *gin.Context) {
h.mu.RLock()
defer h.mu.RUnlock()
// 解析分页参数
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
}
}
// 解析搜索参数
searchTerm := c.Query("search")
searchTermLower := ""
if searchTerm != "" {
searchTermLower = strings.ToLower(searchTerm)
}
// 解析角色参数,用于过滤工具并标注启用状态
roleName := c.Query("role")
var roleToolsSet map[string]bool // 角色配置的工具集合
var roleUsesAllTools bool = true // 角色是否使用所有工具(默认角色)
if roleName != "" && roleName != "默认" && h.config.Roles != nil {
if role, exists := h.config.Roles[roleName]; exists && role.Enabled {
if len(role.Tools) > 0 {
// 角色配置了工具列表,只使用这些工具
roleToolsSet = make(map[string]bool)
for _, toolKey := range role.Tools {
roleToolsSet[toolKey] = true
}
roleUsesAllTools = false
}
}
}
// 获取所有内部工具并应用搜索过滤
configToolMap := make(map[string]bool)
allTools := make([]ToolConfigInfo, 0, len(h.config.Security.Tools))
for _, tool := range h.config.Security.Tools {
configToolMap[tool.Name] = true
toolInfo := ToolConfigInfo{
Name: tool.Name,
Description: h.pickToolDescription(tool.ShortDescription, tool.Description),
Enabled: tool.Enabled,
IsExternal: false,
}
// 根据角色配置标注工具状态
if roleName != "" {
if roleUsesAllTools {
// 角色使用所有工具标注启用的工具为role_enabled=true
if tool.Enabled {
roleEnabled := true
toolInfo.RoleEnabled = &roleEnabled
} else {
roleEnabled := false
toolInfo.RoleEnabled = &roleEnabled
}
} else {
// 角色配置了工具列表,检查工具是否在列表中
// 内部工具使用工具名称作为key
if roleToolsSet[tool.Name] {
roleEnabled := tool.Enabled // 工具必须在角色列表中且本身启用
toolInfo.RoleEnabled = &roleEnabled
} else {
// 不在角色列表中标记为false
roleEnabled := false
toolInfo.RoleEnabled = &roleEnabled
}
}
}
// 如果有关键词,进行搜索过滤
if searchTermLower != "" {
nameLower := strings.ToLower(toolInfo.Name)
descLower := strings.ToLower(toolInfo.Description)
if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) {
continue // 不匹配,跳过
}
}
allTools = append(allTools, toolInfo)
}
// 从MCP服务器获取所有已注册的工具包括直接注册的工具如知识检索工具
if h.mcpServer != nil {
mcpTools := h.mcpServer.GetAllTools()
for _, mcpTool := range mcpTools {
// 跳过已经在配置文件中的工具(避免重复)
if configToolMap[mcpTool.Name] {
continue
}
description := mcpTool.ShortDescription
if description == "" {
description = mcpTool.Description
}
if len(description) > 10000 {
description = description[:10000] + "..."
}
toolInfo := ToolConfigInfo{
Name: mcpTool.Name,
Description: description,
Enabled: true, // 直接注册的工具默认启用
IsExternal: false,
}
// 根据角色配置标注工具状态
if roleName != "" {
if roleUsesAllTools {
// 角色使用所有工具,直接注册的工具默认启用
roleEnabled := true
toolInfo.RoleEnabled = &roleEnabled
} else {
// 角色配置了工具列表,检查工具是否在列表中
// 内部工具使用工具名称作为key
if roleToolsSet[mcpTool.Name] {
roleEnabled := true // 在角色列表中且工具本身启用
toolInfo.RoleEnabled = &roleEnabled
} else {
// 不在角色列表中标记为false
roleEnabled := false
toolInfo.RoleEnabled = &roleEnabled
}
}
}
// 如果有关键词,进行搜索过滤
if searchTermLower != "" {
nameLower := strings.ToLower(toolInfo.Name)
descLower := strings.ToLower(toolInfo.Description)
if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) {
continue // 不匹配,跳过
}
}
allTools = append(allTools, toolInfo)
}
}
// 获取外部MCP工具
if h.externalMCPMgr != nil {
// 创建context用于获取外部工具
ctx := context.Background()
externalTools := h.getExternalMCPTools(ctx)
// 应用搜索过滤和角色配置
for _, toolInfo := range externalTools {
// 搜索过滤
if searchTermLower != "" {
nameLower := strings.ToLower(toolInfo.Name)
descLower := strings.ToLower(toolInfo.Description)
if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) {
continue // 不匹配,跳过
}
}
// 根据角色配置标注工具状态
if roleName != "" {
if roleUsesAllTools {
// 角色使用所有工具标注启用的工具为role_enabled=true
roleEnabled := toolInfo.Enabled
toolInfo.RoleEnabled = &roleEnabled
} else {
// 角色配置了工具列表,检查工具是否在列表中
// 外部工具使用 "mcpName::toolName" 格式作为key
externalToolKey := fmt.Sprintf("%s::%s", toolInfo.ExternalMCP, toolInfo.Name)
if roleToolsSet[externalToolKey] {
roleEnabled := toolInfo.Enabled // 工具必须在角色列表中且本身启用
toolInfo.RoleEnabled = &roleEnabled
} else {
// 不在角色列表中标记为false
roleEnabled := false
toolInfo.RoleEnabled = &roleEnabled
}
}
}
allTools = append(allTools, toolInfo)
}
}
// 如果角色配置了工具列表,过滤工具(只保留列表中的工具,但保留其他工具并标记为禁用)
// 注意:这里我们不直接过滤掉工具,而是保留所有工具,但通过 role_enabled 字段标注状态
// 这样前端可以显示所有工具,并标注哪些工具在当前角色中可用
total := len(allTools)
// 统计已启用的工具数(在角色中的启用工具数)
totalEnabled := 0
for _, tool := range allTools {
if tool.RoleEnabled != nil && *tool.RoleEnabled {
totalEnabled++
} else if tool.RoleEnabled == nil && tool.Enabled {
// 如果未指定角色,统计所有启用的工具
totalEnabled++
}
}
totalPages := (total + pageSize - 1) / pageSize
if totalPages == 0 {
totalPages = 1
}
// 计算分页范围
offset := (page - 1) * pageSize
end := offset + pageSize
if end > total {
end = total
}
var tools []ToolConfigInfo
if offset < total {
tools = allTools[offset:end]
} else {
tools = []ToolConfigInfo{}
}
c.JSON(http.StatusOK, GetToolsResponse{
Tools: tools,
Total: total,
TotalEnabled: totalEnabled,
Page: page,
PageSize: pageSize,
TotalPages: totalPages,
})
}
// UpdateConfigRequest 更新配置请求
type UpdateConfigRequest struct {
OpenAI *config.OpenAIConfig `json:"openai,omitempty"`
FOFA *config.FofaConfig `json:"fofa,omitempty"`
MCP *config.MCPConfig `json:"mcp,omitempty"`
Tools []ToolEnableStatus `json:"tools,omitempty"`
Agent *config.AgentConfig `json:"agent,omitempty"`
Knowledge *config.KnowledgeConfig `json:"knowledge,omitempty"`
Robots *config.RobotsConfig `json:"robots,omitempty"`
MultiAgent *config.MultiAgentAPIUpdate `json:"multi_agent,omitempty"`
}
// ToolEnableStatus 工具启用状态
type ToolEnableStatus struct {
Name string `json:"name"`
Enabled bool `json:"enabled"`
IsExternal bool `json:"is_external,omitempty"` // 是否为外部MCP工具
ExternalMCP string `json:"external_mcp,omitempty"` // 外部MCP名称如果是外部工具
}
// UpdateConfig 更新配置
func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
var req UpdateConfigRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
return
}
h.mu.Lock()
defer h.mu.Unlock()
// 更新OpenAI配置
if req.OpenAI != nil {
h.config.OpenAI = *req.OpenAI
h.logger.Info("更新OpenAI配置",
zap.String("base_url", h.config.OpenAI.BaseURL),
zap.String("model", h.config.OpenAI.Model),
)
}
// 更新FOFA配置
if req.FOFA != nil {
h.config.FOFA = *req.FOFA
h.logger.Info("更新FOFA配置", zap.String("email", h.config.FOFA.Email))
}
// 更新MCP配置
if req.MCP != nil {
h.config.MCP = *req.MCP
h.logger.Info("更新MCP配置",
zap.Bool("enabled", h.config.MCP.Enabled),
zap.String("host", h.config.MCP.Host),
zap.Int("port", h.config.MCP.Port),
)
}
// 更新Agent配置
if req.Agent != nil {
h.config.Agent = *req.Agent
h.logger.Info("更新Agent配置",
zap.Int("max_iterations", h.config.Agent.MaxIterations),
)
}
// 更新Knowledge配置
if req.Knowledge != nil {
// 保存旧的嵌入模型配置(用于检测变更)
if h.config.Knowledge.Enabled {
h.lastEmbeddingConfig = &config.EmbeddingConfig{
Provider: h.config.Knowledge.Embedding.Provider,
Model: h.config.Knowledge.Embedding.Model,
BaseURL: h.config.Knowledge.Embedding.BaseURL,
APIKey: h.config.Knowledge.Embedding.APIKey,
}
}
h.config.Knowledge = *req.Knowledge
h.logger.Info("更新Knowledge配置",
zap.Bool("enabled", h.config.Knowledge.Enabled),
zap.String("base_path", h.config.Knowledge.BasePath),
zap.String("embedding_model", h.config.Knowledge.Embedding.Model),
zap.Int("retrieval_top_k", h.config.Knowledge.Retrieval.TopK),
zap.Float64("similarity_threshold", h.config.Knowledge.Retrieval.SimilarityThreshold),
zap.Float64("hybrid_weight", h.config.Knowledge.Retrieval.HybridWeight),
)
}
// 更新机器人配置
if req.Robots != nil {
h.config.Robots = *req.Robots
h.logger.Info("更新机器人配置",
zap.Bool("wecom_enabled", h.config.Robots.Wecom.Enabled),
zap.Bool("dingtalk_enabled", h.config.Robots.Dingtalk.Enabled),
zap.Bool("lark_enabled", h.config.Robots.Lark.Enabled),
)
}
// 多代理标量sub_agents 等仍由 config.yaml 维护)
if req.MultiAgent != nil {
h.config.MultiAgent.Enabled = req.MultiAgent.Enabled
dm := strings.TrimSpace(req.MultiAgent.DefaultMode)
if dm == "multi" || dm == "single" {
h.config.MultiAgent.DefaultMode = dm
}
h.config.MultiAgent.RobotUseMultiAgent = req.MultiAgent.RobotUseMultiAgent
h.config.MultiAgent.BatchUseMultiAgent = req.MultiAgent.BatchUseMultiAgent
h.logger.Info("更新多代理配置",
zap.Bool("enabled", h.config.MultiAgent.Enabled),
zap.String("default_mode", h.config.MultiAgent.DefaultMode),
zap.Bool("robot_use_multi_agent", h.config.MultiAgent.RobotUseMultiAgent),
zap.Bool("batch_use_multi_agent", h.config.MultiAgent.BatchUseMultiAgent),
)
}
// 更新工具启用状态
if req.Tools != nil {
// 分离内部工具和外部工具
internalToolMap := make(map[string]bool)
// 外部工具状态MCP名称 -> 工具名称 -> 启用状态
externalMCPToolMap := make(map[string]map[string]bool)
for _, toolStatus := range req.Tools {
if toolStatus.IsExternal && toolStatus.ExternalMCP != "" {
// 外部工具:保存每个工具的独立状态
mcpName := toolStatus.ExternalMCP
if externalMCPToolMap[mcpName] == nil {
externalMCPToolMap[mcpName] = make(map[string]bool)
}
externalMCPToolMap[mcpName][toolStatus.Name] = toolStatus.Enabled
} else {
// 内部工具
internalToolMap[toolStatus.Name] = toolStatus.Enabled
}
}
// 更新内部工具状态
for i := range h.config.Security.Tools {
if enabled, ok := internalToolMap[h.config.Security.Tools[i].Name]; ok {
h.config.Security.Tools[i].Enabled = enabled
h.logger.Info("更新工具启用状态",
zap.String("tool", h.config.Security.Tools[i].Name),
zap.Bool("enabled", enabled),
)
}
}
// 更新外部MCP工具状态
if h.externalMCPMgr != nil {
for mcpName, toolStates := range externalMCPToolMap {
// 更新配置中的工具启用状态
if h.config.ExternalMCP.Servers == nil {
h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig)
}
cfg, exists := h.config.ExternalMCP.Servers[mcpName]
if !exists {
h.logger.Warn("外部MCP配置不存在", zap.String("mcp", mcpName))
continue
}
// 初始化ToolEnabled map
if cfg.ToolEnabled == nil {
cfg.ToolEnabled = make(map[string]bool)
}
// 更新每个工具的启用状态
for toolName, enabled := range toolStates {
cfg.ToolEnabled[toolName] = enabled
h.logger.Info("更新外部工具启用状态",
zap.String("mcp", mcpName),
zap.String("tool", toolName),
zap.Bool("enabled", enabled),
)
}
// 检查是否有任何工具启用如果有则启用MCP
hasEnabledTool := false
for _, enabled := range cfg.ToolEnabled {
if enabled {
hasEnabledTool = true
break
}
}
// 如果MCP之前未启用但现在有工具启用则启用MCP
// 如果MCP之前已启用保持启用状态允许部分工具禁用
if !cfg.ExternalMCPEnable && hasEnabledTool {
cfg.ExternalMCPEnable = true
h.logger.Info("自动启用外部MCP因为有工具启用", zap.String("mcp", mcpName))
}
h.config.ExternalMCP.Servers[mcpName] = cfg
}
// 同步更新 externalMCPMgr 中的配置,确保 GetConfigs() 返回最新配置
// 在循环外部统一更新,避免重复调用
h.externalMCPMgr.LoadConfigs(&h.config.ExternalMCP)
// 处理MCP连接状态异步启动避免阻塞
for mcpName := range externalMCPToolMap {
cfg := h.config.ExternalMCP.Servers[mcpName]
// 如果MCP需要启用确保客户端已启动
if cfg.ExternalMCPEnable {
// 启动外部MCP如果未启动- 异步执行,避免阻塞
client, exists := h.externalMCPMgr.GetClient(mcpName)
if !exists || !client.IsConnected() {
go func(name string) {
if err := h.externalMCPMgr.StartClient(name); err != nil {
h.logger.Warn("启动外部MCP失败",
zap.String("mcp", name),
zap.Error(err),
)
} else {
h.logger.Info("启动外部MCP",
zap.String("mcp", name),
)
}
}(mcpName)
}
}
}
}
}
// 保存配置到文件
if err := h.saveConfig(); 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": "配置已更新"})
}
// ApplyConfig 应用配置(重新加载并重启相关服务)
func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
// 先检查是否需要动态初始化知识库(在锁外执行,避免阻塞其他请求)
var needInitKnowledge bool
var knowledgeInitializer KnowledgeInitializer
h.mu.RLock()
needInitKnowledge = h.config.Knowledge.Enabled && h.knowledgeToolRegistrar == nil && h.knowledgeInitializer != nil
if needInitKnowledge {
knowledgeInitializer = h.knowledgeInitializer
}
h.mu.RUnlock()
// 如果需要动态初始化知识库,在锁外执行(这是耗时操作)
if needInitKnowledge {
h.logger.Info("检测到知识库从禁用变为启用,开始动态初始化知识库组件")
if _, err := knowledgeInitializer(); err != nil {
h.logger.Error("动态初始化知识库失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": "初始化知识库失败: " + err.Error()})
return
}
h.logger.Info("知识库动态初始化完成,工具已注册")
}
// 检查嵌入模型配置是否变更(需要在锁外执行,避免阻塞)
var needReinitKnowledge bool
var reinitKnowledgeInitializer KnowledgeInitializer
h.mu.RLock()
if h.config.Knowledge.Enabled && h.knowledgeInitializer != nil && h.lastEmbeddingConfig != nil {
// 检查嵌入模型配置是否变更
currentEmbedding := h.config.Knowledge.Embedding
if currentEmbedding.Provider != h.lastEmbeddingConfig.Provider ||
currentEmbedding.Model != h.lastEmbeddingConfig.Model ||
currentEmbedding.BaseURL != h.lastEmbeddingConfig.BaseURL ||
currentEmbedding.APIKey != h.lastEmbeddingConfig.APIKey {
needReinitKnowledge = true
reinitKnowledgeInitializer = h.knowledgeInitializer
h.logger.Info("检测到嵌入模型配置变更,需要重新初始化知识库组件",
zap.String("old_model", h.lastEmbeddingConfig.Model),
zap.String("new_model", currentEmbedding.Model),
zap.String("old_base_url", h.lastEmbeddingConfig.BaseURL),
zap.String("new_base_url", currentEmbedding.BaseURL),
)
}
}
h.mu.RUnlock()
// 如果需要重新初始化知识库(嵌入模型配置变更),在锁外执行
if needReinitKnowledge {
h.logger.Info("开始重新初始化知识库组件(嵌入模型配置已变更)")
if _, err := reinitKnowledgeInitializer(); err != nil {
h.logger.Error("重新初始化知识库失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": "重新初始化知识库失败: " + err.Error()})
return
}
h.logger.Info("知识库组件重新初始化完成")
}
// 现在获取写锁,执行快速的操作
h.mu.Lock()
defer h.mu.Unlock()
// 如果重新初始化了知识库,更新嵌入模型配置记录
if needReinitKnowledge && h.config.Knowledge.Enabled {
h.lastEmbeddingConfig = &config.EmbeddingConfig{
Provider: h.config.Knowledge.Embedding.Provider,
Model: h.config.Knowledge.Embedding.Model,
BaseURL: h.config.Knowledge.Embedding.BaseURL,
APIKey: h.config.Knowledge.Embedding.APIKey,
}
h.logger.Info("已更新嵌入模型配置记录")
}
// 重新注册工具(根据新的启用状态)
h.logger.Info("重新注册工具")
// 清空MCP服务器中的工具
h.mcpServer.ClearTools()
// 重新注册安全工具
h.executor.RegisterTools(h.mcpServer)
// 重新注册漏洞记录工具(内置工具,必须注册)
if h.vulnerabilityToolRegistrar != nil {
h.logger.Info("重新注册漏洞记录工具")
if err := h.vulnerabilityToolRegistrar(); err != nil {
h.logger.Error("重新注册漏洞记录工具失败", zap.Error(err))
} else {
h.logger.Info("漏洞记录工具已重新注册")
}
}
// 重新注册 WebShell 工具(内置工具,必须注册)
if h.webshellToolRegistrar != nil {
h.logger.Info("重新注册 WebShell 工具")
if err := h.webshellToolRegistrar(); err != nil {
h.logger.Error("重新注册 WebShell 工具失败", zap.Error(err))
} else {
h.logger.Info("WebShell 工具已重新注册")
}
}
// 重新注册Skills工具内置工具必须注册
if h.skillsToolRegistrar != nil {
h.logger.Info("重新注册Skills工具")
if err := h.skillsToolRegistrar(); err != nil {
h.logger.Error("重新注册Skills工具失败", zap.Error(err))
} else {
h.logger.Info("Skills工具已重新注册")
}
}
// 如果知识库启用,重新注册知识库工具
if h.config.Knowledge.Enabled && h.knowledgeToolRegistrar != nil {
h.logger.Info("重新注册知识库工具")
if err := h.knowledgeToolRegistrar(); err != nil {
h.logger.Error("重新注册知识库工具失败", zap.Error(err))
} else {
h.logger.Info("知识库工具已重新注册")
}
}
// 更新Agent的OpenAI配置
if h.agent != nil {
h.agent.UpdateConfig(&h.config.OpenAI)
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
h.logger.Info("Agent配置已更新")
}
// 更新AttackChainHandler的OpenAI配置
if h.attackChainHandler != nil {
h.attackChainHandler.UpdateConfig(&h.config.OpenAI)
h.logger.Info("AttackChainHandler配置已更新")
}
// 更新检索器配置(如果知识库启用)
if h.config.Knowledge.Enabled && h.retrieverUpdater != nil {
retrievalConfig := &knowledge.RetrievalConfig{
TopK: h.config.Knowledge.Retrieval.TopK,
SimilarityThreshold: h.config.Knowledge.Retrieval.SimilarityThreshold,
HybridWeight: h.config.Knowledge.Retrieval.HybridWeight,
}
h.retrieverUpdater.UpdateConfig(retrievalConfig)
h.logger.Info("检索器配置已更新",
zap.Int("top_k", retrievalConfig.TopK),
zap.Float64("similarity_threshold", retrievalConfig.SimilarityThreshold),
zap.Float64("hybrid_weight", retrievalConfig.HybridWeight),
)
}
// 更新嵌入模型配置记录(如果知识库启用)
if h.config.Knowledge.Enabled {
h.lastEmbeddingConfig = &config.EmbeddingConfig{
Provider: h.config.Knowledge.Embedding.Provider,
Model: h.config.Knowledge.Embedding.Model,
BaseURL: h.config.Knowledge.Embedding.BaseURL,
APIKey: h.config.Knowledge.Embedding.APIKey,
}
}
// 重启钉钉/飞书长连接,使前端修改的机器人配置立即生效(无需重启服务)
if h.robotRestarter != nil {
h.robotRestarter.RestartRobotConnections()
h.logger.Info("已触发机器人连接重启(钉钉/飞书)")
}
h.logger.Info("配置已应用",
zap.Int("tools_count", len(h.config.Security.Tools)),
)
c.JSON(http.StatusOK, gin.H{
"message": "配置已应用",
"tools_count": len(h.config.Security.Tools),
})
}
// saveConfig 保存配置到文件
func (h *ConfigHandler) 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)
}
updateAgentConfig(root, h.config.Agent.MaxIterations)
updateMCPConfig(root, h.config.MCP)
updateOpenAIConfig(root, h.config.OpenAI)
updateFOFAConfig(root, h.config.FOFA)
updateKnowledgeConfig(root, h.config.Knowledge)
updateRobotsConfig(root, h.config.Robots)
updateMultiAgentConfig(root, h.config.MultiAgent)
// 更新外部MCP配置使用external_mcp.go中的函数同一包中可直接调用
// 读取原始配置以保持向后兼容
originalConfigs := make(map[string]map[string]bool)
externalMCPNode := findMapValue(root, "external_mcp")
if externalMCPNode != nil && externalMCPNode.Kind == yaml.MappingNode {
serversNode := findMapValue(externalMCPNode, "servers")
if serversNode != nil && serversNode.Kind == yaml.MappingNode {
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)
if enabledVal := findBoolInMap(serverNode, "enabled"); enabledVal != nil {
originalConfigs[serverName]["enabled"] = *enabledVal
}
if disabledVal := findBoolInMap(serverNode, "disabled"); disabledVal != nil {
originalConfigs[serverName]["disabled"] = *disabledVal
}
}
}
}
}
updateExternalMCPConfig(root, h.config.ExternalMCP, originalConfigs)
if err := writeYAMLDocument(h.configPath, root); err != nil {
return fmt.Errorf("保存配置文件失败: %w", err)
}
// 更新工具配置文件中的enabled状态
if h.config.Security.ToolsDir != "" {
configDir := filepath.Dir(h.configPath)
toolsDir := h.config.Security.ToolsDir
if !filepath.IsAbs(toolsDir) {
toolsDir = filepath.Join(configDir, toolsDir)
}
for _, tool := range h.config.Security.Tools {
toolFile := filepath.Join(toolsDir, tool.Name+".yaml")
// 检查文件是否存在
if _, err := os.Stat(toolFile); os.IsNotExist(err) {
// 尝试.yml扩展名
toolFile = filepath.Join(toolsDir, tool.Name+".yml")
if _, err := os.Stat(toolFile); os.IsNotExist(err) {
h.logger.Warn("工具配置文件不存在", zap.String("tool", tool.Name))
continue
}
}
toolDoc, err := loadYAMLDocument(toolFile)
if err != nil {
h.logger.Warn("解析工具配置失败", zap.String("tool", tool.Name), zap.Error(err))
continue
}
setBoolInMap(toolDoc.Content[0], "enabled", tool.Enabled)
if err := writeYAMLDocument(toolFile, toolDoc); err != nil {
h.logger.Warn("保存工具配置文件失败", zap.String("tool", tool.Name), zap.Error(err))
continue
}
h.logger.Info("更新工具配置", zap.String("tool", tool.Name), zap.Bool("enabled", tool.Enabled))
}
}
h.logger.Info("配置已保存", zap.String("path", h.configPath))
return nil
}
func loadYAMLDocument(path string) (*yaml.Node, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, err
}
if len(bytes.TrimSpace(data)) == 0 {
return newEmptyYAMLDocument(), nil
}
var doc yaml.Node
if err := yaml.Unmarshal(data, &doc); err != nil {
return nil, err
}
if doc.Kind != yaml.DocumentNode || len(doc.Content) == 0 {
return newEmptyYAMLDocument(), nil
}
if doc.Content[0].Kind != yaml.MappingNode {
root := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
doc.Content = []*yaml.Node{root}
}
return &doc, nil
}
func newEmptyYAMLDocument() *yaml.Node {
root := &yaml.Node{
Kind: yaml.DocumentNode,
Content: []*yaml.Node{{Kind: yaml.MappingNode, Tag: "!!map"}},
}
return root
}
func writeYAMLDocument(path string, doc *yaml.Node) error {
var buf bytes.Buffer
encoder := yaml.NewEncoder(&buf)
encoder.SetIndent(2)
if err := encoder.Encode(doc); err != nil {
return err
}
if err := encoder.Close(); err != nil {
return err
}
return os.WriteFile(path, buf.Bytes(), 0644)
}
func updateAgentConfig(doc *yaml.Node, maxIterations int) {
root := doc.Content[0]
agentNode := ensureMap(root, "agent")
setIntInMap(agentNode, "max_iterations", maxIterations)
}
func updateMCPConfig(doc *yaml.Node, cfg config.MCPConfig) {
root := doc.Content[0]
mcpNode := ensureMap(root, "mcp")
setBoolInMap(mcpNode, "enabled", cfg.Enabled)
setStringInMap(mcpNode, "host", cfg.Host)
setIntInMap(mcpNode, "port", cfg.Port)
}
func updateOpenAIConfig(doc *yaml.Node, cfg config.OpenAIConfig) {
root := doc.Content[0]
openaiNode := ensureMap(root, "openai")
setStringInMap(openaiNode, "api_key", cfg.APIKey)
setStringInMap(openaiNode, "base_url", cfg.BaseURL)
setStringInMap(openaiNode, "model", cfg.Model)
}
func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) {
root := doc.Content[0]
fofaNode := ensureMap(root, "fofa")
setStringInMap(fofaNode, "base_url", cfg.BaseURL)
setStringInMap(fofaNode, "email", cfg.Email)
setStringInMap(fofaNode, "api_key", cfg.APIKey)
}
func updateKnowledgeConfig(doc *yaml.Node, cfg config.KnowledgeConfig) {
root := doc.Content[0]
knowledgeNode := ensureMap(root, "knowledge")
setBoolInMap(knowledgeNode, "enabled", cfg.Enabled)
setStringInMap(knowledgeNode, "base_path", cfg.BasePath)
// 更新嵌入配置
embeddingNode := ensureMap(knowledgeNode, "embedding")
setStringInMap(embeddingNode, "provider", cfg.Embedding.Provider)
setStringInMap(embeddingNode, "model", cfg.Embedding.Model)
if cfg.Embedding.BaseURL != "" {
setStringInMap(embeddingNode, "base_url", cfg.Embedding.BaseURL)
}
if cfg.Embedding.APIKey != "" {
setStringInMap(embeddingNode, "api_key", cfg.Embedding.APIKey)
}
// 更新检索配置
retrievalNode := ensureMap(knowledgeNode, "retrieval")
setIntInMap(retrievalNode, "top_k", cfg.Retrieval.TopK)
setFloatInMap(retrievalNode, "similarity_threshold", cfg.Retrieval.SimilarityThreshold)
setFloatInMap(retrievalNode, "hybrid_weight", cfg.Retrieval.HybridWeight)
// 更新索引配置
indexingNode := ensureMap(knowledgeNode, "indexing")
setIntInMap(indexingNode, "chunk_size", cfg.Indexing.ChunkSize)
setIntInMap(indexingNode, "chunk_overlap", cfg.Indexing.ChunkOverlap)
setIntInMap(indexingNode, "max_chunks_per_item", cfg.Indexing.MaxChunksPerItem)
setIntInMap(indexingNode, "max_rpm", cfg.Indexing.MaxRPM)
setIntInMap(indexingNode, "rate_limit_delay_ms", cfg.Indexing.RateLimitDelayMs)
setIntInMap(indexingNode, "max_retries", cfg.Indexing.MaxRetries)
setIntInMap(indexingNode, "retry_delay_ms", cfg.Indexing.RetryDelayMs)
}
func updateRobotsConfig(doc *yaml.Node, cfg config.RobotsConfig) {
root := doc.Content[0]
robotsNode := ensureMap(root, "robots")
wecomNode := ensureMap(robotsNode, "wecom")
setBoolInMap(wecomNode, "enabled", cfg.Wecom.Enabled)
setStringInMap(wecomNode, "token", cfg.Wecom.Token)
setStringInMap(wecomNode, "encoding_aes_key", cfg.Wecom.EncodingAESKey)
setStringInMap(wecomNode, "corp_id", cfg.Wecom.CorpID)
setStringInMap(wecomNode, "secret", cfg.Wecom.Secret)
setIntInMap(wecomNode, "agent_id", int(cfg.Wecom.AgentID))
dingtalkNode := ensureMap(robotsNode, "dingtalk")
setBoolInMap(dingtalkNode, "enabled", cfg.Dingtalk.Enabled)
setStringInMap(dingtalkNode, "client_id", cfg.Dingtalk.ClientID)
setStringInMap(dingtalkNode, "client_secret", cfg.Dingtalk.ClientSecret)
larkNode := ensureMap(robotsNode, "lark")
setBoolInMap(larkNode, "enabled", cfg.Lark.Enabled)
setStringInMap(larkNode, "app_id", cfg.Lark.AppID)
setStringInMap(larkNode, "app_secret", cfg.Lark.AppSecret)
setStringInMap(larkNode, "verify_token", cfg.Lark.VerifyToken)
}
func updateMultiAgentConfig(doc *yaml.Node, cfg config.MultiAgentConfig) {
root := doc.Content[0]
maNode := ensureMap(root, "multi_agent")
setBoolInMap(maNode, "enabled", cfg.Enabled)
setStringInMap(maNode, "default_mode", cfg.DefaultMode)
setBoolInMap(maNode, "robot_use_multi_agent", cfg.RobotUseMultiAgent)
setBoolInMap(maNode, "batch_use_multi_agent", cfg.BatchUseMultiAgent)
}
func ensureMap(parent *yaml.Node, path ...string) *yaml.Node {
current := parent
for _, key := range path {
value := findMapValue(current, key)
if value == nil {
keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}
mapNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
current.Content = append(current.Content, keyNode, mapNode)
value = mapNode
}
if value.Kind != yaml.MappingNode {
value.Kind = yaml.MappingNode
value.Tag = "!!map"
value.Style = 0
value.Content = nil
}
current = value
}
return current
}
func findMapValue(mapNode *yaml.Node, key string) *yaml.Node {
if mapNode == nil || mapNode.Kind != yaml.MappingNode {
return nil
}
for i := 0; i < len(mapNode.Content); i += 2 {
if mapNode.Content[i].Value == key {
return mapNode.Content[i+1]
}
}
return nil
}
func ensureKeyValue(mapNode *yaml.Node, key string) (*yaml.Node, *yaml.Node) {
if mapNode == nil || mapNode.Kind != yaml.MappingNode {
return nil, nil
}
for i := 0; i < len(mapNode.Content); i += 2 {
if mapNode.Content[i].Value == key {
return mapNode.Content[i], mapNode.Content[i+1]
}
}
keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}
valueNode := &yaml.Node{}
mapNode.Content = append(mapNode.Content, keyNode, valueNode)
return keyNode, valueNode
}
func setStringInMap(mapNode *yaml.Node, key, value string) {
_, valueNode := ensureKeyValue(mapNode, key)
valueNode.Kind = yaml.ScalarNode
valueNode.Tag = "!!str"
valueNode.Style = 0
valueNode.Value = value
}
func setIntInMap(mapNode *yaml.Node, key string, value int) {
_, valueNode := ensureKeyValue(mapNode, key)
valueNode.Kind = yaml.ScalarNode
valueNode.Tag = "!!int"
valueNode.Style = 0
valueNode.Value = fmt.Sprintf("%d", value)
}
func findBoolInMap(mapNode *yaml.Node, key string) *bool {
if mapNode == nil || mapNode.Kind != yaml.MappingNode {
return nil
}
for i := 0; i < len(mapNode.Content); i += 2 {
if i+1 >= len(mapNode.Content) {
break
}
keyNode := mapNode.Content[i]
valueNode := mapNode.Content[i+1]
if keyNode.Kind == yaml.ScalarNode && keyNode.Value == key {
if valueNode.Kind == yaml.ScalarNode {
if valueNode.Value == "true" {
result := true
return &result
} else if valueNode.Value == "false" {
result := false
return &result
}
}
return nil
}
}
return nil
}
func setBoolInMap(mapNode *yaml.Node, key string, value bool) {
_, valueNode := ensureKeyValue(mapNode, key)
valueNode.Kind = yaml.ScalarNode
valueNode.Tag = "!!bool"
valueNode.Style = 0
if value {
valueNode.Value = "true"
} else {
valueNode.Value = "false"
}
}
func setFloatInMap(mapNode *yaml.Node, key string, value float64) {
_, valueNode := ensureKeyValue(mapNode, key)
valueNode.Kind = yaml.ScalarNode
valueNode.Tag = "!!float"
valueNode.Style = 0
// 对于0.0到1.0之间的值如hybrid_weight使用%.1f确保0.0被明确序列化为"0.0"
// 对于其他值,使用%g自动选择最合适的格式
if value >= 0.0 && value <= 1.0 {
valueNode.Value = fmt.Sprintf("%.1f", value)
} else {
valueNode.Value = fmt.Sprintf("%g", value)
}
}
// getExternalMCPTools 获取外部MCP工具列表公共方法
// 返回 ToolConfigInfo 列表,已处理启用状态和描述信息
func (h *ConfigHandler) getExternalMCPTools(ctx context.Context) []ToolConfigInfo {
var result []ToolConfigInfo
if h.externalMCPMgr == nil {
return result
}
// 使用较短的超时时间5秒进行快速失败避免阻塞页面加载
timeoutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
externalTools, err := h.externalMCPMgr.GetAllTools(timeoutCtx)
if err != nil {
// 记录警告但不阻塞,继续返回已缓存的工具(如果有)
h.logger.Warn("获取外部MCP工具失败可能连接断开尝试返回缓存的工具",
zap.Error(err),
zap.String("hint", "如果外部MCP工具未显示请检查连接状态或点击刷新按钮"),
)
}
// 如果获取到了工具(即使有错误),继续处理
if len(externalTools) == 0 {
return result
}
externalMCPConfigs := h.externalMCPMgr.GetConfigs()
for _, externalTool := range externalTools {
// 解析工具名称mcpName::toolName
mcpName, actualToolName := h.parseExternalToolName(externalTool.Name)
if mcpName == "" || actualToolName == "" {
continue // 跳过格式不正确的工具
}
// 计算启用状态
enabled := h.calculateExternalToolEnabled(mcpName, actualToolName, externalMCPConfigs)
// 处理描述信息
description := h.pickToolDescription(externalTool.ShortDescription, externalTool.Description)
result = append(result, ToolConfigInfo{
Name: actualToolName,
Description: description,
Enabled: enabled,
IsExternal: true,
ExternalMCP: mcpName,
})
}
return result
}
// parseExternalToolName 解析外部工具名称格式mcpName::toolName
func (h *ConfigHandler) parseExternalToolName(fullName string) (mcpName, toolName string) {
idx := strings.Index(fullName, "::")
if idx > 0 {
return fullName[:idx], fullName[idx+2:]
}
return "", ""
}
// calculateExternalToolEnabled 计算外部工具的启用状态
func (h *ConfigHandler) calculateExternalToolEnabled(mcpName, toolName string, configs map[string]config.ExternalMCPServerConfig) bool {
cfg, exists := configs[mcpName]
if !exists {
return false
}
// 首先检查外部MCP是否启用
if !cfg.ExternalMCPEnable && !(cfg.Enabled && !cfg.Disabled) {
return false // MCP未启用所有工具都禁用
}
// MCP已启用检查单个工具的启用状态
// 如果ToolEnabled为空或未设置该工具默认为启用向后兼容
if cfg.ToolEnabled == nil {
// 未设置工具状态,默认为启用
} else if toolEnabled, exists := cfg.ToolEnabled[toolName]; exists {
// 使用配置的工具状态
if !toolEnabled {
return false
}
}
// 工具未在配置中,默认为启用
// 最后检查外部MCP是否已连接
client, exists := h.externalMCPMgr.GetClient(mcpName)
if !exists || !client.IsConnected() {
return false // 未连接时视为禁用
}
return true
}
// pickToolDescription 根据 security.tool_description_mode 选择 short 或 full 描述并限制长度
func (h *ConfigHandler) pickToolDescription(shortDesc, fullDesc string) string {
useFull := strings.TrimSpace(strings.ToLower(h.config.Security.ToolDescriptionMode)) == "full"
description := shortDesc
if useFull {
description = fullDesc
} else if description == "" {
description = fullDesc
}
if len(description) > 10000 {
description = description[:10000] + "..."
}
return description
}