module Dert

# File lib/dert/dns.rb, line 44
def self.query(domain, method, list=nil)
  case method
    when CONSTANTS::ARIN
      return ARIN.query(domain)
    when CONSTANTS::AXFR
      return AXFR.query(domain)
    when CONSTANTS::BRT
      return BRT.query(domain, list, Dnsruby::Types.A)
    when CONSTANTS::IPV6
      return BRT.query(domain, list, Dnsruby::Types.AAAA)
    when CONSTANTS::RVL
      return RVL.query(domain)
    when CONSTANTS::SRV
      return SRV.query(domain)
    when CONSTANTS::STD
      return STD.query(domain)
    when CONSTANTS::TLD
      return TLD.query(domain)
    when CONSTANTS::GTLD
      return GTLD.query(domain)
  end
end

# File lib/dert/dns.rb, line 162
def self.run(options)
  type = 0

  # Validate settings for brute force
  if %w(ipv6 brt).include? options[:type]
    if options[:threads] == nil or options[:domain] == nil or options[:wordlist] == nil
      puts "Usage #{File.basename($0)} -e <brt|ipv6> -d <domain> -w <wordlist> -t <threads>"
      exit
    end
  end

  # RVL requires domain or a word list
  if options[:type] == 'rvl'
    if options[:domain] == nil and options[:wordlist] == nil
      puts "Usage #{File.basename($0)} -e rvl -d IP"
      puts "Usage #{File.basename($0)} -e rvl -w IPLIST"
      exit
    end
  end

  # Validate wordlist
  if options[:wordlist]
    unless File.exist?(options[:wordlist])
      puts 'Word List not found.'
      exit
    end
  end

  # Validate threads
  if options[:threads]
    if options[:threads] > 100 or options[:threads] < 1
      puts 'Thread count must be between 1 and 100'
      exit
    end
  else
    options[:threads] = 1
  end

  # Validate Output
  if options[:output]
    unless Dir.exists?(File.dirname(options[:output]))
      puts 'Output directory does not exists.'
      exit
    end
  end

  # Convert string type to integer type
  case options[:type].downcase
    when 'arin'
      type = 1
    when 'axfr'
      type = 2
    when 'brt'
      type = 3
    when 'ipv6'
      type = 4
    when 'rvl'
      type = 5
    when 'srv'
      type = 6
    when 'std'
      type = 7
    when 'tld'
      type = 8
    when 'gtld'
      type = 9
    else
      puts 'Wrong enumeration type. Try --help to view accepted enumeration inputs.'
      exit
  end

  # Start Enumeration
  results = self.start(options[:domain], type, options[:threads], options[:wordlist])

  # Save results to a file if specified
  if options[:output]
    File.open(options[:output], 'w') do |f|
      f.write(JSON.pretty_generate(results))
    end
  end

  # Print output to terminal unless silent
  unless options[:silent]
    puts 'Results:'
    if type == 1
      results.each do |x|
        puts "  Range: #{x[:cidr]}"
        puts "  Handle: #{x[:handle]}"
        puts "  Customer: #{x[:customer]}"
        puts "  Zip Code: #{x[:zip]}"
        puts ''
      end
    else
      results.each do |x|
        puts "  Hostname: #{x[:hostname]}"
        puts "    IP: #{x[:address]}"
        puts "    Type: #{x[:type]}"
      end
    end
  end

  # Return results as a hash
  results
end

# File lib/dert/dns.rb, line 76
def self.start(domain, method, threads = nil, word_list = nil, output = nil)

  results = []

  # Process for Brute Force DNS Enumeration
  if method == CONSTANTS::BRT or method == CONSTANTS::IPV6 or (method == CONSTANTS::RVL and word_list)

    # Count words/ips in list.
    count = File.foreach(word_list).inject(0) { |c, line| c+1 }
    # Words/IPs per thread
    per = (count / threads) + 1
    # Array of words/ips
    arr = []
    # Array of sets of words/ips per thread
    lists_per_thread = []

    thread_container = []

    # Parse words/ips from word list
    File.open(word_list).each_line do |x|
      if method == CONSTANTS::RVL
        tmp = Rex::Socket::RangeWalker.new(x.chomp.strip)
        if tmp.valid?
          tmp.each do |y|
            arr << y
          end
        end
      else
        arr << x.chomp
      end
    end

    # If word list count is greater than 50, use multiple threads.
    if arr.count > 50
      arr.each_slice(per) { |a| lists_per_thread << a }
    else
      lists_per_thread = [arr]
    end

    # Iterate through sets of words.
    lists_per_thread.each do |x|
      # Create a new thread and add it to a container
      thread_container << Thread.new {
        # Check if RVL, else BRT or IPV
        if method == CONSTANTS::RVL
          # Iterate through IP addresses
          ret = []
          x.each do |y|
            ret.concat(self.query(y, method))
          end
        else
          # Send a single set of words to brute force
          ret = self.query(domain, method, x)
        end
        # Grab thread output
        Thread.current[:output] = ret
      }
    end

    # Join all threads and grab their outputs
    thread_container.each do |t|
      t.join
      results += t[:output] unless t[:output].empty?
    end

    # Process for Single Enumeration
  else
    results = self.query(domain, method)
  end

  # Write output to file if specified
  if output
    File.open(output, 'w') do |f|
      f.write(JSON.pretty_generate(results))
    end
  end

  # Return Results for Console output
  results
end