From d622bf3e5ee4568e6322db9fa7fd348896fa9cc8 Mon Sep 17 00:00:00 2001
From: soh_cah_toa <soh_cah_toa@cybercrud.net>
Date: Fri, 10 May 2013 23:01:10 -0400
Subject: [PATCH] New DNS entries can now be added dynamically without a server
 restart. Database is checked every five seconds and adds new rules if there
 were any changes.

---
 extensions/dns/dns.rb          | 41 ++++++---------------
 extensions/dns/ruby/rubydns.rb | 67 ++++++++++++++++++++++++++++++----
 2 files changed, 72 insertions(+), 36 deletions(-)

diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb
index aa665cc53..b3ceaa642 100644
--- a/extensions/dns/dns.rb
+++ b/extensions/dns/dns.rb
@@ -16,16 +16,6 @@ module DNS
         RubyDNS::run_server(:listen => [[:udp, address, port]]) do
           upstream = RubyDNS::Resolver.new([[:udp, '8.8.8.8', 53], [:tcp, '8.8.8.8', 53]])
 
-          BeEF::Core::Models::DNS.each do |record|
-            name  = record.name
-            type  = BeEF::Extension::DNS::DNS.parse_type(record.type)
-            value = record.value
-
-            match(name, type) do |transaction|
-              transaction.respond!(value)
-            end
-          end
-
           otherwise do |transaction|
             transaction.passthrough!(upstream)
           end
@@ -34,28 +24,21 @@ module DNS
     end
 
     def add_rule(name, type, value)
-      d = BeEF::Core::Models::DNS.new(
-        :name  => name,
-        :type  => type,
-        :value => value
-      ).save
+      catch(:match) do
+        BeEF::Core::Models::DNS.each do |rule|
+          n = rule.name
+          t = rule.type
+          v = rule.value
 
-      type = BeEF::Extension::DNS::DNS.parse_type(type)
+          throw :match if [n, t, v] == [name, type, value]
+        end
 
-      RubyDNS::stop_server
-      run_server
-    end
-
-    # XXX Why must this be a class method? As a private instance method,
-    #     it throws NoMethodError.
-    def self.parse_type(type)
-      resolv = 'Resolv::DNS::Resource'
-
-      if type =~ /(A|AAAA|SRV|WKS)/
-        resolv += '::IN'
+        BeEF::Core::Models::DNS.create(
+          :name  => name,
+          :type  => type,
+          :value => value
+        )
       end
-
-      eval "#{resolv}::#{type}"
     end
 
   end
diff --git a/extensions/dns/ruby/rubydns.rb b/extensions/dns/ruby/rubydns.rb
index 112316667..39ab1af93 100644
--- a/extensions/dns/ruby/rubydns.rb
+++ b/extensions/dns/ruby/rubydns.rb
@@ -3,9 +3,14 @@
 # Browser Exploitation Framework (BeEF) - http://beefproject.com
 # See the file 'doc/COPYING' for copying permission
 #
+
+require 'rubygems'
+require 'rubydns'
+
 module RubyDNS
 
-  # Behaves exactly the same, minus the output
+  # Behaves exactly the same, except without any output and an added periodic
+  # timer that checks for new DNS rules every five seconds
   def self.run_server(options = {}, &block)
     server = RubyDNS::Server.new(&block)
 
@@ -16,22 +21,21 @@ module RubyDNS
 
       options[:listen].each do |spec|
         if spec[0] == :udp
-          @signature = EventMachine.open_datagram_socket(spec[1], spec[2], UDPHandler, server)
+          EventMachine.open_datagram_socket(spec[1], spec[2], UDPHandler, server)
         elsif spec[0] == :tcp
-          @signature = EventMachine.start_server(spec[1], spec[2], TCPHandler, server)
+          EventMachine.start_server(spec[1], spec[2], TCPHandler, server)
         end
       end
 
+      server.load_rules
+      EventMachine.add_periodic_timer(5) { server.check_rules }
+
       server.fire(:start)
     end
 
     server.fire(:stop)
   end
 
-  def self.stop_server
-    EventMachine.stop_server(@signature)
-  end
-
   class Transaction
 
     # Behaves exactly the same, except using debug logger instead of info
@@ -52,4 +56,53 @@ module RubyDNS
 
   end
 
+  class Server
+
+    # Reads current DNS entries in database and adds them as new rules
+    def load_rules
+      rules = get_rules
+      @rule_count = rules.count
+
+      rules.each do |rule|
+        match(rule[0], parse_type(rule[1])) do |transaction|
+          transaction.respond!(rule[2])
+        end
+      end
+    end
+
+    # Re-loads ruleset if new entries have been added to database
+    def check_rules
+      load_rules if get_rules.count != @rule_count
+    end
+
+    private
+
+    # Returns an AoA where each element is a rule of the form [name, type, value]
+    def get_rules
+      rules = []
+
+      BeEF::Core::Models::DNS.each do |record|
+        name  = record.name
+        type  = record.type
+        value = record.value
+
+        rules << [name, type, value]
+      end
+
+      rules
+    end
+
+    # Convenience method for fully-qualifying Resolv::DNS::Resource types
+    def parse_type(type)
+      resolv = 'Resolv::DNS::Resource'
+
+      if type =~ /(A|AAAA|SRV|WKS)/
+        resolv += '::IN'
+      end
+
+      eval "#{resolv}::#{type}"
+    end
+
+  end
+
 end