add query type and class option

Author: wafuwafu13

check-dns/lib/check_dns.go

@@ -5,17 +5,20 @@ import (
 	"net"
 	"os"
 	"strconv"
+	"strings"
 
 	"github.com/jessevdk/go-flags"
 	"github.com/mackerelio/checkers"
 	"github.com/miekg/dns"
 )
 
 type dnsOpts struct {
-	Host   string `short:"H" long:"host" required:"true" description:"The name or address you want to query"`
-	Server string `short:"s" long:"server" description:"DNS server you want to use for the lookup"`
-	Port   int    `short:"p" long:"port" default:"53" description:"Port number you want to use"`
-	Norec  bool   `long:"norec" description:"Set not recursive mode"`
+	Host       string `short:"H" long:"host" required:"true" description:"The name or address you want to query"`
+	Server     string `short:"s" long:"server" description:"DNS server you want to use for the lookup"`
+	Port       int    `short:"p" long:"port" default:"53" description:"Port number you want to use"`
+	QueryType  string `short:"q" long:"querytype" default:"A" description:"DNS record query type where TYPE =(A, AAAA, SRV, TXT, MX, ANY)"`
+	QueryClass string `short:"c" long:"queryclass" default:"IN" description:"DNS record class type where TYPE =(IN, CS, CH, HS, NONE, ANY)"`
+	Norec      bool   `long:"norec" description:"Set not recursive mode"`
 }
 
 // Do the plugin
@@ -48,13 +51,22 @@ func (opts *dnsOpts) run() *checkers.Checker {
 	}
 	nameserver = net.JoinHostPort(nameserver, strconv.Itoa(opts.Port))
 
+	queryType, ok := dns.StringToType[strings.ToUpper(opts.QueryType)]
+	if !ok {
+		return checkers.Critical(fmt.Sprintf("%s is invalid queryType", opts.QueryType))
+	}
+	queryClass, ok := dns.StringToClass[strings.ToUpper(opts.QueryClass)]
+	if !ok {
+		return checkers.Critical(fmt.Sprintf("%s is invalid queryClass", opts.QueryClass))
+	}
+
 	c := new(dns.Client)
 	m := &dns.Msg{
 		MsgHdr: dns.MsgHdr{
 			RecursionDesired: !opts.Norec,
 			Opcode:           dns.OpcodeQuery,
 		},
-		Question: []dns.Question{{Name: dns.Fqdn(opts.Host), Qtype: dns.TypeA, Qclass: uint16(dns.ClassINET)}},
+		Question: []dns.Question{{Name: dns.Fqdn(opts.Host), Qtype: queryType, Qclass: uint16(queryClass)}},
 	}
 	m.Id = dns.Id()
 

check-dns/lib/check_dns_test.go

@@ -26,47 +26,67 @@ func TestCheckDns(t *testing.T) {
 	tests := []struct {
 		args        []string
 		want_status checkers.Status
-		want_msg    string
+		want_msg    []string
 	}{
 		{
 			[]string{"-H", "example.com"},
 			checkers.OK,
-			"status: NOERROR",
+			[]string{"status: NOERROR"},
 		},
 		{
 			[]string{"-H", "example.com", "--norec"},
 			checkers.OK,
-			"status: NOERROR",
+			[]string{"status: NOERROR"},
 		},
 		{
 			[]string{"-H", "exampleeeee.com"},
 			checkers.CRITICAL,
-			"status: NXDOMAIN",
+			[]string{"status: NXDOMAIN"},
 		},
 		{
 			[]string{"-H", "example.com", "-s", "8.8.8.8"},
 			checkers.OK,
-			"status: NOERROR",
+			[]string{"status: NOERROR"},
 		},
 		{
 			[]string{"-H", "exampleeeee.com", "-s", "8.8.8.8"},
 			checkers.CRITICAL,
-			"status: NXDOMAIN",
+			[]string{"status: NXDOMAIN"},
 		},
 		{
 			[]string{"-H", "exampleeeee.com", "-s", "8.8.8"},
 			checkers.CRITICAL,
-			"timeout",
+			[]string{"timeout"},
 		},
 		{
 			[]string{"-H", "jprs.co.jp", "-s", "202.11.16.49", "--norec"},
 			checkers.OK,
-			"status: NOERROR",
+			[]string{"status: NOERROR"},
 		},
 		{
 			[]string{"-H", "www.google.com", "-s", "202.11.16.49", "--norec"},
 			checkers.CRITICAL,
-			"status: REFUSED",
+			[]string{"status: REFUSED"},
+		},
+		{
+			[]string{"-H", "example.com", "-s", "8.8.8.8", "-q", "AAAA"},
+			checkers.OK,
+			[]string{"status: NOERROR", "AAAA"},
+		},
+		{
+			[]string{"-H", "example.com", "-s", "8.8.8.8", "-q", "AAA"},
+			checkers.CRITICAL,
+			[]string{"AAA is invalid queryType"},
+		},
+		{
+			[]string{"-H", "example.com", "-s", "8.8.8.8", "-c", "IN"},
+			checkers.OK,
+			[]string{"status: NOERROR"},
+		},
+		{
+			[]string{"-H", "example.com", "-s", "8.8.8.8", "-c", "INN"},
+			checkers.CRITICAL,
+			[]string{"INN is invalid queryClass"},
 		},
 	}
 
@@ -84,8 +104,10 @@ func TestCheckDns(t *testing.T) {
 
 		assert.Equal(t, tt.want_status, ckr.Status)
 
-		if !strings.Contains(ckr.Message, tt.want_msg) {
-			t.Errorf("%s is not incleded in message", tt.want_msg)
+		for _, want := range tt.want_msg {
+			if !strings.Contains(ckr.Message, want) {
+				t.Errorf("%s is not incleded in message", want)
+			}
 		}
 	}
 }