勉强适配nexttrace v3 api

go.mod

@@ -14,11 +14,15 @@ require (
 	github.com/saracen/go7z v0.0.0-20191010121135-9c09b6bd7fda
 	github.com/spf13/cobra v1.7.0
 	github.com/spf13/viper v1.15.0
+	github.com/stretchr/testify v1.8.4
+	github.com/tsosunchia/powclient v0.1.0
+	github.com/xgadget-lab/nexttrace v1.1.7-0.20230601074004-64371fb41ab4
 	golang.org/x/text v0.9.0
 	gopkg.in/yaml.v2 v2.4.0
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
@@ -28,6 +32,7 @@ require (
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/oschwald/maxminddb-golang v1.10.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.7 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/saracen/go7z-fixtures v0.0.0-20190623165746-aa6b8fba1d2f // indirect
 	github.com/saracen/solidblock v0.0.0-20190426153529-45df20abab6f // indirect
 	github.com/spf13/afero v1.9.5 // indirect

go.sum

@@ -197,12 +197,17 @@ github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/subosito/gotenv v1.4.2 h1:X1TuBLAMDFbaTAChgCBLu3DU3UPyELpnF2jjJ2cz/S8=
 github.com/subosito/gotenv v1.4.2/go.mod h1:ayKnFf/c6rvx/2iiLrJUk1e6plDbT3edrFNGqEflhK0=
+github.com/tsosunchia/powclient v0.1.0 h1:yUCa1Zd0KENOCiWJ3NOYhV1uVWUjy1ujy4w8gap3j4M=
+github.com/tsosunchia/powclient v0.1.0/go.mod h1:Pm4MP3QqN74SfNskPpFIEyT+NQrcABGoXkkeRwjlMEE=
 github.com/ulikunitz/xz v0.5.11 h1:kpFauv27b6ynzBNT/Xy+1k+fK4WswhN/6PN5WhFAGw8=
 github.com/ulikunitz/xz v0.5.11/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14=
+github.com/xgadget-lab/nexttrace v1.1.7-0.20230601074004-64371fb41ab4 h1:mqXG5rRwjZKnH6zRZnmeNvq0qU0twhGBN63xPnaqAtM=
+github.com/xgadget-lab/nexttrace v1.1.7-0.20230601074004-64371fb41ab4/go.mod h1:n/RsjvHSaeZS0shqYWzugI9ZyX/9G4LU4kJju+59B9o=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

internal/db/db.go

@@ -16,6 +16,8 @@ import (
 	"github.com/zu1k/nali/pkg/zxipv6wry"
 )
 
+var tokenCache = make(map[string]string)
+
 func GetDB(typ dbif.QueryType) (db dbif.DB) {
 	if db, found := dbTypeCache[typ]; found {
 		return db
@@ -82,7 +84,16 @@ func Find(typ dbif.QueryType, query string) string {
 	if nali == "1" {
 		result, err = GetDB(typ).Find(query)
 	} else {
-		result, err = leomoeapi.Find(query)
+		token, ok := tokenCache["api.leo.moe"]
+		//fmt.Println("token:", token)
+		if !ok {
+			token = ""
+			//fmt.Println("token not found")
+		} else {
+			//fmt.Println("token found")
+		}
+		result, token, err = leomoeapi.Find(query, token)
+		tokenCache["api.leo.moe"] = token
 	}
 	if err != nil {
 		return ""

pkg/leomoeapi/leomoeapi.go

@@ -1,10 +1,17 @@
 package leomoeapi
 
 import (
+	"crypto/tls"
 	"encoding/json"
 	"fmt"
 	"github.com/gorilla/websocket"
+	"github.com/xgadget-lab/nexttrace/util"
+	"github.com/zu1k/nali/pow"
+	"log"
 	"net"
+	"net/http"
+	"net/url"
+	"os"
 	"strings"
 )
 
@@ -29,28 +36,90 @@ type IPGeoData struct {
 	Source    string              `json:"source"`
 }
 
-func FetchIPInfo(ip string) (*IPGeoData, error) {
-	c, _, err := websocket.DefaultDialer.Dial("wss://api.leo.moe/v2/ipGeoWs", nil)
-	if err != nil {
-		return nil, fmt.Errorf("websocket dial: %w", err)
+var conn *websocket.Conn
+
+func FetchIPInfo(ip string, token string) (*IPGeoData, string, error) {
+	var host, port, fastIp string
+	host, port = util.GetHostAndPort()
+	// 如果 host 是一个 IP 使用默认域名
+	if valid := net.ParseIP(host); valid != nil {
+		host = "api.leo.moe"
+	} else {
+		// 默认配置完成，开始寻找最优 IP
+		fastIp = util.GetFastIP(host, port, false)
+	}
+	//host, port, fastIp = "103.120.18.35", "api.leo.moe", "443"
+	envToken := util.EnvToken
+	jwtToken := "Bearer " + token
+	ua := []string{"Privileged Client"}
+	if token == "" {
+		jwtToken = envToken
+		err := error(nil)
+		if envToken == "" {
+			jwtToken, err = pow.GetToken(fastIp, host, port)
+			if err != nil {
+				log.Println(err)
+				os.Exit(1)
+			}
+			ua = []string{pow.UserAgent}
+		}
+		jwtToken = "Bearer " + jwtToken
+	}
+
+	requestHeader := http.Header{
+		"Host":          []string{host},
+		"User-Agent":    ua,
+		"Authorization": []string{jwtToken},
+	}
+	jwtToken = strings.TrimPrefix(jwtToken, "Bearer ")
+	dialer := websocket.DefaultDialer
+	dialer.TLSClientConfig = &tls.Config{
+		ServerName: host,
+	}
+	u := url.URL{Scheme: "wss", Host: fastIp + ":" + port, Path: "/v3/ipGeoWs"}
+
+	var c *websocket.Conn
+	var err error
+
+	if conn == nil {
+		c, _, err = websocket.DefaultDialer.Dial(u.String(), requestHeader)
+		if err != nil {
+			return nil, "", fmt.Errorf("websocket dial: %w", err)
+		}
+		c.SetCloseHandler(func(code int, text string) error {
+			conn = nil // 将全局的 conn 设为 nil
+			return nil
+		})
+		conn = c
+	} else {
+		c = conn
 	}
-	defer c.Close()
+
+	//defer func(c *websocket.Conn) {
+	//	err := c.Close()
+	//	if err != nil {
+	//		log.Println(err)
+	//	}
+	//}(c)
+	// TODO: 现在是一直不关闭，以后想办法在程序退出时关闭
+	// 在这种情况下，你可以考虑使用Go的 os/signal 包来监听操作系统发出的终止信号。当程序收到这样的信号时，
+	// 比如 SIGINT（即 Ctrl+C）或 SIGTERM，你可以优雅地关闭你的 WebSocket 连接。
 
 	if err := c.WriteMessage(websocket.TextMessage, []byte(ip)); err != nil {
-		return nil, fmt.Errorf("write message: %w", err)
+		return nil, "", fmt.Errorf("write message: %w", err)
 	}
 
 	_, message, err := c.ReadMessage()
 	if err != nil {
-		return nil, fmt.Errorf("read message: %w", err)
+		return nil, "", fmt.Errorf("read message: %w", err)
 	}
 
 	var data IPGeoData
 	if err := json.Unmarshal(message, &data); err != nil {
-		return nil, fmt.Errorf("json unmarshal: %w", err)
+		return nil, "", fmt.Errorf("json unmarshal: %w", err)
 	}
 
-	return &data, nil
+	return &data, jwtToken, nil
 }
 
 type Result struct {
@@ -119,17 +188,26 @@ func isPrivateOrReserved(ip net.IP) bool {
 	return false
 }
 
-func Find(query string, params ...string) (result fmt.Stringer, err error) {
+func Find(query string, token string) (result fmt.Stringer, retToken string, err error) {
 	if net.ParseIP(query) == nil {
-		return Result{""}, nil // 如果 query 不是一个有效的 IP 地址，返回空字符串
+		return Result{""}, token, nil // 如果 query 不是一个有效的 IP 地址，返回空字符串
 	}
 	if isPrivateOrReserved(net.ParseIP(query)) {
-		return Result{""}, nil // 如果 query 是一个私有或保留地址，返回空字符串
+		return Result{""}, token, nil // 如果 query 是一个私有或保留地址，返回空字符串
 	}
-	res, err := FetchIPInfo(query)
-	if err != nil {
-		return nil, err
+	i := 0
+	var res *IPGeoData
+	for i = 0; i < 3; i++ {
+		res, token, err = FetchIPInfo(query, token)
+		if err != nil {
+			continue
+		}
+		break
 	}
+	if i == 3 {
+		return nil, "", err
+	}
+
 	result = Result{
 		Data: strings.Join(func() []string {
 			dataSlice := make([]string, 0, 7)
@@ -154,5 +232,5 @@ func Find(query string, params ...string) (result fmt.Stringer, err error) {
 		}(), ";"),
 	}
 
-	return result, nil
+	return result, token, nil
 }

pow/pow.go

@@ -0,0 +1,41 @@
+package pow
+
+import (
+	"fmt"
+	"github.com/tsosunchia/powclient"
+	"github.com/zu1k/nali/internal/constant"
+	"net/url"
+	"os"
+	"runtime"
+)
+
+const (
+	baseURL = "/v3/challenge"
+)
+
+var UserAgent = fmt.Sprintf("Nali-NextTrace %s/%s/%s", constant.Version, runtime.GOOS, runtime.GOARCH)
+
+func GetToken(fastIp string, host string, port string) (string, error) {
+	getTokenParams := powclient.NewGetTokenParams()
+	u := url.URL{Scheme: "https", Host: fastIp + ":" + port, Path: baseURL}
+	getTokenParams.BaseUrl = u.String()
+	getTokenParams.SNI = host
+	getTokenParams.Host = host
+	getTokenParams.UserAgent = UserAgent
+	var err error
+	// 尝试三次RetToken，如果都失败了，异常退出
+	for i := 0; i < 3; i++ {
+		token, err := powclient.RetToken(getTokenParams)
+		//fmt.Println(token, err)
+		if err != nil {
+			continue
+		}
+		return token, nil
+	}
+	if err != nil {
+		fmt.Println(err)
+	}
+	fmt.Println("RetToken failed 3 times, exit")
+	os.Exit(1)
+	return "", nil
+}

pow/pow_test.go

@@ -0,0 +1,13 @@
+package pow
+
+import (
+	"fmt"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestGetToken(t *testing.T) {
+	token, err := GetToken("103.120.18.35", "api.leo.moe", "443")
+	fmt.Println(token, err)
+	assert.NoError(t, err, "GetToken() returned an error")
+}