Preliminary fix for #26

Signed-off-by: Sam Therapy <sam@samtherapy.net>

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
+				}
+			}
 		}
 	}
 

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)
 	}

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