From dd209deeb310948cbf7f00b25ae75c7f2a4b9dab Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Thu, 2 Nov 2023 19:44:04 +1100
Subject: [PATCH 01/14] Initial syntax handling for DNS command

---
 lib/msf/core/feature_manager.rb              |   7 +
 lib/msf/ui/console/command_dispatcher/dns.rb | 151 +++++++++++++++++++
 lib/msf/ui/console/driver.rb                 |   3 +-
 3 files changed, 160 insertions(+), 1 deletion(-)
 create mode 100755 lib/msf/ui/console/command_dispatcher/dns.rb

diff --git a/lib/msf/core/feature_manager.rb b/lib/msf/core/feature_manager.rb
index 167e41d549a6..d58443ac3c61 100644
--- a/lib/msf/core/feature_manager.rb
+++ b/lib/msf/core/feature_manager.rb
@@ -20,6 +20,7 @@ class FeatureManager
     MANAGER_COMMANDS = 'manager_commands'
     METASPLOIT_PAYLOAD_WARNINGS = 'metasploit_payload_warnings'
     DEFER_MODULE_LOADS = 'defer_module_loads'
+    DNS_FEATURE = 'dns_feature'
     DEFAULTS = [
       {
         name: WRAPPED_TABLES,
@@ -53,6 +54,12 @@ class FeatureManager
         description: 'When enabled will not eagerly load all modules',
         requires_restart: true,
         default_value: false
+      }.freeze,
+      {
+        name: DNS_FEATURE,
+        description: 'When enabled, allows configuration of DNS resolution behaviour in Metasploit',
+        requires_restart: false,
+        default_value: false
       }.freeze
     ].freeze
 
diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
new file mode 100755
index 000000000000..05fa6a648f54
--- /dev/null
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -0,0 +1,151 @@
+# -*- coding: binary -*-
+
+class Msf::Ui::Console::CommandDispatcher::DNS
+
+  include Msf::Ui::Console::CommandDispatcher
+
+  @@add_opts = Rex::Parser::Arguments.new(
+    ['-r', '--rule'] => [true, 'Set a DNS wildcard entry to match against' ],
+    ['-s', '--session'] => [true, 'Force the DNS request to occur over a particular channel (override routing rules)' ],
+  )
+
+  def initialize(driver)
+    super
+  end
+
+  def name
+    'DNS'
+  end
+
+  def commands
+    commands = {}
+
+    if framework.features.enabled?(Msf::FeatureManager::DNS_FEATURE)
+      commands = {
+        'dns'        => "Manage Metasploit's DNS resolving behaviour"
+      }
+    end
+    commands
+  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 [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 "  get - display the DNS server(s) and communication channel that would be used for a given target"
+    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 "    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 "    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
+    print_line "  Set the DNS server to be used for all requests that match no rules"
+    print_line "    route add 8.8.8.8 8.8.4.4"
+    print_line
+    print_line "  Display the DNS server that would be used for the given domain name"
+    print_line "    route get subdomain.metasploit.com"
+    print_line
+  end
+
+  #
+  # Manage Metasploit's DNS resolution rules
+  #
+  def cmd_dns(*args)
+    args << 'print' if args.length == 0
+    # Short-circuit help
+    if args.delete("-h") || args.delete("--help")
+      cmd_dns_help
+      return
+    end
+
+    action = args.shift
+    begin
+      case action
+      when "add"
+        add_dns(*args)
+      when "remove", "del"
+        remove_dns(*args)
+      when "purge"
+        purge_dns(*args)
+      when "print"
+        print_dns(*args)
+      when "help"
+        cmd_dns_help
+      end
+    rescue ::ArgumentError => e
+      print_error(e.message)
+    end
+  end
+
+  def add_dns(*args)
+    rules = []
+    comm = nil
+    servers = []
+    @@add_opts.parse(args) do |opt, idx, val|
+      unless servers.empty? || opt.nil?
+        raise ::ArgumentError.new("Invalid command near #{opt}")
+      end
+      case opt
+      when '--rule', '-r'
+        if val.nil?
+          raise ::ArgumentError.new('No rule specified')
+        end
+
+        rules << val
+      when '--session', '-s'
+        if val.nil?
+          raise ::ArgumentError.new('No session specified')
+        end
+
+        unless comm.nil?
+          raise ::ArgumentError.new('Only one session can be specified')
+        end
+
+        comm = val
+      when nil
+        servers << val
+      else
+        raise ::ArgumentError.new("Unknown flag: #{opt}")
+      end
+    end
+
+    # The remaining args should be the DNS servers
+
+    if servers.length < 1
+      raise ::ArgumentError.new("You must specify at least one DNS server")
+    end
+
+    servers.each do |host|
+      unless host =~ Resolv::IPv4::Regex || 
+             host =~ Resolv::IPv6::Regex
+        raise ::ArgumentError.new("Invalid DNS server: #{host}")
+      end
+    end
+  end
+
+  def remove_dns(*args)
+  end
+
+  def purge_dns(*args)
+  end
+
+  def print_dns(*args)
+  end
+end
diff --git a/lib/msf/ui/console/driver.rb b/lib/msf/ui/console/driver.rb
index 3f97edca68d6..9edca71ee62c 100644
--- a/lib/msf/ui/console/driver.rb
+++ b/lib/msf/ui/console/driver.rb
@@ -29,7 +29,8 @@ class Driver < Msf::Ui::Driver
     CommandDispatcher::Resource,
     CommandDispatcher::Db,
     CommandDispatcher::Creds,
-    CommandDispatcher::Developer
+    CommandDispatcher::Developer,
+    CommandDispatcher::DNS
   ]
 
   #

From a7c4b297488534b99087c99e38281183b622d352 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Tue, 7 Nov 2023 06:32:03 +1100
Subject: [PATCH 02/14] Register nameserver with framework

---
 lib/msf/core/framework.rb                     | 16 ++++
 lib/msf/ui/console/command_dispatcher/dns.rb  | 81 ++++++++++++++++--
 .../proto/dns/custom_nameserver_provider.rb   | 85 +++++++++++++++++++
 lib/rex/proto/dns/resolver.rb                 | 26 ++++--
 4 files changed, 195 insertions(+), 13 deletions(-)
 create mode 100755 lib/rex/proto/dns/custom_nameserver_provider.rb

diff --git a/lib/msf/core/framework.rb b/lib/msf/core/framework.rb
index 5aaa76e3b00b..6cb122341fe5 100644
--- a/lib/msf/core/framework.rb
+++ b/lib/msf/core/framework.rb
@@ -81,6 +81,7 @@ def initialize(options={})
     # Configure the SSL certificate generator
     require 'msf/core/cert_provider'
     Rex::Socket::Ssl.cert_provider = Msf::Ssl::CertProvider
+    initialize_dns_resolver
 
     subscriber = FrameworkEventSubscriber.new(self)
     events.add_exploit_subscriber(subscriber)
@@ -90,6 +91,16 @@ def initialize(options={})
     events.add_ui_subscriber(subscriber)
   end
 
+  def initialize_dns_resolver
+    self.dns_resolver = Rex::Proto::DNS::CachedResolver.new
+    self.dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    Rex::Socket._install_global_resolver(self.dns_resolver)
+  end
+
+  def dns_resolver
+    self.dns_resolver
+  end
+
   def inspect
     "#<Framework (#{sessions.length} sessions, #{jobs.length} jobs, #{plugins.length} plugins#{db.active ? ", #{db.driver} database active" : ""})>"
   end
@@ -147,6 +158,10 @@ def version
     Version
   end
 
+  #
+  # DNS resolver for the framework
+  #
+  attr_reader   :dns_resolver
   #
   # Event management interface for registering event handler subscribers and
   # for interacting with the correlation engine.
@@ -278,6 +293,7 @@ def eicar_corrupted?
   #   @return [Hash]
   attr_accessor :options
 
+  attr_writer   :dns_resolver #:nodoc:
   attr_writer   :events # :nodoc:
   attr_writer   :modules # :nodoc:
   attr_writer   :datastore # :nodoc:
diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index 05fa6a648f54..3d1c6ec3897a 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -1,6 +1,11 @@
 # -*- coding: binary -*-
 
-class Msf::Ui::Console::CommandDispatcher::DNS
+module Msf
+module Ui
+module Console
+module CommandDispatcher
+
+class DNS
 
   include Msf::Ui::Console::CommandDispatcher
 
@@ -83,9 +88,9 @@ def cmd_dns(*args)
       when "remove", "del"
         remove_dns(*args)
       when "purge"
-        purge_dns(*args)
+        purge_dns
       when "print"
-        print_dns(*args)
+        print_dns
       when "help"
         cmd_dns_help
       end
@@ -133,19 +138,81 @@ def add_dns(*args)
     end
 
     servers.each do |host|
-      unless host =~ Resolv::IPv4::Regex || 
-             host =~ Resolv::IPv6::Regex
+      unless Rex::Socket.is_ip_addr?(host)
         raise ::ArgumentError.new("Invalid DNS server: #{host}")
       end
     end
+
+    unless comm.nil?
+      raise ::ArgumentError.new("Not a valid number: #{comm}") unless comm =~ /^\d+$/
+      comm_int = comm.to_i
+      raise ::ArgumentError.new("Session does not exist: #{comm}") unless driver.framework.sessions.include?(comm_int)
+
+    end
+
+    # Split each DNS server entry up into a separate entry
+    servers.each do |server|
+      driver.framework.dns_resolver.add_nameserver(rules, server, comm_int)
+    end
   end
 
   def remove_dns(*args)
   end
 
-  def purge_dns(*args)
+  def purge_dns
+    driver.framework.dns_resolver.purge
   end
 
-  def print_dns(*args)
+  def print_dns
+    results = driver.framework.dns_resolver.nameserver_entries
+    columns = ['ID','Rule(s)', 'DNS Server(s)', '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])]})
+
+    # Default nameservers don't include a rule
+    columns = ['ID', 'DNS Server(s)', 'Comm channel']
+    print_dns_set('Default nameservers', columns, results[1].map {|hash| [hash[:id], hash[:dns_server], prettify_comm(hash[:comm], hash[:dns_server])]})
   end
+
+  private
+
+  def prettify_comm(comm, dns_server)
+    if comm.nil?
+      channel = Rex::Socket::SwitchBoard.best_comm(dns_server)
+      if channel.nil?
+        comm_text = nil
+      else
+        comm_text = "Session #{channel.sid} (auto)"
+      end
+    else
+      if driver.framework.sessions.include?(comm)
+        comm_text = "Session #{comm}"
+      else
+        comm_text = "Broken session (#{comm})"
+      end
+    end
+  end
+
+  def print_dns_set(heading, columns, result_set)
+    tbl = Table.new(
+        Table::Style::Default,
+        'Header'  => heading,
+        'Prefix'  => "\n",
+        'Postfix' => "\n",
+        'Columns' => columns
+        )
+    result_set.each do |row|
+      tbl << row
+    end
+
+    print(tbl.to_s) if tbl.rows.length > 0
+  end
+
+  def resolver
+    self.driver.framework.dns_resolver
+  end
+end
+
+end
+end
 end
+end
\ No newline at end of file
diff --git a/lib/rex/proto/dns/custom_nameserver_provider.rb b/lib/rex/proto/dns/custom_nameserver_provider.rb
new file mode 100755
index 000000000000..3b7ec6f63db8
--- /dev/null
+++ b/lib/rex/proto/dns/custom_nameserver_provider.rb
@@ -0,0 +1,85 @@
+module Rex
+module Proto
+module DNS
+
+  ##
+  # Provides a DNS resolver the ability to use different nameservers
+  # for different requests, based on the domain being queried.
+  ##
+  module CustomNameserverProvider
+
+    def init
+      self.entries_with_rules = []
+      self.entries_without_rules = []
+      self.next_id = 0
+    end
+
+    # Add a custom nameserver entry to the custom provider
+    # @param [wildcard_rules] Array<String> The wildcard rules to match a DNS request against
+    # @param [dns_server] Array<String> The list of IP addresses that would be used for this custom rule
+    # @param comm [Integer] The communication channel to be used for these DNS requests
+    def add_nameserver(wildcard_rules, dns_server, comm)
+      entry = {
+        :wildcard_rules => wildcard_rules,
+        :dns_server => dns_server,
+        :comm => comm,
+        :id => self.next_id
+      }
+      self.next_id += 1
+      if wildcard_rules.empty?
+        entries_without_rules << entry
+      else
+        entries_with_rules << entry
+      end
+    end
+
+    #
+    # The custom nameserver entries that have been configured
+    # @return [Array<Array>] An array containing two elements: The entries with rules, and the entries without rules
+    def nameserver_entries
+      [entries_with_rules, entries_without_rules]
+    end
+
+    def purge
+      init
+    end
+
+    def nameservers_for_packet(packet)
+      name = packet.question.qName
+      dns_servers = []
+
+      self.entries_with_rules.each do |entry|
+        entry[:wildcard_rules].each do |rule|
+          if matches(name, rule)
+            dns_servers.concat([entry[:dns_server], entry[:comm]])
+            break
+          end
+        end
+      end
+
+      # Only look at the rule-less entries if no rules were found (avoids DNS leaks)
+      if dns_servers.empty?
+        self.entries_without_rules.each do |entry|
+          dns_servers.concat([entry[:dns_server], entry[:comm]])
+        end
+      end
+      dns_servers.uniq!
+    end
+
+    def self.extended(mod)
+      mod.init
+    end
+
+    private
+
+    def matches(domain, pattern)
+      true
+    end
+
+    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
+  end
+end
+end
+end
\ No newline at end of file
diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index 8551ae331d61..6f7355e1fd6e 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/lib/rex/proto/dns/resolver.rb
@@ -110,6 +110,15 @@ def proxies=(prox, timeout_added = 250)
       end
     end
 
+    #
+    # Find the nameservers to use for a given DNS request
+    # @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)
+      @config[:nameservers].map {|ns| [ns, {}]}
+    end
+
     #
     # Send DNS request over appropriate transport and process response
     #
@@ -119,10 +128,6 @@ def proxies=(prox, timeout_added = 250)
     #
     # @return [Dnsruby::Message] DNS response
     def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
-      if @config[:nameservers].size == 0
-        raise ResolverError, "No nameservers specified!"
-      end
-
       method = self.use_tcp? ? :send_tcp : :send_udp
 
       case argument
@@ -136,6 +141,10 @@ def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
         packet = Rex::Proto::DNS::Packet.encode_drb(net_packet)
       end
 
+      if nameservers_for_packet(packet).size == 0
+        raise ResolverError, "No nameservers specified!"
+      end
+
       # Store packet_data for performance improvements,
       # so methods don't keep on calling Packet#encode
       packet_data = packet.encode
@@ -195,7 +204,8 @@ def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
     def send_tcp(packet,packet_data,prox = @config[:proxies])
       ans = nil
       length = [packet_data.size].pack("n")
-      @config[:nameservers].each do |ns|
+      nameservers = nameservers_for_packet(packet)
+      nameservers.each do |ns, socket_options|
         begin
           socket = nil
           @config[:tcp_timeout].timeout do
@@ -208,6 +218,8 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
                   'Context' => @config[:context],
                   'Comm' => @config[:comm]
                 }
+                config.update(socket_options)
+
                 if @config[:source_port] > 0
                   config['LocalPort'] = @config[:source_port]
                 end
@@ -289,7 +301,8 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
     def send_udp(packet,packet_data)
       ans = nil
       response = ""
-      @config[:nameservers].each do |ns|
+      nameservers = nameservers_for_packet(packet)
+      nameservers.each do |ns, socket_options|
         begin
           @config[:udp_timeout].timeout do
             begin
@@ -299,6 +312,7 @@ def send_udp(packet,packet_data)
                 'Context' => @config[:context],
                 'Comm' => @config[:comm]
               }
+              config.update(socket_options)
               if @config[:source_port] > 0
                 config['LocalPort'] = @config[:source_port]
               end

From 21f3335c315730ea1c51188d4de1f8900289a168 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Tue, 7 Nov 2023 11:10:37 +1100
Subject: [PATCH 03/14] Fully integrated Rex-socket-friendly DNS

---
 lib/msf/ui/console/command_dispatcher/dns.rb  | 49 +++++++++++---
 .../proto/dns/custom_nameserver_provider.rb   | 64 +++++++++++++++----
 lib/rex/proto/dns/resolver.rb                 | 30 ++++++---
 3 files changed, 110 insertions(+), 33 deletions(-)

diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index 3d1c6ec3897a..c37963ab3854 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -14,6 +14,10 @@ class DNS
     ['-s', '--session'] => [true, 'Force the DNS request to occur over a particular channel (override routing rules)' ],
   )
 
+  @@remove_opts = Rex::Parser::Arguments.new(
+    ['-i'] => [true, 'Index to remove']
+  )
+
   def initialize(driver)
     super
   end
@@ -109,9 +113,8 @@ def add_dns(*args)
       end
       case opt
       when '--rule', '-r'
-        if val.nil?
-          raise ::ArgumentError.new('No rule specified')
-        end
+        raise ::ArgumentError.new('No rule specified') if val.nil?
+        raise ::ArgumentError.new("Invalid rule: #{val}") unless valid_rule(val)
 
         rules << val
       when '--session', '-s'
@@ -143,26 +146,50 @@ def add_dns(*args)
       end
     end
 
+    comm_obj = nil
+
     unless comm.nil?
       raise ::ArgumentError.new("Not a valid number: #{comm}") unless comm =~ /^\d+$/
       comm_int = comm.to_i
       raise ::ArgumentError.new("Session does not exist: #{comm}") unless driver.framework.sessions.include?(comm_int)
-
+      comm_obj = driver.framework.sessions[comm_int]
     end
 
     # Split each DNS server entry up into a separate entry
     servers.each do |server|
-      driver.framework.dns_resolver.add_nameserver(rules, server, comm_int)
+      driver.framework.dns_resolver.add_nameserver(rules, server, comm_obj)
     end
   end
 
+  #
+  # Is the given wildcard DNS entry valid?
+  def valid_rule(rule)
+    rule =~ /^(\*\.)?([a-z\d][a-z\d-]*[a-z\d]\.)+[a-z]+$/
+  end
+
+  #
+  # Remove all matching user-configured DNS entries
   def remove_dns(*args)
+    remove_ids = []
+    @@remove_opts.parse(args) do |opt, idx, val|
+      case opt
+      when '-i'
+        raise ::ArgumentError.new("Not a valid number: #{val}") unless val =~ /^\d+$/
+        remove_ids << val.to_i
+      end
+    end
+
+    driver.framework.dns_resolver.remove_ids(remove_ids)
   end
 
+  #
+  # Delete all user-configured DNS settings
   def purge_dns
     driver.framework.dns_resolver.purge
   end
 
+  #
+  # Display the user-configured DNS settings
   def print_dns
     results = driver.framework.dns_resolver.nameserver_entries
     columns = ['ID','Rule(s)', 'DNS Server(s)', 'Comm channel']
@@ -175,19 +202,21 @@ def print_dns
 
   private
 
+  #
+  # Get user-friendly text for displaying the session that this entry would go through
   def prettify_comm(comm, dns_server)
     if comm.nil?
       channel = Rex::Socket::SwitchBoard.best_comm(dns_server)
       if channel.nil?
-        comm_text = nil
+        nil
       else
-        comm_text = "Session #{channel.sid} (auto)"
+        "Session #{channel.sid} (route)"
       end
     else
-      if driver.framework.sessions.include?(comm)
-        comm_text = "Session #{comm}"
+      if comm.alive
+        "Session #{comm.sid}"
       else
-        comm_text = "Broken session (#{comm})"
+        "Closed session (#{comm.sid})"
       end
     end
   end
diff --git a/lib/rex/proto/dns/custom_nameserver_provider.rb b/lib/rex/proto/dns/custom_nameserver_provider.rb
index 3b7ec6f63db8..15e91209fd7c 100755
--- a/lib/rex/proto/dns/custom_nameserver_provider.rb
+++ b/lib/rex/proto/dns/custom_nameserver_provider.rb
@@ -33,6 +33,16 @@ def add_nameserver(wildcard_rules, dns_server, comm)
       end
     end
 
+    #
+    # Remove entries with the given IDs
+    # Ignore entries that are not found
+    def remove_ids(ids)
+      ids.each do |id|
+        self.entries_with_rules.delete_if {|entry| entry[:id] == id}
+        self.entries_without_rules.delete_if {|entry| entry[:id] == id}
+      end
+    end
+
     #
     # The custom nameserver entries that have been configured
     # @return [Array<Array>] An array containing two elements: The entries with rules, and the entries without rules
@@ -45,25 +55,47 @@ def purge
     end
 
     def nameservers_for_packet(packet)
-      name = packet.question.qName
-      dns_servers = []
+      # 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
+      # will throw an error here.
+      results_from_all_questions = []
+      packet.question.each do |question|
+        name = question.qname.to_s
+        dns_servers = []
 
-      self.entries_with_rules.each do |entry|
-        entry[:wildcard_rules].each do |rule|
-          if matches(name, rule)
-            dns_servers.concat([entry[:dns_server], entry[:comm]])
-            break
+        self.entries_with_rules.each do |entry|
+          entry[:wildcard_rules].each do |rule|
+            socket_options = {}
+            socket_options['Comm'] = entry[:comm] unless entry[:comm].nil?
+            if matches(name, rule)
+              dns_servers.append([entry[:dns_server], socket_options])
+              break
+            end
+          end
+        end
+
+        # Only look at the rule-less entries if no rules were found (avoids DNS leaks)
+        if dns_servers.empty?
+          self.entries_without_rules.each do |entry|
+            socket_options = {}
+            socket_options['Comm'] = entry[:comm] unless entry[:comm].nil?
+            dns_servers.append([entry[:dns_server], socket_options])
           end
         end
-      end
 
-      # Only look at the rule-less entries if no rules were found (avoids DNS leaks)
-      if dns_servers.empty?
-        self.entries_without_rules.each do |entry|
-          dns_servers.concat([entry[:dns_server], entry[:comm]])
+        if dns_servers.empty?
+          # Fall back to default nameservers
+          dns_servers = super
         end
+        results_from_all_questions << dns_servers.uniq
+      end
+      results_from_all_questions.uniq!
+      if results_from_all_questions.size != 1
+        raise ResolverError.new('Inconsistent nameserver entries attempted to be sent in the one packet')
       end
-      dns_servers.uniq!
+
+      results_from_all_questions[0]
     end
 
     def self.extended(mod)
@@ -73,7 +105,11 @@ def self.extended(mod)
     private
 
     def matches(domain, pattern)
-      true
+      if pattern.start_with?('*.')
+        domain.downcase.end_with?(pattern[1..-1].downcase)
+      else
+        domain.casecmp?(pattern)
+      end
     end
 
     attr_accessor :entries_with_rules # Set of custom nameserver entries that specify a rule
diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index 6f7355e1fd6e..3310a3ddbe0d 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/lib/rex/proto/dns/resolver.rb
@@ -14,7 +14,7 @@ module DNS
   class Resolver < Net::DNS::Resolver
 
     Defaults = {
-      :config_file => "/dev/null", # default can lead to info leaks
+      :config_file => "/etc/resolv.conf",
       :log_file => "/dev/null", # formerly $stdout, should be tied in with our loggers
       :port => 53,
       :searchlist => [],
@@ -141,7 +141,8 @@ def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
         packet = Rex::Proto::DNS::Packet.encode_drb(net_packet)
       end
 
-      if nameservers_for_packet(packet).size == 0
+      nameservers = nameservers_for_packet(packet)
+      if nameservers.size == 0
         raise ResolverError, "No nameservers specified!"
       end
 
@@ -210,6 +211,7 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
           socket = nil
           @config[:tcp_timeout].timeout do
             catch(:next_ns) do
+              suffix = ''
               begin
                 config = {
                   'PeerHost' => ns.to_s,
@@ -219,7 +221,12 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
                   'Comm' => @config[:comm]
                 }
                 config.update(socket_options)
+                unless config['Comm'].nil? || config['Comm'].alive?
+                  @logger.warn("Session #{config['Comm'].sid} not active, and cannot be used to resolve DNS")
+                  throw :next_ns
+                end
 
+                suffix = " over session #{@config['Comm'].sid}" unless @config['Comm'].nil?
                 if @config[:source_port] > 0
                   config['LocalPort'] = @config[:source_port]
                 end
@@ -228,11 +235,11 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
                 end
                 socket = Rex::Socket::Tcp.create(config)
               rescue
-                @logger.warn "TCP Socket could not be established to #{ns}:#{@config[:port]} #{@config[:proxies]}"
+                @logger.warn "TCP Socket could not be established to #{ns}:#{@config[:port]} #{@config[:proxies]}#{suffix}"
                 throw :next_ns
               end
               next unless socket #
-              @logger.info "Contacting nameserver #{ns} port #{@config[:port]}"
+              @logger.info "Contacting nameserver #{ns} port #{@config[:port]}#{suffix}"
               socket.write(length+packet_data)
               got_something = false
               loop do
@@ -241,7 +248,7 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
                 begin
                   ans = socket.recv(2)
                 rescue Errno::ECONNRESET
-                  @logger.warn "TCP Socket got Errno::ECONNRESET from #{ns}:#{@config[:port]} #{@config[:proxies]}"
+                  @logger.warn "TCP Socket got Errno::ECONNRESET from #{ns}:#{@config[:port]} #{@config[:proxies]}#{suffix}"
                   attempts -= 1
                   retry if attempts > 0
                 end
@@ -249,7 +256,7 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
                   if got_something
                     break #Proper exit from loop
                   else
-                    @logger.warn "Connection reset to nameserver #{ns}, trying next."
+                    @logger.warn "Connection reset to nameserver #{ns}#{suffix}, trying next."
                     throw :next_ns
                   end
                 end
@@ -259,7 +266,7 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
                 @logger.info "Receiving #{len} bytes..."
 
                 if len.nil? or len == 0
-                  @logger.warn "Receiving 0 length packet from nameserver #{ns}, trying next."
+                  @logger.warn "Receiving 0 length packet from nameserver #{ns}#{suffix}, trying next."
                   throw :next_ns
                 end
 
@@ -270,7 +277,7 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
                 end
 
                 unless buffer.size == len
-                  @logger.warn "Malformed packet from nameserver #{ns}, trying next."
+                  @logger.warn "Malformed packet from nameserver #{ns}#{suffix}, trying next."
                   throw :next_ns
                 end
                 if block_given?
@@ -282,7 +289,7 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
             end
           end
         rescue Timeout::Error
-          @logger.warn "Nameserver #{ns} not responding within TCP timeout, trying next one"
+          @logger.warn "Nameserver #{ns}#{suffix} not responding within TCP timeout, trying next one"
           next
         ensure
           socket.close if socket
@@ -313,6 +320,11 @@ def send_udp(packet,packet_data)
                 'Comm' => @config[:comm]
               }
               config.update(socket_options)
+              unless config['Comm'].nil? || config['Comm'].alive?
+                @logger.warn("Session #{config['Comm'].sid} not active, and cannot be used to resolve DNS")
+                throw :next_ns
+              end
+
               if @config[:source_port] > 0
                 config['LocalPort'] = @config[:source_port]
               end

From 7442655ab9a2918f6f62d464f72ab0c254bc2c02 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Tue, 7 Nov 2023 15:33:10 +1100
Subject: [PATCH 04/14] Override to TCP when encountering UDP-unfriendly comms

---
 lib/msf/base/sessions/meterpreter.rb            |  4 ++++
 lib/msf/base/sessions/ssh_command_shell_bind.rb |  4 ++++
 lib/msf/core/session/comm.rb                    |  7 +++++++
 lib/rex/proto/dns/resolver.rb                   | 16 +++++++++++++++-
 4 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/lib/msf/base/sessions/meterpreter.rb b/lib/msf/base/sessions/meterpreter.rb
index 166011c13a66..0640b93199fc 100644
--- a/lib/msf/base/sessions/meterpreter.rb
+++ b/lib/msf/base/sessions/meterpreter.rb
@@ -593,6 +593,10 @@ def create(param)
     sock
   end
 
+  def supports_udp?
+    true
+  end
+
   #
   # Get a string representation of the current session platform
   #
diff --git a/lib/msf/base/sessions/ssh_command_shell_bind.rb b/lib/msf/base/sessions/ssh_command_shell_bind.rb
index 6b6a86622bde..a6bc1dd97b33 100644
--- a/lib/msf/base/sessions/ssh_command_shell_bind.rb
+++ b/lib/msf/base/sessions/ssh_command_shell_bind.rb
@@ -287,6 +287,10 @@ def create(params)
       sock
     end
 
+    def supports_udp?
+      false
+    end
+
     def create_server_channel(params)
       msf_channel = nil
       mutex = Mutex.new
diff --git a/lib/msf/core/session/comm.rb b/lib/msf/core/session/comm.rb
index 64501e6f33f7..0340442953ff 100644
--- a/lib/msf/core/session/comm.rb
+++ b/lib/msf/core/session/comm.rb
@@ -22,6 +22,13 @@ module Comm
   def create(param)
     raise NotImplementedError
   end
+
+  #
+  # Does the Comm support sending UDP messages?
+  #
+  def supports_udp?
+    raise NotImplementedError
+  end
 end
 
 end
diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index 3310a3ddbe0d..66832d5ca213 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/lib/rex/proto/dns/resolver.rb
@@ -159,6 +159,9 @@ def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
         if use_tcp? or !(proxies.nil? or proxies.empty?) # User requested TCP
           @logger.info "Sending #{packet_size} bytes using TCP due to tcp flag"
           method = :send_tcp
+        elsif !supports_udp?(nameservers)
+          @logger.info "Sending #{packet_size} bytes using TCP due to the presence of a non-UDP-compatible comm channel"
+          method = :send_tcp
         else # Finally use UDP
           @logger.info "Sending #{packet_size} bytes using UDP"
           method = :send_udp unless method == :send_tcp
@@ -334,7 +337,7 @@ def send_udp(packet,packet_data)
               socket = Rex::Socket::Udp.create(config)
             rescue
               @logger.warn "UDP Socket could not be established to #{ns}:#{@config[:port]}"
-              return nil
+              throw :next_ds
             end
             @logger.info "Contacting nameserver #{ns} port #{@config[:port]}"
             #socket.sendto(packet_data, ns.to_s, @config[:port].to_i, 0)
@@ -403,6 +406,17 @@ def query(name, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
       return send(name,type,cls)
 
     end
+
+    private
+
+    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) }}
+        next if comm.nil?
+        return false unless comm.supports_udp?
+      end
+      true
+    end
   end # Resolver
 
 end

From 1a7eefd972b56a035c678aca49a46ee0075b8020 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Wed, 8 Nov 2023 07:38:12 +1100
Subject: [PATCH 05/14] Support saving and loading DNS in the MSF config file

---
 lib/msf/core/framework.rb                     |  1 +
 lib/msf/ui/console/command_dispatcher/core.rb |  1 +
 lib/msf/ui/console/command_dispatcher/dns.rb  | 13 +--
 .../proto/dns/custom_nameserver_provider.rb   | 91 ++++++++++++++++++-
 4 files changed, 97 insertions(+), 9 deletions(-)

diff --git a/lib/msf/core/framework.rb b/lib/msf/core/framework.rb
index 6cb122341fe5..133b1ac8eb0f 100644
--- a/lib/msf/core/framework.rb
+++ b/lib/msf/core/framework.rb
@@ -94,6 +94,7 @@ def initialize(options={})
   def initialize_dns_resolver
     self.dns_resolver = Rex::Proto::DNS::CachedResolver.new
     self.dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    self.dns_resolver.load_config
     Rex::Socket._install_global_resolver(self.dns_resolver)
   end
 
diff --git a/lib/msf/ui/console/command_dispatcher/core.rb b/lib/msf/ui/console/command_dispatcher/core.rb
index 117b9d0cdb1f..6fbfacbe86fb 100644
--- a/lib/msf/ui/console/command_dispatcher/core.rb
+++ b/lib/msf/ui/console/command_dispatcher/core.rb
@@ -1350,6 +1350,7 @@ def cmd_save(*args)
       # Save the framework's datastore
       begin
         framework.save_config
+        driver.framework.dns_resolver.save_config
 
         if active_module
           active_module.save_config
diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index c37963ab3854..4c184dbf9d93 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -114,7 +114,6 @@ def add_dns(*args)
       case opt
       when '--rule', '-r'
         raise ::ArgumentError.new('No rule specified') if val.nil?
-        raise ::ArgumentError.new("Invalid rule: #{val}") unless valid_rule(val)
 
         rules << val
       when '--session', '-s'
@@ -161,14 +160,9 @@ def add_dns(*args)
     end
   end
 
-  #
-  # Is the given wildcard DNS entry valid?
-  def valid_rule(rule)
-    rule =~ /^(\*\.)?([a-z\d][a-z\d-]*[a-z\d]\.)+[a-z]+$/
-  end
-
   #
   # Remove all matching user-configured DNS entries
+  #
   def remove_dns(*args)
     remove_ids = []
     @@remove_opts.parse(args) do |opt, idx, val|
@@ -184,12 +178,14 @@ def remove_dns(*args)
 
   #
   # Delete all user-configured DNS settings
+  #
   def purge_dns
     driver.framework.dns_resolver.purge
   end
 
   #
   # Display the user-configured DNS settings
+  #
   def print_dns
     results = driver.framework.dns_resolver.nameserver_entries
     columns = ['ID','Rule(s)', 'DNS Server(s)', 'Comm channel']
@@ -204,6 +200,7 @@ def print_dns
 
   #
   # Get user-friendly text for displaying the session that this entry would go through
+  #
   def prettify_comm(comm, dns_server)
     if comm.nil?
       channel = Rex::Socket::SwitchBoard.best_comm(dns_server)
@@ -213,7 +210,7 @@ def prettify_comm(comm, dns_server)
         "Session #{channel.sid} (route)"
       end
     else
-      if comm.alive
+      if comm.alive?
         "Session #{comm.sid}"
       else
         "Closed session (#{comm.sid})"
diff --git a/lib/rex/proto/dns/custom_nameserver_provider.rb b/lib/rex/proto/dns/custom_nameserver_provider.rb
index 15e91209fd7c..d64df2cef6b4 100755
--- a/lib/rex/proto/dns/custom_nameserver_provider.rb
+++ b/lib/rex/proto/dns/custom_nameserver_provider.rb
@@ -7,6 +7,28 @@ module DNS
   # for different requests, based on the domain being queried.
   ##
   module CustomNameserverProvider
+    CONFIG_KEY = 'framework/dns'
+
+    #
+    # A Comm implementation that always reports as dead, so should never
+    # be used. This is used to prevent DNS leaks of saved DNS rules that 
+    # were attached to a specific channel.
+    ##
+    class CommSink
+      include Msf::Session::Comm
+      def alive?
+        false
+      end
+
+      def supports_udp?
+        # It won't be used anyway, so let's just say we support it
+        true
+      end
+
+      def sid
+        'previous MSF session'
+      end
+    end
 
     def init
       self.entries_with_rules = []
@@ -14,11 +36,71 @@ def init
       self.next_id = 0
     end
 
+    def save_config
+      new_config = {}
+      [self.entries_with_rules, self.entries_without_rules].each do |entry_set|
+        entry_set.each do |entry|
+          key = entry[:id].to_s
+          val = [entry[:wildcard_rules].join(','),
+                 entry[:dns_server],
+                 (!entry[:comm].nil?).to_s
+                ].join(';')
+          new_config[key] = val
+        end
+      end
+
+      Msf::Config.save(CONFIG_KEY => new_config)
+    end
+
+    def load_config
+      config = Msf::Config.load
+
+      with_rules = []
+      without_rules = []
+      next_id = 0
+
+      dns_settings = config.fetch(CONFIG_KEY, {}).each do |name, value|
+        id = name.to_i
+        wildcard_rules, dns_server, uses_comm = value.split(';')
+        wildcard_rules = wildcard_rules.split(',')
+
+        raise Msf::Config::ConfigError.new('DNS parsing failed: Comm must be true or false') unless ['true','false'].include?(uses_comm)
+        raise Msf::Config::ConfigError.new('Invalid DNS config: Invalid DNS server') unless Rex::Socket.is_ip_addr?(dns_server)
+        raise Msf::Config::ConfigError.new('Invalid DNS config: Invalid rule') unless wildcard_rules.all? {|rule| valid_rule?(rule)}
+
+        comm = uses_comm == 'true' ? CommSink.new : nil
+        entry = {
+          :wildcard_rules => wildcard_rules,
+          :dns_server => dns_server,
+          :comm => comm,
+          :id => id
+        }
+
+        if wildcard_rules.empty?
+          without_rules << entry
+        else
+          with_rules << entry
+        end
+
+        next_id = [id, next_id].max
+      end
+
+      # Now that config has successfully read, update the global values
+      self.entries_with_rules = with_rules
+      self.entries_without_rules = without_rules
+      self.next_id = next_id
+    end
+
     # Add a custom nameserver entry to the custom provider
     # @param [wildcard_rules] Array<String> The wildcard rules to match a DNS request against
     # @param [dns_server] Array<String> The list of IP addresses that would be used for this custom rule
-    # @param comm [Integer] The communication channel to be used for these DNS requests
+    # @param comm [Msf::Session::Comm] The communication channel to be used for these DNS requests
     def add_nameserver(wildcard_rules, dns_server, comm)
+      raise ::ArgumentError.new("Invalid DNS server: #{dns_server}") unless Rex::Socket.is_ip_addr?(dns_server)
+      wildcard_rules.each do |rule|
+        raise ::ArgumentError.new("Invalid rule: #{rule}") unless valid_rule?(rule)
+      end
+
       entry = {
         :wildcard_rules => wildcard_rules,
         :dns_server => dns_server,
@@ -103,6 +185,13 @@ def self.extended(mod)
     end
 
     private
+    #
+    # Is the given wildcard DNS entry valid?
+    #
+    def valid_rule?(rule)
+      rule =~ /^(\*\.)?([a-z\d][a-z\d-]*[a-z\d]\.)+[a-z]+$/
+    end
+
 
     def matches(domain, pattern)
       if pattern.start_with?('*.')

From 00f508170cf8de864ee28622ce2e545882ada374 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Wed, 8 Nov 2023 10:16:22 +1100
Subject: [PATCH 06/14] Implemented tab completion for DNS command

---
 lib/msf/ui/console/command_dispatcher/dns.rb | 55 ++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index 4c184dbf9d93..e0ca9ec7e8ce 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -37,6 +37,61 @@ def commands
     commands
   end
 
+  #
+  # 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
+  def cmd_dns_tabs(str, words)
+    if words.length == 1
+      options = ['add','del','remove','flush','print']
+      return options.select { |opt| opt.start_with?(str) }
+    end
+
+    cmd = words[1]
+    case cmd
+    when 'flush','print'
+      # These commands don't have any arguments
+      return
+    when 'add'
+      # We expect a repeating pattern of tag (e.g. -r) and then a value (e.g. *.metasploit.com)
+      # Once this pattern is violated, we're just specifying DNS servers at that point.
+      tag_is_expected = true
+      if words.length > 2
+        words[2..-1].each do |word|
+          if tag_is_expected && !word.start_with?('-')
+            return # They're trying to specify a DNS server - we can't help them from here on out
+          end
+          tag_is_expected = !tag_is_expected
+        end
+      end
+
+      case words[-1]
+      when '-s', '--session'
+        session_ids = driver.framework.sessions.keys.map { |k| k.to_s }
+        return session_ids.select { |id| id.start_with?(str) }
+      when '-r', '--rule'
+        # Hard to auto-complete a rule with any meaningful value; just return
+        return
+      when /^-/
+        # Unknown tag
+        return
+      end
+
+      options = @@add_opts.option_keys.select { |opt| opt.start_with?(str) }
+      options << '' # Prevent tab-completion of a dash, given they could provide an IP address at this point
+      return options
+    when 'del','remove'
+      if words[-1] == '-i'
+        ids = driver.framework.dns_resolver.nameserver_entries.flatten.map { |entry| entry[:id].to_s }
+        return ids.select { |id| id.start_with? str }
+      else
+        return @@remove_opts.option_keys.select { |opt| opt.start_with?(str) }
+      end
+    end
+  end
+
   def cmd_dns_help
     print_line 'Usage: dns'
     print_line

From 8ce328022cc0d42d90020eaf5abd8cbfa4b3ee88 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Wed, 8 Nov 2023 11:15:01 +1100
Subject: [PATCH 07/14] Clearer dns feature results

---
 lib/msf/ui/console/command_dispatcher/dns.rb   | 18 +++++++++++++++---
 .../proto/dns/custom_nameserver_provider.rb    | 18 +++++++++++++++---
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index e0ca9ec7e8ce..42d67d55cd5d 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -209,10 +209,15 @@ def add_dns(*args)
       comm_obj = driver.framework.sessions[comm_int]
     end
 
+    rules.each do |rule|
+      print_warning("DNS rule #{rule} does not contain wildcards, so will not match subdomains") unless rule.include?('*')
+    end
+
     # Split each DNS server entry up into a separate entry
     servers.each do |server|
       driver.framework.dns_resolver.add_nameserver(rules, server, comm_obj)
     end
+    print_good("DNS #{servers.length > 1 ? 'entries' : 'entry'} added")
   end
 
   #
@@ -228,7 +233,13 @@ def remove_dns(*args)
       end
     end
 
-    driver.framework.dns_resolver.remove_ids(remove_ids)
+    removed = driver.framework.dns_resolver.remove_ids(remove_ids)
+    difference = remove_ids.difference(removed.map { |entry| entry[:id] })
+    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])]})
+    end
   end
 
   #
@@ -236,6 +247,7 @@ def remove_dns(*args)
   #
   def purge_dns
     driver.framework.dns_resolver.purge
+    print_good('DNS entries purged')
   end
 
   #
@@ -243,11 +255,11 @@ def purge_dns
   #
   def print_dns
     results = driver.framework.dns_resolver.nameserver_entries
-    columns = ['ID','Rule(s)', 'DNS Server(s)', 'Comm channel']
+    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])]})
 
     # Default nameservers don't include a rule
-    columns = ['ID', 'DNS Server(s)', 'Comm channel']
+    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])]})
   end
 
diff --git a/lib/rex/proto/dns/custom_nameserver_provider.rb b/lib/rex/proto/dns/custom_nameserver_provider.rb
index d64df2cef6b4..d0f284b5ad72 100755
--- a/lib/rex/proto/dns/custom_nameserver_provider.rb
+++ b/lib/rex/proto/dns/custom_nameserver_provider.rb
@@ -82,7 +82,7 @@ def load_config
           with_rules << entry
         end
 
-        next_id = [id, next_id].max
+        next_id = [id + 1, next_id].max
       end
 
       # Now that config has successfully read, update the global values
@@ -118,11 +118,23 @@ def add_nameserver(wildcard_rules, dns_server, comm)
     #
     # Remove entries with the given IDs
     # Ignore entries that are not found
+    # @param [Array<Integer>] The IDs to removed
+    # @return [Array<Hash>] The removed entries
+    #
     def remove_ids(ids)
+      removed= []
       ids.each do |id|
-        self.entries_with_rules.delete_if {|entry| entry[:id] == id}
-        self.entries_without_rules.delete_if {|entry| entry[:id] == id}
+        removed_with, remaining_with = self.entries_with_rules.partition {|entry| entry[:id] == id}
+        self.entries_with_rules.replace(remaining_with)
+
+        removed_without, remaining_without = self.entries_without_rules.partition {|entry| entry[:id] == id}
+        self.entries_without_rules.replace(remaining_without)
+
+        removed.concat(removed_with)
+        removed.concat(removed_without)
       end
+
+      removed
     end
 
     #

From d0585e0df55fca069a44b6084f464e1b66930a64 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Thu, 9 Nov 2023 14:45:06 +1100
Subject: [PATCH 08/14] Added unit tests

---
 lib/rex/proto/dns/cached_resolver.rb          |   3 +-
 .../dns/custom_nameserver_provider_spec.rb    | 114 ++++++++++++++++++
 2 files changed, 116 insertions(+), 1 deletion(-)
 create mode 100755 spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb

diff --git a/lib/rex/proto/dns/cached_resolver.rb b/lib/rex/proto/dns/cached_resolver.rb
index 02c85ccc9fae..1adbf0f34fa1 100644
--- a/lib/rex/proto/dns/cached_resolver.rb
+++ b/lib/rex/proto/dns/cached_resolver.rb
@@ -21,6 +21,7 @@ class CachedResolver < Resolver
     #
     # @return [nil]
     def initialize(config = {})
+      dns_cache_no_start = config.delete(:dns_cache_no_start)
       super(config)
       self.cache = Rex::Proto::DNS::Cache.new
       # Read hostsfile into cache
@@ -72,7 +73,7 @@ def initialize(config = {})
         end
       end
       # TODO: inotify or similar on hostsfile for live updates? Easy-button functionality
-      self.cache.start unless config[:dns_cache_no_start]
+      self.cache.start unless dns_cache_no_start
       return
     end
 
diff --git a/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb b/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb
new file mode 100755
index 000000000000..853f39912f75
--- /dev/null
+++ b/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb
@@ -0,0 +1,114 @@
+# -*- coding:binary -*-
+require 'spec_helper'
+require 'net/dns'
+
+
+RSpec.describe Rex::Proto::DNS::CustomNameserverProvider do
+  def packet_for(name)
+    packet = Net::DNS::Packet.new(name, Net::DNS::A, Net::DNS::IN)
+    Rex::Proto::DNS::Packet.encode_drb(packet)
+  end
+
+  let(:base_nameserver) do
+    '1.2.3.4'
+  end
+
+  let(:ruleless_nameserver) do
+    '1.2.3.5'
+  end
+
+  let(:ruled_nameserver) do
+    '1.2.3.6'
+  end
+
+  let(:ruled_nameserver2) do
+    '1.2.3.7'
+  end
+
+  let(:ruled_nameserver3) do
+    '1.2.3.8'
+  end
+
+  let (:config) do
+    {:dns_cache_no_start => true}
+  end
+
+  subject(:many_ruled_provider) do
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.nameservers = [base_nameserver]
+    dns_resolver.add_nameserver([], ruleless_nameserver, nil)
+    dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver, nil)
+    dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver2, nil)
+    dns_resolver.add_nameserver(['*.notmetasploit.com'], ruled_nameserver3, nil)
+
+    dns_resolver
+  end
+
+  subject(:ruled_provider) do
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.nameservers = [base_nameserver]
+    dns_resolver.add_nameserver([], ruleless_nameserver, nil)
+    dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver, nil)
+
+    dns_resolver
+  end
+
+  subject(:ruleless_provider) do
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.nameservers = [base_nameserver]
+    dns_resolver.add_nameserver([], ruleless_nameserver, nil)
+
+    dns_resolver
+  end
+
+  subject(:empty_provider) do
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.nameservers = [base_nameserver]
+
+    dns_resolver
+  end
+
+  context 'When no nameserver is configured' do
+   it 'The resolver base is returned' do
+     packet = packet_for('subdomain.metasploit.com')
+     ns = empty_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[base_nameserver, {}]])
+   end
+  end
+
+  context 'When a base nameserver is configured' do
+   it 'The base nameserver is returned' do
+     packet = packet_for('subdomain.metasploit.com')
+     ns = ruleless_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[ruleless_nameserver, {}]])
+   end
+  end
+
+  context 'When a nameserver rule is configured and a rule entry matches' do
+   it 'The correct nameserver is returned' do
+     packet = packet_for('subdomain.metasploit.com')
+     ns = ruled_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[ruled_nameserver, {}]])
+    end
+  end
+
+  context 'When a nameserver rule is configured and no rule entry is applicable' do
+   it 'The base nameserver is returned when no rule entry' do
+     packet = packet_for('subdomain.notmetasploit.com')
+     ns = ruled_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[ruleless_nameserver, {}]])
+   end
+  end
+
+  context 'When many rules are configured' do
+   it 'Returns multiple entries if multiple rules match' do
+     packet = packet_for('subdomain.metasploit.com')
+     ns = many_ruled_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[ruled_nameserver, {}], [ruled_nameserver2, {}]])
+   end
+  end
+end
\ No newline at end of file

From f351d7b5e1aac2daf2dc5ee2f95a52769c06d5ec Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Thu, 9 Nov 2023 16:26:53 +1100
Subject: [PATCH 09/14] Don't create DNS cached resolver on every test, because
 it'll exhaust the OS's resources

---
 lib/msf/core/framework.rb                           | 13 +++++--------
 lib/msf/ui/console/command_dispatcher/dns.rb        |  4 ++++
 lib/msf/ui/console/driver.rb                        | 10 +++++++++-
 lib/rex/proto/dns/resolver.rb                       |  2 +-
 .../proto/dns/custom_nameserver_provider_spec.rb    | 10 ++++++++++
 5 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/lib/msf/core/framework.rb b/lib/msf/core/framework.rb
index 133b1ac8eb0f..b983fffb8327 100644
--- a/lib/msf/core/framework.rb
+++ b/lib/msf/core/framework.rb
@@ -81,7 +81,11 @@ def initialize(options={})
     # Configure the SSL certificate generator
     require 'msf/core/cert_provider'
     Rex::Socket::Ssl.cert_provider = Msf::Ssl::CertProvider
-    initialize_dns_resolver
+
+    if options.include?('CustomDnsResolver')
+      self.dns_resolver = options['CustomDnsResolver']
+      Rex::Socket._install_global_resolver(self.dns_resolver)
+    end
 
     subscriber = FrameworkEventSubscriber.new(self)
     events.add_exploit_subscriber(subscriber)
@@ -91,13 +95,6 @@ def initialize(options={})
     events.add_ui_subscriber(subscriber)
   end
 
-  def initialize_dns_resolver
-    self.dns_resolver = Rex::Proto::DNS::CachedResolver.new
-    self.dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
-    self.dns_resolver.load_config
-    Rex::Socket._install_global_resolver(self.dns_resolver)
-  end
-
   def dns_resolver
     self.dns_resolver
   end
diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index 42d67d55cd5d..086eb5290e2d 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -44,6 +44,8 @@ def commands
   # @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
   def cmd_dns_tabs(str, words)
+    return if driver.framework.dns_resolver.nil?
+
     if words.length == 1
       options = ['add','del','remove','flush','print']
       return options.select { |opt| opt.start_with?(str) }
@@ -132,6 +134,8 @@ def cmd_dns_help
   # Manage Metasploit's DNS resolution rules
   #
   def cmd_dns(*args)
+    return if driver.framework.dns_resolver.nil?
+
     args << 'print' if args.length == 0
     # Short-circuit help
     if args.delete("-h") || args.delete("--help")
diff --git a/lib/msf/ui/console/driver.rb b/lib/msf/ui/console/driver.rb
index 9edca71ee62c..8586b4ea8a2b 100644
--- a/lib/msf/ui/console/driver.rb
+++ b/lib/msf/ui/console/driver.rb
@@ -80,8 +80,16 @@ def initialize(prompt = DefaultPrompt, prompt_char = DefaultPromptChar, opts = {
 
     # Initialize attributes
 
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.load_config
+
     # Defer loading of modules until paths from opts can be added below
-    framework_create_options = opts.merge('DeferModuleLoads' => true)
+    framework_create_options = opts.merge({
+                                            'DeferModuleLoads' => true,
+                                            'CustomDnsResolver' => dns_resolver
+                                          }
+                                         )
     self.framework = opts['Framework'] || Msf::Simple::Framework.create(framework_create_options)
 
     if self.framework.datastore['Prompt']
diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index 66832d5ca213..8047342f10b3 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/lib/rex/proto/dns/resolver.rb
@@ -15,7 +15,7 @@ class Resolver < Net::DNS::Resolver
 
     Defaults = {
       :config_file => "/etc/resolv.conf",
-      :log_file => "/dev/null", # formerly $stdout, should be tied in with our loggers
+      :log_file => File::NULL, # formerly $stdout, should be tied in with our loggers
       :port => 53,
       :searchlist => [],
       :nameservers => [IPAddr.new("127.0.0.1")],
diff --git a/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb b/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb
index 853f39912f75..2b4f7aca5d31 100755
--- a/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb
+++ b/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb
@@ -111,4 +111,14 @@ def packet_for(name)
      expect(ns).to eq([[ruled_nameserver, {}], [ruled_nameserver2, {}]])
    end
   end
+
+  context 'When a packet contains multiple questions that have different nameserver results' do
+   it 'Throws an error' do
+     packet = packet_for('subdomain.metasploit.com')
+     q = Dnsruby::Question.new('subdomain.notmetasploit.com', Dnsruby::Types::A, Dnsruby::Classes::IN)
+
+     packet.question.append(q)
+     expect {many_ruled_provider.nameservers_for_packet(packet)}.to raise_error(ResolverError)
+   end
+  end
 end
\ No newline at end of file

From 987bed6972791955cd47f7f30c1c2cf2db644e73 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Fri, 10 Nov 2023 15:01:45 +1100
Subject: [PATCH 10/14] Remove unimplemented command

---
 lib/msf/ui/console/command_dispatcher/dns.rb | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index 086eb5290e2d..8b0dd0b6a72e 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -109,7 +109,6 @@ def cmd_dns_help
     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 "  get - display the DNS server(s) and communication channel that would be used for a given target"
     print_line "  print - show all active DNS resolution entries"
     print_line
     print_line "Examples:"
@@ -125,9 +124,6 @@ def cmd_dns_help
     print_line "  Set the DNS server to be used for all requests that match no rules"
     print_line "    route add 8.8.8.8 8.8.4.4"
     print_line
-    print_line "  Display the DNS server that would be used for the given domain name"
-    print_line "    route get subdomain.metasploit.com"
-    print_line
   end
 
   #

From 1a07ab5aee39c31772d4f6b9d31c8c45511dd6ec Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Fri, 10 Nov 2023 15:56:45 +1100
Subject: [PATCH 11/14] Improved comments on functions

---
 .../proto/dns/custom_nameserver_provider.rb    | 18 +++++++++++++++---
 lib/rex/proto/dns/resolver.rb                  |  3 ++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/lib/rex/proto/dns/custom_nameserver_provider.rb b/lib/rex/proto/dns/custom_nameserver_provider.rb
index d0f284b5ad72..4bc947550d4d 100755
--- a/lib/rex/proto/dns/custom_nameserver_provider.rb
+++ b/lib/rex/proto/dns/custom_nameserver_provider.rb
@@ -36,6 +36,9 @@ def init
       self.next_id = 0
     end
 
+    #
+    # Save the custom settings to the MSF config file
+    #
     def save_config
       new_config = {}
       [self.entries_with_rules, self.entries_without_rules].each do |entry_set|
@@ -52,6 +55,9 @@ def save_config
       Msf::Config.save(CONFIG_KEY => new_config)
     end
 
+    #
+    # Load the custom settings from the MSF config file
+    #
     def load_config
       config = Msf::Config.load
 
@@ -92,8 +98,8 @@ def load_config
     end
 
     # Add a custom nameserver entry to the custom provider
-    # @param [wildcard_rules] Array<String> The wildcard rules to match a DNS request against
-    # @param [dns_server] Array<String> The list of IP addresses that would be used for this custom rule
+    # @param wildcard_rules [Array<String>] The wildcard rules to match a DNS request against
+    # @param dns_server [Array<String>] The list of IP addresses that would be used for this custom rule
     # @param comm [Msf::Session::Comm] The communication channel to be used for these DNS requests
     def add_nameserver(wildcard_rules, dns_server, comm)
       raise ::ArgumentError.new("Invalid DNS server: #{dns_server}") unless Rex::Socket.is_ip_addr?(dns_server)
@@ -118,7 +124,7 @@ def add_nameserver(wildcard_rules, dns_server, comm)
     #
     # Remove entries with the given IDs
     # Ignore entries that are not found
-    # @param [Array<Integer>] The IDs to removed
+    # @param ids [Array<Integer>] The IDs to removed
     # @return [Array<Hash>] The removed entries
     #
     def remove_ids(ids)
@@ -140,6 +146,7 @@ def remove_ids(ids)
     #
     # The custom nameserver entries that have been configured
     # @return [Array<Array>] An array containing two elements: The entries with rules, and the entries without rules
+    #
     def nameserver_entries
       [entries_with_rules, entries_without_rules]
     end
@@ -148,6 +155,11 @@ def purge
       init
     end
 
+    # The nameservers that match the given packet
+    # @param packet [Dnsruby::Message] The DNS packet to be sent
+    # @raise [ResolveError] If the packet contains multiple questions, which would end up sending to a different set of nameservers
+    # @return [Array<Array>] A list of nameservers, each with Rex::Socket options
+    #
     def nameservers_for_packet(packet)
       # Leaky abstraction: a packet could have multiple question entries,
       # and each of these could have different nameservers, or travel via
diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index 8047342f10b3..ed7105e2c00c 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/lib/rex/proto/dns/resolver.rb
@@ -115,6 +115,7 @@ def proxies=(prox, timeout_added = 250)
     # @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)
       @config[:nameservers].map {|ns| [ns, {}]}
     end
@@ -125,8 +126,8 @@ def nameservers_for_packet(dns_message)
     # @param argument [Object] An object holding the DNS message to be processed.
     # @param type [Fixnum] Type of record to look up
     # @param cls [Fixnum] Class of question to look up
-    #
     # @return [Dnsruby::Message] DNS response
+    #
     def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
       method = self.use_tcp? ? :send_tcp : :send_udp
 

From 34bd661d3f19f83d8f1e7b7a24009eba62b399fd Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Wed, 22 Nov 2023 09:06:06 +1100
Subject: [PATCH 12/14] Fall back to other server if first one fails

---
 lib/msf/core/framework.rb     |  4 ---
 lib/rex/proto/dns/resolver.rb | 64 ++++++++++++++++++-----------------
 2 files changed, 33 insertions(+), 35 deletions(-)

diff --git a/lib/msf/core/framework.rb b/lib/msf/core/framework.rb
index b983fffb8327..f9e6ae9a5ba8 100644
--- a/lib/msf/core/framework.rb
+++ b/lib/msf/core/framework.rb
@@ -95,10 +95,6 @@ def initialize(options={})
     events.add_ui_subscriber(subscriber)
   end
 
-  def dns_resolver
-    self.dns_resolver
-  end
-
   def inspect
     "#<Framework (#{sessions.length} sessions, #{jobs.length} jobs, #{plugins.length} plugins#{db.active ? ", #{db.driver} database active" : ""})>"
   end
diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index ed7105e2c00c..0006591608f9 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/lib/rex/proto/dns/resolver.rb
@@ -314,41 +314,43 @@ def send_udp(packet,packet_data)
       response = ""
       nameservers = nameservers_for_packet(packet)
       nameservers.each do |ns, socket_options|
-        begin
-          @config[:udp_timeout].timeout do
-            begin
-              config = {
-                'PeerHost' => ns.to_s,
-                'PeerPort' => @config[:port].to_i,
-                'Context' => @config[:context],
-                'Comm' => @config[:comm]
-              }
-              config.update(socket_options)
-              unless config['Comm'].nil? || config['Comm'].alive?
-                @logger.warn("Session #{config['Comm'].sid} not active, and cannot be used to resolve DNS")
-                throw :next_ns
-              end
+        catch(:next_ns) do
+          begin
+            @config[:udp_timeout].timeout do
+              begin
+                config = {
+                  'PeerHost' => ns.to_s,
+                  'PeerPort' => @config[:port].to_i,
+                  'Context' => @config[:context],
+                  'Comm' => @config[:comm]
+                }
+                config.update(socket_options)
+                unless config['Comm'].nil? || config['Comm'].alive?
+                  @logger.warn("Session #{config['Comm'].sid} not active, and cannot be used to resolve DNS")
+                  throw :next_ns
+                end
 
-              if @config[:source_port] > 0
-                config['LocalPort'] = @config[:source_port]
-              end
-              if @config[:source_host] != IPAddr.new('0.0.0.0')
-                config['LocalHost'] = @config[:source_host] unless @config[:source_host].nil?
+                if @config[:source_port] > 0
+                  config['LocalPort'] = @config[:source_port]
+                end
+                if @config[:source_host] != IPAddr.new('0.0.0.0')
+                  config['LocalHost'] = @config[:source_host] unless @config[:source_host].nil?
+                end
+                socket = Rex::Socket::Udp.create(config)
+              rescue
+                @logger.warn "UDP Socket could not be established to #{ns}:#{@config[:port]}"
+                throw :next_ns
               end
-              socket = Rex::Socket::Udp.create(config)
-            rescue
-              @logger.warn "UDP Socket could not be established to #{ns}:#{@config[:port]}"
-              throw :next_ds
+              @logger.info "Contacting nameserver #{ns} port #{@config[:port]}"
+              #socket.sendto(packet_data, ns.to_s, @config[:port].to_i, 0)
+              socket.write(packet_data)
+              ans = socket.recvfrom(@config[:packet_size])
             end
-            @logger.info "Contacting nameserver #{ns} port #{@config[:port]}"
-            #socket.sendto(packet_data, ns.to_s, @config[:port].to_i, 0)
-            socket.write(packet_data)
-            ans = socket.recvfrom(@config[:packet_size])
+            break if ans
+          rescue Timeout::Error
+            @logger.warn "Nameserver #{ns} not responding within UDP timeout, trying next one"
+            throw :next_ds
           end
-          break if ans
-        rescue Timeout::Error
-          @logger.warn "Nameserver #{ns} not responding within UDP timeout, trying next one"
-          next
         end
       end
       return ans

From ef9a165d2289ccf277ec654acd84f671d2a77d00 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Wed, 22 Nov 2023 15:45:03 +1100
Subject: [PATCH 13/14] Changes from code review

---
 lib/msf/core/framework.rb                     |  1 +
 lib/msf/ui/console/command_dispatcher/dns.rb  | 65 ++++++++++++-------
 .../proto/dns/custom_nameserver_provider.rb   | 12 +++-
 lib/rex/proto/dns/resolver.rb                 | 16 ++---
 .../dns/custom_nameserver_provider_spec.rb    | 18 +++++
 5 files changed, 80 insertions(+), 32 deletions(-)

diff --git a/lib/msf/core/framework.rb b/lib/msf/core/framework.rb
index f9e6ae9a5ba8..6c104f368d15 100644
--- a/lib/msf/core/framework.rb
+++ b/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
 
diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index 8b0dd0b6a72e..4824c939a97d 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/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,7 +295,14 @@ 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,
@@ -293,8 +310,12 @@ def print_dns_set(heading, columns, result_set)
         '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
diff --git a/lib/rex/proto/dns/custom_nameserver_provider.rb b/lib/rex/proto/dns/custom_nameserver_provider.rb
index 4bc947550d4d..bab8c55390dc 100755
--- a/lib/rex/proto/dns/custom_nameserver_provider.rb
+++ b/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
@@ -172,9 +175,9 @@ def nameservers_for_packet(packet)
 
         self.entries_with_rules.each do |entry|
           entry[:wildcard_rules].each do |rule|
-            socket_options = {}
-            socket_options['Comm'] = entry[:comm] unless entry[:comm].nil?
             if matches(name, rule)
+              socket_options = {}
+              socket_options['Comm'] = entry[:comm] unless entry[:comm].nil?
               dns_servers.append([entry[:dns_server], socket_options])
               break
             end
@@ -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
diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index 0006591608f9..31be7a4826bd 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/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
 
@@ -174,7 +174,7 @@ def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
         method = :send_tcp
       end
 
-      ans = self.__send__(method, packet, packet_data)
+      ans = self.__send__(method, packet, packet_data, nameservers)
 
       unless (ans and ans[0].length > 0)
         @logger.fatal "No response from nameservers list: aborting"
@@ -203,13 +203,13 @@ def send(argument, type = Dnsruby::Types::A, cls = Dnsruby::Classes::IN)
     #
     # @param packet [Net::DNS::Packet] Packet associated with packet_data
     # @param packet_data [String] Data segment of DNS request packet
+    # @param nameservers [Array<[String,Hash]>] List of nameservers to use for this request, and their associated socket options
     # @param prox [String] Proxy configuration for TCP socket
     #
     # @return ans [String] Raw DNS reply
-    def send_tcp(packet,packet_data,prox = @config[:proxies])
+    def send_tcp(packet, packet_data, nameservers, prox = @config[:proxies])
       ans = nil
       length = [packet_data.size].pack("n")
-      nameservers = nameservers_for_packet(packet)
       nameservers.each do |ns, socket_options|
         begin
           socket = nil
@@ -307,12 +307,12 @@ def send_tcp(packet,packet_data,prox = @config[:proxies])
     #
     # @param packet [Net::DNS::Packet] Packet associated with packet_data
     # @param packet_data [String] Data segment of DNS request packet
+    # @param nameservers [Array<[String,Hash]>] List of nameservers to use for this request, and their associated socket options
     #
     # @return ans [String] Raw DNS reply
-    def send_udp(packet,packet_data)
+    def send_udp(packet,packet_data, nameservers)
       ans = nil
       response = ""
-      nameservers = nameservers_for_packet(packet)
       nameservers.each do |ns, socket_options|
         catch(:next_ns) do
           begin
@@ -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
diff --git a/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb b/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb
index 2b4f7aca5d31..89313095c35b 100755
--- a/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb
+++ b/spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb
@@ -33,6 +33,20 @@ def packet_for(name)
     {:dns_cache_no_start => true}
   end
 
+  let (:framework_with_dns_enabled) do
+    framework = Object.new
+    def framework.features
+      f = Object.new
+      def f.enabled?(_name)
+        true
+      end
+
+      f
+    end
+
+    framework
+  end
+
   subject(:many_ruled_provider) do
     dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
     dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
@@ -41,6 +55,7 @@ def packet_for(name)
     dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver, nil)
     dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver2, nil)
     dns_resolver.add_nameserver(['*.notmetasploit.com'], ruled_nameserver3, nil)
+    dns_resolver.set_framework(framework_with_dns_enabled)
 
     dns_resolver
   end
@@ -51,6 +66,7 @@ def packet_for(name)
     dns_resolver.nameservers = [base_nameserver]
     dns_resolver.add_nameserver([], ruleless_nameserver, nil)
     dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver, nil)
+    dns_resolver.set_framework(framework_with_dns_enabled)
 
     dns_resolver
   end
@@ -60,6 +76,7 @@ def packet_for(name)
     dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
     dns_resolver.nameservers = [base_nameserver]
     dns_resolver.add_nameserver([], ruleless_nameserver, nil)
+    dns_resolver.set_framework(framework_with_dns_enabled)
 
     dns_resolver
   end
@@ -68,6 +85,7 @@ def packet_for(name)
     dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
     dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
     dns_resolver.nameservers = [base_nameserver]
+    dns_resolver.set_framework(framework_with_dns_enabled)
 
     dns_resolver
   end

From 473ded345ba24cc325ef1ba5dffc3a5dce546c72 Mon Sep 17 00:00:00 2001
From: Ashley Donaldson <smashery@gmail.com>
Date: Thu, 23 Nov 2023 09:54:47 +1100
Subject: [PATCH 14/14] Fix UDP detection when DNS resolution is not on

---
 lib/msf/ui/console/command_dispatcher/dns.rb |  4 +-
 lib/net/dns/resolver.rb                      | 13 +++---
 lib/rex/proto/dns/resolver.rb                | 42 ++++++++++----------
 3 files changed, 31 insertions(+), 28 deletions(-)

diff --git a/lib/msf/ui/console/command_dispatcher/dns.rb b/lib/msf/ui/console/command_dispatcher/dns.rb
index 4824c939a97d..ec538a7c14e9 100755
--- a/lib/msf/ui/console/command_dispatcher/dns.rb
+++ b/lib/msf/ui/console/command_dispatcher/dns.rb
@@ -160,6 +160,8 @@ def cmd_dns(*args)
         print_dns
       when "help"
         cmd_dns_help
+      else
+        print_error("Invalid command. To view help: dns -h")
       end
     rescue ::ArgumentError => e
       print_error(e.message)
@@ -298,7 +300,7 @@ def prettify_comm(comm, dns_server)
   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']
+      columns = ['ID', 'Rules(s)', 'DNS Server', 'Comm channel']
     else
       columns = ['ID', 'DNS Server', 'Commm channel']
     end
diff --git a/lib/net/dns/resolver.rb b/lib/net/dns/resolver.rb
index af5f1a356124..f89690bcf91a 100644
--- a/lib/net/dns/resolver.rb
+++ b/lib/net/dns/resolver.rb
@@ -975,7 +975,7 @@ def send(argument,type=Net::DNS::A,cls=Net::DNS::IN)
           end
         end
 
-        ans = self.old_send(method,packet,packet_data)
+        ans = self.old_send(method,packet,packet_data, nameservers.map {|ns| [ns, {}]})
 
         unless ans
           @logger.fatal "No response from nameservers list: aborting"
@@ -1027,7 +1027,8 @@ def axfr(name,cls=Net::DNS::IN)
 
         answers = []
         soa = 0
-        self.old_send(method, packet, packet_data) do |ans|
+        nameservers_and_hash = nameservers.map {|ns| [ns, {}]}
+        self.old_send(method, packet, packet_data, nameservers_and_hash) do |ans|
           @logger.info "Received #{ans[0].size} bytes from #{ans[1][2]+":"+ans[1][1].to_s}"
 
           begin
@@ -1161,12 +1162,12 @@ def make_query_packet(string,type,cls)
 
       end
 
-      def send_tcp(packet,packet_data)
+      def send_tcp(packet,packet_data, nameservers)
 
         ans = nil
         length = [packet_data.size].pack("n")
 
-        @config[:nameservers].each do |ns|
+        nameservers.each do |ns, _unused|
           begin
             socket = Socket.new(Socket::AF_INET,Socket::SOCK_STREAM,0)
             socket.bind(Socket.pack_sockaddr_in(@config[:source_port],@config[:source_address].to_s))
@@ -1233,13 +1234,13 @@ def send_tcp(packet,packet_data)
         return nil
       end
 
-      def send_udp(packet,packet_data)
+      def send_udp(packet, packet_data, nameservers)
         socket = UDPSocket.new
         socket.bind(@config[:source_address].to_s,@config[:source_port])
 
         ans = nil
         response = ""
-        @config[:nameservers].each do |ns|
+        nameservers.each do |ns, _unused|
           begin
             @config[:udp_timeout].timeout do
               @logger.info "Contacting nameserver #{ns} port #{@config[:port]}"
diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index 31be7a4826bd..8e4aee0fd5a2 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/lib/rex/proto/dns/resolver.rb
@@ -117,7 +117,7 @@ def proxies=(prox, timeout_added = 250)
     # @return [Array<Array>] A list of nameservers, each with Rex::Socket options
     #
     def nameservers_for_packet(_dns_message)
-      @config[:nameservers].map {|ns| [ns, {}]}
+      @config[:nameservers].map {|ns| [ns.to_s, {}]}
     end
 
     #
@@ -213,30 +213,30 @@ def send_tcp(packet, packet_data, nameservers, prox = @config[:proxies])
       nameservers.each do |ns, socket_options|
         begin
           socket = nil
+          config = {
+            'PeerHost' => ns.to_s,
+            'PeerPort' => @config[:port].to_i,
+            'Proxies' => prox,
+            'Context' => @config[:context],
+            'Comm' => @config[:comm]
+          }
+          config.update(socket_options)
+          unless config['Comm'].nil? || config['Comm'].alive?
+            @logger.warn("Session #{config['Comm'].sid} not active, and cannot be used to resolve DNS")
+            throw :next_ns
+          end
+
+          suffix = " over session #{@config['Comm'].sid}" unless @config['Comm'].nil?
+          if @config[:source_port] > 0
+            config['LocalPort'] = @config[:source_port]
+          end
+          if @config[:source_host].to_s != '0.0.0.0'
+            config['LocalHost'] = @config[:source_host] unless @config[:source_host].nil?
+          end
           @config[:tcp_timeout].timeout do
             catch(:next_ns) do
               suffix = ''
               begin
-                config = {
-                  'PeerHost' => ns.to_s,
-                  'PeerPort' => @config[:port].to_i,
-                  'Proxies' => prox,
-                  'Context' => @config[:context],
-                  'Comm' => @config[:comm]
-                }
-                config.update(socket_options)
-                unless config['Comm'].nil? || config['Comm'].alive?
-                  @logger.warn("Session #{config['Comm'].sid} not active, and cannot be used to resolve DNS")
-                  throw :next_ns
-                end
-
-                suffix = " over session #{@config['Comm'].sid}" unless @config['Comm'].nil?
-                if @config[:source_port] > 0
-                  config['LocalPort'] = @config[:source_port]
-                end
-                if @config[:source_host].to_s != '0.0.0.0'
-                  config['LocalHost'] = @config[:source_host] unless @config[:source_host].nil?
-                end
                 socket = Rex::Socket::Tcp.create(config)
               rescue
                 @logger.warn "TCP Socket could not be established to #{ns}:#{@config[:port]} #{@config[:proxies]}#{suffix}"