Merge pull request #19756 from smashery/dns_reorder

Add the ability to reorder DNS entries

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

@@ -37,6 +37,11 @@ class DNS
     ['-f'] => [true,  'Address family - IPv4 or IPv6 (default IPv4)']
   )
 
+  REORDER_USAGE = 'dns [reorder] -i <existing ID> <new ID>'.freeze
+  @@reorder_opts = Rex::Parser::Arguments.new(
+    ['-i', '--index'] => [true, 'Index of the rule to move']
+  )
+
   def initialize(driver)
     super
   end
@@ -108,7 +113,7 @@ def cmd_dns_tabs(str, words)
     when 'help'
       # These commands don't have any arguments
       return subcommands.select { |sc| sc.start_with?(str) }
-    when 'remove','delete'
+    when 'remove','delete','reorder'
       if words[-1] == '-i'
         return
       else
@@ -156,6 +161,7 @@ def cmd_dns_help(*args)
     print_line "  #{ADD_STATIC_USAGE}"
     print_line "  #{REMOVE_USAGE}"
     print_line "  #{REMOVE_STATIC_USAGE}"
+    print_line "  #{REORDER_USAGE}"
     print_line "  dns [flush-cache]"
     print_line "  dns [flush-entries]"
     print_line "  dns [flush-static]"
@@ -171,10 +177,11 @@ def cmd_dns_help(*args)
     print_line "  flush-entries - Remove all configured DNS resolution entries"
     print_line "  flush-static  - Remove all statically defined hostnames"
     print_line "  print         - Show all configured DNS resolution entries"
-    print_line "  remove        - Delete a DNS resolution entry"
+    print_line "  remove        - Delete one or more DNS resolution entries"
     print_line "  remove-static - Delete a statically defined hostname"
     print_line "  reset-config  - Reset the DNS configuration"
     print_line "  resolve       - Resolve a hostname"
+    print_line "  reorder       - Reorder one or more rules"
     print_line
     print_line "EXAMPLES:"
     print_line "  Display help information for the 'add' subcommand"
@@ -222,6 +229,8 @@ def cmd_dns(*args)
         cmd_dns_help(*args)
       when "print"
         print_dns
+      when 'reorder'
+        reorder_dns(*args)
       when "remove", "rm", "delete", "del"
         remove_dns(*args)
       when "remove-static"
@@ -491,6 +500,54 @@ def remove_dns_help
     print_line
   end
 
+  def reorder_dns(*args)
+    reorder_ids = []
+    new_id = -1
+    @@remove_opts.parse(args) do |opt, idx, val|
+      case opt
+      when '-i', '--index'
+        raise ::ArgumentError.new("Not a valid index: #{val}") unless val.to_i > 0
+        raise ::ArgumentError.new("Duplicate index: #{val}") if reorder_ids.include?(val.to_i - 1)
+
+        reorder_ids << val.to_i - 1
+      when nil
+        raise ::ArgumentError.new("Not a valid index: #{val}") unless (val.to_i > 0 || val.to_i == -1)
+        new_id = val.to_i
+        new_id -= 1 unless new_id == -1
+      end
+    end
+
+    if reorder_ids.empty?
+      raise ::ArgumentError.new('At least one index to reorder must be provided')
+    end
+
+    reordered = resolver.reorder_ids(reorder_ids, new_id)
+    print_warning('Some entries were not reordered') unless reordered.length == reorder_ids.length
+    if reordered.length > 0
+      print_good("#{reordered.length} DNS #{reordered.length > 1 ? 'entries' : 'entry'} reordered")
+      print_resolver_rules
+    end
+  end
+
+  def reorder_dns_help
+    print_line "USAGE:"
+    print_line "  #{REORDER_USAGE}"
+    print_line "If providing multiple IDs, they will be inserted at the given index in the order you provide."
+    print_line(@@reorder_opts.usage)
+    print_line "EXAMPLES:"
+    print_line "  Move the third DNS entry to the top of the resolution order"
+    print_line "    dns reorder -i 3  1"
+    print_line
+    print_line "  Move the third and fifth DNS entries just below the first entry (i.e. becoming the second and third entries, respectively)"
+    print_line "    dns reorder -i 3 -i 5  2"
+    print_line
+    print_line "  Move the second and third DNS entries to the bottom of the resolution order"
+    print_line "    dns reorder -i 2 -i 3  -1"
+    print_line "  Alternatively, assuming there are 6 entries in the list"
+    print_line "    dns reorder -i 2 -i 3  7"
+    print_line
+  end
+
   def remove_static_dns(*args)
     if args.length < 1
       raise ::ArgumentError.new('A hostname must be provided')
@@ -604,6 +661,15 @@ def flush_static_dns
     print_good('DNS static hostnames flushed')
   end
 
+  def print_resolver_rules
+    upstream_rules = resolver.upstream_rules
+    print_dns_set('Resolver rule entries', upstream_rules, ids: (1..upstream_rules.length).to_a)
+    if upstream_rules.empty?
+      print_line
+      print_error('No DNS nameserver entries configured')
+    end
+  end
+
   #
   # Display the user-configured DNS settings
   #
@@ -628,12 +694,7 @@ def print_dns
     end
     print_line("Current cache size:    #{resolver.cache.records.length}")
 
-    upstream_rules = resolver.upstream_rules
-    print_dns_set('Resolver rule entries', upstream_rules, ids: (1..upstream_rules.length).to_a)
-    if upstream_rules.empty?
-      print_line
-      print_error('No DNS nameserver entries configured')
-    end
+    print_resolver_rules
 
     tbl = Table.new(
       Table::Style::Default,

lib/rex/proto/dns/custom_nameserver_provider.rb

@@ -132,9 +132,9 @@ def add_upstream_rule(resolvers, comm: nil, wildcard: '*', index: -1)
     end
 
     #
-    # Remove upstream rules with the given indexes
+    # Remove upstream rules with the given indices
     # Ignore entries that are not found
-    # @param ids [Array<Integer>] The IDs to removed
+    # @param ids [Array<Integer>] The IDs to remove
     # @return [Array<UpstreamRule>] The removed entries
     def remove_ids(ids)
       removed = []
@@ -146,6 +146,47 @@ def remove_ids(ids)
       removed.reverse
     end
 
+    #
+    # Move upstream rules with the given indices into the location provided.
+    # If multiple IDs are provided, they will all be inserted into the provided location, 
+    # in the order provided.
+    # Ignore entries that are not found
+    # @param ids [Array<Integer>] The IDs to move
+    # @param insertion_id [Integer] The ID to insert the entries at (in the order provided), or -1 to insert at the end
+    # @return [Array<UpstreamRule>] The moved entries
+    def reorder_ids(ids, new_id)
+      if new_id == -1
+        new_id = @upstream_rules.length
+      end
+      if new_id > @upstream_rules.length
+        raise ::ArgumentError.new("Insertion ID is past the end of the ruleset")
+      end
+      to_move = []
+      to_subtract = 0
+      # Get the entries before we delete (gets too complicated with indices changing otherwise)
+      ids.each do |id|
+        upstream_rule = @upstream_rules[id]
+        unless upstream_rule.nil?
+          to_move << upstream_rule
+          if new_id > id
+            to_subtract += 1 # Adjust for the fact that are about to delete one, so the indices would be off-by-one after that index is deleted
+          end
+        end
+      end
+
+      new_id -= to_subtract
+
+      ids.sort.reverse.each do |id|
+        @upstream_rules.delete_at(id)
+      end
+
+      to_move.reverse.each do |rule|
+        @upstream_rules.insert(new_id, rule)
+      end
+
+      to_move
+    end
+
     def flush
       @upstream_rules.clear
     end