diff --git a/cli.go b/cli.go index d12d11c..e810161 100644 --- a/cli.go +++ b/cli.go @@ -148,7 +148,7 @@ func prepareCLI() *cli.App { } // Parse the wildcard arguments, drill style -func parseArgs(args []string) (query.Answers, error) { +func parseArgs(args []string, opts query.Options) (query.Answers, error) { var ( resp query.Response err error @@ -193,7 +193,22 @@ func parseArgs(args []string) (query.Answers, error) { if err != nil { // Query Google by default resp.Answers.Server = "8.8.4.4" } else { - resp.Answers.Server = resolv.Servers[rand.Intn(len(resolv.Servers))] + for _, srv := range resolv.Servers { + if opts.IPv4 { + if strings.Contains(srv, ".") { + resp.Answers.Server = srv + break + } + } else if opts.IPv6 { + if strings.Contains(srv, ":") { + resp.Answers.Server = srv + break + } + } else { + resp.Answers.Server = resolv.Servers[rand.Intn(len(resolv.Servers))] + break + } + } } } diff --git a/cli_test.go b/cli_test.go index 393d070..797ebf3 100644 --- a/cli_test.go +++ b/cli_test.go @@ -37,7 +37,7 @@ func TestArgParse(t *testing.T) { }, } for _, test := range tests { - act, err := parseArgs(test.in) + act, err := parseArgs(test.in, query.Options{}) assert.Nil(t, err) assert.Equal(t, test.want, act) } diff --git a/query.go b/query.go index 40f654a..feb7b4e 100644 --- a/query.go +++ b/query.go @@ -42,7 +42,7 @@ func doQuery(c *cli.Context) error { Reverse: c.Bool("reverse"), Debug: c.Bool("debug"), } - Options.Answers, err = parseArgs(c.Args().Slice()) + Options.Answers, err = parseArgs(c.Args().Slice(), Options) if err != nil { Options.Logger.Error("Unable to parse args") return err