From 1f051459bb738abb6efa7d5fb2b6c84ca3426876 Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Thu, 31 Mar 2022 22:45:26 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=92=A9=20=E4=BC=98=E5=8C=96=E8=8E=B7?=
 =?UTF-8?q?=E5=8F=96=20GeoIP?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/agent/monitor/myip.go      | 68 +++++++++++++++++++---------------
 cmd/agent/monitor/myip_test.go | 30 +++++++++++++++
 2 files changed, 68 insertions(+), 30 deletions(-)
 create mode 100644 cmd/agent/monitor/myip_test.go

diff --git a/cmd/agent/monitor/myip.go b/cmd/agent/monitor/myip.go
index c41b774..9872084 100644
--- a/cmd/agent/monitor/myip.go
+++ b/cmd/agent/monitor/myip.go
@@ -3,6 +3,7 @@ package monitor
 import (
 	"fmt"
 	"io/ioutil"
+	"math/rand"
 	"net/http"
 	"strings"
 	"time"
@@ -19,11 +20,11 @@ type geoIP struct {
 var (
 	geoIPApiList = []string{
 		"https://api.ip.sb/geoip",
-		"https://ip.seeip.org/geoip",
 		"https://ipapi.co/json",
 		"https://freegeoip.app/json/",
 		"http://ip-api.com/json/",
 		"https://extreme-ip-lookup.com/json/",
+		// "https://ip.seeip.org/geoip",
 	}
 	cachedIP, cachedCountry string
 	httpClientV4            = utils.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, false)
@@ -56,35 +57,42 @@ func fetchGeoIP(servers []string, isV6 bool) geoIP {
 	var ip geoIP
 	var resp *http.Response
 	var err error
-	for i := 0; i < len(servers); i++ {
-		if isV6 {
-			resp, err = httpClientV6.Get(servers[i])
-		} else {
-			resp, err = httpClientV4.Get(servers[i])
-		}
-		if err == nil {
-			body, err := ioutil.ReadAll(resp.Body)
-			if err != nil {
-				continue
-			}
-			resp.Body.Close()
-			err = utils.Json.Unmarshal(body, &ip)
-			if err != nil {
-				continue
-			}
-			if ip.IP == "" && ip.Query != "" {
-				ip.IP = ip.Query
-			}
-			// 没取到 v6 IP
-			if isV6 && !strings.Contains(ip.IP, ":") {
-				continue
-			}
-			// 没取到 v4 IP
-			if !isV6 && !strings.Contains(ip.IP, ".") {
-				continue
-			}
-			return ip
-		}
+	if isV6 {
+		resp, err = httpGetWithUA(httpClientV6, servers[rand.Intn(len(servers))])
+	} else {
+		resp, err = httpGetWithUA(httpClientV4, servers[rand.Intn(len(servers))])
+	}
+	if err != nil {
+		return ip
+	}
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return ip
+	}
+	resp.Body.Close()
+	err = utils.Json.Unmarshal(body, &ip)
+	if err != nil {
+		return ip
+	}
+	if ip.IP == "" && ip.Query != "" {
+		ip.IP = ip.Query
+	}
+	// 没取到 v6 IP
+	if isV6 && !strings.Contains(ip.IP, ":") {
+		return ip
+	}
+	// 没取到 v4 IP
+	if !isV6 && !strings.Contains(ip.IP, ".") {
+		return ip
 	}
 	return ip
 }
+
+func httpGetWithUA(client *http.Client, url string) (*http.Response, error) {
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Add("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/99.0.4844.84 Safari/537.36")
+	return client.Do(req)
+}
diff --git a/cmd/agent/monitor/myip_test.go b/cmd/agent/monitor/myip_test.go
new file mode 100644
index 0000000..f7c0e1b
--- /dev/null
+++ b/cmd/agent/monitor/myip_test.go
@@ -0,0 +1,30 @@
+package monitor
+
+import (
+	"io/ioutil"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/naiba/nezha/pkg/utils"
+)
+
+func TestGeoIPApi(t *testing.T) {
+	for i := 0; i < len(geoIPApiList); i++ {
+		resp, err := httpGetWithUA(httpClientV4, geoIPApiList[i])
+		assert.Nil(t, err)
+		body, err := ioutil.ReadAll(resp.Body)
+		assert.Nil(t, err)
+		resp.Body.Close()
+		var ip geoIP
+		err = utils.Json.Unmarshal(body, &ip)
+		assert.Nil(t, err)
+		t.Logf("%s %s %s %s", geoIPApiList[i], ip.CountryCode, utils.IPDesensitize(ip.IP), utils.IPDesensitize(ip.Query))
+		assert.True(t, ip.IP != "" || ip.Query != "")
+	}
+}
+
+func TestFetchGeoIP(t *testing.T) {
+	ip := fetchGeoIP(geoIPApiList, false)
+	assert.NotEmpty(t, ip.IP)
+}