#!/usr/bin/env ruby
# Copyright, 2009, 2012, by Samuel G. D. Williams. <http://www.codeotaku.com>
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

require 'rubydns/version'

require 'resolv'
require 'optparse'

OPTIONS = {
	:Domains => [],
	:Nameservers => [],
	:Timeout => 0.5,

	:Threads => 10,
	:Requests => 20
}

ARGV.options do |o|
	script_name = File.basename($0)

	o.banner = "Usage: #{script_name} [options] nameserver [nameserver]"

	o.on("-d [path]", String, "Specify a file that contains a list of domains") do |path|
		OPTIONS[:Domains] += File.readlines(path).collect { |name| name.strip.downcase }
	end

	o.on("-t [timeout]", Float, "Queries that take longer than this will be printed") do |timeout|
		OPTIONS[:Timeout] = timeout.to_f
	end

	o.on("--threads [count]", Integer, "Number of threads to resolve names concurrently") do |count|
		OPTIONS[:Threads] = count.to_i
	end

	o.on("--requests [count]", Integer, "Number of requests to perform per thread") do |count|
		OPTIONS[:Requests] = count.to_i
	end

	o.on_tail("--copy", "Display copyright information") do
		puts "#{script_name} v#{RubyDNS::VERSION::STRING}. Copyright (c) 2009, 2011 Samuel Williams. Released under the MIT license."
		puts "See http://www.oriontransfer.co.nz/ for more information."

		exit
	end

	o.on_tail("-h", "--help", "Show this help message.") { puts o; exit }
end.parse!

OPTIONS[:Nameservers] = ARGV

if OPTIONS[:Nameservers].size > 0
	$R = Resolv::DNS.new(:nameserver => ARGV)
else
	$R = Resolv::DNS.new
end

$TG = ThreadGroup.new
$M = Mutex.new
$STATUS = {}
$TOTAL = [0.0, 0]

if OPTIONS[:Domains].size == 0
	OPTIONS[:Domains] += ["www.google.com", "www.amazon.com", "www.apple.com", "www.microsoft.com"]
	OPTIONS[:Domains] += ["www.rubygems.org", "www.ruby-lang.org", "www.slashdot.org", "www.lucidsystems.org"]
	OPTIONS[:Domains] += ["www.facebook.com", "www.twitter.com", "www.myspace.com", "www.youtube.com"]
	OPTIONS[:Domains] += ["www.oriontransfer.co.nz", "www.digg.com"]
end

def random_domain
	d = OPTIONS[:Domains]

	d[rand(d.size - 1)]
end

puts "Starting test with #{OPTIONS[:Domains].size} domains..."
puts "Using nameservers: " + OPTIONS[:Nameservers].join(", ")
puts "Only long running queries will be printed..."

def resolve_domain
	s = Time.now
	result = nil
	name = random_domain

	begin
		result = [$R.getaddress(name)]
	rescue Resolv::ResolvError
		$M.synchronize do
			puts "Name #{name} failed to resolve!"
			$STATUS[name] ||= []
			$STATUS[name] << :failure

			if $STATUS[name].include?(:success)
				puts "Name #{name} has had previous successes!"
			end
		end

		return
	end

	result.unshift(name)
	result.unshift(Time.now - s)

	$M.synchronize do	
		$TOTAL[0] += result[0]
		$TOTAL[1] += 1

		if result[0] > OPTIONS[:Timeout]
			puts "\t\t%0.2fs: %s => %s" % result
		end

		$STATUS[name] ||= []
		$STATUS[name] << :success

		if $STATUS[name].include?(:failure)
			puts "Name #{name} has had previous failures!"
		end
	end
end

puts "Starting threads..."
Thread.abort_on_exception = true

OPTIONS[:Threads].times do
	th = Thread.new do
		OPTIONS[:Requests].times do
			resolve_domain
		end
	end

	$TG.add th
end

$TG.list.each { |thr| thr.join }

$STATUS.each do |name, results|
	if results.include?(:failure)
		puts "Name #{name} failed at least once!"
	end
end

puts
puts "Requests: #{$TOTAL[1]} Average time: #{$TOTAL[0] / $TOTAL[1]}"
