Manually merged DNS extension code (pull request 967 from @soh-cah-toa)

Gemfile

@@ -32,6 +32,8 @@ gem "erubis"
 gem "dm-migrations"
 gem "msfrpc-client"
 gem "rubyzip", "~> 1.0.0"
+gem "rubydns"
+gem "sourcify"
 
 # notifications
 gem "twitter", ">= 5.0.0"

config.yaml

@@ -120,3 +120,5 @@ beef:
                 enable: false
         ipec:
             enable: true
+        dns:
+            enable: true

core/loader.rb

@@ -15,6 +15,8 @@ require 'ipaddr'
 require 'base64'
 require 'xmlrpc/client'
 require 'openssl'
+require 'rubydns'
+require 'sourcify'
 
 # @note Include the filters
 require 'core/filters'
@@ -29,4 +31,4 @@ require 'core/api'
 require 'core/settings'
 
 # @note Include the core of BeEF
-require 'core/core'
+require 'core/core'

extensions/dns/api.rb

@@ -0,0 +1,48 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+module BeEF
+module Extension
+module Dns
+module API
+
+  module NameserverHandler
+
+    BeEF::API::Registrar.instance.register(
+      BeEF::Extension::Dns::API::NameserverHandler,
+      BeEF::API::Server,
+      'pre_http_start'
+    )
+
+    BeEF::API::Registrar.instance.register(
+      BeEF::Extension::Dns::API::NameserverHandler,
+      BeEF::API::Server,
+      'mount_handler'
+    )
+
+    # Begins main DNS server run-loop at BeEF startup
+    def self.pre_http_start(http_hook_server)
+      dns_config = BeEF::Core::Configuration.instance.get('beef.extension.dns')
+
+      address = dns_config['address']
+      port    = dns_config['port']
+
+      dns = BeEF::Extension::Dns::Server.instance
+      dns.run_server(address, port)
+
+      print_info "DNS Server: #{address}:#{port}"
+    end
+
+    # Mounts handler for processing RESTful API calls
+    def self.mount_handler(beef_server)
+      beef_server.mount('/api/dns', BeEF::Extension::Dns::DnsRest.new)
+    end
+
+  end
+
+end
+end
+end
+end

extensions/dns/config.yaml

@@ -0,0 +1,13 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+beef:
+    extension:
+        dns:
+            enable: true
+            name: 'DNS Server'
+            authors: ['soh_cah_toa']
+            address: '127.0.0.1'
+            port: 5300

extensions/dns/dns.rb

@@ -0,0 +1,151 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+module BeEF
+module Extension
+module Dns
+
+  # This class is responsible for providing a DNS nameserver that can be dynamically
+  # configured by other modules and extensions. It is particularly useful for
+  # performing DNS spoofing, hijacking, tunneling, etc.
+  #
+  # Only a single instance will exist during runtime (known as the "singleton pattern").
+  # This makes it easier to coordinate actions across the various BeEF systems.
+  class Server
+
+    include Singleton
+
+    attr_reader :address, :port
+
+    # @!method self.instance
+    #  Returns the singleton instance. Use this in place of {#initialize}.
+
+    # @note This method cannot be invoked! Use {.instance} instead.
+    # @see ::instance
+    def initialize
+      @lock = Mutex.new
+      @server = nil
+    end
+
+    # Starts the main DNS server run-loop in a new thread.
+    #
+    # @param address [String] interface address server should run on
+    # @param port [Integer] desired server port number
+    def run_server(address = '0.0.0.0', port = 5300)
+      @address = address
+      @port = port
+
+      @lock.synchronize do
+        Thread.new do
+          # @note Calling #sleep is a quick fix that prevents race conditions
+          #       with WebSockets. A better solution is needed; perhaps a
+          #       global EventMachine mutex.
+          sleep(1)
+
+          if EventMachine.reactor_running?
+            EventMachine.next_tick { run_server_block(@address, @port) }
+          else
+            run_server_block(@address, @port)
+          end
+        end
+      end
+    end
+
+    # Adds a new DNS rule or "resource record". Does nothing if rule is already present.
+    #
+    # @example Adds an A record for foobar.com with the value 1.2.3.4
+    #
+    #   dns = BeEF::Extension::Dns::Server.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.)
+    #
+    # @note When parameter 'pattern' is a literal Regexp object, it must NOT be passed
+    #       using the /.../ literal syntax. Instead use either %r{...} or Regexp::new.
+    #       This does not apply if 'pattern' is a variable.
+    #
+    # @yield callback to invoke when pattern is matched
+    # @yieldparam transaction [RubyDNS::Transaction] details of query question and response
+    #
+    # @return [String] unique 7-digit hex identifier for use with {#remove_rule}
+    #
+    # @see #remove_rule
+    # @see http://rubydoc.info/gems/rubydns/RubyDNS/Transaction
+    def add_rule(pattern, type, &block)
+      @lock.synchronize { @server.match(pattern, type, block) }
+    end
+
+    # Removes the given DNS rule. Any future queries for it will be passed through.
+    #
+    # @param id [Integer] id returned from {#add_rule}
+    #
+    # @return [Boolean] true on success, false on failure
+    #
+    # @see #add_rule
+    def remove_rule(id)
+      @lock.synchronize { @server.remove_rule(id) }
+    end
+
+    # Retrieves a specific rule given its id
+    #
+    # @param id [Integer] unique identifier for rule
+    #
+    # @return [Hash] hash representation of rule
+    def get_rule(id)
+      @lock.synchronize { @server.get_rule(id) }
+    end
+
+    # Returns an AoH representing the entire current DNS ruleset.
+    #
+    # Each element is a hash with the following keys:
+    #
+    # * <code>:id</code>
+    # * <code>:pattern</code>
+    # * <code>:type</code>
+    # * <code>:response</code>
+    #
+    # @return [Array<Hash>] DNS ruleset (empty if no rules are currently loaded)
+    def get_ruleset
+      @lock.synchronize { @server.get_ruleset }
+    end
+
+    # Clears the entire DNS ruleset.
+    #
+    # Requests made after doing so will be passed through to the root nameservers.
+    #
+    # @return [Boolean] true on success, false on failure
+    def remove_ruleset
+      @lock.synchronize { @server.remove_ruleset }
+    end
+
+    private
+
+    # Common code needed by {#run_server} to start DNS server.
+    #
+    # @param address [String] interface address server should run on
+    # @param port [Integer] desired server port number
+    def run_server_block(address, port)
+      RubyDNS.run_server(:listen => [[:udp, address, port]]) do
+        server = self
+        BeEF::Extension::Dns::Server.instance.instance_eval { @server = server }
+
+        # Pass unmatched queries upstream to root nameservers
+        otherwise do |transaction|
+          transaction.passthrough!(
+            RubyDNS::Resolver.new([[:udp, '8.8.8.8', 53], [:tcp, '8.8.8.8', 53]])
+          )
+        end
+      end
+    end
+
+  end
+
+end
+end
+end

extensions/dns/extension.rb

@@ -0,0 +1,25 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+module BeEF
+module Extension
+module Dns
+
+  extend BeEF::API::Extension
+
+  @short_name  = 'dns'
+  @full_name   = 'DNS Server'
+  @description = 'A configurable DNS nameserver for performing DNS spoofing, ' +
+                 'hijacking, and other related attacks against hooked zombies'
+
+end
+end
+end
+
+require 'extensions/dns/api'
+require 'extensions/dns/dns'
+require 'extensions/dns/model'
+require 'extensions/dns/rest/dns'
+require 'extensions/dns/ruby'

extensions/dns/model.rb

@@ -0,0 +1,27 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+module BeEF
+module Core
+module Models
+module Dns
+
+  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
+
+  end
+
+end
+end
+end
+end

extensions/dns/rest/dns.rb

@@ -0,0 +1,245 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+module BeEF
+module Extension
+module Dns
+
+  # This class handles the routing of RESTful API requests that query BeEF's DNS server
+  class DnsRest < BeEF::Core::Router::Router
+
+    # Filters out bad requests before performing any routing
+    before do
+      config = BeEF::Core::Configuration.instance
+
+      # Require a valid API token from a valid IP address
+      halt 401 unless params[:token] == config.get('beef.api_token')
+      halt 403 unless BeEF::Core::Rest.permitted_source?(request.ip)
+
+      headers 'Content-Type' => 'application/json; charset=UTF-8',
+              'Pragma' => 'no-cache',
+              'Cache-Control' => 'no-cache',
+              'Expires' => '0'
+    end
+
+    # Returns the entire current DNS ruleset
+    get '/ruleset' do
+      begin
+        ruleset = BeEF::Extension::Dns::Server.instance.get_ruleset
+        count = ruleset.length
+
+        result = {}
+        result[:count] = count
+        result[:ruleset] = ruleset
+        result.to_json
+      rescue StandardError => e
+        print_error "Internal error while retrieving DNS ruleset (#{e.message})"
+        halt 500
+      end
+    end
+
+    # Returns a specific rule given its id
+    get '/rule/:id' do
+      begin
+        id = params[:id]
+
+        unless BeEF::Filters.alphanums_only?(id)
+          raise InvalidParamError, 'Invalid "id" parameter passed to endpoint /api/dns/rule/:id'
+        end
+
+        result = BeEF::Extension::Dns::Server.instance.get_rule(id)
+        halt 404 if result.length == 0
+
+        result.to_json
+      rescue InvalidParamError => e
+        print_error e.message
+        halt 400
+      rescue StandardError => e
+        print_error "Internal error while retrieving DNS rule with id #{id} (#{e.message})"
+        halt 500
+      end
+    end
+
+    # Adds a new DNS rule
+    post '/rule' do
+      begin
+        body = JSON.parse(request.body.read)
+
+        pattern = body['pattern']
+        type = body['type']
+        response = body['response']
+
+        # Validate required JSON keys
+        unless [pattern, type, response].include?(nil)
+          # Determine whether 'pattern' is a String or Regexp
+          begin
+            pattern_test = eval pattern
+            pattern = pattern_test if pattern_test.class == Regexp
+          rescue => e; end
+
+          if response.class == Array
+              if response.length == 0
+                raise InvalidJsonError, 'Empty "reponse" key passed to endpoint /api/dns/rule'
+              end
+          else
+            raise InvalidJsonError, 'Non-array "reponse" key passed to endpoint /api/dns/rule'
+          end
+
+          unless BeEF::Filters.is_non_empty_string?(pattern)
+            raise InvalidJsonError, 'Empty "pattern" key passed to endpoint /api/dns/rule'
+          end
+
+          unless BeEF::Filters.is_non_empty_string?(type)
+            raise InvalidJsonError, 'Empty "type" key passed to endpoint /api/dns/rule'
+          end
+
+          id = ''
+
+          block_src = format_response(type, response)
+          type_obj  = eval "Resolv::DNS::Resource::IN::#{type}"
+
+          # Bypass #add_rule so that 'block_src' can be passed as a String
+          BeEF::Extension::Dns::Server.instance.instance_eval do
+            id = @server.match(pattern, type_obj, block_src)
+          end
+
+          result = {}
+          result['success'] = true
+          result['id'] = id
+          result.to_json
+        end
+      rescue InvalidJsonError => e
+        print_error e.message
+        halt 400
+      rescue StandardError => e
+        print_error "Internal error while adding DNS rule (#{e.message})"
+        halt 500
+      end
+    end
+
+    # Removes a rule given its id
+    delete '/rule/:id' do
+      begin
+        id = params[:id]
+
+        unless BeEF::Filters.alphanums_only?(id)
+          raise InvalidParamError, 'Invalid "id" parameter passed to endpoint /api/dns/rule/:id'
+        end
+
+        result = {}
+        result['success'] = BeEF::Extension::Dns::Server.instance.remove_rule(id)
+        result.to_json
+      rescue InvalidParamError => e
+        print_error e.message
+        halt 400
+      rescue StandardError => e
+        print_error "Internal error while removing DNS rule with id #{id} (#{e.message})"
+        halt 500
+      end
+    end
+
+    private
+
+    # Generates a formatted string representation of the callback to invoke as a response.
+    #
+    # @param [String] type resource record type (e.g. A, CNAME, NS, etc.)
+    # @param [Array] rdata record data to include in response
+    #
+    # @return [String] string representation of response callback
+    def format_response(type, rdata)
+      src = 'proc { |t| t.respond!(%s) }'
+
+      args = case type
+             when 'A'
+               data = { :address => rdata[0] }
+               sprintf "'%<address>s'", data
+             when 'AAAA'
+               data = { :address => rdata[0] }
+               sprintf "'%<address>s'", data
+             when 'CNAME'
+               data = { :cname => rdata[0] }
+               sprintf "Resolv::DNS::Name.create('%<cname>s')", data
+             when 'HINFO'
+               data = { :cpu => rdata[0], :os => rdata[1] }
+               sprintf "'%<cpu>s', '%<os>s'", data
+             when 'MINFO'
+               data = { :rmailbx => rdata[0], :emailbx => rdata[1] }
+
+               sprintf "Resolv::DNS::Name.create('%<rmailbx>s'), " +
+                       "Resolv::DNS::Name.create('%<emailbx>s')",
+                       data
+             when 'MX'
+               data = { :preference => rdata[0], :exchange => rdata[1] }
+               sprintf "%<preference>d, Resolv::DNS::Name.create('%<exchange>s')", data
+             when 'NS'
+               data = { :nsdname => rdata[0] }
+               sprintf "Resolv::DNS::Name.create('%<nsdname>s')", data
+             when 'PTR'
+               data = { :ptrdname => rdata[0] }
+               sprintf "Resolv::DNS::Name.create('%<ptrdname>s')", data
+             when 'SOA'
+               data = {
+                 :mname   => rdata[0],
+                 :rname   => rdata[1],
+                 :serial  => rdata[2],
+                 :refresh => rdata[3],
+                 :retry   => rdata[4],
+                 :expire  => rdata[5],
+                 :minimum => rdata[6]
+               }
+
+               sprintf "Resolv::DNS::Name.create('%<mname>s'), " +
+                       "Resolv::DNS::Name.create('%<rname>s'), " +
+                       '%<serial>d, ' +
+                       '%<refresh>d, ' +
+                       '%<retry>d, ' +
+                       '%<expire>d, ' +
+                       '%<minimum>d',
+                       data
+             when 'TXT'
+               data = { :txtdata => rdata[0] }
+               sprintf "'%<txtdata>s'", data
+             when 'WKS'
+               data = {
+                 :address  => rdata[0],
+                 :protocol => rdata[1],
+                 :bitmap   => rdata[2]
+               }
+
+               sprintf "'%<address>s', %<protocol>d, %<bitmap>d", data
+             else
+               raise InvalidJsonError, 'Unknown "type" key passed to endpoint /api/dns/rule'
+             end
+
+      sprintf(src, args)
+    end
+
+    # Raised when invalid JSON input is passed to an /api/dns handler.
+    class InvalidJsonError < StandardError
+
+      DEFAULT_MESSAGE = 'Invalid JSON input passed to /api/dns handler'
+
+      def initialize(message = nil)
+        super(message || DEFAULT_MESSAGE)
+      end
+
+    end
+
+    # Raised when an invalid named parameter is passed to an /api/dns handler.
+    class InvalidParamError < StandardError
+
+      DEFAULT_MESSAGE = 'Invalid parameter passed to /api/dns handler'
+
+      def initialize(message = nil)
+        super(message || DEFAULT_MESSAGE)
+      end
+
+    end
+
+  end
+
+end
+end
+end

extensions/dns/ruby.rb

@@ -0,0 +1,7 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+require 'extensions/dns/ruby/logger'
+require 'extensions/dns/ruby/rubydns'

extensions/dns/ruby/logger.rb

@@ -0,0 +1,27 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+
+# Overrives the logger used by RubyDNS to use BeEF's {#print_info} and friends.
+class Logger
+
+  def debug(msg)
+    print_debug "DNS Server: #{msg}"
+  end
+
+  def info(msg)
+    print_info "DNS Server: #{msg}"
+  end
+
+  def error(msg)
+    print_error "DNS Server: #{msg}"
+  end
+
+  def warn(msg)
+    print_error "DNS Server: #{msg}"
+  end
+
+end
+

extensions/dns/ruby/rubydns.rb

@@ -0,0 +1,242 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+
+# This module is a modified version of RubyDNS built to be compatible with BeEF.
+# For the most part, it will behave exactly the same except where otherwise noted.
+#
+# Additional features include database support, BeEF logger, assignment of unique
+# identifiers to rules, rule removal, and more.
+#
+# The core functionality of BeEF's DNS server is implemented here, whereas
+# BeEF::Extension::Dns::Server is simply a small wrapper around it.
+#
+# @see http://rubydoc.info/gems/rubydns/frames
+module RubyDNS
+
+  # Behaves exactly the same, except without any logger output
+  def self.run_server(options = {}, &block)
+    server = RubyDNS::Server.new(&block)
+
+    options[:listen] ||= [[:udp, '0.0.0.0', 53], [:tcp, '0.0.0.0', 53]]
+
+    EventMachine.run do
+      server.fire(:setup)
+
+      options[:listen].each do |spec|
+        if spec[0] == :udp
+          EventMachine.open_datagram_socket(spec[1], spec[2], UDPHandler, server)
+        elsif spec[0] == :tcp
+          EventMachine.start_server(spec[1], spec[2], TCPHandler, server)
+        end
+      end
+
+      server.load_rules
+      server.fire(:start)
+    end
+
+    server.fire(:stop)
+  end
+
+  class Server
+
+    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
+
+    # New method that loads all rules from the database at server startup
+    def load_rules
+      BeEF::Core::Models::Dns::Rule.each do |rule|
+        id = rule.id
+        pattern = [rule.pattern, rule.type]
+        block = eval rule.block
+
+        regex = pattern[0]
+        pattern[0] = Regexp.new(regex) if regex =~ /^\(\?-mix:/
+
+        @rules << Rule.new(id, pattern, block)
+      end
+    end
+
+    # Now includes BeEF database support and checks for already present rules
+    def match(*pattern, block)
+      id = ''
+
+      catch :match do
+        begin
+          # Sourcify block (already a string only for RESTful API calls)
+          block_src = case block
+                      when String then block
+                      when Proc   then block.to_source
+                      end
+
+          # Break out and return id if rule is already present
+          BeEF::Core::Models::Dns::Rule.each do |rule|
+            if pattern[0] == rule.pattern &&
+               pattern[1] == rule.type &&
+               block_src  == rule.block
+
+              id = rule.id
+              throw :match
+            end
+          end
+
+          id = generate_id
+
+          case block
+          when String
+            @rules << Rule.new(id, pattern, eval(block_src))
+          when Proc
+            @rules << Rule.new(id, pattern, block)
+          end
+
+          BeEF::Core::Models::Dns::Rule.create(
+            :id => id,
+            :pattern => pattern[0].to_s,
+            :type => pattern[1],
+            :block => block_src
+          )
+        rescue Sourcify::CannotHandleCreatedOnTheFlyProcError,
+               Sourcify::CannotParseEvalCodeError,
+               Sourcify::MultipleMatchingProcsPerLineError,
+               Sourcify::NoMatchingProcError,
+               Sourcify::ParserInternalError
+
+          @logger.error "Failed to sourcify block for DNS rule '#{id}'"
+          raise
+        end
+      end
+
+      id
+    end
+
+    # New method that removes a rule given its id and returns boolean result
+    def remove_rule(id)
+      @rules.delete_if { |rule| rule.id == id }
+
+      rule = BeEF::Core::Models::Dns::Rule.get(id)
+
+      rule != nil ? rule.destroy : false
+    end
+
+    # New method that returns a hash representing the given rule
+    def get_rule(id)
+      result = {}
+
+      begin
+        rule = BeEF::Core::Models::Dns::Rule.get!(id)
+
+        result[:id] = rule.id
+        result[:pattern] = rule.pattern
+        result[:type] = rule.type.to_s.split('::')[-1]
+        result[:response] = parse_response(rule.block)
+      rescue DataMapper::ObjectNotFoundError => e
+        @logger.error(e.message)
+      end
+
+      result
+    end
+
+    # New method that returns the entire DNS ruleset as an AoH
+    def get_ruleset
+      result = []
+
+      BeEF::Core::Models::Dns::Rule.each do |rule|
+        element = {}
+
+        element[:id] = rule.id
+        element[:pattern] = rule.pattern
+        element[:type] = rule.type.to_s.split('::')[-1]
+        element[:response] = parse_response(rule.block)
+
+        result << element
+      end
+
+      result
+    end
+
+    # New method that removes the entire DNS ruleset
+    def remove_ruleset
+      @rules = []
+      BeEF::Core::Models::Dns::Rule.destroy
+    end
+
+    private
+
+    # New method that generates a unique id for a rule
+    def generate_id
+      begin
+        id = BeEF::Core::Crypto.secure_token.byteslice(0..6)
+
+        # Make sure id isn't already in use
+        BeEF::Core::Models::Dns::Rule.each { |rule| throw StandardError if id == rule.id }
+      rescue StandardError
+        retry
+      end
+
+      id
+    end
+
+    # New method that parses response callback and returns RDATA as an array
+    def parse_response(block)
+      # Extract response arguments into an array
+      methods = '(respond|failure)'
+      args = /(?<=\.#{methods}!\().*(?=\))/.match(block).to_s.split(/,\s*/)
+
+      result = []
+
+      # Determine whether each argument is a domain name, integer, or IP address
+      args.each do |elem|
+        arg = nil
+
+        if /Name\.create\((.*)\)/.match(elem)
+          arg = $1
+        elsif /:(NoError|FormErr|ServFail|NXDomain|NotImp|Refused|NotAuth)/.match(elem)
+          arg = $1.upcase
+        else
+          int_test = elem.to_i
+          arg = (int_test != 0 ? int_test : elem)
+        end
+
+        arg.gsub!(/['"]/, '') unless arg.is_a?(Integer)
+
+        result << arg
+      end
+
+      result
+    end
+
+  end
+
+  class Transaction
+
+    # Behaves exactly the same, except using debug logger instead of info
+    def respond!(*data)
+      options = data.last.kind_of?(Hash) ? data.pop : {}
+      resource_class = options[:resource_class] || @resource_class
+
+      if resource_class == nil
+        raise ArgumentError, "Could not instantiate resource #{resource_class}!"
+      end
+
+      @server.logger.debug("Resource class: #{resource_class.inspect}")
+      resource = resource_class.new(*data)
+      @server.logger.debug("Resource: #{resource.inspect}")
+
+      append!(resource, options)
+    end
+
+  end
+
+end

extensions/social_engineering/rest/socialengineering.rb

@@ -20,10 +20,15 @@ module BeEF
                   'Expires' => '0'
         end
 
-        #Example: curl -H "Content-Type: application/json; charset=UTF-8"
+        # Example: curl -H "Content-Type: application/json; charset=UTF-8" -d json_body
-        #-d '{"url":"https://accounts.google.com/ServiceLogin?service=mail&passive=true&rm=false&continue=
+        # -X POST http://127.0.0.1:3000/api/seng/clone_page?token=851a937305f8773ee82f5259e792288cdcb01cd7
-        #https://mail.google.com/mail/&ss=1&scc=1&ltmpl=default&ltmplcache=2", "mount":"/url"}'
+        #
-        #-X POST http://127.0.0.1:3000/api/seng/clone_page?token=851a937305f8773ee82f5259e792288cdcb01cd7
+        # Example json_body:
+        # {
+        #     "url": "https://accounts.google.com/ServiceLogin?service=mail&continue=https://mail.google.com/mail/"
+        #     "mount": "/gmail",
+        #     "dns_spoof": true
+        # }
         post '/clone_page' do
           request.body.rewind
           begin
@@ -31,6 +36,7 @@ module BeEF
             uri = body["url"]
             mount = body["mount"]
             use_existing = body["use_existing"]
+            dns_spoof = body["dns_spoof"]
 
             if uri != nil && mount != nil
               if (uri =~ URI::regexp).nil? #invalid URI
@@ -44,7 +50,8 @@ module BeEF
               end
 
               web_cloner = BeEF::Extension::SocialEngineering::WebCloner.instance
-              success = web_cloner.clone_page(uri,mount,use_existing)
+              success = web_cloner.clone_page(uri, mount, use_existing, dns_spoof)
+
               if success
                 result = {
                     "success" => true,
@@ -125,4 +132,4 @@ module BeEF
       end
     end
   end
-end
+end

extensions/social_engineering/web_cloner/web_cloner.rb

@@ -7,9 +7,9 @@ module BeEF
   module Extension
     module SocialEngineering
       class WebCloner
+        require 'socket'
         include Singleton
 
-
         def initialize
           @http_server = BeEF::Core::Server.instance
           @config = BeEF::Core::Configuration.instance
@@ -17,7 +17,7 @@ module BeEF
           @beef_hook = "http://#{@config.get('beef.http.host')}:#{@config.get('beef.http.port')}#{@config.get('beef.http.hook_file')}"
         end
 
-        def clone_page(url, mount, use_existing)
+        def clone_page(url, mount, use_existing, dns_spoof)
           print_info "Cloning page at URL #{url}"
           uri = URI(url)
           output = uri.host
@@ -113,6 +113,18 @@ module BeEF
             @http_server.mount("#{mount}", interceptor.new)
             print_info "Mounting cloned page on URL [#{mount}]"
             @http_server.remap
+
+            # Add a DNS record spoofing the address of the cloned webpage as the BeEF server
+            if dns_spoof
+              dns = BeEF::Extension::Dns::Server.instance
+              ip = Socket.ip_address_list.detect {|i| !(i.ipv4_loopback? || i.ipv6_loopback?)}
+              domain = url.gsub(%r{^http://}, '')
+
+              id = dns.add_rule(domain, Resolv::DNS::Resource::IN::A) do |transaction|
+                transaction.respond!(ip.ip_address)
+              end
+            end
+
             success = true
           else
             print_error "Error cloning #{url}. Be sure that you don't have errors while retrieving the page with 'wget'."

test/common/test_constants.rb

@@ -19,4 +19,6 @@ BEEF_PASSWD = "beef"
 RESTAPI_HOOKS = "http://" + ATTACK_DOMAIN + ":3000/api/hooks"
 RESTAPI_LOGS = "http://" + ATTACK_DOMAIN + ":3000/api/logs"
 RESTAPI_MODULES = "http://" + ATTACK_DOMAIN + ":3000/api/modules"
-RESTAPI_ADMIN = "http://" + ATTACK_DOMAIN + ":3000/api/admin"
+RESTAPI_DNS = "http://" + ATTACK_DOMAIN + ":3000/api/dns"
+RESTAPI_SENG = "http://" + ATTACK_DOMAIN + ":3000/api/seng"
+RESTAPI_ADMIN = "http://" + ATTACK_DOMAIN + ":3000/api/admin"

test/integration/tc_dns_rest.rb

@@ -0,0 +1,388 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+require 'test/unit'
+require 'rest_client'
+require 'json'
+require '../common/test_constants'
+
+class TC_DnsRest < Test::Unit::TestCase
+
+  class << self
+
+    def startup
+      json = {:username => BEEF_USER, :password => BEEF_PASSWD}.to_json
+      @@headers = {:content_type => :json, :accept => :json}
+
+      response = RestClient.post("#{RESTAPI_ADMIN}/login",
+                                 json,
+                                 @@headers)
+
+      result  = JSON.parse(response.body)
+      @@token = result['token']
+
+      $root_dir = '../../'
+      $:.unshift($root_dir)
+
+      require 'core/loader'
+
+      BeEF::Core::Configuration.new(File.join($root_dir, 'config.yaml'))
+      BeEF::Core::Configuration.instance.load_extensions_config
+
+      @@config = BeEF::Core::Configuration.instance
+    end
+
+    def shutdown
+      $root_dir = nil
+    end
+
+  end
+
+  # Tests POST /api/dns/rule handler with valid input
+  def test_1_add_rule_good
+    pattern = 'foo.bar'
+    type = 'A'
+    dns_response = ['1.2.3.4']
+
+    json = {:pattern => pattern, :type => type, :response => dns_response}.to_json
+
+    rest_response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                    json,
+                                    @@headers)
+
+    check_rest_response(rest_response)
+
+    result = JSON.parse(rest_response.body)
+    first_id = result['id']
+
+    rest_response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                    json,
+                                    @@headers)
+
+    # Verify that adding an existing rule returns its id
+    check_rest_response(rest_response)
+
+    result = JSON.parse(rest_response.body)
+    second_id = result['id']
+
+    assert_equal(first_id, second_id)
+  end
+
+  # Tests POST /api/dns/rule handler with invalid input
+  def test_2_add_rule_bad
+    pattern = ''
+    type = 'A'
+    dns_response = ['1.1.1.1']
+
+    hash = {:pattern => pattern, :type => type, :response => dns_response}
+
+    # Test that an empty "pattern" key returns 400
+    assert_raise RestClient::BadRequest do
+      rest_response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                      hash.to_json,
+                                      @@headers)
+    end
+
+    hash['pattern'] = 'foo.bar.baz'
+    hash['type'] = ''
+
+    # Test that an empty "type" key returns 400
+    assert_raise RestClient::BadRequest do
+      rest_response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                      hash.to_json,
+                                      @@headers)
+    end
+
+    hash['type'] = 'A'
+    hash['response'] = []
+
+    # Test that an empty "response" key returns 400
+    assert_raise RestClient::BadRequest do
+      rest_response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                      hash.to_json,
+                                      @@headers)
+    end
+
+    hash['response'] = 42
+
+    # Test that a non-array "response" key returns 400
+    assert_raise RestClient::BadRequest do
+      rest_response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                      hash.to_json,
+                                      @@headers)
+    end
+  end
+
+  # Tests POST /api/dns/rule handler with each supported RR type
+  def test_3_add_rule_types
+    pattern = 'be.ef'
+    type = 'AAAA'
+    response = ['2001:db8:ac10:fe01::']
+
+    # Test AAAA type
+    rule = {'pattern' => pattern, 'type' => type, 'response' => response}
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      #{rule['response'][0]}$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test CNAME type
+    rule['type'] = 'CNAME'
+    rule['response'] = ['fe.eb.']
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      #{rule['response'][0]}$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test HINFO type
+    rule['type'] = 'HINFO'
+    rule['response'] = ['M6800', 'VMS']
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      "#{rule['response'][0]}"\s+
+      "#{rule['response'][1]}"$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test MINFO type
+    rule['type'] = 'MINFO'
+    rule['response'] = ['rmail.be.ef.', 'email.be.ef.']
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      #{rule['response'][0]}\s+
+      #{rule['response'][1]}$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test MX type
+    rule['type'] = 'MX'
+    rule['response'] = [10, 'mail.be.ef.']
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      #{rule['response'][0]}\s+
+      #{rule['response'][1]}$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test NS type
+    rule['type'] = 'NS'
+    rule['response'] = ['ns.be.ef.']
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      #{rule['response'][0]}$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test PTR type
+    rule['type'] = 'PTR'
+    rule['response'] = ['4.3.2.1.in-addr.arpa.']
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      #{rule['response'][0]}$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test SOA type
+    rule['type'] = 'SOA'
+    rule['response'] = [
+      "ns.#{rule['pattern']}.",
+      "mail.#{rule['pattern']}.",
+      2012031500,
+      10800,
+      3600,
+      604800,
+      3600
+    ]
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      .*
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test TXT type
+    rule['type'] = 'TXT'
+    rule['response'] = ['b33f_is_s0_l33t']
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      "#{rule['response'][0]}"$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test WKS type
+    rule['type'] = 'WKS'
+    rule['response'] = ['9.9.9.9', 6, 0]
+
+    regex = %r{
+      ^#{rule['pattern']}\.\t+
+      \d+\t+
+      IN\t+
+      #{rule['type']}\t+
+      #{rule['response'][0]}\s
+      0\s5\s6$
+    }x
+
+    add_rule(rule)
+    check_dns_response(regex, rule['type'], rule['pattern'])
+
+    # Test that an invalid RR returns 400
+    rule['type'] = 'BeEF'
+
+    assert_raise RestClient::BadRequest do
+      rest_response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                      rule.to_json,
+                                      @@headers)
+    end
+  end
+
+  # Tests GET /api/dns/rule/:id handler with valid input
+  def test_4_get_rule_good
+    pattern = 'wheres.the.beef'
+    type = 'A'
+    dns_response = ['4.2.4.2']
+
+    json = {:pattern => pattern, :type => type, :response => dns_response}.to_json
+
+    rest_response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                    json,
+                                    @@headers)
+
+    check_rest_response(rest_response)
+    result = JSON.parse(rest_response.body)
+    id = result['id']
+
+    rest_response = RestClient.get("#{RESTAPI_DNS}/rule/#{id}", :params => {:token => @@token})
+
+    assert_not_nil(rest_response.body)
+    assert_equal(200, rest_response.code)
+
+    result = JSON.parse(rest_response.body)
+
+    assert_equal(id, result['id'])
+    assert_equal(pattern, result['pattern'])
+    assert_equal(type, result['type'])
+    assert_equal(dns_response, result['response'])
+  end
+
+  # Tests GET /api/dns/rule/:id handler with invalid input
+  def test_4_get_rule_bad
+    id = 42
+
+    assert_raise RestClient::ResourceNotFound do
+      response = RestClient.get("#{RESTAPI_DNS}/rule/#{id}", :params => {:token => @@token})
+    end
+
+    id = '(*_*)'
+
+    assert_raise RestClient::BadRequest do
+      RestClient.get("#{RESTAPI_DNS}/rule/#{id}", :params => {:token => @@token})
+    end
+  end
+
+  # Tests GET /api/dns/ruleset handler
+  def test_4_get_ruleset
+    rest_response = RestClient.get("#{RESTAPI_DNS}/ruleset", :params => {:token => @@token})
+
+    assert_not_nil(rest_response.body)
+    assert_equal(200, rest_response.code)
+
+    result = JSON.parse(rest_response.body)
+    assert_equal(15, result['count'])
+
+    result['ruleset'].each do |rule|
+      assert(rule['id'])
+      assert(rule['pattern'])
+      assert(rule['type'])
+      assert(rule['response'].length != 0)
+    end
+  end
+
+  private
+
+  # Adds a new DNS rule
+  def add_rule(params)
+    response = RestClient.post("#{RESTAPI_DNS}/rule?token=#{@@token}",
+                                    params.to_json,
+                                    @@headers)
+
+    check_rest_response(response)
+  end
+
+  # Standard assertions for verifying response from RESTful API
+  def check_rest_response(response)
+    assert_not_nil(response.body)
+    assert_equal(200, response.code)
+
+    result = JSON.parse(response.body)
+
+    assert(result['success'])
+    assert(result['id'])
+  end
+
+  # Compares output of dig command against regex
+  def check_dns_response(regex, type, pattern)
+    address = @@config.get('beef.extension.dns.address')
+    port    = @@config.get('beef.extension.dns.port')
+
+    dig_output = `dig @#{address} -p #{port} -t #{type} #{pattern}`
+    assert_match(regex, dig_output)
+  end
+
+end

test/integration/tc_social_engineering_rest.rb

@@ -0,0 +1,91 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+require 'test/unit'
+require 'rest_client'
+require 'json'
+require '../common/test_constants'
+
+# @todo RESTful API for the social engineering extension lacks some serious test coverage.
+class TC_SocialEngineeringRest < Test::Unit::TestCase
+
+  class << self
+
+    # Login to API before performing any tests
+    def startup
+      json = {:username => BEEF_USER, :password => BEEF_PASSWD}.to_json
+      @@headers = {:content_type => :json, :accept => :json}
+
+      response = RestClient.post("#{RESTAPI_ADMIN}/login",
+                                 json,
+                                 @@headers)
+
+      result  = JSON.parse(response.body)
+      @@token = result['token']
+
+      $root_dir = '../../'
+      $:.unshift($root_dir)
+
+      require 'core/loader'
+
+      BeEF::Core::Configuration.new(File.join($root_dir, 'config.yaml'))
+      BeEF::Core::Configuration.instance.load_extensions_config
+
+      @@config = BeEF::Core::Configuration.instance
+    end
+
+    def shutdown
+      $root_dir = nil
+    end
+
+  end
+
+  # Tests DNS spoofing of cloned webpages
+  def test_1_dns_spoof
+    url = 'http://beefproject.com'
+    mount = '/beefproject'
+    dns_spoof = true
+
+    json = {:url => url, :mount => mount, :dns_spoof => dns_spoof}.to_json
+
+    response = RestClient.post("#{RESTAPI_SENG}/clone_page?token=#{@@token}",
+                                    json,
+                                    @@headers)
+
+    check_response(response)
+
+    ip = Socket.ip_address_list.detect {|i| !(i.ipv4_loopback? || i.ipv6_loopback?)}
+    domain = url.gsub(%r{^http://}, '')
+
+    regex = %r{
+      ^#{domain}\.\t+
+      \d+\t+
+      IN\t+
+      A\t+
+      #{ip.ip_address}$
+    }x
+
+    # Send DNS request to server to verify that a new rule was added
+    dns_address = @@config.get('beef.extension.dns.address')
+    dns_port    = @@config.get('beef.extension.dns.port')
+
+    dig_output = `dig @#{dns_address} -p #{dns_port} -t A #{domain}`
+    assert_match(regex, dig_output)
+  end
+
+  private
+
+  # Assertions for verifying a response from the RESTful API
+  def check_response(response)
+    assert_not_nil(response.body)
+    assert_equal(200, response.code)
+
+    result = JSON.parse(response.body)
+
+    assert(result['success'])
+    assert(result['mount'])
+  end
+
+end

test/integration/ts_dns_rest.rb

@@ -0,0 +1,20 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+require '../common/ts_common'
+require './tc_dns_rest'
+
+class TS_DnsIntegrationTests
+
+  def self.suite
+    suite = Test::Unit::TestSuite.new(name="BeEF DNS Integration Test Suite")
+    suite << TC_DnsRest.suite
+
+    return suite
+  end
+
+end
+
+Test::Unit::UI::Console::TestRunner.run(TS_DnsIntegrationTests)

test/integration/ts_integration.rb

@@ -16,6 +16,8 @@ require './check_environment' # Basic log in and log out tests
 require './tc_debug_modules' # RESTful API tests (as well as debug modules)
 require './tc_login' # Basic log in and log out tests
 require './tc_jools' # Basic tests for jools
+require './tc_dns_rest' # Basic tests for DNS RESTful API interface
+require './tc_social_engineering_rest' # Basic tests for social engineering RESTful API interface
 
 class TS_BeefIntegrationTests
   def self.suite
@@ -25,6 +27,8 @@ class TS_BeefIntegrationTests
     suite << TC_login.suite
     suite << TC_DebugModules.suite
     suite << TC_Jools.suite
+    suite << TC_DnsRest.suite
+    suite << TC_SocialEngineeringRest.suite
 
     return suite
   end

test/unit/extensions/tc_dns.rb

@@ -0,0 +1,310 @@
+#
+# Copyright (c) 2006-2013 Wade Alcorn - wade@bindshell.net
+# Browser Exploitation Framework (BeEF) - http://beefproject.com
+# See the file 'doc/COPYING' for copying permission
+#
+require 'test/unit'
+require 'resolv'
+
+class TC_Dns < Test::Unit::TestCase
+
+  IN = Resolv::DNS::Resource::IN
+
+  class << self
+
+    def startup
+      $root_dir = '../../'
+      $:.unshift(File.expand_path($root_dir))
+
+      require 'extensions/dns/extension'
+
+      BeEF::Core::Configuration.new(File.join($root_dir, 'config.yaml'))
+      config = BeEF::Core::Configuration.instance
+      config.load_extensions_config
+
+      @@dns_config = config.get('beef.extension.dns')
+    end
+
+    def shutdown
+      $root_dir = nil
+    end
+
+  end
+
+  # Connects to in-memory database (does not test anything)
+  def test_01_database
+    DataMapper.setup(:default, 'sqlite3::memory:')
+    DataMapper.auto_migrate!
+  end
+
+  # Checks for required settings in config file
+  def test_02_config
+    assert(@@dns_config.has_key?('address'))
+    assert(@@dns_config.has_key?('port'))
+  end
+
+  # Verifies public interface
+  def test_03_interface
+    @@dns = BeEF::Extension::Dns::Server.instance
+
+    assert_respond_to(@@dns, :run_server)
+    assert_respond_to(@@dns, :add_rule)
+    assert_respond_to(@@dns, :remove_rule)
+    assert_respond_to(@@dns, :get_rule)
+    assert_respond_to(@@dns, :get_ruleset)
+    assert_respond_to(@@dns, :remove_ruleset)
+  end
+
+  # Tests that DNS server runs correctly on desired address and port
+  def test_04_run_server
+    address = @@dns_config['address']
+    port    = @@dns_config['port']
+
+    @@dns.run_server(address, port)
+    sleep(3)
+
+    assert_equal(address, @@dns.address)
+    assert_equal(port, @@dns.port)
+  end
+
+  # Tests procedure for properly adding new DNS rules
+  def test_05_add_rule_good
+    id1 = nil
+    id2 = nil
+
+    assert_nothing_raised do
+      id1 = @@dns.add_rule('foo.bar', IN::A) do |transaction|
+        transaction.respond!('1.2.3.4')
+      end
+    end
+
+    assert_not_nil(id1)
+
+    assert_nothing_raised do
+      id2 = @@dns.add_rule(%r{i\.(love|hate)\.beef\.com?}, IN::A) do |transaction|
+        transaction.respond!('9.9.9.9')
+      end
+    end
+
+    assert_not_nil(id2)
+
+    domain1 = 'i.hate.beef.com'
+    domain2 = 'i.love.beef.com'
+    domain3 = 'i.love.beef.co'
+    domain4 = 'i.love.beef.co'
+
+    [domain1, domain2, domain3, domain4].each do |domain|
+      regex = /^#{domain}\.\t+\d+\t+IN\t+A\t+9\.9\.9\.9$/
+      check_dns_response(regex, 'A', domain)
+    end
+  end
+
+  # Tests addition of new rules with invalid parameters
+  def test_06_add_rule_bad
+    id = nil
+    same_id = nil
+
+    # Add the same rule twice
+    assert_nothing_raised do
+      id = @@dns.add_rule('j.random.hacker', IN::A) do |transaction|
+        transaction.respond!('4.2.4.2')
+      end
+    end
+
+    assert_nothing_raised do
+      same_id = @@dns.add_rule('j.random.hacker', IN::A) do |transaction|
+        transaction.respond!('4.2.4.2')
+      end
+    end
+
+    assert_equal(id, same_id)
+
+    # Use /.../ literal syntax to throw Sourcify exception
+    assert_raise do
+      id = @@dns.add_rule(/.*/, IN::A) do |transaction|
+        transaction.respond!('5.1.5.0')
+      end
+    end
+  end
+
+  # Verifies the proper format for rule identifiers
+  def test_07_id_format
+    id = @@dns.add_rule('dead.beef', IN::A) do |transaction|
+      transaction.respond!('2.2.2.2')
+    end
+
+    assert_equal(7, id.length)
+    assert_not_nil(id =~ /^\h{7}$/)
+  end
+
+  # Tests retrieval of valid DNS rules
+  def test_08_get_rule_good
+    id = @@dns.add_rule('be.ef', IN::A) do |transaction|
+      transaction.respond!('1.1.1.1')
+    end
+
+    rule = @@dns.get_rule(id)
+
+    assert_equal(Hash, rule.class)
+    assert(rule.length > 0)
+
+    assert(rule.has_key?(:id))
+    assert(rule.has_key?(:pattern))
+    assert(rule.has_key?(:type))
+    assert(rule.has_key?(:response))
+
+    assert_equal(id, rule[:id])
+    assert_equal('be.ef', rule[:pattern])
+    assert_equal('A', rule[:type])
+
+    response = rule[:response]
+
+    assert_equal(Array, response.class)
+    assert(response.length > 0)
+    assert_equal('1.1.1.1', response[0])
+  end
+
+  # Tests retrieval of invalid DNS rules
+  def test_09_get_rule_bad
+    rule = @@dns.get_rule(42)
+
+    assert_equal(Hash, rule.class)
+    assert_equal(0, rule.length)
+  end
+
+  # Tests the removal of existing DNS rules
+  def test_10_remove_rule_good
+    id = @@dns.add_rule('hack.the.gibson', IN::A) do |transaction|
+      transaction.respond!('1.9.9.5')
+    end
+
+    removed = @@dns.remove_rule(id)
+
+    assert(removed)
+  end
+
+  # Tests the removal of unknown DNS rules
+  def test_11_remove_rule_bad
+    removed = @@dns.remove_rule(42)
+
+    assert(!removed)
+  end
+
+  # Tests the retrieval of the entire DNS ruleset
+  def test_12_get_ruleset
+    ruleset = @@dns.get_ruleset
+    ruleset.sort! {|a, b| a[:pattern] <=> b[:pattern] }
+
+    assert_equal(Array, ruleset.class)
+    assert_equal(5, ruleset.length)
+
+    check_rule(ruleset[0], {:pattern=>'(?-mix:i\\.(love|hate)\\.beef\\.com?)',
+                            :type => 'A',
+                            :response => '9.9.9.9'})
+
+    check_rule(ruleset[1], {:pattern => 'be.ef', :type => 'A', :response => '1.1.1.1'})
+    check_rule(ruleset[2], {:pattern => 'dead.beef', :type => 'A', :response => '2.2.2.2'})
+    check_rule(ruleset[3], {:pattern => 'foo.bar', :type => 'A', :response => '1.2.3.4'})
+    check_rule(ruleset[4], {:pattern => 'j.random.hacker', :type => 'A', :response => '4.2.4.2'})
+  end
+
+  # Tests the removal of the entire DNS ruleset
+  def test_13_remove_ruleset
+    removed = @@dns.remove_ruleset
+    ruleset = @@dns.get_ruleset
+
+    assert(removed)
+    assert_equal(0, ruleset.length)
+  end
+
+  # Tests each supported type of query failure
+  def test_14_failure_types
+    begin
+      id = @@dns.add_rule('noerror.beef.com', IN::A) do |transaction|
+        transaction.failure!(:NoError)
+      end
+
+      check_failure_status(id, :NoError)
+    end
+
+    begin
+      id = @@dns.add_rule('formerr.beef.com', IN::A) do |transaction|
+        transaction.failure!(:FormErr)
+      end
+
+      check_failure_status(id, :FormErr)
+    end
+
+    begin
+      id = @@dns.add_rule('servfail.beef.com', IN::A) do |transaction|
+        transaction.failure!(:ServFail)
+      end
+
+      check_failure_status(id, :ServFail)
+    end
+
+    begin
+      id = @@dns.add_rule('nxdomain.beef.com', IN::A) do |transaction|
+        transaction.failure!(:NXDomain)
+      end
+
+      check_failure_status(id, :NXDomain)
+    end
+
+    begin
+      id = @@dns.add_rule('notimp.beef.com', IN::A) do |transaction|
+        transaction.failure!(:NotImp)
+      end
+
+      check_failure_status(id, :NotImp)
+    end
+
+    begin
+      id = @@dns.add_rule('refused.beef.com', IN::A) do |transaction|
+        transaction.failure!(:Refused)
+      end
+
+      check_failure_status(id, :Refused)
+    end
+
+    begin
+      id = @@dns.add_rule('notauth.beef.com', IN::A) do |transaction|
+        transaction.failure!(:NotAuth)
+      end
+
+      check_failure_status(id, :NotAuth)
+    end
+  end
+
+  private
+
+  # Compares each key in hash 'rule' with the respective key in hash 'expected'
+  def check_rule(rule, expected = {})
+    assert_equal(expected[:pattern], rule[:pattern])
+    assert_equal(expected[:type], rule[:type])
+    assert_equal(expected[:response], rule[:response][0])
+  end
+
+  # Confirms that a query for the rule given in 'id' returns a 'type' failure status
+  def check_failure_status(id, type)
+    rule   = @@dns.get_rule(id)
+    status = type.to_s.force_encoding('UTF-8').upcase
+    assert_equal(status, rule[:response][0])
+
+    check_dns_response(/status: #{status}/, rule[:type], rule[:pattern])
+  end
+
+  # Compares output of dig command against regex
+  def check_dns_response(regex, type, pattern)
+    dig_output = `dig @#{@@dns.address} -p #{@@dns.port} -t #{type} #{pattern}`
+    assert_match(regex, dig_output)
+  end
+
+end
+
+# Suppresses unnecessary output from RubyDNS
+module Kernel
+
+  def puts(*args); end
+
+end

test/unit/tc_grep.rb

@@ -12,12 +12,13 @@ class TC_Grep < Test::Unit::TestCase
       File.open( path ) do |f|
         next if /tc_grep.rb/.match(path) # skip this file
         next if /\/msf-test\//.match(path) # skip this file
+        next if /extensions\/dns/.match(path) # skip this file
+
         f.grep( /\Weval\W/im ) do |line|
           assert(false, "Illegal use of 'eval' in framework: " + path + ':' + line)
         end
       end
     end
-    
   end
 
 end

test/unit/ts_unit.rb

@@ -27,6 +27,7 @@ require './extensions/tc_hooks'
 require './extensions/tc_proxy'
 require './extensions/tc_requester'
 require './extensions/tc_event_logger'
+require './extensions/tc_dns'
 require './tc_grep'
 require './tc_filesystem'
 
@@ -55,6 +56,7 @@ class TS_BeefTests
     suite << TC_EventLogger.suite
 	suite << TC_Hooks.suite
     suite << TC_Redirector.suite
+    suite << TC_Dns.suite
 
     return suite
   end