#
# 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)

    BeEF::Extension::Dns::Server.instance.set_server(server)

    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]
        # antisnatchor: this would be unsafe, but input gets validated in extensions/dns/rest/dns.rb (lines 95 to 105)
        # in this case input comes from the DB, but that data stored in the DB was originally coming from the now safe code
        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

          if @rules == nil
            @rules = []
          end

          case block
            when String
              # antisnatchor: this would be unsafe, but input gets validated in extensions/dns/rest/dns.rb (lines 95 to 105)
              @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