lib/ronin/cli/commands/dns.rb in ronin-2.0.3 vs lib/ronin/cli/commands/dns.rb in ronin-2.0.4

@@ -48,26 +48,11 @@
 
         usage '[options] {HOST ... | --file FILE}'
 
         option :type, short: '-t',
                       value: {
-                        type: {
-                          A:     :a,
-                          AAAA:  :aaaa,
-                          ANY:   :any,
-                          CNAME: :cname,
-                          HINFO: :hinfo,
-                          LOC:   :loc,
-                          MINFO: :minfo,
-                          MX:    :mx,
-                          NS:    :ns,
-                          PTR:   :ptr,
-                          SOA:   :soa,
-                          SRV:   :srv,
-                          TXT:   :txt,
-                          WKS:   :wks
-                        }
+                        type: RECORD_TYPES
                       },
                       desc: 'Queries a specific type of DNS record'
 
         argument :host, required: true,
                         desc:     'The host name to query'
@@ -81,9 +66,26 @@
         #
         # @param [String] host
         #
         def process_value(host)
           print_records(query_records(host))
+        end
+
+        #
+        # Queries the records for the given host name.
+        #
+        # @param [String] host
+        #   The host name to query.
+        #
+        # @return [Array<Resolv::DNS::Resource>]
+        #   The returned DNS resource records.
+        #
+        def query_records(host)
+          if options[:type]
+            resolver.get_records(host,options[:type].downcase)
+          else
+            resolver.get_a_records(host) + resolver.get_aaaa_records(host)
+          end
         end
 
       end
     end
   end