#!/usr/bin/env ruby
$:.unshift(File.dirname(__FILE__) + "/../lib")

require 'bitcoin'
require 'socket'
require 'json'
require 'rubydns'
require 'optparse'

Bitcoin::Network::Node # autoload to include Array#weighted_sample

DNS_OPTS = {
  :name => nil,
  :connect => "127.0.0.1:9999",
  :port => 8333,
  :cache => 1024,
  :send => 16,
  :interval => 10,
  :ttl => nil,
  :hosts => nil,
}
optparse = OptionParser.new do|opts|
  opts.banner = "Usage: bitcoin_dns_seed [options] <domain>"

  opts.on("-c", "--connect [HOST:PORT]",
    "Connect to server (default: 127.0.0.1:9999)") do |connect|
    DNS_OPTS[:connect] = connect
  end

  opts.on("-p", "--port [PORT]",
    "Network default port each addr must match (default: 8333)") do |port|
    DNS_OPTS[:port] = port.to_i
  end

  opts.on("-c", "--cache [SIZE]", "Cache addresses (default: 1024)") do |size|
    DNS_OPTS[:cache] = size.to_i
  end

  opts.on("-s", "--send [SIZE]", "Send addresses (default: 16)") do |size|
    DNS_OPTS[:send] = size.to_i
  end

  opts.on("-i", "--interval [SECONDS]", "Refresh addresses (default: 10)") do |s|
    DNS_OPTS[:interval] = s.to_i
  end

  opts.on("-t", "--ttl [SECONDS]", "TTL to set on served records (default: relative to addr age)") do |t|
    DNS_OPTS[:ttl] = t.to_i
  end

  opts.on("--hosts [HOSTS]", "Additional hosts file with name => addr mappings") do |h|
    DNS_OPTS[:hosts] = h
  end

  opts.on( '-h', '--help', 'Display this screen' ) do
    puts opts
    exit
  end
end
optparse.parse!

DNS_OPTS[:domain] = ARGV.shift
unless DNS_OPTS[:domain]
  puts optparse
  exit
end


class RubyDNS::Server
  attr_accessor :addrs
end

class AddrFetcher < EM::Connection
  def initialize(dns)
    @dns = dns
    @dns.addrs ||= []
    @buf = BufferedTokenizer.new("\x00")
    send_data(["addrs", DNS_OPTS[:cache].to_s].to_json)
  end

  def receive_data(data)
    @buf.extract(data).each do |packet|
      json = JSON.load(packet)[1]
      @dns.logger.info { "received #{json.size} new addrs" }
      @dns.addrs += json.select{|a| a[1] == DNS_OPTS[:port]}
      @dns.addrs.uniq! {|a| a[0]}
      pop = @dns.addrs.size - DNS_OPTS[:cache]
      @dns.addrs.pop(pop)  if pop > 0
      @dns.logger.debug { "addr pool size now #{@dns.addrs.size}" }
      close_connection
    end
  end

  def self.run(host, port, dns, interval = 30)
    EM.connect(host, port, self, dns)
    EM.add_periodic_timer(interval) do
      EM.connect(host, port, self, dns)
    end
  end
end

dns = RubyDNS::Server.new
dns.match(DNS_OPTS[:domain], :A) do |transaction|
  addrs = dns.addrs.weighted_sample(DNS_OPTS[:send]) {|addr| 10800 - addr[2] }.uniq
  addrs.each do |addr|
    begin
      ttl = DNS_OPTS[:ttl] || ((10800 - addr[2]) / 30)
      transaction.respond!(addr[0], :ttl => ttl)
    rescue
      # ignore faulty addrs ("localhost" etc)
    end
  end
end

if DNS_OPTS[:hosts]
  dns.logger.info { "serving additional addrs from file" }
  File.read(DNS_OPTS[:hosts]).each_line do |line|
    dns.logger.debug { line.strip }
    addr, *names = line.split(/\s+/)
    names.each {|n| dns.match(n, :A) {|t| t.respond!(addr) } }
  end
end

EM.run do
  EventMachine.open_datagram_socket("0.0.0.0", 53, RubyDNS::UDPHandler, dns)
  EventMachine.start_server("0.0.0.0", 53, RubyDNS::TCPHandler, dns)
  dns.logger.info { "DNS Server listening on 0.0.0.0:53" }

  timer = AddrFetcher.run(*DNS_OPTS[:connect].split(":"), dns, DNS_OPTS[:interval])
end