From 2bd5c6f3da1800065638f0ef9e9dbccfcfd2ab75 Mon Sep 17 00:00:00 2001 From: buyuxiang <347586493@qq.com> Date: Thu, 28 Jul 2022 21:15:03 +0800 Subject: [PATCH] feat: support DNS resolution --- go.mod | 1 + hrp/cmd/dial.go | 30 ++++- hrp/internal/dial/dns.go | 250 ++++++++++++++++++++++++++++++++++++++ hrp/internal/dial/ping.go | 1 + 4 files changed, 281 insertions(+), 1 deletion(-) create mode 100644 hrp/internal/dial/dns.go diff --git a/go.mod b/go.mod index d82800ae..51da4a79 100644 --- a/go.mod +++ b/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 diff --git a/hrp/cmd/dial.go b/hrp/cmd/dial.go index d460d76f..5f394db9 100644 --- a/hrp/cmd/dial.go +++ b/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") } diff --git a/hrp/internal/dial/dns.go b/hrp/internal/dial/dns.go new file mode 100644 index 00000000..7c169742 --- /dev/null +++ b/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 +} diff --git a/hrp/internal/dial/ping.go b/hrp/internal/dial/ping.go index 8b920a3a..29c24295 100644 --- a/hrp/internal/dial/ping.go +++ b/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)