From bc31032ca85f7f3ef97bd1bdaf85d702f9f147a2 Mon Sep 17 00:00:00 2001
From: UUBulb <35923940+uubulb@users.noreply.github.com>
Date: Tue, 26 Nov 2024 21:29:00 +0800
Subject: [PATCH] feat: custom ip api (#100)

* feat: custom ip api

feat: add switch for nat and query
Co-authored-by: ChrisKim <1515748548@qq.com>

* remove dash.cloudflare

---------

Co-authored-by: ChrisKim <1515748548@qq.com>
Co-authored-by: naiba <hi@nai.ba>
---
 cmd/agent/main.go   | 21 +++++++++++++++++++++
 model/config.go     |  3 +++
 pkg/monitor/myip.go | 16 +++++++++++++---
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index 1df8237..f2e38c3 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -155,6 +155,7 @@ func preRun(configPath string) error {
 	}
 
 	monitor.InitConfig(&agentConfig)
+	monitor.CustomEndpoints = agentConfig.CustomIPApi
 
 	if agentConfig.ClientSecret == "" {
 		return errors.New("ClientSecret 不能为空")
@@ -570,6 +571,11 @@ func handleUpgradeTask(*pb.Task, *pb.TaskResult) {
 }
 
 func handleTcpPingTask(task *pb.Task, result *pb.TaskResult) {
+	if agentConfig.DisableSendQuery {
+		result.Data = "This server has disabled query sending"
+		return
+	}
+
 	host, port, err := net.SplitHostPort(task.GetData())
 	if err != nil {
 		result.Data = err.Error()
@@ -596,6 +602,11 @@ func handleTcpPingTask(task *pb.Task, result *pb.TaskResult) {
 }
 
 func handleIcmpPingTask(task *pb.Task, result *pb.TaskResult) {
+	if agentConfig.DisableSendQuery {
+		result.Data = "This server has disabled query sending"
+		return
+	}
+
 	ipAddr, err := lookupIP(task.GetData())
 	if err != nil {
 		result.Data = err.Error()
@@ -623,6 +634,11 @@ func handleIcmpPingTask(task *pb.Task, result *pb.TaskResult) {
 }
 
 func handleHttpGetTask(task *pb.Task, result *pb.TaskResult) {
+	if agentConfig.DisableSendQuery {
+		result.Data = "This server has disabled query sending"
+		return
+	}
+
 	start := time.Now()
 	taskUrl := task.GetData()
 	resp, err := httpClient.Get(taskUrl)
@@ -836,6 +852,11 @@ func handleTerminalTask(task *pb.Task) {
 }
 
 func handleNATTask(task *pb.Task) {
+	if agentConfig.DisableNat {
+		println("This server has disabled NAT traversal")
+		return
+	}
+
 	var nat model.TaskNAT
 	err := util.Json.Unmarshal([]byte(task.GetData()), &nat)
 	if err != nil {
diff --git a/model/config.go b/model/config.go
index 083b5cb..5c1500d 100644
--- a/model/config.go
+++ b/model/config.go
@@ -29,7 +29,10 @@ type AgentConfig struct {
 	InsecureTLS                 bool            `mapstructure:"insecure_tls" json:"insecure_tls"`                       // 是否禁用证书检查
 	UseIPv6CountryCode          bool            `mapstructure:"use_ipv6_country_code" json:"use_ipv6_country_code"`     // 默认优先展示IPv6旗帜
 	UseGiteeToUpgrade           bool            `mapstructure:"use_gitee_to_upgrade" json:"use_gitee_to_upgrade"`       // 强制从Gitee获取更新
+	DisableNat                  bool            `mapstructure:"disable_nat" json:"disable_nat"`                         // 关闭内网穿透
+	DisableSendQuery            bool            `mapstructure:"disable_send_query" json:"disable_send_query"`           // 关闭发送TCP/ICMP/HTTP请求
 	IPReportPeriod              uint32          `mapstructure:"ip_report_period" json:"ip_report_period"`               // IP上报周期
+	CustomIPApi                 []string        `mapstructure:"custom_ip_api" json:"custom_ip_api"`                     // 自定义 IP API
 
 	v *viper.Viper
 }
diff --git a/pkg/monitor/myip.go b/pkg/monitor/myip.go
index 1e31a34..de093d7 100644
--- a/pkg/monitor/myip.go
+++ b/pkg/monitor/myip.go
@@ -15,9 +15,11 @@ import (
 var (
 	cfList = []string{
 		"https://blog.cloudflare.com/cdn-cgi/trace",
-		"https://dash.cloudflare.com/cdn-cgi/trace",
+		"https://hostinger.com/cdn-cgi/trace",
+		"https://ahrefs.com/cdn-cgi/trace",
 		"https://developers.cloudflare.com/cdn-cgi/trace",
 	}
+	CustomEndpoints               []string
 	GeoQueryIP, CachedCountryCode string
 	GeoQueryIPChanged             bool = true
 	httpClientV4                       = util.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, false)
@@ -32,11 +34,19 @@ func FetchIP(useIPv6CountryCode bool) *pb.GeoIP {
 	var ipv4, ipv6 string
 	go func() {
 		defer wg.Done()
-		ipv4 = fetchIP(cfList, false)
+		if len(CustomEndpoints) > 0 {
+			ipv4 = fetchIP(CustomEndpoints, false)
+		} else {
+			ipv4 = fetchIP(cfList, false)
+		}
 	}()
 	go func() {
 		defer wg.Done()
-		ipv6 = fetchIP(cfList, true)
+		if len(CustomEndpoints) > 0 {
+			ipv6 = fetchIP(CustomEndpoints, true)
+		} else {
+			ipv6 = fetchIP(cfList, true)
+		}
 	}()
 	wg.Wait()