# -*- coding: binary -*-
require 'msf/core'
require 'rex/proto/dns'


module Msf

###
#
# This module exposes methods for querying a remote DNS service
#
###
module Exploit::Remote::DNS
module Server
  include Common
  include Exploit::Remote::SocketServer

  #
  # Initializes an exploit module that serves DNS requests
  #
  def initialize(info = {})
    super

    register_options(
      [
        OptPort.new('SRVPORT', [true, 'The local port to listen on.', 53]),
        OptString.new('STATIC_ENTRIES', [ false, "DNS domain search list (hosts file or space/semicolon separate entries)"]),
        OptBool.new('DISABLE_RESOLVER', [ false, "Disable DNS request forwarding", false]),
        OptBool.new('DISABLE_NS_CACHE', [ false, "Disable DNS response caching", false])
      ], Exploit::Remote::DNS::Server
    )

    register_advanced_options(
      [
        OptBool.new('DnsServerUdp', [true, "Serve UDP DNS requests", true]),
        OptBool.new('DnsServerTcp', [true, "Serve TCP DNS requests", false])
      ], Exploit::Remote::DNS::Server
    )
  end

  attr_accessor :service

  #
  # Process static entries
  #
  # @param entries [String] Filename or String containing static entries
  # @param type [String] Type of record for which to add static entries
  #
  # @return [Array] List of static entries in the cache
  def add_static_hosts(entries = datastore['STATIC_ENTRIES'], type = 'A')
    return if entries.nil? or entries.empty?
    if File.file?(File.expand_path(entries))
      data = File.read(File.expand_path(entries)).split("\n")
    else
      data = entries.split(';')
    end
    data.each do |entry|
      next if entry.gsub(/\s/,'').empty?
      addr, names = entry.split(' ', 2)
      names.split.each do |name|
        name << '.' unless name[-1] == '.' or name == '*'
        service.cache.add_static(name, addr, type)
      end
    end
    service.cache.records.select {|r,e| e == 0}
  end

  #
  # Flush all static entries
  #
  def flush_static_hosts
    data.cache.records.select {|r,e| e == 0}.each do |flush|
      data.cache.delete(flush)
    end
  end

  #
  # Flush cache entries
  # @param static [TrueClass, FalseClass] flush static hosts
  def flush_cache(static = false)
    self.service.cache.stop(true)
    flush_static_hosts if static
    self.service.cache.start
  end

  #
  # Handle incoming requests
  # Override this method in modules to take flow control
  #
  def on_dispatch_request(cli, data)
    service.default_dispatch_request(cli,data)
  end

  #
  # Handle incoming requests
  # Override this method in modules to take flow control
  #
  def on_send_response(cli, data)
    cli.write(data)
  end

  #
  # Starts the server
  #
  def start_service
    options.validate(datastore) # This is a hack, DS values should not be Strings prior to this
    if !datastore['DISABLE_RESOLVER'] and self.respond_to?(:setup_resolver)
      setup_resolver
    end
    begin

      comm = _determine_server_comm
      self.service = Rex::Proto::DNS::Server.new(
        datastore['SRVHOST'],
        datastore['SRVPORT'],
        datastore['DnsServerUdp'],
        datastore['DnsServerTcp'],
        (datastore['DISABLE_RESOLVER'] ? false : @dns_resolver),
        comm,
        {'Msf' => framework, 'MsfExploit' => self}
      ) if self.service.nil?

      self.service.dispatch_request_proc = Proc.new do |cli, data|
        on_dispatch_request(cli,data)
      end
      self.service.send_response_proc = Proc.new do |cli, data|
        on_send_response(cli,data)
      end
      
      add_static_hosts
      self.service.start(!datastore['DISABLE_NS_CACHE'])

    rescue ::Errno::EACCES => e
      if (srvport.to_i < 1024)
        print_line(" ")
        print_error("Could not start the DNS server: #{e}.")
        print_error(
          "This module is configured to use a privileged port (#{srvport}). " +
          "On Unix systems, only the root user account is allowed to bind to privileged ports." +
          "Please run the framework as root to use this module."
        )
        print_error(
          "On Microsoft Windows systems, this error is returned when a process attempts to "+
          "listen on a host/port combination that is already in use. For example, Windows XP "+
          "will return this error if a process attempts to bind() over the system SMB/NetBIOS services."
        )
        print_line(" ")
      end
      raise e
    end
  end

  #
  # Stops the server
  # @param destroy [TrueClass,FalseClass] Dereference the server object
  def stop_service(destroy = false)
    self.service.stop unless self.service.nil?
    if destroy
      @dns_resolver = nil
      self.service = nil
    end
  end

  #
  # Resets the DNS server
  #
  def reset_service
    stop_service(true)
    start_service
  end

end
end
end