feat: refactor flag

cmd/cmd.go

@@ -0,0 +1,233 @@
+package cmd
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net"
+	"os"
+	"os/signal"
+	"runtime"
+	"strings"
+	"time"
+
+	"github.com/akamensky/argparse"
+	"github.com/syndtr/gocapability/capability"
+	fastTrace "github.com/xgadget-lab/nexttrace/fast_trace"
+	"github.com/xgadget-lab/nexttrace/ipgeo"
+	"github.com/xgadget-lab/nexttrace/printer"
+	"github.com/xgadget-lab/nexttrace/reporter"
+	"github.com/xgadget-lab/nexttrace/trace"
+	"github.com/xgadget-lab/nexttrace/tracelog"
+	"github.com/xgadget-lab/nexttrace/tracemap"
+	"github.com/xgadget-lab/nexttrace/util"
+	"github.com/xgadget-lab/nexttrace/wshandle"
+)
+
+func Excute() {
+	parser := argparse.NewParser("nexttrace", "An open source visual route tracking CLI tool")
+	// Create string flag
+	tcp := parser.Flag("T", "tcp", &argparse.Options{Help: "Use TCP SYN for tracerouting (default port is 80)"})
+	udp := parser.Flag("U", "udp", &argparse.Options{Help: "Use UDP SYN for tracerouting (default port is 53)"})
+	fast_trace := parser.Flag("F", "fast-trace", &argparse.Options{Help: "One-Key Fast Trace to China ISPs"})
+	port := parser.Int("p", "port", &argparse.Options{Help: "Set the destination port to use. It is either initial udp port value for \"default\"" +
+		"method (incremented by each probe, default is 33434), or initial seq for \"icmp\" (incremented as well, default from 1), or some constant" +
+		"destination port for other methods (with default of 80 for \"tcp\", 53 for \"udp\", etc.)"})
+	numMeasurements := parser.Int("q", "queries", &argparse.Options{Default: 3, Help: "Set the number of probes per each hop"})
+	parallelRequests := parser.Int("", "parallel-requests", &argparse.Options{Default: 18, Help: "Set ParallelRequests number. It should be 1 when there is a multi-routing."})
+	maxHops := parser.Int("m", "max-hops", &argparse.Options{Default: 30, Help: "Set the max number of hops (max TTL to be reached)"})
+	dataOrigin := parser.Selector("d", "data-provider", []string{"IP.SB", "IPInfo", "IPInsight", "IPAPI.com"}, &argparse.Options{Default: "LeoMoeAPI",
+		Help: "Choose IP Geograph Data Provider [LeoMoeAPI,IP.SB, IPInfo, IPInsight, IPAPI.com]"})
+	noRdns := parser.Flag("n", "no-rdns", &argparse.Options{Help: " Do not resolve IP addresses to their domain names"})
+	routePath := parser.Flag("r", "route-path", &argparse.Options{Help: "Print traceroute hop path by ASN and location"})
+	output := parser.Flag("o", "output", &argparse.Options{Help: "Write trace result to file (RealTimePrinter ONLY)"})
+	tablePrint := parser.Flag("t", "table", &argparse.Options{Help: "Output trace results as table"})
+	classicPrint := parser.Flag("c", "classic", &argparse.Options{Help: "Classic Output trace results like BestTrace"})
+	beginHop := parser.Int("f", "first", &argparse.Options{Default: 1, Help: "Start from the first_ttl hop (instead from 1)"})
+	maptrace := parser.Flag("M", "map", &argparse.Options{Help: "Print Trace Map. This will return a Trace Map URL"})
+	ver := parser.Flag("v", "version", &argparse.Options{Help: "Print version info and exit"})
+	src_addr := parser.String("s", "source", &argparse.Options{Help: "Use source src_addr for outgoing packets"})
+	src_dev := parser.String("D", "dev", &argparse.Options{Help: "Use the following Network Devices as the source address in outgoing packets"})
+	router := parser.Flag("R", "route", &argparse.Options{Help: "Show Routing Table [Provided By BGP.Tools]"})
+	str := parser.StringPositional(&argparse.Options{Help: "IP Address or domain name"})
+
+	err := parser.Parse(os.Args)
+	if err != nil {
+		// In case of error print error and print usage
+		// This can also be done by passing -h or --help flags
+		fmt.Print(parser.Usage(err))
+	}
+
+	if *ver {
+		printer.CopyRight()
+		os.Exit(0)
+	}
+
+	domain := *str
+
+	if domain == "" {
+		fmt.Print(parser.Usage(err))
+		return
+	}
+
+	if *fast_trace {
+		fastTrace.FastTest(*tcp, *output)
+		if *output {
+			fmt.Println("您的追踪日志已经存放在 /tmp/trace.log 中")
+		}
+
+		os.Exit(0)
+	}
+
+	capabilities_check()
+	// return
+
+	var ip net.IP
+
+	if runtime.GOOS == "windows" && (*tcp || *udp) {
+		fmt.Println("NextTrace 基于 Windows 的路由跟踪还在早期开发阶段，目前还存在诸多问题，TCP/UDP SYN 包请求可能不能正常运行")
+	}
+
+	if *tcp || *udp {
+		ip = util.DomainLookUp(domain, true)
+	} else {
+		ip = util.DomainLookUp(domain, false)
+	}
+
+	if *src_dev != "" {
+		dev, _ := net.InterfaceByName(*src_dev)
+
+		if addrs, err := dev.Addrs(); err == nil {
+			for _, addr := range addrs {
+				if (addr.(*net.IPNet).IP.To4() == nil) == (ip.To4() == nil) {
+					*src_addr = addr.(*net.IPNet).IP.String()
+				}
+			}
+		}
+	}
+
+	if strings.ToUpper(*dataOrigin) == "LEOMOEAPI" {
+		w := wshandle.New()
+		w.Interrupt = make(chan os.Signal, 1)
+		signal.Notify(w.Interrupt, os.Interrupt)
+		defer func() {
+			w.Conn.Close()
+		}()
+	}
+
+	printer.PrintTraceRouteNav(ip, domain, *dataOrigin)
+
+	var m trace.Method = ""
+
+	switch {
+	case *tcp:
+		m = trace.TCPTrace
+	case *udp:
+		m = trace.UDPTrace
+	default:
+		m = trace.ICMPTrace
+	}
+
+	if !*tcp && *port == 80 {
+		*port = 53
+	}
+
+	var conf = trace.Config{
+		SrcAddr:          *src_addr,
+		BeginHop:         *beginHop,
+		DestIP:           ip,
+		DestPort:         *port,
+		MaxHops:          *maxHops,
+		NumMeasurements:  *numMeasurements,
+		ParallelRequests: *parallelRequests,
+		RDns:             !*noRdns,
+		IPGeoSource:      ipgeo.GetSource(*dataOrigin),
+		Timeout:          1 * time.Second,
+	}
+
+	if !*tablePrint {
+		if *classicPrint {
+			conf.RealtimePrinter = printer.ClassicPrinter
+		} else {
+			if *output {
+				conf.RealtimePrinter = tracelog.RealtimePrinter
+			} else if *router {
+				conf.RealtimePrinter = printer.RealtimePrinterWithRouter
+				fmt.Println("路由表数据源由 BGP.Tools 提供，在此特表感谢")
+			} else {
+				conf.RealtimePrinter = printer.RealtimePrinter
+			}
+		}
+	} else {
+		conf.AsyncPrinter = printer.TracerouteTablePrinter
+	}
+
+	res, err := trace.Traceroute(m, conf)
+
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	if *tablePrint {
+		printer.TracerouteTablePrinter(res)
+	}
+
+	if *routePath {
+		r := reporter.New(res, ip.String())
+		r.Print()
+	}
+
+	if *maptrace {
+		r, _ := json.Marshal(res)
+		tracemap.GetMapUrl(string(r))
+	}
+}
+
+func capabilities_check() {
+
+	// Windows 判断放在前面，防止遇到一些奇奇怪怪的问题
+	if runtime.GOOS == "windows" {
+		// Running on Windows, skip checking capabilities
+		return
+	}
+
+	uid := os.Getuid()
+	if uid == 0 {
+		// Running as root, skip checking capabilities
+		return
+	}
+
+	/***
+	* 检查当前进程是否有两个关键的权限
+	==== 看不到我 ====
+	* 没办法啦
+	* 自己之前承诺的坑补全篇
+	* 被迫填坑系列 qwq
+	==== 看不到我 ====
+	***/
+
+	// NewPid 已经被废弃了，这里改用 NewPid2 方法
+	caps, err := capability.NewPid2(0)
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+
+	// load 获取全部的 caps 信息
+	err = caps.Load()
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+
+	// 判断一下权限有木有
+	if caps.Get(capability.EFFECTIVE, capability.CAP_NET_RAW) && caps.Get(capability.EFFECTIVE, capability.CAP_NET_ADMIN) {
+		// 有权限啦
+		return
+	} else {
+		// 没权限啦
+		log.Println("您正在以普通用户权限运行 NextTrace，但 NextTrace 未被赋予监听网络套接字的ICMP消息包、修改IP头信息（TTL）等路由跟踪所需的权限")
+		log.Println("请使用管理员用户执行 `sudo setcap cap_net_raw,cap_net_admin+eip ${your_nexttrace_path}/nexttrace` 命令，赋予相关权限后再运行~")
+		log.Fatalln("什么？为什么 ping 普通用户执行不要 root 权限？因为这些工具在管理员安装时就已经被赋予了一些必要的权限，具体请使用 `getcap /usr/bin/ping` 查看")
+	}
+}

cmd/cmd_test.go

@@ -0,0 +1,7 @@
+package cmd
+
+import "testing"
+
+func TestCMD(t *testing.T) {
+	Excute()
+}

go.mod

@@ -3,6 +3,7 @@ module github.com/xgadget-lab/nexttrace
 go 1.19
 
 require (
+	github.com/akamensky/argparse v1.4.0
 	github.com/google/gopacket v1.1.19
 	github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635
 	golang.org/x/net v0.5.0

go.sum

@@ -1,3 +1,5 @@
+github.com/akamensky/argparse v1.4.0 h1:YGzvsTqCvbEZhL8zZu2AiA5nq805NZh75JNj4ajn1xc=
+github.com/akamensky/argparse v1.4.0/go.mod h1:S5kwC7IuDcEr5VeXtGPRVZ5o/FdhcMlQz4IZQuw64xA=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=

ipgeo/ipgeo_test.go

@@ -1,11 +1,72 @@
 package ipgeo
 
+import (
+	"errors"
+	"fmt"
+	"log"
+	"os"
+	"strconv"
+	"testing"
+)
+
 // import (
 // 	"testing"
 
 // 	"github.com/stretchr/testify/assert"
 // )
 
+func TestXxx(t *testing.T) {
+	const ID_FIXED_HEADER = "10"
+	var processID = fmt.Sprintf("%07b", os.Getpid()&0x7f) //取进程ID的前7位
+	var ttl = fmt.Sprintf("%06b", 95)                     //取TTL的后6位
+	fmt.Println(os.Getpid()&0x7f, 95)
+
+	var parity int
+	id := ID_FIXED_HEADER + processID + ttl
+	for _, c := range id {
+		if c == '1' {
+			parity++
+		}
+	}
+	if parity%2 == 0 {
+		id += "1"
+	} else {
+		id += "0"
+	}
+	process_id, ttl_r, _ := reverseID(id)
+	log.Println(process_id, ttl_r)
+}
+
+func reverseID(id string) (int64, int64, error) {
+	ttl, _ := strconv.ParseInt(id[9:15], 2, 32)
+	//process ID
+	processID, _ := strconv.ParseInt(id[2:9], 2, 32)
+
+	parity := 0
+	for i := 0; i < len(id)-1; i++ {
+		if id[i] == '1' {
+			parity++
+		}
+	}
+
+	if parity%2 == 1 {
+		if id[len(id)-1] == '0' {
+			fmt.Println("Parity check passed.")
+		} else {
+			fmt.Println("Parity check failed.")
+			return 0, 0, errors.New("err")
+		}
+	} else {
+		if id[len(id)-1] == '1' {
+			fmt.Println("Parity check passed.")
+		} else {
+			fmt.Println("Parity check failed.")
+			return 0, 0, errors.New("err")
+		}
+	}
+	return processID, ttl, nil
+}
+
 // func TestLeoIP(t *testing.T) {
 // 	// res, err := LeoIP("1.1.1.1")
 // 	// assert.Nil(t, err)