Changes from code review

lib/msf/core/framework.rb

@@ -84,6 +84,7 @@ def initialize(options={})
 
     if options.include?('CustomDnsResolver')
       self.dns_resolver = options['CustomDnsResolver']
+      self.dns_resolver.set_framework(self)
       Rex::Socket._install_global_resolver(self.dns_resolver)
     end
 

lib/msf/ui/console/command_dispatcher/dns.rb

@@ -41,19 +41,19 @@ def commands
   # Tab completion for the dns command
   #
   # @param str [String] the string currently being typed before tab was hit
-  # @param words [Array<String>] the previously completed words on the command line.  words is always
-  # at least 1 when tab completion has reached this stage since the command itself has been completed
+  # @param words [Array<String>] the previously completed words on the command line. The array
+  # contains at least one entry when tab completion has reached this stage since the command itself has been completed
   def cmd_dns_tabs(str, words)
     return if driver.framework.dns_resolver.nil?
 
     if words.length == 1
-      options = ['add','del','remove','flush','print']
+      options = ['add','del','remove','purge','print']
       return options.select { |opt| opt.start_with?(str) }
     end
 
     cmd = words[1]
     case cmd
-    when 'flush','print'
+    when 'purge','print'
       # These commands don't have any arguments
       return
     when 'add'
@@ -95,33 +95,41 @@ def cmd_dns_tabs(str, words)
   end
 
   def cmd_dns_help
-    print_line 'Usage: dns'
-    print_line
     print_line "Manage Metasploit's DNS resolution behaviour"
     print_line
     print_line "Usage:"
-    print_line "  dns [add/remove] [--session <session_id>] [--rule <wildcard DNS entry>] <IP Address> <IP Address> ..."
-    print_line "  dns [get] <hostname>"
-    print_line "  dns [flush]"
+    print_line "  dns [add] [--session <session_id>] [--rule <wildcard DNS entry>] <IP Address> <IP Address> ..."
+    print_line "  dns [remove/del] -i <entry id> [-i <entry id> ...]"
+    print_line "  dns [purge]"
     print_line "  dns [print]"
     print_line
     print_line "Subcommands:"
     print_line "  add - add a DNS resolution entry to resolve certain domain names through a particular DNS server"
     print_line "  remove - delete a DNS resolution entry; 'del' is an alias"
-    print_line "  flush - remove all DNS resolution entries"
+    print_line "  purge - remove all DNS resolution entries"
     print_line "  print - show all active DNS resolution entries"
     print_line
     print_line "Examples:"
-    print_line "  Set the DNS server to be used for *.metasploit.com to 192.168.1.10"
+    print_line "  Display all current DNS nameserver entries"
+    print_line "    dns"
+    print_line "    dns print"
+    print_line
+    print_line "  Set the DNS server(s) to be used for *.metasploit.com to 192.168.1.10"
     print_line "    route add --rule *.metasploit.com 192.168.1.10"
     print_line
-    print_line "  Set the DNS server to be used for *.metasploit.com to 192.168.1.10, but specifically to go through session 2"
+    print_line "  Add multiple entries at once"
+    print_line "    route add --rule *.metasploit.com --rule *.google.com 192.168.1.10 192.168.1.11"
+    print_line
+    print_line "  Set the DNS server(s) to be used for *.metasploit.com to 192.168.1.10, but specifically to go through session 2"
     print_line "    route add --session 2 --rule *.metasploit.com 192.168.1.10"
     print_line
-    print_line "  Delete the above DNS resolution rule"
-    print_line "    route remove --session 2 --rule *.metasploit.com 192.168.1.10"
+    print_line "  Delete the DNS resolution rule with ID 3"
+    print_line "    route remove -i 3"
     print_line
-    print_line "  Set the DNS server to be used for all requests that match no rules"
+    print_line "  Delete multiple entries in one command"
+    print_line "    route remove -i 3 -i 4 -i 5"
+    print_line
+    print_line "  Set the DNS server(s) to be used for all requests that match no rules"
     print_line "    route add 8.8.8.8 8.8.4.4"
     print_line
   end
@@ -217,7 +225,7 @@ def add_dns(*args)
     servers.each do |server|
       driver.framework.dns_resolver.add_nameserver(rules, server, comm_obj)
     end
-    print_good("DNS #{servers.length > 1 ? 'entries' : 'entry'} added")
+    print_good("#{servers.length} DNS #{servers.length > 1 ? 'entries' : 'entry'} added")
   end
 
   #
@@ -238,7 +246,7 @@ def remove_dns(*args)
     print_warning("Some entries were not removed: #{difference.join(', ')}") unless difference.empty?
     if removed.length > 0
       print_good("#{removed.length} DNS #{removed.length > 1 ? 'entries' : 'entry'} removed") 
-      print_dns_set('Deleted entries', ['ID', 'Rules(s)', 'DNS Server', 'Commm channel'], removed.map {|hash| [hash[:id], hash[:wildcard_rules].join(','), hash[:dns_server], prettify_comm(hash[:comm], hash[:dns_server])]})
+      print_dns_set('Deleted entries', removed)
     end
   end
 
@@ -256,11 +264,13 @@ def purge_dns
   def print_dns
     results = driver.framework.dns_resolver.nameserver_entries
     columns = ['ID','Rule(s)', 'DNS Server', 'Comm channel']
-    print_dns_set('Custom nameserver rules', columns, results[0].map {|hash| [hash[:id], hash[:wildcard_rules].join(','), hash[:dns_server], prettify_comm(hash[:comm], hash[:dns_server])]})
+    print_dns_set('Custom nameserver rules', results[0])
 
     # Default nameservers don't include a rule
     columns = ['ID', 'DNS Server', 'Comm channel']
-    print_dns_set('Default nameservers', columns, results[1].map {|hash| [hash[:id], hash[:dns_server], prettify_comm(hash[:comm], hash[:dns_server])]})
+    print_dns_set('Default nameservers', results[1])
+
+    print_line('No custom DNS nameserver entries configured') if results[0].length + results[1].length == 0
   end
 
   private
@@ -285,16 +295,27 @@ def prettify_comm(comm, dns_server)
     end
   end
 
-  def print_dns_set(heading, columns, result_set)
+  def print_dns_set(heading, result_set)
+    return if result_set.length == 0
+    if result_set[0][:wildcard_rules].any?
+      columns = ['ID', 'Rules(s)', 'DNS Server', 'Commm channel']
+    else
+      columns = ['ID', 'DNS Server', 'Commm channel']
+    end
+
     tbl = Table.new(
         Table::Style::Default,
         'Header'  => heading,
         'Prefix'  => "\n",
         'Postfix' => "\n",
         'Columns' => columns
         )
-    result_set.each do |row|
-      tbl << row
+    result_set.each do |hash|
+      if columns.size == 4
+        tbl << [hash[:id], hash[:wildcard_rules].join(','), hash[:dns_server], prettify_comm(hash[:comm], hash[:dns_server])]
+      else
+        tbl << [hash[:id], hash[:dns_server], prettify_comm(hash[:comm], hash[:dns_server])]
+      end
     end
 
     print(tbl.to_s) if tbl.rows.length > 0

lib/rex/proto/dns/custom_nameserver_provider.rb

@@ -161,6 +161,9 @@ def purge
     # @return [Array<Array>] A list of nameservers, each with Rex::Socket options
     #
     def nameservers_for_packet(packet)
+      unless feature_set.enabled?(Msf::FeatureManager::DNS_FEATURE)
+        return super
+      end
       # Leaky abstraction: a packet could have multiple question entries,
       # and each of these could have different nameservers, or travel via
       # different comm channels. We can't allow DNS leaks, so for now, we
@@ -208,6 +211,10 @@ def self.extended(mod)
       mod.init
     end
 
+    def set_framework(framework)
+      self.feature_set = framework.features
+    end
+
     private
     #
     # Is the given wildcard DNS entry valid?
@@ -228,6 +235,7 @@ def matches(domain, pattern)
     attr_accessor :entries_with_rules # Set of custom nameserver entries that specify a rule
     attr_accessor :entries_without_rules # Set of custom nameserver entries that do not include a rule
     attr_accessor :next_id # The next ID to have been allocated to an entry
+    attr_accessor :feature_set
   end
 end
 end

lib/rex/proto/dns/resolver.rb

@@ -112,11 +112,11 @@ def proxies=(prox, timeout_added = 250)
 
     #
     # Find the nameservers to use for a given DNS request
-    # @param dns_message [Dnsruby::Message] The DNS message to be sent
+    # @param _dns_message [Dnsruby::Message] The DNS message to be sent
     #
     # @return [Array<Array>] A list of nameservers, each with Rex::Socket options
     #
-    def nameservers_for_packet(dns_message)
+    def nameservers_for_packet(_dns_message)
       @config[:nameservers].map {|ns| [ns, {}]}
     end
 
@@ -414,7 +414,7 @@ def query(name, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
 
     def supports_udp?(nameserver_results)
       nameserver_results.each do |nameserver, socket_options|
-        comm = socket_options.fetch('Comm') { @config.fetch(:comm) { Rex::Socket::SwitchBoard.best_comm(ns) }}
+        comm = socket_options.fetch('Comm') { @config[:comm] || Rex::Socket::SwitchBoard.best_comm(nameserver) }
         next if comm.nil?
         return false unless comm.supports_udp?
       end