package data import ( "context" "encoding/json" "fmt" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo/options" "fwh-toolbox-server/cmd/mcp/tools/helper" "fwh-toolbox-server/database" ) // mongoAllowedCollections 允许查询的 collection 白名单 var mongoAllowedCollections = []string{ // 站点信息管理 "SiteInfoManageData", "SiteInfoManageRecord", // SIM卡管理 "SimCardManage", "SimCardMonthlyData", // 批量下载 "BatchDownloadConfig", // 站点分析 "SiteAnalysisConfig", // 统计 "StatisticsShutdown", "StatisticsAlarm", // 站点监控 "SiteMonitorConfig", // 自动升级 "AutoUpdateConfigOption", "AutoUpdateConfig", // 注意:UserInfo(内部用户身份数据)与 AiChatSessions(历史对话自由文本, // 字段级脱敏无效)已从白名单移除,禁止通过 data_mongo_* 直查以防隐私泄露。 } func registerMongoTools(s *server.MCPServer) { s.AddTool( mcp.NewTool("data_mongo_find", mcp.WithDescription("直接查询 MongoDB 集合数据(只读)。用于 service 层无法满足的灵活查询场景。"), mcp.WithString("collection", mcp.Required(), mcp.Description( fmt.Sprintf("集合名,可选值: %v", mongoAllowedCollections), )), mcp.WithString("filter", mcp.Description("查询条件,JSON 格式的 MongoDB filter(如 {\"sn\": \"xxx\"}),不传则查全部")), mcp.WithNumber("limit", mcp.Description("返回数量限制,默认10,最大100")), mcp.WithNumber("skip", mcp.Description("跳过数量(分页偏移),默认0")), mcp.WithString("sort", mcp.Description("排序条件,JSON 格式(如 {\"createdAt\": -1})")), ), handleMongoFind, ) s.AddTool( mcp.NewTool("data_mongo_count", mcp.WithDescription("统计 MongoDB 集合中符合条件的文档数量"), mcp.WithString("collection", mcp.Required(), mcp.Description( fmt.Sprintf("集合名,可选值: %v", mongoAllowedCollections), )), mcp.WithString("filter", mcp.Description("查询条件,JSON 格式的 MongoDB filter,不传则统计全部")), ), handleMongoCount, ) } func handleMongoFind(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { collection, err := req.RequireString("collection") if err != nil { return helper.ParamError("collection"), nil } if !isAllowedCollection(collection) { return helper.ParamInvalidError("collection", fmt.Sprintf("不允许查询此集合,可选: %v", mongoAllowedCollections)), nil } // 解析 filter filterStr := req.GetString("filter", "{}") var filter bson.M if err := json.Unmarshal([]byte(filterStr), &filter); err != nil { return helper.ParamInvalidError("filter", fmt.Sprintf("JSON 解析失败: %v", err)), nil } // limit limit := int64(req.GetFloat("limit", 10)) if limit > 100 { limit = 100 } if limit < 1 { limit = 10 } // skip skip := int64(req.GetFloat("skip", 0)) if skip < 0 { skip = 0 } // sort var sort bson.M sortStr := req.GetString("sort", "") if sortStr != "" { if err := json.Unmarshal([]byte(sortStr), &sort); err != nil { return helper.ParamInvalidError("sort", fmt.Sprintf("JSON 解析失败: %v", err)), nil } } // 查询 coll := database.GetCollection(collection) findOpts := options.Find().SetLimit(limit).SetSkip(skip) if sort != nil { findOpts.SetSort(sort) } cursor, err := coll.Find(ctx, filter, findOpts) if err != nil { return helper.ServiceError("MongoDB查询", err), nil } defer cursor.Close(ctx) var results []bson.M if err := cursor.All(ctx, &results); err != nil { return helper.ServiceError("读取MongoDB结果", err), nil } if len(results) == 0 { return helper.EmptyResult("查询结果为空"), nil } return helper.JSONResultOrFileAuto(results), nil } func handleMongoCount(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { collection, err := req.RequireString("collection") if err != nil { return helper.ParamError("collection"), nil } if !isAllowedCollection(collection) { return helper.ParamInvalidError("collection", fmt.Sprintf("不允许查询此集合,可选: %v", mongoAllowedCollections)), nil } filterStr := req.GetString("filter", "{}") var filter bson.M if err := json.Unmarshal([]byte(filterStr), &filter); err != nil { return helper.ParamInvalidError("filter", fmt.Sprintf("JSON 解析失败: %v", err)), nil } coll := database.GetCollection(collection) count, err := coll.CountDocuments(ctx, filter) if err != nil { return helper.ServiceError("MongoDB统计", err), nil } return helper.JSONResultOrFileAuto(map[string]any{ "collection": collection, "count": count, }), nil } func isAllowedCollection(name string) bool { for _, c := range mongoAllowedCollections { if c == name { return true } } return false }