bin/check-dns.rb in sensu-plugins-dns-1.0.0 vs bin/check-dns.rb in sensu-plugins-dns-1.1.0

@@ -29,10 +29,11 @@
 #   for details.
 #
 
 require 'sensu-plugin/check/cli'
 require 'dnsruby'
+require 'ipaddr'
 #
 # DNS
 #
 class DNS < Sensu::Plugin::Check::CLI
   option :domain,
@@ -44,15 +45,27 @@
          description: 'Record type to resolve (A, AAAA, TXT, etc) use PTR for reverse lookup',
          short: '-t RECORD',
          long: '--type RECORD',
          default: 'A'
 
+  option :class,
+         description: 'Record class to resolve (IN, CH, HS, ANY)',
+         short: '-c CLASS',
+         long: '--class CLASS',
+         default: 'IN'
+
   option :server,
          description: 'Server to use for resolution',
          short: '-s SERVER',
          long: '--server SERVER'
 
+  option :port,
+         description: 'Port to use for resolution',
+         short: '-p PORT',
+         long: '--port PORT',
+         proc: proc(&:to_i)
+
   option :result,
          description: 'A positive result entry',
          short: '-r RESULT',
          long: '--result RESULT'
 
@@ -76,14 +89,32 @@
          description: 'Validate dnssec responses',
          short: '-v',
          long: '--validate',
          boolean: true
 
+  option :use_tcp,
+         description: 'Use tcp for resolution',
+         short: '-T',
+         long: '--use-tcp',
+         boolean: true
+
+  option :timeout,
+         description: 'Set timeout for query',
+         short: '-T TIMEOUT',
+         long: '--timeout TIMEOUT',
+         proc: proc(&:to_i),
+         default: 5
+
   def resolve_domain
-    resolv = config[:server].nil? ? Dnsruby::Resolver.new : Dnsruby::Resolver.new(nameserver: [config[:server]])
+    dnsruby_config = {}
+    dnsruby_config[:nameserver] = [config[:server]] unless config[:server].nil?
+    dnsruby_config[:port] = config[:port] unless config[:port].nil?
+    dnsruby_config[:use_tcp] = config[:use_tcp] unless config[:use_tcp].nil?
+    resolv = Dnsruby::Resolver.new(dnsruby_config)
     resolv.do_validation = true if config[:validate]
-    entries = resolv.query(config[:domain], config[:type])
+    entries = resolv.query(config[:domain], config[:type], config[:class])
+    resolv.query_timeout = config[:timeout]
     puts "Entries: #{entries}" if config[:debug]
 
     entries
   end
 
@@ -100,38 +131,57 @@
       end
     end # b.each()
     critical "Resolved #{config[:domain]} #{config[:type]} did not match #{regex}"
   end
 
+  def check_ips(entries)
+    ips = entries.answer.rrsets(config[:type]).flat_map(&:rrs).map(&:address).map(&:to_s)
+    result = IPAddr.new config[:result]
+    if ips.any? { |ip| (IPAddr.new ip) == result }
+      ok "Resolved #{entries.security_level} #{config[:domain]} #{config[:type]} included #{config[:result]}"
+    else
+      critical "Resolved #{config[:domain]} #{config[:type]} did not include #{config[:result]}"
+    end
+  end
+
   def run
     unknown 'No domain specified' if config[:domain].nil?
 
     begin
       entries = resolve_domain
     rescue Dnsruby::NXDomain
       output = "Could not resolve #{config[:domain]} #{config[:type]} record"
       critical(output)
       return
     rescue => e
-      output = "Couldn not resolve  #{config[:domain]}: #{e}"
+      output = "Could not resolve  #{config[:domain]}: #{e}"
       config[:warn_only] ? warning(output) : critical(output)
       return
     end
     puts entries.answer if config[:debug]
     if entries.answer.length.zero?
       output = "Could not resolve #{config[:domain]} #{config[:type]} record"
       config[:warn_only] ? warning(output) : critical(output)
     elsif config[:result]
-      b = if entries.answer.count > 1
-            entries.answer.rrsets(config[:type].to_s).to_s
-          else
-            entries.answer.first.to_s
-          end
-      if b.include?(config[:result])
-        ok "Resolved #{entries.security_level} #{config[:domain]} #{config[:type]} included #{config[:result]}"
+      # special logic for checking ipaddresses with result
+      # mostly for ipv6 but decided to use the same logic for
+      # consistency reasons
+      if config[:type] == 'A' || config[:type] == 'AAAA'
+        check_ips(entries)
+      # non ip type
       else
-        critical "Resolved #{config[:domain]} #{config[:type]} did not include #{config[:result]}"
+        b = if entries.answer.count > 1
+              entries.answer.rrsets(config[:type].to_s).to_s
+            else
+              entries.answer.first.to_s
+            end
+        if b.include?(config[:result])
+          ok "Resolved #{entries.security_level} #{config[:domain]} #{config[:type]} included #{config[:result]}"
+        else
+          critical "Resolved #{config[:domain]} #{config[:type]} did not include #{config[:result]}"
+        end
       end
+
     elsif config[:regex]
       check_against_regex(entries, Regexp.new(config[:regex]))
 
     elsif config[:validate]
       if entries.security_level != 'SECURE'