require "ipaddr"
require "socket"
require "optparse"
require "json"
require "securerandom"

require "xmpp4r"

puts ARGV

options = {}
OptionParser.new do |opts|
  opts.banner = "usage: ${$0} [opts]"

  opts.on("-i", "--interface=INTERFACE",
          "Publicly-accessible interface") do |interface|
    options[:interface] = interface
  end

  opts.on("-d", "--domain=DOMAIN",
          "Domain on which we wish to set the new ip") do |domain|
    options[:domain] = domain
  end

  opts.on("-s", "--server=SERVER",
          "Backplane DNS XMPP server") do |server|
    options[:server] = server
  end

  opts.on("-p", "--password-file=/path/to/file",
          "File containing password for XMPP server") do |pw_file|
    options[:pw_file] = pw_file
  end

  opts.on("-4", "--ipv4",
          "Check for a public IPv4 and register with the backplane.") do
    options[:ipv4] = true
  end

  opts.on("-6", "--ipv6",
          "Check for a public IPv6 and register with the backplane.") do
    options[:ipv6] = true
  end

  opts.on("-f", "--sshfp=FILE",
          "Register host SSH key fingerprints with the backplane.") do |file| 
    options[:sshfp] = [] if not options[:sshfp]
    options[:sshfp] = options[:sshfp] + [file]
  end
end.parse!

def error(msg)
  puts msg
  throw msg
end

error("domain is required") if not options[:domain]
error("server is required") if not options[:server]
error("password file is required") if not options[:pw_file]
error("at least one of -4 or -6 required") if not (options[:ipv4] or options[:ipv6])

if not File::readable?(options[:pw_file])
  error("file does not exist or is not readable")
end
password = File::open(options[:pw_file]) { |f| f.gets.strip }

class XMPPClient
  def initialize(domain, hostname, server, password)
    @jid = "host-#{hostname}@#{server}"
    @service_jid = "service-dns@#{server}"
    @server = server
    @domain = domain
    @password = password
    @responses = Queue.new
    @responses_lock = Mutex.new
  end

  def connect
    disconnect if connected?
    @client = Jabber::Client::new(@jid)
    @client.connect # will use SRV records
    error("failed to initialize TLS connection") if not @client.is_tls?
    @client.auth(@password)
    register_response_callback
  end

  def connected?
    @client ||= nil
    @client.respond_to?(:is_connected?) and @client.is_connected?
  end

  def disconnect
    if @client.respond_to?(:is_connected?) && @client.is_connected?
      begin
        @client.close
      rescue Errno::EPIPE, IOError => e
        nil
      end
    end
    @client = nil
  end

  def send(msg_content)
    msg_id = SecureRandom::uuid
    encoded_payload = payload(msg_content, msg_id).to_json
    puts "payload: #{encoded_payload}"
    msg = Jabber::Message.new(@service_jid, encoded_payload)
    msg.type = :chat
    @client.send(msg)
    response = receive_response(msg_id)
    puts "response: #{response}"
    response and response["status"] == "OK"
  end

  def send_ip(ip)
    send(ip_payload(ip))
  end

  def send_sshfp(fps)
    send(sshfp_payload(fps))
  end

  def payload(req, msg_id)
    {
      version: 1,
      service: :dns,
      msgid:   msg_id,
      payload: req
    }
  end

  def ip_payload(ip)
    {
      request: ip.ipv4? ? :change_ipv4 : :change_ipv6,
      domain:  @domain,
      ip:      ip.to_s
    }
  end

  def sshfp_payload(fp)
    {
      request: :change_sshfp,
      domain:  @domain,
      sshfp:   fp
    }
  end

  def register_response_callback
    @client.add_message_callback do |msg|
      enqueue_message(JSON.parse(msg.body))
    end
  end

  def enqueue_message(msg)
    @responses << msg
  end

  def receive_response(msg_id)
    msg = @responses.pop
    return msg if (msg and (msg["msgid"] == msg_id.to_s))
    raise "failed to receive message: #{msg}"
  end
end

RESERVED_V4_NETWORKS = [
  "0.0.0.0/8",
  "10.0.0.0/8",
  "100.64.0.0/10",
  "127.0.0.0/8",
  "169.254.0.0/16",
  "172.16.0.0/12",
  "192.0.0.0/24",
  "192.0.2.0/24",
  "192.88.99.0/24",
  "192.168.0.0/16",
  "198.18.0.0/15",
  "198.51.100.0/24",
  "203.0.113.0/24",
  "224.0.0.0/4",
  "240.0.0.0/4",
  "255.255.255.255/32"
].map { |ip| IPAddr.new(ip) }

def public_ip?(ip)
  if (ip.ipv4?)
    not RESERVED_V4_NETWORKS.any? { |network| network.include? ip }
  elsif (ip.ipv6?)
    not (ip.link_local? or ip.loopback? or ip.private?)
  else
    false
  end
end

def to_ipaddr(addrinfo)
  if addrinfo.ipv4?
    IPAddr.new addrinfo.ip_address
  else
    IPAddr.new(addrinfo.ip_address.split("%")[0])
  end
end

def local_addresses
  Socket::ip_address_list.map do |addrinfo|
    to_ipaddr(addrinfo)
  end.select { |ip| public_ip?(ip) }
end

def interface_addresses(interface)
  Socket::getifaddrs.select do |ifaddr|
    ifaddr.name == interface
  end.select do |ifaddr|
    ifaddr.addr.ip? and (ifaddr.flags & Socket::IFF_MULTICAST != 0)
  end.map do |ifaddr|
    to_ipaddr(ifaddr.addr)
  end.filter do |ip|
    public_ip? ip
  end
end

def host_sshfp(keys)
  keys.flat_map { |keyfile|
    `ssh-keygen -r hostname -f #{keyfile}`.split("\n")
  }.map { |fp|
    fp.match(/[0-9] [0-9] [a-fA-F0-9]{32,64}$/)[0]
  }.compact
end

def hostname
  Socket.gethostname.split(".").first
end

client = XMPPClient::new(options[:domain],
                         hostname,
                         options[:server],
                         password)

success = true

begin
  client.connect

  addrs = if options[:interface]
            interface_addresses(options[:interface])
          else
            local_addresses
          end

  if options[:ipv4]
    ipv4 = addrs.find { |ip| ip.ipv4? }
    if ipv4
      puts "#{options[:server]}: #{hostname}.#{options[:domain]} IN A => #{ipv4.to_s}"
      if client.send_ip(ipv4)
        puts "OK"
      else
        puts "ERROR"
        success = false
      end
    else
      puts "#{options[:server]}: no valid public IPv4 found on the local host"
    end
  end

  if options[:ipv6]
    ipv6 = addrs.find { |ip| ip.ipv6? }
    if ipv6
      puts "#{options[:server]}: #{hostname}.#{options[:domain]} IN AAAA => #{ipv6.to_s}"
      if client.send_ip(ipv6)
        puts "OK"
      else
        puts "ERROR"
        success = false
      end
    else
      puts "#{options[:server]}: no valid public IPv6 found on the local host"
    end
  end

  if options[:sshfp]
    fps = host_sshfp(options[:sshfp])
    if not fps.empty?
      puts "#{options[:server]}: #{hostname}.#{options[:domain]} IN SSHFP => #{fps}"
      if client.send_sshfp(fps)
        puts "OK"
      else
        puts "ERROR"
        success = false
      end
    else
      puts "#{options[:server]}: no valid sshfps found"
    end
  end
ensure
  client.disconnect
end

exit success ? 0 : 1