feat: support DNS resolution

go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/jmespath/go-jmespath v0.4.0
 	github.com/json-iterator/go v1.1.12
 	github.com/maja42/goval v1.2.1
+	github.com/miekg/dns v1.0.14
 	github.com/mitchellh/mapstructure v1.4.1
 	github.com/olekukonko/tablewriter v0.0.5
 	github.com/pkg/errors v0.9.1

hrp/cmd/dial.go

@@ -3,12 +3,16 @@ package cmd
 import (
 	"time"
 
+	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 
 	"github.com/httprunner/httprunner/v4/hrp/internal/dial"
 )
 
-var pingOptions dial.PingOptions
+var (
+	pingOptions dial.PingOptions
+	dnsOptions  dial.DnsOptions
+)
 
 var pingCmd = &cobra.Command{
 	Use:   "ping $url",
@@ -22,10 +26,34 @@ var pingCmd = &cobra.Command{
 	},
 }
 
+var dnsCmd = &cobra.Command{
+	Use:   "dns $url",
+	Short: "DNS resolution for different source and record types",
+	Args:  cobra.ExactArgs(1),
+	PreRun: func(cmd *cobra.Command, args []string) {
+		setLogLevel(logLevel)
+	},
+	RunE: func(cmd *cobra.Command, args []string) error {
+		if dnsOptions.DnsSourceType != dial.DnsSourceTypeLocal && dnsOptions.DnsServer != "" {
+			log.Warn().Msg("DNS server not supported for non-local DNS source, ignored")
+		}
+		if dnsOptions.DnsSourceType == dial.DnsSourceTypeHttp && dnsOptions.DnsRecordType == dial.DnsRecordTypeCNAME {
+			log.Warn().Msg("CNAME record not supported for http DNS source, using default record type(A)")
+		}
+		return dial.DoDns(&dnsOptions, args)
+	},
+}
+
 func init() {
 	rootCmd.AddCommand(pingCmd)
 	pingCmd.Flags().IntVarP(&pingOptions.Count, "count", "c", 10, "Stop after sending (and receiving) N packets")
 	pingCmd.Flags().DurationVarP(&pingOptions.Timeout, "timeout", "t", 20*time.Second, "Ping exits after N seconds")
 	pingCmd.Flags().DurationVarP(&pingOptions.Interval, "interval", "i", 1*time.Second, "Wait N seconds between sending each packet")
 	pingCmd.Flags().BoolVar(&pingOptions.SaveTests, "save-tests", false, "Save ping results json")
+
+	rootCmd.AddCommand(dnsCmd)
+	dnsCmd.Flags().IntVar(&dnsOptions.DnsSourceType, "dns-source", 0, "DNS source type\n0: local DNS\n1: http DNS\n2: google DNS")
+	dnsCmd.Flags().IntVar(&dnsOptions.DnsRecordType, "dns-record", 1, "DNS record type\n1: A\n28: AAAA\n5: CNAME")
+	dnsCmd.Flags().StringVar(&dnsOptions.DnsServer, "dns-server", "", "DNS server, only available for local DNS source")
+	dnsCmd.Flags().BoolVar(&dnsOptions.SaveTests, "save-tests", false, "Save DNS resolution results json")
 }

hrp/internal/dial/dns.go

@@ -0,0 +1,250 @@
+package dial
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"net/url"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/miekg/dns"
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
+
+	"github.com/httprunner/httprunner/v4/hrp/internal/builtin"
+)
+
+const (
+	httpDnsUrl   = "https://dig.bdurl.net/q"
+	googleDnsUrl = "https://dns.google/resolve"
+)
+
+const (
+	DnsSourceTypeLocal = iota
+	DnsSourceTypeHttp
+	DnsSourceTypeGoogle
+)
+
+const (
+	DnsRecordTypeA     = 1
+	DnsRecordTypeAAAA  = 28
+	DnsRecordTypeCNAME = 5
+)
+
+type DnsOptions struct {
+	DnsSourceType int
+	DnsRecordType int
+	DnsServer     string
+	SaveTests     bool
+}
+
+type DnsResult struct {
+	DnsList       []string `json:"dnsList"`
+	DnsSource     int      `json:"dnsType"`
+	DnsRecordType int      `json:"dnsRecordType"`
+	DnsServer     string   `json:"dnsServer,omitempty"`
+	Ttl           int      `json:"ttl"`
+	Suc           bool     `json:"suc"`
+	ErrMsg        string   `json:"errMsg"`
+}
+
+type googleDnsResp struct {
+	Answer []googleDnsAnswer `json:"Answer"`
+}
+
+type httpDnsResp struct {
+	Ips []string `json:"ips"`
+	Ttl int      `json:"ttl"`
+}
+
+type googleDnsAnswer struct {
+	Name string `json:"name"`
+	Type int    `json:"type"`
+	TTL  int    `json:"TTL"`
+	Data string `json:"data"`
+}
+
+func ParseIP(s string) (net.IP, int) {
+	ip := net.ParseIP(s)
+	if ip == nil {
+		return nil, 0
+	}
+	for i := 0; i < len(s); i++ {
+		switch s[i] {
+		case '.':
+			return ip, 4
+		case ':':
+			return ip, 6
+		}
+	}
+	return nil, 0
+}
+
+func localDns(src string, dnsRecordType int, dnsServer string) (dnsResult DnsResult, err error) {
+	dnsResult.DnsSource = DnsSourceTypeLocal
+	dnsResult.DnsRecordType = dnsRecordType
+
+	if dnsServer == "" {
+		config, _ := dns.ClientConfigFromFile("/etc/resolv.conf")
+		dnsServer = config.Servers[0]
+	} else {
+		dnsResult.DnsServer = dnsServer
+	}
+
+	_, ipType := ParseIP(dnsServer)
+	if ipType == 4 {
+		dnsServer += ":53"
+	}
+
+	c := dns.Client{
+		Timeout: 5 * time.Second,
+	}
+	m := dns.Msg{}
+
+	m.SetQuestion(src+".", uint16(dnsRecordType))
+	r, _, err := c.Exchange(&m, dnsServer)
+	if err != nil {
+		return
+	}
+	for _, ans := range r.Answer {
+		switch dnsRecordType {
+		case DnsRecordTypeA:
+			record, isType := ans.(*dns.A)
+			if isType {
+				dnsResult.Ttl = int(record.Hdr.Ttl)
+				dnsResult.DnsList = append(dnsResult.DnsList, record.A.String())
+			}
+		case DnsRecordTypeAAAA:
+			record, isType := ans.(*dns.AAAA)
+			if isType {
+				dnsResult.Ttl = int(record.Hdr.Ttl)
+				dnsResult.DnsList = append(dnsResult.DnsList, record.AAAA.String())
+			}
+		case DnsRecordTypeCNAME:
+			record, isType := ans.(*dns.CNAME)
+			if isType {
+				dnsResult.Ttl = int(record.Hdr.Ttl)
+				dnsResult.DnsList = append(dnsResult.DnsList, record.Target)
+			}
+		}
+	}
+	return
+}
+
+func httpDns(url string, dnsRecordType int) (dnsResult DnsResult, err error) {
+	target := httpDnsUrl + "?host=" + url
+	if dnsRecordType == DnsRecordTypeAAAA {
+		target += "&aid=13&f=2"
+	}
+	resp, err := http.Get(target)
+
+	dnsResult.DnsSource = DnsSourceTypeHttp
+	dnsResult.DnsRecordType = dnsRecordType
+
+	if err != nil {
+		return
+	} else {
+		defer resp.Body.Close()
+		var buf []byte
+		buf, err = ioutil.ReadAll(resp.Body)
+		if err != nil {
+			return
+		}
+		var result httpDnsResp
+		err = json.Unmarshal(buf, &result)
+		if err != nil {
+			return
+		}
+		dnsResult.DnsList = result.Ips
+		dnsResult.Ttl = result.Ttl
+	}
+	return
+}
+
+func googleDns(url string, dnsRecordType int) (dnsResult DnsResult, err error) {
+	resp, err := http.Get(googleDnsUrl + "?name=" + url + "&type=" + strconv.Itoa(dnsRecordType))
+
+	dnsResult.DnsSource = DnsSourceTypeGoogle
+	dnsResult.DnsRecordType = dnsRecordType
+
+	if err != nil {
+		return
+	} else {
+		defer resp.Body.Close()
+		var buf []byte
+		buf, err = ioutil.ReadAll(resp.Body)
+		if err != nil {
+			return
+		}
+		var result googleDnsResp
+		err = json.Unmarshal(buf, &result)
+		if err != nil {
+			return
+		}
+		if len(result.Answer) == 0 {
+			return
+		}
+		for _, answer := range result.Answer {
+			if answer.Type == dnsRecordType {
+				dnsResult.Ttl = answer.TTL
+				dnsResult.DnsList = append(dnsResult.DnsList, answer.Data)
+			}
+		}
+
+	}
+	return
+}
+
+func DoDns(dnsOptions *DnsOptions, args []string) (err error) {
+	if len(args) != 1 {
+		return errors.New("there should be one argument")
+	}
+
+	var dnsResult DnsResult
+	defer func() {
+		if dnsOptions.SaveTests {
+			dir, _ := os.Getwd()
+			dnsResultName := fmt.Sprintf("dns_result_%v.json", time.Now().Format("20060102150405"))
+			dnsResultPath := filepath.Join(dir, dnsResultName)
+			err = builtin.Dump2JSON(dnsResult, dnsResultPath)
+			if err != nil {
+				log.Error().Err(err).Msg("save ping result failed")
+			}
+		}
+	}()
+
+	dnsTarget := args[0]
+
+	parsedURL, err := url.Parse(dnsTarget)
+	if err == nil && parsedURL.Host != "" {
+		log.Info().Msgf("parse input url %v and extract host %v", dnsTarget, parsedURL.Host)
+		dnsTarget = strings.Split(parsedURL.Host, ":")[0]
+	}
+	log.Info().Msgf("resolve DNS for %v", dnsTarget)
+	dnsRecordType := dnsOptions.DnsRecordType
+	dnsServer := dnsOptions.DnsServer
+	switch dnsOptions.DnsSourceType {
+	case DnsSourceTypeLocal:
+		dnsResult, err = localDns(dnsTarget, dnsRecordType, dnsServer)
+	case DnsSourceTypeHttp:
+		dnsResult, err = httpDns(dnsTarget, dnsRecordType)
+	case DnsSourceTypeGoogle:
+		dnsResult, err = googleDns(dnsTarget, dnsRecordType)
+	}
+	if err != nil {
+		dnsResult.Suc = false
+		dnsResult.ErrMsg = err.Error()
+		log.Error().Err(err).Msgf("fail to do DNS for %s", dnsTarget, err)
+	} else {
+		dnsResult.Suc = true
+		dnsResult.ErrMsg = ""
+		fmt.Printf("\nDNS resolution done, result IP list: %v\n", dnsResult.DnsList)
+	}
+	return
+}

hrp/internal/dial/ping.go

@@ -98,6 +98,7 @@ func DoPing(pingOptions *PingOptions, args []string) (err error) {
 		pingResult.DebugLog = err.Error()
 		return
 	}
+	fmt.Print(pingResult.DebugLog)
 	stats := pinger.Statistics() // get send/receive/rtt stats
 	pingResult.Ip = pinger.IPAddr().String()
 	pingResult.AvgCost = int(stats.AvgRtt / time.Millisecond)