增加配置文件

Author: keac

common/config.go

@@ -199,7 +199,11 @@ var ScanTypeKeys = map[string]string{
 type ENConfig struct {
 	Version   float64 `yaml:"version"`
 	UserAgent string  `yaml:"user_agent"` // 自定义 User-Agent
-	Cookies   struct {
+	Api       struct {
+		Api string `yaml:"api"`
+		Mcp string `yaml:"mcp"`
+	}
+	Cookies struct {
 		Aldzs       string `yaml:"aldzs"`
 		Xlb         string `yaml:"xlb"`
 		Aiqicha     string `yaml:"aiqicha"`
@@ -221,11 +225,14 @@ type ENConfig struct {
 }
 
 var cfgYName = filepath.Join(utils.GetConfigPath(), "config.yaml")
-var cfgYV = 0.6
-var configYaml = `version: 0.6
+var cfgYV = 0.7
+var configYaml = `version: 0.7
 user_agent: ""			# 自定义 User-Agent（可设置为获取Cookie的浏览器）
 app:
   miit_api: ''          # HG-ha的ICP_Query (非狼组维护 https://github.com/HG-ha/ICP_Query) 
+api:
+  api: ':31000'    # API监听地址
+  mcp: 'http://localhost:8080'    # MCP SSE监听地址
 cookies:
   aiqicha: ''           # 爱企查   Cookie
   tianyancha: ''        # 天眼查   Cookie

common/utils/utils.go

@@ -12,6 +12,8 @@ import (
 	"math"
 	"math/big"
 	mrand "math/rand"
+	"net"
+	"net/url"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -268,3 +270,18 @@ func RandomElement(c string) string {
 	// 返回该索引对应的元素
 	return slice[index]
 }
+
+func ExtractPortString(rawURL string) (string, error) {
+	// 解析 URL
+	u, err := url.Parse(rawURL)
+	if err != nil {
+		return "", fmt.Errorf("无效的 URL: %v", err)
+	}
+	// 提取 host:port 部分
+	_, port, err := net.SplitHostPort(u.Host)
+	if err != nil {
+		return "", fmt.Errorf("未指定端口且协议 %q 无默认端口", u.Scheme)
+	}
+
+	return port, nil // port 已经是字符串
+}

runner/api.go

@@ -40,7 +40,6 @@ func api(options *common.ENOptions) {
 			if data, ok := enApiData[w.OrgName]; !ok || w.IsRefresh {
 				enApiStatus[w.OrgName] = true
 				rdata = enTask.getInfoByApi(w)
-				fmt.Println(rdata)
 				message = "查询成功"
 				enApiStatus[w.OrgName] = false
 				enApiData[w.OrgName] = rdata
@@ -88,7 +87,7 @@ func api(options *common.ENOptions) {
 
 	}
 	enTask.StartENWorkers()
-	err := r.Run(":31000")
+	err := r.Run(options.ENConfig.Api.Api)
 	if err != nil {
 		gologger.Error().Msgf("API服务启动失败！")
 		gologger.Fatal().Msgf(err.Error())

runner/mcpServer.go

@@ -7,6 +7,7 @@ import (
 	"github.com/mark3labs/mcp-go/server"
 	"github.com/wgpsec/ENScan/common"
 	"github.com/wgpsec/ENScan/common/gologger"
+	"github.com/wgpsec/ENScan/common/utils"
 	"golang.org/x/net/context"
 	"log"
 	"strconv"
@@ -179,9 +180,14 @@ func McpServer(options *common.ENOptions) {
 		),
 	), getInfoByKeyword(enTask))
 	enTask.StartENWorkers()
-	sseServer := server.NewSSEServer(s, server.WithBaseURL("http://localhost:8080"))
-	gologger.Info().Msgf("SSE server listening on :8080")
-	if err := sseServer.Start(":8080"); err != nil {
+	sseServer := server.NewSSEServer(s, server.WithBaseURL(options.ENConfig.Api.Mcp))
+	port, err := utils.ExtractPortString(options.ENConfig.Api.Mcp)
+	if err != nil {
+		gologger.Error().Msgf("MCP服务启动失败！")
+		gologger.Fatal().Msgf(err.Error())
+	}
+	gologger.Info().Msgf("SSE server listening on :" + port)
+	if err := sseServer.Start(":" + port); err != nil {
 		log.Fatalf("Server error: %v", err)
 	}
 }