add: 新增配置文件模块

config/config.go

@@ -0,0 +1,31 @@
+package config
+
+import "os"
+
+type tracerConfig struct {
+    Token `yaml:"Token"`
+    Preference `yaml:"Preference"`
+}
+
+type Token struct {
+	LeoMoeAPI string `yaml:"LeoMoeAPI"`
+	IPInfo    string `yaml:"IPInfo"`
+}
+
+type Preference struct {
+    AlwaysRoutePath bool `yaml:"AlwaysRoutePath"`
+}
+
+type configPath func() (string, error)
+
+func configFromRunDir() (string, error) {
+    return "./", nil
+}
+
+func configFromUserHomeDir() (string, error) {
+    dir, err := os.UserHomeDir()
+    if err != nil {
+        return "", err
+    }
+    return dir + "/.nexttrace/", nil
+}

config/config_test.go

@@ -0,0 +1,18 @@
+package config
+
+import (
+	"testing"
+	"log"
+)
+
+func TestReadConfig(t *testing.T) {
+	if res, err := Read(); err != nil {
+		log.Println(err)
+	} else {
+		log.Println(res)
+	}
+}
+
+func TestGenerateConfig(t *testing.T) {
+	Generate()
+}

config/generate_config.go

@@ -0,0 +1,94 @@
+package config
+
+import (
+	"os"
+	"io/ioutil"
+	"fmt"
+	"gopkg.in/yaml.v2"
+)
+
+func pathExists(path string) (bool, error) {
+	_, err := os.Stat(path)
+	if err == nil {
+		return true, nil
+	}
+	if os.IsNotExist(err) {
+		return false, nil
+	}
+	return false, err
+}
+
+func writeFile(content []byte) error {
+	var err error
+	var path string
+	path, err = configFromUserHomeDir()
+	if err != nil {
+		path, err = configFromRunDir()
+		if err != nil {
+			return err
+		}
+	}
+
+	if exist, _ := pathExists(path); !exist {
+		err := os.Mkdir(path, os.ModePerm)
+		if err != nil {
+			return err
+		}
+	}
+
+	if err = ioutil.WriteFile(path + "ntraceConfig.yml", []byte(content), 0644); err != nil{
+		return err
+	}
+
+	return nil
+}
+
+func Generate() (*tracerConfig, error) {
+	var leotoken string
+	var iPInfoToken string
+	var routePathEnable string
+	
+	fmt.Println("这是一个配置向导，我们会帮助您生成配置文件，它是一次性的，除非您主动要求重新生成，否则它将不会再出现")
+
+	fmt.Println("请输入您的LeoMoeAPI Token，如果您没有，请到 Telegram Bot @NextTraceBot 获取一个")
+	fmt.Scanln(&leotoken)
+	if leotoken == "" {
+		fmt.Println("检测到您的输入为空，您将使用公共Token。这意味着您将和所有使用此Token的客户端共用每分钟150次IP查询的额度")
+		leotoken = "NextTraceDemo"
+	}
+
+	fmt.Println("请输入您的IPInfo Token，如果您不需要使用IPInfo，可以直接回车")
+	fmt.Scanln(&iPInfoToken)
+
+	token := Token{
+		LeoMoeAPI: leotoken,
+		IPInfo:    iPInfoToken,
+	}
+
+	var preference Preference
+	fmt.Print("您是否希望在每次Traceroute结束后显示Route-Path图? (y/n)")
+	fmt.Scanln(&routePathEnable)
+	if routePathEnable == "n" || routePathEnable == "N" || routePathEnable == "no" || routePathEnable == "No" || routePathEnable == "NO" {
+		preference = Preference{AlwaysRoutePath: false}
+	} else {
+		preference = Preference{AlwaysRoutePath: true}
+	}
+
+	finalConfig := tracerConfig{
+		Token:      token,
+		Preference: preference,
+	}
+
+	yamlData, err := yaml.Marshal(&finalConfig)
+
+	if err != nil {
+        return nil, err
+    }
+
+	if err = writeFile(yamlData); err != nil {
+		return nil, err
+	} else {
+		fmt.Println("配置文件创建成功")
+		return &finalConfig, nil
+	}
+}

config/read_config.go

@@ -0,0 +1,47 @@
+package config
+
+import (
+    "io/ioutil"
+    "log"
+    "gopkg.in/yaml.v2"
+)
+
+func (c *tracerConfig) Parse(data []byte) error {
+    return yaml.Unmarshal(data, c)
+}
+
+func readFile(cp configPath) ([]byte, error) {
+    var content []byte
+    path, err := cp()
+    if err != nil {
+        log.Println(err)
+        return nil, err
+    }
+    content, err = ioutil.ReadFile(path + "ntraceConfig.yml")
+    if err != nil {
+        return nil, err
+    }
+    return content, nil
+}
+
+func Read() (*tracerConfig, error) {
+    var data []byte
+    var err  error
+    
+    data, err = readFile(configFromRunDir)
+
+    if err != nil {
+        data, err = readFile(configFromUserHomeDir)
+
+        if err != nil {
+            return nil, err
+        }
+    }
+
+    var config tracerConfig
+    if err := config.Parse(data); err != nil {
+        return nil, err
+    }
+    
+    return &config, err
+}

go.mod

@@ -11,6 +11,7 @@ require (
 require (
 	github.com/mattn/go-colorable v0.1.9 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
 )
 
 require (

go.sum

@@ -49,6 +49,8 @@ golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapK
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

ipgeo/ipgeo.go

@@ -26,7 +26,9 @@ func GetSource(s string) Source {
 		return IPInSight
 	case "IPAPI.COM":
 		return IPApiCom
+	case "IPINFO":
+		return IPInfo
 	default:
-		return nil
+		return LeoIP
 	}
 }

ipgeo/ipinfo.go

@@ -8,7 +8,6 @@ import (
 )
 
 func IPInfo(ip string) (*IPGeoData, error) {
-
 	resp, err := http.Get("https://ipinfo.io/" + ip + "?token=" + token.ipinfo)
 	if err != nil {
 		return nil, err

ipgeo/tokens.go

@@ -1,5 +1,7 @@
 package ipgeo
 
+import "github.com/xgadget-lab/nexttrace/config"
+
 type tokenData struct {
 	ipinsight string
 	ipinfo    string
@@ -11,3 +13,9 @@ var token = tokenData{
 	ipinfo:    "",
 	ipleo:     "NextTraceDemo",
 }
+
+
+func SetToken(c config.Token) {
+	token.ipleo = c.LeoMoeAPI
+	token.ipinfo = c.IPInfo
+}

main.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/xgadget-lab/nexttrace/config"
 	"github.com/xgadget-lab/nexttrace/ipgeo"
 	"github.com/xgadget-lab/nexttrace/printer"
 	"github.com/xgadget-lab/nexttrace/reporter"
@@ -18,7 +19,7 @@ import (
 var fSet = flag.NewFlagSet("", flag.ExitOnError)
 var tcpSYNFlag = fSet.Bool("T", false, "Use TCP SYN for tracerouting (default port is 80)")
 var udpPackageFlag = fSet.Bool("U", false, "Use UDP Package for tracerouting (default port is 53 in UDP)")
-var port = flag.Int("p", 80, "Set SYN Traceroute Port")
+var port = fSet.Int("p", 80, "Set SYN Traceroute Port")
 var numMeasurements = fSet.Int("q", 3, "Set the number of probes per each hop.")
 var parallelRequests = fSet.Int("r", 18, "Set ParallelRequests number. It should be 1 when there is a multi-routing.")
 var maxHops = fSet.Int("m", 30, "Set the max number of hops (max TTL to be reached).")
@@ -65,6 +66,16 @@ func main() {
 		log.Fatalln("Traceroute requires root/sudo privileges.")
 	}
 
+	configData, err := config.Read();
+
+	if err != nil {
+		if configData, err = config.Generate(); err != nil {
+			log.Fatal(err)
+		}
+	}
+
+	ipgeo.SetToken(configData.Token)
+
 	ip := util.DomainLookUp(domain)
 	printer.PrintTraceRouteNav(ip, domain, *dataOrigin)
 
@@ -110,7 +121,9 @@ func main() {
 		printer.TraceroutePrinter(res)
 	}
 
-	if *routePath {
+	p := configData.Preference
+
+	if *routePath || p.AlwaysRoutePath {
 		r := reporter.New(res, ip.String())
 		r.Print()
 	}