From f3858fbce0de1f617b0495b7488bf445b6758f66 Mon Sep 17 00:00:00 2001
From: buyuxiang <347586493@qq.com>
Date: Tue, 2 Aug 2022 10:55:45 +0800
Subject: [PATCH] using original traceroute command for dial

---
 docs/CHANGELOG.md               |  1 +
 hrp/cmd/dial.go                 |  2 +
 hrp/internal/dial/traceroute.go | 92 +++++++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index 5b1128b6..6ee78e52 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -14,6 +14,7 @@
 - fix: optimize websocket step initialization
 - fix: reuse plugin instance if already initialized
 - fix: deep copy api step to avoid data racing
+- feat: support ping/dns/traceroute for dial test
 
 ## v4.1.6 (2022-07-04)
 
diff --git a/hrp/cmd/dial.go b/hrp/cmd/dial.go
index 68ffc3d5..1eeb0881 100644
--- a/hrp/cmd/dial.go
+++ b/hrp/cmd/dial.go
@@ -71,5 +71,7 @@ func init() {
 	dnsCmd.Flags().BoolVar(&dnsOptions.SaveTests, "save-tests", false, "Save DNS resolution result as json")
 
 	rootCmd.AddCommand(traceRouteCmd)
+	traceRouteCmd.Flags().IntVarP(&traceRouteOptions.MaxTTL, "max-hops", "m", 30, "Set the max number of hops (max TTL to be reached)")
+	traceRouteCmd.Flags().IntVarP(&traceRouteOptions.Queries, "queries", "q", 1, "Set the number of probes per each hop")
 	traceRouteCmd.Flags().BoolVar(&traceRouteOptions.SaveTests, "save-tests", false, "Save traceroute result as json")
 }
diff --git a/hrp/internal/dial/traceroute.go b/hrp/internal/dial/traceroute.go
index 5f19c4a1..f9d6550c 100644
--- a/hrp/internal/dial/traceroute.go
+++ b/hrp/internal/dial/traceroute.go
@@ -1,10 +1,14 @@
 package dial
 
 import (
+	"bufio"
 	"fmt"
 	"net/url"
 	"os"
+	"os/exec"
 	"path/filepath"
+	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
@@ -17,7 +21,16 @@ import (
 	"github.com/httprunner/httprunner/v4/hrp/internal/json"
 )
 
+var (
+	regexIPAddr            = regexp.MustCompile(`([\d.]+)`)
+	regexElapsedTime       = regexp.MustCompile(`(\d+\.\d+)`)
+	regexTraceroutePass    = regexp.MustCompile(fmt.Sprintf(`(\d+)[\s*]+(\S+)\s+\(%s\)\s+%s\s+ms`, regexIPAddr, regexElapsedTime))
+	regexTracerouteFailure = regexp.MustCompile(`(\d+)[\s*]+$`)
+)
+
 type TraceRouteOptions struct {
+	MaxTTL    int
+	Queries   int
 	SaveTests bool
 }
 
@@ -48,6 +61,85 @@ func DoTraceRoute(traceRouteOptions *TraceRouteOptions, args []string) (err erro
 	if len(args) != 1 {
 		return errors.New("there should be one argument")
 	}
+	var traceRouteResult TraceRouteResult
+	defer func() {
+		if traceRouteOptions.SaveTests {
+			dir, _ := os.Getwd()
+			traceRouteResultName := fmt.Sprintf("traceroute_result_%v.json", time.Now().Format("20060102150405"))
+			traceRouteResultPath := filepath.Join(dir, traceRouteResultName)
+			err = builtin.Dump2JSON(traceRouteResult, traceRouteResultPath)
+			if err != nil {
+				log.Error().Err(err).Msg("save traceroute result failed")
+			}
+		}
+	}()
+
+	traceRouteTarget := args[0]
+	parsedURL, err := url.Parse(traceRouteTarget)
+	if err == nil && parsedURL.Host != "" {
+		log.Info().Msgf("parse input url %v and extract host %v", traceRouteTarget, parsedURL.Host)
+		traceRouteTarget = strings.Split(parsedURL.Host, ":")[0]
+	}
+
+	cmd := exec.Command("traceroute", "-m", strconv.Itoa(traceRouteOptions.MaxTTL),
+		"-q", strconv.Itoa(traceRouteOptions.Queries), traceRouteTarget)
+	stdout, _ := cmd.StdoutPipe()
+
+	startT := time.Now()
+	defer func() {
+		log.Info().Msgf("for target %s, traceroute costs %v", traceRouteTarget, time.Since(startT))
+	}()
+
+	log.Info().Msgf("start to traceroute %v", traceRouteTarget)
+	err = cmd.Start()
+	if err != nil {
+		traceRouteResult.Suc = false
+		traceRouteResult.ErrMsg = "execute traceroute failed"
+		log.Error().Err(err).Msg("start command failed")
+		return
+	}
+
+	scanner := bufio.NewScanner(stdout)
+	for scanner.Scan() {
+		hopLine := scanner.Text()
+		fmt.Println(hopLine)
+		passLine := regexTraceroutePass.FindStringSubmatch(hopLine)
+		if len(passLine) == 5 {
+			hopID, _ := strconv.Atoi(passLine[1])
+			traceRouteResult.Details = append(traceRouteResult.Details, TraceRouteResultNode{
+				Id:   hopID,
+				Ip:   passLine[3],
+				Time: passLine[4],
+			})
+			traceRouteResult.Suc = true
+			continue
+		}
+		failureLine := regexTracerouteFailure.FindStringSubmatch(hopLine)
+		if len(failureLine) == 2 {
+			hopID, _ := strconv.Atoi(failureLine[1])
+			traceRouteResult.Details = append(traceRouteResult.Details, TraceRouteResultNode{
+				Id: hopID,
+			})
+			continue
+		}
+	}
+	hopCount := len(traceRouteResult.Details)
+	traceRouteResult.IP = traceRouteResult.Details[hopCount-1].Ip
+	err = cmd.Wait()
+	if err != nil {
+		traceRouteResult.Suc = false
+		traceRouteResult.ErrMsg = "wait traceroute finish failed"
+		log.Error().Err(err).Msg("wait command failed")
+		return
+	}
+	return
+}
+
+// DoTraceRouteSDK with golang SDK, which needs root privilege
+func DoTraceRouteSDK(traceRouteOptions *TraceRouteOptions, args []string) (err error) {
+	if len(args) != 1 {
+		return errors.New("there should be one argument")
+	}
 
 	var traceRouteResult TraceRouteResult
 	defer func() {