From 6e86ac6e1b223755360f9a3a89b6bda4d1eec555 Mon Sep 17 00:00:00 2001
From: RageLtMan <rageltman [at] sempervictus>
Date: Wed, 24 Feb 2016 16:58:19 -0500
Subject: [PATCH] Tweak Server and Resolver

Create default_dispatch_request method in Server to allow an
intercepted dispatch request to fall back into default exec flow.

Add attr_reader to the records hash in Cache

Provide Resolver and Server with comm option for their sockets.
---
 lib/rex/proto/dns/resolver.rb |  8 +++-
 lib/rex/proto/dns/server.rb   | 69 +++++++++++++++++++----------------
 2 files changed, 44 insertions(+), 33 deletions(-)

diff --git a/lib/rex/proto/dns/resolver.rb b/lib/rex/proto/dns/resolver.rb
index 8dc665260a..387bff9046 100644
--- a/lib/rex/proto/dns/resolver.rb
+++ b/lib/rex/proto/dns/resolver.rb
@@ -27,7 +27,9 @@ module DNS
       :ignore_truncated => false,
       :packet_size => 512,
       :tcp_timeout => TcpTimeout.new(30),
-      :udp_timeout => UdpTimeout.new(30)
+      :udp_timeout => UdpTimeout.new(30),
+      :context => {},
+      :comm => nil
     }
 
     #
@@ -191,6 +193,7 @@ module DNS
                   'PeerPort' => @config[:port].to_i,
                   'Proxies' => prox,
                   'Context' => @config[:context]
+                  'Comm' => @config[:comm]
                 }
                 if @config[:source_port] > 0
                   config['LocalPort'] = @config[:source_port]
@@ -272,7 +275,8 @@ module DNS
               config = {
                 'PeerHost' => ns.to_s,
                 'PeerPort' => @config[:port].to_i,
-                'Context' => @config[:context]
+                'Context' => @config[:context],
+                'Comm' => @config[:comm]
               }
               if @config[:source_port] > 0
                 config['LocalPort'] = @config[:source_port]
diff --git a/lib/rex/proto/dns/server.rb b/lib/rex/proto/dns/server.rb
index d132bc4e61..f50fe41d02 100644
--- a/lib/rex/proto/dns/server.rb
+++ b/lib/rex/proto/dns/server.rb
@@ -10,6 +10,7 @@ module DNS
 class Server
 
   class Cache
+    attr_reader :records
     include Rex::Proto::DNS::Constants
     # class DNSRecordError < ::Exception
     #
@@ -138,17 +139,19 @@ class Server
   #
   # @return [Rex::Proto::DNS::Server] DNS Server object
   attr_accessor :fwd_res, :dispatch_block, :send_block, :cache
-  def initialize(lhost = '0.0.0.0', lport = 53, udp = true, tcp = false, res = nil, ctx = {}, dblock = nil, sblock = nil)
+  def initialize(lhost = '0.0.0.0', lport = 53, udp = true, tcp = false, res = nil, comm = nil, ctx = {}, dblock = nil, sblock = nil)
     
     @udp_sock = udp ? Rex::Socket::Udp.create(
       'LocalHost' => lhost,
       'LocalPort' => lport,
-      'Context'   => ctx
+      'Context'   => ctx,
+      'Comm'      => comm
     ) : nil
     @tcp_sock = tcp ? Rex::Socket::TcpServer.create(
       'LocalHost' => lhost,
       'LocalPort' => lport,
-      'Context'   => ctx
+      'Context'   => ctx,
+      'Comm'      => comm
     ) : nil
     @fwd_res = res.nil? ? Rex::Proto::DNS::Resolver.new : res
     @udp_mon = nil
@@ -226,37 +229,41 @@ class Server
     if @dispatch_block
       @dispatch_block.call(cli,data)
     else
-      req = Net::DNS::Packet.parse(data)
-      forward = req.dup
-      # Find cached items, remove request from forwarded packet
-      req.question.each do |ques|
-        cached = @cache.find(ques.qName, ques.qType.to_s)
-        if cached.empty?
-          next
-        else
-          req.answer = req.answer + cached
-          forward.question.delete(ques)
-        end
-      end
-      # Forward remaining requests, cache responses
-      if forward.question.count > 0
-        forwarded = @fwd_res.send(validate_packet(forward))
-        req.answer = req.answer + forwarded.answer 
-        forwarded.answer.each do |ans|
-          @cache.cache_record(ans)
-        end
-        req.header.ra = 1 # Set recursion bit
-      end
-      # Finalize answers in response
-      # Check for empty response prior to sending
-      if req.answer.size < 1
-        req.header.rCode = Net::DNS::Header::RCode.new(3)
-      end
-      req.header.qr = 1 # Set response bit
-      send_response(cli, validate_packet(req).data)
+      default_dispatch_request(cli,data)
     end
   end
 
+  def default_dispatch_request(cli,data)
+    req = Net::DNS::Packet.parse(data)
+    forward = req.dup
+    # Find cached items, remove request from forwarded packet
+    req.question.each do |ques|
+      cached = @cache.find(ques.qName, ques.qType.to_s)
+      if cached.empty?
+        next
+      else
+        req.answer = req.answer + cached
+        forward.question.delete(ques)
+      end
+    end
+    # Forward remaining requests, cache responses
+    if forward.question.count > 0
+      forwarded = @fwd_res.send(validate_packet(forward))
+      req.answer = req.answer + forwarded.answer 
+      forwarded.answer.each do |ans|
+        @cache.cache_record(ans)
+      end
+      req.header.ra = 1 # Set recursion bit
+    end
+    # Finalize answers in response
+    # Check for empty response prior to sending
+    if req.answer.size < 1
+      req.header.rCode = Net::DNS::Header::RCode.new(3)
+    end
+    req.header.qr = 1 # Set response bit
+    send_response(cli, validate_packet(req).data)
+  end
+
   #
   # Send response to client, handled with @send_block if set
   #