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

module Msf

###
#
# This module exposes methods for querying a remote DNS service
#
###
module Exploit::Remote::DNS
module Server
  include Exploit::Remote::DNS::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 # :nodoc:

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

      comm = _determine_server_comm
      self.service = Rex::ServiceManager.start(
        Rex::Proto::DNS::Server,
        datastore['SRVHOST'],
        datastore['SRVPORT'],
        datastore['DnsServerUdp'],
        datastore['DnsServerTcp'],
        (use_resolver? ? setup_resolver : false),
        comm,
        {'Msf' => framework, 'MsfExploit' => self}
      )

      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
      raise Rex::BindFailed.new(e.message)
    end
  end

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

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

  #
  # Determines if resolver is available and configured for use
  #
  def use_resolver?
    !datastore['DISABLE_RESOLVER'] and self.respond_to?(:setup_resolver)
  end

end
end
end