# -*- 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 Client

  include Common
  include Exploit::Remote::Udp
  include Exploit::Remote::Tcp

  #
  # Initializes an exploit module that interacts with a DNS server.
  #
  def initialize(info = {})
    super

    deregister_options('RHOST')
    register_options(
      [
        Opt::RPORT(53),
        Opt::Proxies,
        OptString.new('DOMAIN', [ false, "The target domain name"]),
        OptString.new('NS', [ false, "Specify the nameservers to use for queries, space separated" ]),
        OptString.new('SEARCHLIST', [ false, "DNS domain search list, comma separated"]),
        OptInt.new('THREADS', [true, "Number of threads to use in threaded queries", 1])
      ], Exploit::Remote::DNS::Client
    )

    register_advanced_options(
      [
        OptString.new('DnsClientDefaultNS', [ false, "Specify the default to use for queries, space separated", '8.8.8.8 8.8.4.4' ]),   
        OptInt.new('DnsClientRetry', [ false, "Number of times to try to resolve a record if no response is received", 2]),
        OptInt.new('DnsClientRetryInterval', [ false, "Number of seconds to wait before doing a retry", 2]),
        OptBool.new('DnsClientReportARecords', [false, "Add hosts found via BRT and RVL to DB", true]),
        OptBool.new('DnsClientRVLExistingOnly', [false, "Only perform lookups on hosts in DB", true]),
        OptBool.new('DnsClientTcpDns', [false, "Run queries over TCP", false]),
        OptPath.new('DnsClientResolvconf', [true, "Resolvconf formatted configuration file to use for Resolver", "/dev/null"])
      ], Exploit::Remote::DNS::Client
    )

    register_autofilter_ports([ 53 ]) if respond_to?(:register_autofilter_ports)
    register_autofilter_services(%W{ dns }) if respond_to?(:register_autofilter_services)
  end


  #
  # Convenience wrapper around Resolver's query method - send DNS request
  #
  # @param domain [String] Domain for which to request a record
  # @param type [String] Type of record to request for domain
  #
  # @return [Dnsruby::RR] DNS response
  def query(domain = datastore['DOMAIN'], type = 'A')
    client.query(domain, type)
  end

  #
  # Performs a set of asynchronous lookups for an array of domain,type pairs
  #
  # @param queries [Array] Set of domain,type pairs to pass into #query
  # @param threadmax [Fixnum] Max number of running threads at a time
  # @param block [Proc] Code block to execute with the query result
  #
  # @return [Array] Resulting set of responses or responses processed by passed blocks
  def query_async(queries = [], threadmax = datastore['THREADS'], &block)
    running = []
    while !queries.empty?
      domain, type = queries.shift
      running << framework.threads.spawn("Module(#{self.refname})-#{domain} #{type}", false) do |qat|
        if block
          block.call(query(domain,type))
        else
          query(domain,type)
        end
      end
      while running.select(&:alive?).count >= threadmax
        Rex::ThreadSafe.sleep(1)
      end
    end
    return running.join
  end

  #
  # Switch DNS forwarders in resolver with thread safety
  #
  # @param ns [Array, String] List of (or single) nameservers to use
  def set_nameserver(ns = [])
    if ns.respond_to?(:split)
      ns = [ns]
    end
    @lock.synchronize do
      @dns_resolver.nameserver = ns.flatten
    end
  end

  #
  # Switch nameservers to use explicit NS or SOA for target
  #
  # @param domain [String] Domain for which to find SOA
  def switchdns(domain)
    if datastore['NS'].blank?
      resp_soa = client.query(target, "SOA")
      if (resp_soa)
        (resp_soa.answer.select { |i| i.is_a?(Dnsruby::RR::SOA)}).each do |rr|
          resp_1_soa = client.search(rr.mname)
          if (resp_1_soa and resp_1_soa.answer[0])
            set_nameserver(resp_1_soa.answer.map(&:address).compact.map(&:to_s))
            print_status("Set DNS Server to #{target} NS: #{client.nameserver.join(', ')}")
            break
          end
        end
      end
    else
      vprint_status("Using DNS Server: #{client.nameserver.join(', ')}")
      client.nameserver = process_nameservers
    end
  end

  #
  # Detect if target has wildcards enabled for a record type
  #
  # @param target [String] Domain to test
  # @param type [String] Record type to test
  #
  # @return [String] Address which is returned for wildcard requests
  def wildcard(domain, type = "A")
    addr = false
    rendsub = rand(10000).to_s
    response = query("#{rendsub}.#{target}", type)
    if response.answer.length != 0
      vprint_status("This domain has wildcards enabled!!")
      response.answer.each do |rr|
        print_status("Wildcard IP for #{rendsub}.#{target} is: #{rr.address.to_s}") if rr.class != Dnsruby::RR::CNAME
        addr = rr.address.to_s
      end
    end
    return addr
  end

  #
  # Create and configure Resolver object
  #
  def setup_resolver
    options.validate(datastore) # This is a hack, DS values should not be Strings prior to this
    config = {
      :config_file => datastore['DnsClientResolvconf'],
      :nameservers => process_nameservers,
      :port => datastore['RPORT'],
      :retry_number => datastore['DnsClientRetry'].to_i,
      :retry_interval => datastore['DnsClientRetryInterval'].to_i,
      :use_tcp => datastore['DnsClientTcpDns'],
      :context => {'Msf' => framework, 'MsfExploit' => self}
    }
    if datastore['SEARCHLIST']
      if datastore['SEARCHLIST'].split(',').all? do |search|
        search.match(MATCH_HOSTNAME)
      end
        config[:search_list] = datastore['SEARCHLIST'].split(',')
      else
        raise 'Domain search list must consist of valid domains'
      end
    end
    if datastore['CHOST']
      config[:source_address] = IPAddr.new(datastore['CHOST'].to_s)
    end
    if datastore['CPORT']
      config[:source_port] = datastore['CPORT'] unless datastore['CPORT'] == 0
    end
    if datastore['Proxies']
      vprint_status("Using DNS/TCP resolution for proxy config")
      config[:use_tcp] = true
      config[:proxies] = datastore['Proxies']
    end
    @dns_resolver_lock = Mutex.new unless @dns_resolver_lock
    @dns_resolver = Rex::Proto::DNS::Resolver.new(config)
  end

  #
  # Convenience method for DNS resolver as client
  # Executes setup_resolver if none exists
  #
  def client
    setup_resolver unless @dns_resolver
    @dns_resolver
  end

  #
  # Sets the resolver's nameservers
  # Uses explicitly defined NS option if set
  # Uses RHOSTS if not explicitly defined
  def process_nameservers
    if datastore['NS'].blank?
      nameservers = datastore['DnsClientDefaultNS'].split(/\s|,/)
    else
      nameservers = datastore['NS'].split(/\s|,/)
    end

    invalid = nameservers.select { |ns| !Rex::Socket.dotted_ip?(ns) } 
    if !invalid.empty?
      raise "Nameservers must be IP addresses. The following were invalid: #{invalid.join(", ")}"
    end

    nameservers
  end

end
end
end