Began implementing new method of adding rules without periodic timer.

Also added improved documentation for add_rule() and remove_rule().

extensions/dns/dns.rb

@@ -11,17 +11,29 @@ module DNS
 
     include Singleton
 
-    # Starts DNS server run-loop.
+    # @!method instance
+    #   Returns the singleton instance.
+    def initialize
+      @server  = nil
+      @next_id = 0
+    end
+
+    # Starts the DNS server run-loop.
     #
     # @param address [String] interface address server should run on
     # @param port [Integer] desired server port number
     def run_server(address, port)
       EventMachine::next_tick do
         RubyDNS::run_server(:listen => [[:udp, address, port]]) do
-          upstream = RubyDNS::Resolver.new([[:udp, '8.8.8.8', 53], [:tcp, '8.8.8.8', 53]])
+          server = self
+          BeEF::Extension::DNS::DNS.instance.instance_eval { @server = server }
+          BeEF::Extension::DNS::DNS.instance.load_rules
 
+          # Pass unmatched queries upstream to root nameservers
           otherwise do |transaction|
-            transaction.passthrough!(upstream)
+            transaction.passthrough!(
+              RubyDNS::Resolver.new([[:udp, '8.8.8.8', 53], [:tcp, '8.8.8.8', 53]])
+            )
           end
         end
       end
@@ -29,25 +41,56 @@ module DNS
 
     # Adds a new DNS rule or "resource record". Does nothing if rule is already present.
     #
-    # @param name [String] name of query
-    # @param type [String] query type (e.g. A, CNAME, MX, NS, etc.)
-    # @param value [String] response to send back to resolver
-    def add_rule(name, type, value)
-      catch(:match) do
-        BeEF::Core::Models::DNS.each do |rule|
-          n = rule.name
-          t = rule.type
-          v = rule.value
+    # @example Adds an A record for foobar.com with the value 1.2.3.4
+    #
+    #   dns = BeEF::Extension::DNS::DNS.instance
+    #
+    #   id = dns.add_rule('foobar.com', Resolv::DNS::Resource::IN::A) do |transaction|
+    #     transaction.respond!('1.2.3.4')
+    #   end
+    #
+    # @param pattern [String, Regexp] query pattern to recognize
+    # @param type [Resolv::DNS::Resource::IN] resource record type (e.g. A, CNAME, NS, etc.)
+    #
+    # @yieldparam [RubyDNS::Transaction] details of query question and response
+    #
+    # @return [Integer] unique id for use with {#remove_rule}
+    #
+    # @see #remove_rule
+    # @see http://rubydoc.info/gems/rubydns/RubyDNS/Transaction
+    def add_rule(pattern, type, &block)
+      @next_id += 1
+      @server.match(@next_id, pattern, type, block)
+      @next_id
+    end
 
-          throw :match if [n, t, v] == [name, type, value]
-        end
+    # Removes the given DNS rule. Any future queries for it will be passed through.
+    #
+    # @param id [Integer] id returned from {#add_rule}
+    # @see #add_rule
+    def remove_rule(id)
+      @server.remove_rule(id)
+      @next_id -= 1
+    end
 
-        BeEF::Core::Models::DNS.create(
-          :name  => name,
-          :type  => type,
-          :value => value
-        )
-      end
+    # Loads all rules from the database at server startup
+    def load_rules
+      # TODO Load rules from database
+    end
+
+    private
+
+    # Convenience method for fully-qualifying Resolv::DNS::Resource::IN types
+    def parse_type(type)
+      eval "Resolv::DNS::Resource::IN::#{type}"
+    end
+
+    # Convenience method for generating proper server responses
+    def parse_response(type, value)
+      response = 'value.to_s'
+      response = 'Resolv::DNS::Name.create(value)' if type =~ /(CNAME|NS|PTR)/
+      
+      eval response
     end
 
   end

extensions/dns/ruby/rubydns.rb

@@ -9,8 +9,7 @@ require 'rubydns'
 
 module RubyDNS
 
-  # Behaves exactly the same, except without any output and an added periodic
-  # timer that checks for new DNS rules every five seconds
+  # Behaves exactly the same, except without any logger output
   def self.run_server(options = {}, &block)
     server = RubyDNS::Server.new(&block)
 
@@ -27,15 +26,41 @@ module RubyDNS
         end
       end
 
-      server.load_rules
-      EventMachine.add_periodic_timer(5) { server.check_rules }
-
       server.fire(:start)
     end
 
     server.fire(:stop)
   end
 
+  class Server
+
+    attr_accessor :rules
+
+    class Rule
+
+      attr_accessor :id
+
+      # Now uses an 'id' parameter to uniquely identify rules
+      def initialize(id, pattern, callback)
+        @id = id
+        @pattern  = pattern
+        @callback = callback
+      end
+
+    end
+
+    # Now uses an 'id' parameter to uniquely identify rules
+    def match(id, *pattern, block)
+      @rules << Rule.new(id, pattern, block)
+    end
+
+    # New method that removes a rule given its id
+    def remove_rule(id)
+      @rules.delete_if { |rule| rule.id == id }
+    end
+
+  end
+
   class Transaction
 
     # Behaves exactly the same, except using debug logger instead of info
@@ -56,53 +81,4 @@ 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