From b8772d4cca45b66f96b72eecedd56d7892ded9df Mon Sep 17 00:00:00 2001
From: sjlleo <sjlleo001@gmail.com>
Date: Sun, 12 Jun 2022 12:18:43 +0800
Subject: [PATCH] add: detect hop type & highlight

---
 printer/printer.go          | 43 +++++++++++++-----
 printer/printer_test.go     |  6 +--
 printer/realtime_printer.go | 87 +++++++++++++++++++++++++++++++++++--
 3 files changed, 120 insertions(+), 16 deletions(-)

diff --git a/printer/printer.go b/printer/printer.go
index 75285d6..d034c5b 100644
--- a/printer/printer.go
+++ b/printer/printer.go
@@ -9,18 +9,27 @@ import (
 	"github.com/xgadget-lab/nexttrace/ipgeo"
 )
 
-var dataOrigin string
+// var dataOrigin string
 
-func TraceroutePrinter(res *trace.Result) {
-	for i, hop := range res.Hops {
-		fmt.Print(i + 1)
-		for _, h := range hop {
-			HopPrinter(h)
-		}
-	}
-}
+// func TraceroutePrinter(res *trace.Result) {
+// 	for i, hop := range res.Hops {
+// 		fmt.Print(i + 1)
+// 		for _, h := range hop {
+// 			HopPrinter(h)
+// 		}
+// 	}
+// }
 
-func HopPrinter(h trace.Hop) {
+const (
+	RED_PREFIX    = "\033[1;31m"
+	GREEN_PREFIX  = "\033[1;32m"
+	YELLOW_PREFIX = "\033[1;33m"
+	BLUE_PREFIX   = "\033[1;34m"
+	CYAN_PREFIX   = "\033[1;36m"
+	RESET_PREFIX  = "\033[0m"
+)
+
+func HopPrinter(h trace.Hop, info HopInfo) {
 	if h.Address == nil {
 		fmt.Println("\t*")
 	} else {
@@ -35,8 +44,22 @@ func HopPrinter(h trace.Hop) {
 		if h.Geo != nil {
 			txt += " " + formatIpGeoData(h.Address.String(), h.Geo)
 		}
+		switch info {
+		case IXP:
+			fmt.Print(CYAN_PREFIX)
+		case PoP:
+			fmt.Print(CYAN_PREFIX)
+		case Peer:
+			fmt.Print(YELLOW_PREFIX)
+		case Aboard:
+			fmt.Print(GREEN_PREFIX)
+		}
 
 		fmt.Println(txt)
+
+		if info != General {
+			fmt.Print(RESET_PREFIX)
+		}
 	}
 }
 
diff --git a/printer/printer_test.go b/printer/printer_test.go
index 25460f4..2533c85 100644
--- a/printer/printer_test.go
+++ b/printer/printer_test.go
@@ -90,9 +90,9 @@ var testResult = &trace.Result{
 	},
 }
 
-func TestTraceroutePrinter(t *testing.T) {
-	TraceroutePrinter(testResult)
-}
+// func TestTraceroutePrinter(t *testing.T) {
+// 	TraceroutePrinter(testResult)
+// }
 
 func TestTracerouteTablePrinter(t *testing.T) {
 	TracerouteTablePrinter(testResult)
diff --git a/printer/realtime_printer.go b/printer/realtime_printer.go
index cf52e2b..0701ba6 100644
--- a/printer/realtime_printer.go
+++ b/printer/realtime_printer.go
@@ -2,14 +2,95 @@ package printer
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/xgadget-lab/nexttrace/trace"
 )
 
-func RealtimePrinter(res *trace.Result, ttl int) {
-	fmt.Print(ttl + 1)
+type HopInfo int
+
+const (
+	General HopInfo = 0
+	IXP     HopInfo = 1
+	Peer    HopInfo = 2
+	PoP     HopInfo = 3
+	Aboard  HopInfo = 4
+)
+
+func findLatestAvailableHop(res *trace.Result, ttl int, probesIndex int) int {
+	for ttl > 0 {
+		// 查找上一个跃点是不是有效结果
+		ttl--
+		if res.Hops[ttl][probesIndex].Address != nil {
+			return ttl
+		}
+	}
+	// 没找到
+	return -1
+}
+
+func unifyName(name string) string {
+	if name == "China" || name == "CN" {
+		return "中国"
+	} else if name == "Hong kong" || name == "香港" || name == "Central and Western" {
+		return "中国香港"
+	} else if name == "Taiwan" || name == "台湾" {
+		return "中国台湾"
+	} else {
+		return name
+	}
+}
+
+func chinaISPPeer(hostname string) bool {
+	var keyWords = []string{"china", "ct", "cu", "cm", "cnc", "4134", "4837", "4809", "9929"}
+	for _, k := range keyWords {
+		if strings.Contains(strings.ToLower(hostname), k) {
+			return true
+		}
+	}
+	return false
+}
+
+func chinaMainland(h trace.Hop) bool {
+	if unifyName(h.Geo.Country) == "中国" && unifyName(h.Geo.Prov) != "中国香港" && unifyName(h.Geo.Prov) != "中国台湾" {
+		return true
+	} else {
+		return false
+	}
+}
+
+func makeHopsType(res *trace.Result, ttl int) map[int]HopInfo {
+	// 创建一个字典，存放所有当前TTL的跃点类型集合
+	hopProbesMap := make(map[int]HopInfo)
 	for i := range res.Hops[ttl] {
-		HopPrinter(res.Hops[ttl][i])
+		if res.Hops[ttl][i].Address != nil {
+			if availableTTL := findLatestAvailableHop(res, ttl, i); availableTTL != -1 {
+				switch {
+				case strings.Contains(res.Hops[ttl][i].Geo.District, "IXP") || strings.Contains(strings.ToLower(res.Hops[ttl][i].Hostname), "ix"):
+					hopProbesMap[i] = IXP
+				case strings.Contains(res.Hops[ttl][i].Geo.District, "Peer") || chinaISPPeer(res.Hops[ttl][i].Hostname):
+					hopProbesMap[i] = Peer
+				case strings.Contains(res.Hops[ttl][i].Geo.District, "PoP"):
+					hopProbesMap[i] = PoP
+				// 2个有效跃点必须都为有效数据
+				case res.Hops[availableTTL][i].Geo.Country != "LAN Address" && res.Hops[ttl][i].Geo.Country != "LAN Address" &&
+					res.Hops[availableTTL][i].Geo.Country != "" && res.Hops[ttl][i].Geo.Country != "" &&
+					chinaMainland(res.Hops[availableTTL][i]) != chinaMainland(res.Hops[ttl][i]):
+					hopProbesMap[i] = Aboard
+				}
+			} else {
+				hopProbesMap[i] = General
+			}
+		}
 	}
 
+	return hopProbesMap
+}
+
+func RealtimePrinter(res *trace.Result, ttl int) {
+	fmt.Print(ttl + 1)
+	hopsTypeMap := makeHopsType(res, ttl)
+	for i := range res.Hops[ttl] {
+		HopPrinter(res.Hops[ttl][i], hopsTypeMap[i])
+	}
 }