From 3494542b54fe7c876833c3fce8b44a66e9e7229a Mon Sep 17 00:00:00 2001
From: bcoles <bcoles@gmail.com>
Date: Tue, 4 Mar 2014 00:56:41 +1030
Subject: [PATCH] Make upstream DNS server configurable

---
 extensions/dns/api.rb      | 22 ++++++++++++++++++++++
 extensions/dns/config.yaml |  2 ++
 extensions/dns/dns.rb      | 23 ++++++++++++++++++++++-
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/extensions/dns/api.rb b/extensions/dns/api.rb
index d7d0d050c..5070a3ac2 100644
--- a/extensions/dns/api.rb
+++ b/extensions/dns/api.rb
@@ -33,6 +33,28 @@ module BeEF
             dns.run_server(address, port)
 
             print_info "DNS Server: #{address}:#{port}"
+            servers = []
+            unless dns_config['upstream'].nil?
+              dns_config['upstream'].each do |server|
+                if server[1].nil? or server[2].nil?
+                  next
+                end
+                if server[0] == 'tcp'
+                  servers << ['tcp', server[1], server[2]]
+                elsif server[0] == 'udp'
+                  servers << ['udp', server[1], server[2]]
+                end
+              end
+            end
+            if servers.empty?
+              servers << ['tcp', '8.8.8.8', 53]
+              servers << ['udp', '8.8.8.8', 53]
+            end
+            upstream_servers = ''
+            servers.each do |server|
+              upstream_servers << "Upstream server: #{server[1]}:#{server[2]} (#{server[0]})\n"
+            end
+            print_more upstream_servers
           end
 
           # Mounts handler for processing RESTful API calls
diff --git a/extensions/dns/config.yaml b/extensions/dns/config.yaml
index d90e93730..d382ac70f 100644
--- a/extensions/dns/config.yaml
+++ b/extensions/dns/config.yaml
@@ -11,3 +11,5 @@ beef:
             authors: ['soh_cah_toa']
             address: '127.0.0.1'
             port: 5300
+            upstream:
+                [['tcp', '8.8.8.8', 53], ['udp', '8.8.8.8', 53]]
diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb
index 94c291af8..6ebd71571 100644
--- a/extensions/dns/dns.rb
+++ b/extensions/dns/dns.rb
@@ -136,9 +136,30 @@ module BeEF
             BeEF::Extension::Dns::Server.instance.instance_eval { @server = server }
 
             # Pass unmatched queries upstream to root nameservers
+            server = []
+            dns_config = BeEF::Core::Configuration.instance.get('beef.extension.dns')
+            unless dns_config['upstream'].nil?
+              dns_config['upstream'].each do |server|
+                if server[1].nil? or server[2].nil?
+                  print_error "Invalid server '#{server[1]}:#{server[2]}' specified for upstream DNS server."
+                  next
+                elsif server[0] == 'tcp'
+                  servers << [:tcp, server[1], server[2]]
+                elsif server[0] == 'udp'
+                  servers << [:udp, server[1], server[2]]
+                else
+                  print_error "Invalid protocol '#{server[0]}' specified for upstream DNS server."
+                end
+              end
+            end
+            if servers.empty?
+              print_debug "No upstream DNS servers specified. Using '8.8.8.8'"
+              servers << [:tcp, '8.8.8.8', 53]
+              servers << [:udp, '8.8.8.8', 53]
+            end
             otherwise do |transaction|
               transaction.passthrough!(
-                  RubyDNS::Resolver.new([[:udp, '8.8.8.8', 53], [:tcp, '8.8.8.8', 53]])
+                  RubyDNS::Resolver.new servers
               )
             end
           end