From 861d66207db6247b3dd9a646b5753fd6d0e36844 Mon Sep 17 00:00:00 2001
From: soh_cah_toa <soh_cah_toa@cybercrud.net>
Date: Wed, 23 Apr 2014 11:27:49 -0400
Subject: [PATCH] Implemented new Rule model and #add_rule method.

---
 extensions/dns/dns.rb   | 33 +++++++++++++++++++++++++++++++++
 extensions/dns/model.rb | 37 +++++++++++++++++++++++++++++++------
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/extensions/dns/dns.rb b/extensions/dns/dns.rb
index ac18ce303..8b469505d 100644
--- a/extensions/dns/dns.rb
+++ b/extensions/dns/dns.rb
@@ -22,6 +22,39 @@ module BeEF
           @lock = Mutex.new
         end
 
+        # Adds a new DNS rule. If the rule already exists, its current ID is returned.
+        #
+        # @example Adds an A record for browserhacker.com with the IP address 1.2.3.4
+        #
+        #   dns = BeEF::Extension::Dns::Server.instance
+        #
+        #   id = dns.add_rule(
+        #     :pattern  => 'browserhacker.com',
+        #     :resource => Resolv::DNS::Resource::IN::A,
+        #     :response => '1.2.3.4'
+        #   )
+        #
+        # @param rule [Hash] hash representation of rule
+        # @option rule [String, Regexp] :pattern match criteria
+        # @option rule [Resolv::DNS::Resource::IN] :resource resource record type
+        # @option rule [String, Array] :response server response
+        #
+        # @return [String] unique 8-digit hex identifier
+        def add_rule(rule = {})
+          @lock.synchronize do
+            # Temporarily disable warnings regarding IGNORECASE flag
+            verbose = $VERBOSE
+            $VERBOSE = nil
+            pattern = Regexp.new(rule[:pattern], Regexp::IGNORECASE)
+            $VERBOSE = verbose
+
+            BeEF::Core::Models::Dns::Rule.first_or_create(
+              { :resource => rule[:resource], :pattern => pattern.source },
+              { :response => rule[:response] }
+            ).id
+          end
+        end
+
         # Entry point for processing incoming DNS requests. Attempts to find a matching rule and
         # sends back its associated response.
         #
diff --git a/extensions/dns/model.rb b/extensions/dns/model.rb
index 14f5a2310..062699bdf 100644
--- a/extensions/dns/model.rb
+++ b/extensions/dns/model.rb
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Copyright (c) 2006-2014 Wade Alcorn - wade@bindshell.net
 # Browser Exploitation Framework (BeEF) - http://beefproject.com
 # See the file 'doc/COPYING' for copying permission
 #
@@ -8,16 +8,41 @@ module BeEF
     module Models
       module Dns
 
+        # Represents an individual DNS rule.
         class Rule
-
           include DataMapper::Resource
 
           storage_names[:default] = 'extension_dns_rules'
 
-          property :id, String, :key => true # Unique identifier
-          property :pattern, Object # Query pattern
-          property :type, Object # Resource type
-          property :block, Text # Associated callback
+          property :id, String, :key => true
+          property :pattern, Object, :required => true
+          property :resource, Object, :required => true
+          property :response, Object, :required => true
+          property :callback, String, :required => true
+
+          # Hooks the model's "save" event. Generates a rule identifier and callback.
+          before :save do |rule|
+            rule.callback = validate_response(rule.resource, rule.response)
+            rule.id = generate_id
+          end
+
+        private
+          # Strict validator which ensures that only an appropriate response is given.
+          #
+          # @param resource [Resolv::DNS::Resource::IN] resource record type
+          # @param response [String, Symbol, Array] response to include in callback
+          #
+          # @return [String] string representation of callback that can safely be eval'd
+          def validate_response(resource, response)
+            "t.respond!('1.1.1.1')"
+          end
+
+          # Generates a unique identifier for use as a primary key.
+          #
+          # @return [String] 8-character hex identifier
+          def generate_id
+            '42'
+          end
 
         end