#
# Copyright (c) 2009-2011 RightScale Inc
#
# 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 'ipaddr'
require 'socket'
require 'uri'

module RightSupport::Net

  # IPAddr does not support returning a block of addresses in CIDR notation.
  class IPNet < ::IPAddr
    def cidr_block
      # Ruby Math:: does not have a log2() function :(
      cidr_block = 32 - (Math.log(self.to_range.count) / Math.log(2)).to_i
    end

    # Don't want to override to_s here since we also
    # use this during URI transformation and need a way
    # to get the raw dotted-quad address
    #
    def to_cidr
      "#{self.to_s}/#{self.cidr_block}"
    end
  end

  # A ResolvedEndpoint represents the resolution of a host which can contain
  # multiple IP addresses or entire IP address ranges of varying size.
  class ResolvedEndpoint
    attr_accessor :uri

    def initialize(ips, opts={})
      @ip_addresses = []
      @uri = opts[:uri]

      ips.to_a.each do |address|
        @ip_addresses << IPNet.new(address)
      end

      unless self.all_hosts? || @uri == nil
        raise URI::InvalidURIError, "Cannot resolve URI with CIDR block bigger than a single host" unless self.all_hosts?
      end
    end

    def addresses()
      @ip_addresses.map do |addr|
        if addr.to_range.count == 1 && @uri != nil
            transformed_uri = uri.dup
            transformed_uri.host = addr.to_s
            transformed_uri.to_s
        else
          addr
        end
      end
    end
    alias addrs :addresses

    def blocks()
      @ip_addresses.map {|addr| addr.to_cidr}
    end

    def all_hosts?()
      @ip_addresses.all? {|addr| addr.to_range.count == 1}
    end

    def ==(another_endpoint)
      another_endpoint.addresses.all? {|addr| addresses.member? addr } && another_endpoint.uri == @uri
    end
  end

  module DNS
    DEFAULT_RESOLVE_OPTIONS = {
      :address_family => Socket::AF_INET,
      :socket_type    => Socket::SOCK_STREAM,
      :protocol       => Socket::IPPROTO_TCP,
      :retry          => 3,
      :uri            => nil,
    }.freeze

    STATIC_HOSTNAME_TRANSLATIONS = {
      # This list of CIDR blocks comes directly from Amazon and
      # can be found here: https://forums.aws.amazon.com/ann.jspa?annID=2051
      'cf-mirror.rightscale.com' => [ '54.192.0.0/16',
                                      '54.230.0.0/16',
                                      '54.239.128.0/18',
                                      '54.240.128.0/18',
                                      '204.246.164.0/22',
                                      '204.246.168.0/22',
                                      '204.246.174.0/23',
                                      '204.246.176.0/20',
                                      '205.251.192.0/19',
                                      '205.251.249.0/24',
                                      '205.251.250.0/23',
                                      '205.251.252.0/23',
                                      '205.251.254.0/24',
                                      '216.137.32.0/19' ],
    }.freeze

    # Resolve a set of DNS hostnames to the individual IP addresses to which they map. Only handles
    # IPv4 addresses.
    #
    # @deprecated due to broken error handling - do not use; please use #resolve instead!
    def self.resolve_all_ip_addresses(hostnames)
      ips       = []
      hostnames = [hostnames] unless hostnames.respond_to?(:each)
      hostnames.each do |hostname|
        infos = nil
        begin
          infos = Socket.getaddrinfo(hostname, 443, Socket::AF_INET, Socket::SOCK_STREAM, Socket::IPPROTO_TCP)
        rescue Exception => e
          # NOTE: Need to figure out, which logger can we use here?
          # Log.error "Rescued #{e.class.name} resolving Repose hostnames: #{e.message}; retrying"
          retry
        end

        #Randomly permute the addrinfos of each hostname to help spread load.
        infos.shuffle.each do |info|
          ips << info[3]
        end
      end
      ips
    end

    # Perform DNS resolution on a set of endpoints, where the endpoints may be hostnames or URIs.
    # Expand out the list to include one entry per distinct address that is assigned to a given
    # hostname, but preserve other aspects of the endpoints --. URIs will remain URIs with the
    # same protocol, path-info, and so forth, but the hostname component will be resolved to IP
    # addresses and the URI will be duplicated in the output, once for each distinct IP address.
    #
    # Although this method does accept IPv4 dotted-quad addresses as input, it does not accept
    # IPv6 addresses. However, given hostnames or URIs as input, one _can_ resolve the hostnames
    # to IPv6 addresses by specifying the appropriate address_family in the options.
    #
    # It should never be necessary to specify a different :socket_type or :protocol, but these
    # options are exposed just in case.
    #
    # @param [Array<String>] endpoints a mixed list of hostnames, IPv4 addresses or URIs that contain them
    # @option opts [Integer] :retry number of times to retry SocketError; default is 3
    # @option opts [Integer] :address_family what kind of IP addresses to resolve; default is Socket::AF_INET (IPv4)
    # @option opts [Integer] :socket_type socket-type context to pass to getaddrinfo, default is Socket::SOCK_STREAM
    # @option opts [Integer] :protocol protocol context to pass to getaddrinfo, default is Socket::IPPROTO_TCP
    #
    # @return [Array<String>] larger list of endpoints with all hostnames resolved to IP addresses
    #
    # @raise URI::InvalidURIError if endpoints contains an invalid or URI
    # @raise SocketError if endpoints contains an invalid or unresolvable hostname
    def self.resolve(endpoints, opts={})
      opts = DEFAULT_RESOLVE_OPTIONS.merge(opts)
      endpoints = Array(endpoints)

      retries = 0
      resolved_endpoints = []

      endpoints.each do |endpoint|
        begin
          resolved_endpoint = nil
          if endpoint.include?(':')
            # It contains a colon, therefore it must be a URI -- we don't support IPv6
            uri = URI.parse(endpoint)
            hostname = uri.host
            raise URI::InvalidURIError, "Could not parse host component of URI" unless hostname

            resolved_endpoint = resolve_endpoint(hostname, opts.merge(:uri=>uri))
          else
            resolved_endpoint = resolve_endpoint(endpoint, opts)
          end
          resolved_endpoints << resolved_endpoint
        rescue SocketError => e
          retries += 1
          if retries < opts[:retry]
            retry
          else
            raise
          end
        end
      end

      resolved_endpoints
    end

    # Similar to resolve, but return a hash of { hostnames => [endpoints] }
    #
    # Perform DNS resolution on a set of endpoints, where the endpoints may be hostnames or URIs.
    # Expand out the list to include one entry per distinct address that is assigned to a given
    # hostname, but preserve other aspects of the endpoints --. URIs will remain URIs with the
    # same protocol, path-info, and so forth, but the hostname component will be resolved to IP
    # addresses and the URI will be duplicated in the output, once for each distinct IP address.
    #
    # Although this method does accept IPv4 dotted-quad addresses as input, it does not accept
    # IPv6 addresses. However, given hostnames or URIs as input, one _can_ resolve the hostnames
    # to IPv6 addresses by specifying the appropriate address_family in the options.
    #
    # It should never be necessary to specify a different :socket_type or :protocol, but these
    # options are exposed just in case.
    #
    # @param [Array<String>] endpoints a mixed list of hostnames, IPv4 addresses or URIs that contain them
    # @option opts [Integer] :retry number of times to retry SocketError; default is 3
    # @option opts [Integer] :address_family what kind of IP addresses to resolve; default is Socket::AF_INET (IPv4)
    # @option opts [Integer] :socket_type socket-type context to pass to getaddrinfo, default is Socket::SOCK_STREAM
    # @option opts [Integer] :protocol protocol context to pass to getaddrinfo, default is Socket::IPPROTO_TCP
    #
    # @return [Hash<hostnames => [endpoints]>] Hash with keys of hostnames and values of arrays of all associated IP addresses
    #
    # @raise URI::InvalidURIError if endpoints contains an invalid or URI
    # @raise SocketError if endpoints contains an invalid or unresolvable hostname
    def self.resolve_with_hostnames(endpoints, opts={})
      hostname_hash = {}
      endpoints.each {|endpoint| hostname_hash[endpoint] = resolve(endpoint, opts).first}
      hostname_hash
    end

    private

    # Lookup the address(es) associated with the given endpoints.
    #
    # Perform an address lookup in this order:
    #   1. Hardcoded translation table.
    #   2. System address lookup (e.g. DNS, hosts files...etc)
    #
    # Although this method does accept IPv4 dotted-quad addresses as input, it does not accept
    # IPv6 addresses. However, given hostnames as input, one _can_ resolve the hostnames
    # to IPv6 addresses by specifying the appropriate address_family in the options.
    #
    # It should never be necessary to specify a different :socket_type or :protocol, but these
    # options are exposed just in case.
    #
    # @param [String] endpoint as a hostname or IPv4 address
    # @option opts [Integer] :retry number of times to retry SocketError; default is 3
    # @option opts [Integer] :address_family what kind of IP addresses to resolve; default is Socket::AF_INET (IPv4)
    # @option opts [Integer] :socket_type socket-type context to pass to getaddrinfo, default is Socket::SOCK_STREAM
    # @option opts [Integer] :protocol protocol context to pass to getaddrinfo, default is Socket::IPPROTO_TCP
    #
    # @return [Array<String>] List of resolved IPv4/IPv6 addresses.
    #
    # @raise SocketError if endpoints contains an invalid or unresolvable hostname
    def self.resolve_endpoint(endpoint, opts=DEFAULT_RESOLVE_OPTIONS)
      if STATIC_HOSTNAME_TRANSLATIONS.has_key?(endpoint)
        ResolvedEndpoint.new(STATIC_HOSTNAME_TRANSLATIONS[endpoint], opts)
      else
        infos = Socket.getaddrinfo(endpoint, nil,
                                   opts[:address_family], opts[:socket_type], opts[:protocol])
        ResolvedEndpoint.new(infos.map { |info| info[3] }, opts)
      end
    end
  end
end