Skip to content

Commit 4a08a44

Browse files
keackeac
committed
增加MCP协议功能，方便接入大模型调用
1 parent 6a088cc commit 4a08a44

File tree

9 files changed

+140
-4
lines changed

9 files changed

+140
-4
lines changed

‎README.md

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,24 @@
145145
./enscan -n 小米 -delay 3
146146
```
147147

148+
### 使用MCP
149+
150+
开启MCP服务器，将会监听本地的 http://localhost:8080
151+
152+
```
153+
./enscan --mcp
154+
```
155+
156+
以 Cherry Studio 配置为例
157+
158+
![image-20250329160425571](./README/image-20250329160425571.png)
159+
160+
配置完成完成后开启MCP服务
161+
162+
![image-20250329160556011](./README/image-20250329160556011.png)
163+
164+
配置完成后可以根据自己的需求编写prompt 欢迎 [在此](https://github.com/wgpsec/ENScan_GO/discussions/163) 分享好用的prompt
165+
148166
### Cookie配置
149167

150168
**AQC**

‎common/config.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ type ENOptions struct {
4949
IsNoMerge bool //聚合
5050
OutPutType string // 导出文件类型
5151
IsApiMode bool
52+
IsMCPServer bool
5253
ENConfig *ENConfig
5354
BranchFilter string
5455
}

‎common/flag.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ func Flag(Info *ENOptions) {
4949
//其他设定
5050
flag.BoolVar(&Info.IsGroup, "is-group", false, "查询关键词为集团")
5151
flag.BoolVar(&Info.IsApiMode, "api", false, "API模式运行")
52+
flag.BoolVar(&Info.IsMCPServer, "mcp", false, "MCP模式运行")
5253
flag.BoolVar(&Info.ISKeyPid, "is-pid", false, "批量查询文件是否为公司PID")
5354
flag.IntVar(&Info.DelayTime, "delay", 0, "每个请求延迟（S）-1为随机延迟1-5S")
5455
flag.StringVar(&Info.Proxy, "proxy", "", "设置代理")

‎common/parse.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ func Parse(options *ENOptions) {
5454
gologger.Fatal().Msgf("配置文件当前[V%.1f] 程序需要[V%.1f] 不匹配，请备份配置文件重新运行-v\n", conf.Version, cfgYV)
5555
}
5656

57-
if options.KeyWord == "" && options.CompanyID == "" && options.InputFile == "" && !options.IsApiMode {
57+
if options.KeyWord == "" && options.CompanyID == "" && options.InputFile == "" && !options.IsApiMode && !options.IsMCPServer {
5858
flag.PrintDefaults()
5959
os.Exit(0)
6060
}

‎enscan.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ func main() {
2020
for {
2121
select {
2222
case <-quitSig:
23-
if !enOptions.IsApiMode {
23+
if !enOptions.IsApiMode && !enOptions.IsMCPServer {
2424
gologger.Error().Msgf("任务未完成退出，将自动保存过程文件！")
2525
enDataList := make(map[string][]map[string]string)
2626
close(runner.EnCh)

‎go.mod

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
module github.com/wgpsec/ENScan
22

3-
go 1.21
3+
go 1.23
4+
5+
toolchain go1.23.2
46

57
require (
68
github.com/antchfx/htmlquery v1.3.1
79
github.com/gin-gonic/gin v1.10.0
810
github.com/imroc/req/v3 v3.43.7
11+
github.com/mark3labs/mcp-go v0.17.0
912
github.com/olekukonko/tablewriter v0.0.5
1013
github.com/projectdiscovery/gologger v1.1.12
1114
github.com/robertkrimen/otto v0.0.0-20211024170158-b87d35c0b86f
@@ -36,6 +39,7 @@ require (
3639
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
3740
github.com/golang/snappy v0.0.4 // indirect
3841
github.com/google/pprof v0.0.0-20240227163752-401108e1b7e7 // indirect
42+
github.com/google/uuid v1.6.0 // indirect
3943
github.com/hashicorp/errwrap v1.1.0 // indirect
4044
github.com/hashicorp/go-multierror v1.1.1 // indirect
4145
github.com/json-iterator/go v1.1.12 // indirect
@@ -67,6 +71,7 @@ require (
6771
github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect
6872
github.com/xuri/efp v0.0.0-20231025114914-d1ff6096ae53 // indirect
6973
github.com/xuri/nfp v0.0.0-20230919160717-d98342af3f05 // indirect
74+
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
7075
go.uber.org/mock v0.4.0 // indirect
7176
golang.org/x/arch v0.8.0 // indirect
7277
golang.org/x/crypto v0.23.0 // indirect

‎go.sum

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@ github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN
5959
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
6060
github.com/google/pprof v0.0.0-20240227163752-401108e1b7e7 h1:y3N7Bm7Y9/CtpiVkw/ZWj6lSlDF3F74SfKwfTCer72Q=
6161
github.com/google/pprof v0.0.0-20240227163752-401108e1b7e7/go.mod h1:czg5+yv1E0ZGTi6S6vVK1mke0fV+FaUhNGcd6VRS9Ik=
62+
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
63+
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
6264
github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
6365
github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY2I=
6466
github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
@@ -89,6 +91,8 @@ github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
8991
github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
9092
github.com/logrusorgru/aurora v2.0.3+incompatible h1:tOpm7WcpBTn4fjmVfgpQq0EfczGlG91VSDkswnjF5A8=
9193
github.com/logrusorgru/aurora v2.0.3+incompatible/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=
94+
github.com/mark3labs/mcp-go v0.17.0 h1:5Ps6T7qXr7De/2QTqs9h6BKeZ/qdeUeGrgM5lPzi930=
95+
github.com/mark3labs/mcp-go v0.17.0/go.mod h1:KmJndYv7GIgcPVwEKJjNcbhVQ+hJGJhrCCB/9xITzpE=
9296
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
9397
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
9498
github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/QdE+0=
@@ -176,6 +180,8 @@ github.com/xuri/excelize/v2 v2.8.1 h1:pZLMEwK8ep+CLIUWpWmvW8IWE/yxqG0I1xcN6cVMGu
176180
github.com/xuri/excelize/v2 v2.8.1/go.mod h1:oli1E4C3Pa5RXg1TBXn4ENCXDV5JUMlBluUhG7c+CEE=
177181
github.com/xuri/nfp v0.0.0-20230919160717-d98342af3f05 h1:qhbILQo1K3mphbwKh1vNm4oGezE1eF9fQWmNiIpSfI4=
178182
github.com/xuri/nfp v0.0.0-20230919160717-d98342af3f05/go.mod h1:WwHg+CVyzlv/TX9xqBFXEZAuxOPxn2k1GNHwG41IIUQ=
183+
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
184+
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
179185
github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d/go.mod h1:rHwXgn7JulP+udvsHwJoVG1YGAP6VLg4y9I5dyZdqmA=
180186
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
181187
go.mongodb.org/mongo-driver v1.9.0 h1:f3aLGJvQmBl8d9S40IL+jEyBC6hfLPbJjv9t5hEM9ck=

‎runner/mcpServer.go

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
package runner
2+
3+
import (
4+
"encoding/json"
5+
"errors"
6+
"fmt"
7+
"github.com/mark3labs/mcp-go/mcp"
8+
"github.com/mark3labs/mcp-go/server"
9+
"github.com/tidwall/gjson"
10+
"github.com/wgpsec/ENScan/common"
11+
"github.com/wgpsec/ENScan/common/gologger"
12+
"github.com/wgpsec/ENScan/common/utils"
13+
"github.com/wgpsec/ENScan/internal/aiqicha"
14+
"golang.org/x/net/context"
15+
"log"
16+
)
17+
18+
func helloSearchById(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
19+
arguments := request.Params.Arguments
20+
var enOptions common.ENOptions
21+
enOptions.IsMCPServer = true
22+
pid, ok := arguments["pid"].(string)
23+
if !ok {
24+
return nil, errors.New("pid must be a string")
25+
}
26+
filed, ok := arguments["filed"].(string)
27+
if ok {
28+
enOptions.GetFlags = filed
29+
}
30+
common.Parse(&enOptions)
31+
enOptions.CompanyID = pid
32+
enOptions.IsMergeOut = true
33+
data := RunJob(&enOptions)
34+
r, err := json.Marshal(data)
35+
if err != nil {
36+
return mcp.NewToolResultText(fmt.Sprintf("处理异常！")), err
37+
}
38+
return mcp.NewToolResultText(fmt.Sprintf("%s", r)), nil
39+
}
40+
func helloSearchListByOgrName(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
41+
arguments := request.Params.Arguments
42+
keyWord, ok := arguments["OrgName"].(string)
43+
if !ok {
44+
return nil, errors.New("orgName must be a string")
45+
}
46+
var enOptions common.ENOptions
47+
enOptions.IsMCPServer = true
48+
common.Parse(&enOptions)
49+
enOptions.KeyWord = keyWord
50+
job := &aiqicha.AQC{Options: &enOptions}
51+
enList, err := job.AdvanceFilter()
52+
enMap := job.GetENMap()["enterprise_info"]
53+
if err != nil {
54+
gologger.Error().Msg(err.Error())
55+
return mcp.NewToolResultText(fmt.Sprintf("处理异常！")), err
56+
} else {
57+
utils.TBS(append(enMap.KeyWord[:3], "PID"), append(enMap.Field[:3], enMap.Field[10]), "企业信息", enList)
58+
plList := common.InfoToMap(map[string][]gjson.Result{
59+
"enterprise_info": enList,
60+
}, job.GetENMap(), "")
61+
r, err := json.Marshal(plList["enterprise_info"])
62+
if err != nil {
63+
return mcp.NewToolResultText(fmt.Sprintf("处理异常！")), err
64+
}
65+
return mcp.NewToolResultText(fmt.Sprintf("%s", r)), nil
66+
67+
}
68+
}
69+
70+
func mcpServer(options *common.ENOptions) {
71+
s := server.NewMCPServer(
72+
"EnScan",
73+
"1.0.0",
74+
)
75+
// Add tool
76+
77+
// Add tool handler
78+
s.AddTool(mcp.NewTool("根据PID详细信息",
79+
mcp.WithDescription("根据pid搜索企业的icp备案、微博、微信、app、微博、招聘、微信小程序、版权信息"),
80+
81+
mcp.WithString("pid",
82+
mcp.Required(),
83+
mcp.Description("企业搜索结果的PID"),
84+
),
85+
mcp.WithString("filed",
86+
mcp.Description("获取信息类别多个类别需要以,分隔"),
87+
mcp.Enum("icp", "weibo", "wechat", "app", "weibo", "job", "wx_app", "copyright"),
88+
),
89+
), helloSearchById)
90+
s.AddTool(mcp.NewTool("关键词匹配企业列表",
91+
mcp.WithDescription("根据关键词搜索匹配企业列表"),
92+
mcp.WithString("OrgName",
93+
mcp.Required(),
94+
mcp.Description("企业名称"),
95+
),
96+
), helloSearchListByOgrName)
97+
98+
sseServer := server.NewSSEServer(s, server.WithBaseURL("http://localhost:8080"))
99+
gologger.Info().Msgf("SSE server listening on :8080")
100+
if err := sseServer.Start(":8080"); err != nil {
101+
log.Fatalf("Server error: %v", err)
102+
}
103+
}

‎runner/runner.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,8 @@ func RunEnumeration(options *common.ENOptions) {
8686

8787
} else if options.IsApiMode {
8888
api(options)
89+
} else if options.IsMCPServer {
90+
mcpServer(options)
8991
} else {
9092
RunJob(options)
9193
}
@@ -164,7 +166,7 @@ func RunJob(options *common.ENOptions) map[string][]map[string]string {
164166
}
165167

166168
wg.Wait()
167-
if options.IsMergeOut && options.InputFile == "" && !options.IsApiMode {
169+
if options.IsMergeOut && options.InputFile == "" && !options.IsApiMode && !options.IsMCPServer {
168170
err := common.OutFileByEnInfo(enDataList, options.KeyWord, options.OutPutType, options.Output)
169171
if err != nil {
170172
gologger.Error().Msgf(err.Error())

0 commit comments

Comments
 (0)