#
# 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 'digest/md5'
require 'digest/sha1'
require 'digest/sha2'
require 'openssl'

module RightSupport::Crypto
  # An easy way to compute digital signatures of data contained in a Ruby hash. To work with
  # signed hashes, you must first obtain an asymmetric key pair (any subclass of OpenSSL::PKey);
  # you can generate it from scratch or load it from a file on disk.
  #
  # Signature computation is influenced by four factors:
  #    - The digital signature algorithm and key length
  #    - The encoding used to serialize the hash contents to a byte stream
  #    - The hash algorithm used to compute a message digest of the byte stream
  #    - The OpenSSL API level used (EVP or raw crypto API)
  #
  # You are responsible for providing the PKey object, which determines the signature algorithm
  # and key length. This occasionally constrains your choice of hash algorithm; for instance,
  # a 512-bit RSA key would not be sufficiently long to create signatures of a SHA3-512 hash
  # due to the mathematical underpinnings of the RSA cipher. In practice this is not an issue,
  # because you should be using strong RSA keys (2048 bit or higher) for security reasons,
  # and even the strongest hash algorithms do not exceed 512-bit output. 
  #
  # SignedHash provides reasonable defaults for the other three factors:
  #     - JSON for message encoding (Yajl gem, JSON gem, Oj gem or built-in Ruby 1.9 JSON)
  #     - SHA1 for message digest
  #     - raw crypto API (for compatibility with older RightSupport versions)
  #
  # If you are adopting SignedHash for a new use case, it's best to use the default
  # encoding and message digest, but specify :envelope=>true to use the OpenSSL EVP
  # API! Using an envelope provides better protection against various cryptographic
  # attacks and ensures that the sign and verify operations can't be used.
  #
  # SignedHash defaults to raw-crypto signatures for compatibility reasons, but
  # with RightSupport v3 the raw-crypto will be deprecated and EVP will be used by default.
  #
  # @see OpenSSL::PKey
  # @see Digest
  class SignedHash
    # The default encoding to use when dumping the hash to binary form. Defaults to any available
    # commonly-known JSON library, in the following order of preference:
    #  - Yajl (ruby-yajl gem)
    #  - JSON (json gem, or built-in JSON parser for ruby >= 1.9)
    #  - Oj (oj gem)
    #  - nil (if no JSON library is present)
    # @!parse DefaultEncoding = nil

    if require_succeeds?('yajl')
      DefaultEncoding = ::Yajl
    elsif require_succeeds?('json')
      DefaultEncoding = ::JSON
    elsif require_succeeds?('oj')
      DefaultEncoding = ::Oj
    else
      DefaultEncoding = nil
    end unless defined?(DefaultEncoding)

    DEFAULT_OPTIONS = {
      :digest   => Digest::SHA1,
      :envelope => false,
      :encoding => DefaultEncoding,
    }

    # Mapping of Ruby built-in hash algorithms to their OpenSSL counterparts
    DIGEST_MAP = {
      Digest::MD5  => OpenSSL::Digest::MD5,
      Digest::SHA1 => OpenSSL::Digest::SHA1,
      Digest::SHA2 => OpenSSL::Digest::SHA256,
    }

    # Create a new sign/verify context, passing in a Hash full of data that is to be signed or
    # verified. The new SignedHash will store a reference to the raw data, so be careful not to
    # modify the data hash in a way that will influence the outcome of sign/verify!
    #
    # @param [Hash] hash the actual data that is to be signed
    # @option opts [Class] :digest hash-algorithm class from Ruby's Digest module MD5, SHA1 or SHA2; default SHA1
    # @option opts [true,false] :envelope use the OpenSSL EVP API if true, or raw-crypto API if false; default false
    # @option opts [#dump] :encoding serialization method for dumping hash data; default DefaultEncoding
    # @option opts [OpenSSL::PKey] :public_key key to use when verifying digital signatures
    # @option opts [OpenSSL::PKey] :private_key key to use when computing digital signatures
    #
    # @see DefaultEncoding
    def initialize(hash={}, opts={})
      opts = DEFAULT_OPTIONS.merge(opts)
      @hash        = hash
      @digest      = opts[:digest]
      @encoding    = opts[:encoding]
      @envelope    = !!opts[:envelope]
      @public_key  = opts[:public_key]
      @private_key = opts[:private_key]
      duck_type_check
    end

    # Produce a digital signature of the hash contents, including the expiration timestamp
    # of the signature. The caller must provide the exact same hash and expires_at in order
    # to successfully verify the signature.
    #
    # @param [Time] expires_at
    # @return [String] a binary signature of the hash's contents
    def sign(expires_at)
      raise ArgumentError, "Cannot sign; missing private_key" unless @private_key
      raise ArgumentError, "expires_at must be a Time in the future" unless time_check(expires_at)

      metadata = {:expires_at => expires_at}
      encoded = encode(canonicalize(frame(@hash, metadata)))

      if @envelope
        digest = DIGEST_MAP[@digest].new(encoded)

        @private_key.sign(digest, encoded)
      else
        digest = @digest.new.update(encoded).digest
        @private_key.private_encrypt(digest)
      end
    end

    # Verify a digital signature of the hash's contents. In order for the signature to verify,
    # the expires_at, signature and hash contents must be identical to those used by the
    # signer.
    #
    # @param [String] signature a binary signature to verify
    # @param [Time] expires_at
    # @return [true] always returns true (except when it raises)
    # @raise [ExpiredSignature] if the signature is expired
    # @raise [InvalidSignature] if the signature is invalid
    def verify!(signature, expires_at)
      raise ArgumentError, "Cannot verify; missing public_key" unless @public_key

      metadata  = {:expires_at => expires_at}
      plaintext = encode( canonicalize( frame(@hash, metadata) ) )

      if @envelope
        digest = DIGEST_MAP[@digest].new
        result = @public_key.verify(digest, signature, plaintext)
        raise InvalidSignature, "Signature verification failed" unless true == result
      else
        expected = @digest.new.update(plaintext).digest
        actual = @public_key.public_decrypt(signature)
        raise InvalidSignature, "Signature mismatch: expected #{expected}, got #{actual}" unless actual == expected
      end

      raise ExpiredSignature, "The signature has expired (or expires_at is not a Time)" unless time_check(expires_at)

      true
    end

    # Verify a digital signature of the hash's contents. In order for the signature to verify,
    # the expires_at, signature and hash contents must be identical to those used by the
    # signer.
    #
    # @param [String] signature a binary signature to verify
    # @param [Time] expires_at
    # @return [true] if the signature and expiration verify OK
    # @return [false] if the signature or expiration failed to verify
    def verify(signature, expires_at)
      verify!(signature, expires_at)
    rescue Exception => e
      false
    end

    # Free the inner Hash.
    def method_missing(meth, *args)
      @hash.__send__(meth, *args)
    end

    # Free the inner Hash.
    def respond_to?(meth)
      super || @hash.respond_to?(meth)
    end

    private

    def duck_type_check
      unless DIGEST_MAP.key?(@digest)
        raise ArgumentError, "Digest must be a built-in Ruby Digest class: MD5, SHA1 or SHA2"
      end
      unless @encoding.respond_to?(str_or_symb('dump'))
        raise ArgumentError, "Encoding class/module/object must respond to .dump method"
      end
      if @public_key && !@public_key.respond_to?(str_or_symb('public_decrypt'))
        raise ArgumentError, "Public key must respond to :public_decrypt (e.g. an OpenSSL::PKey instance)"
      end
      if @private_key && !@private_key.respond_to?(str_or_symb('private_encrypt'))
        raise ArgumentError, "Private key must respond to :private_encrypt (e.g. an OpenSSL::PKey instance)"
      end
    end

    def str_or_symb(method)
      RUBY_VERSION > '1.9' ? method.to_sym : method.to_s
    end

    # Ensure that an expiration time is in the future.
    def time_check(t)
      t.is_a?(Time) && (t >= Time.now)
    end

    # Incorporate the hash and its signature metadata into an enclosing hash.
    def frame(data, metadata) # :nodoc:
      {:data => data, :metadata => metadata}
    end

    # Encode a canonicalized representation of the hash.
    def encode(input)
      @encoding.dump(input)
    end

    # Canonicalize the hash (and any nested data) by transforming it deterministically into a
    # structure of arrays-in-arrays whose elements are ordered according to the lexical ordering
    # of hash keys. Canonicalization ensures that the signer and verifier agree on the contents
    # of the thing being signed irrespective of Ruby version, CPU architecture, etc.
    def canonicalize(input) # :nodoc:
      case input
        when Hash
          # Hash is the only complex case. We canonicalize a Hash as an Array of pairs, each of which
          # consists of one key and one value. The ordering of the pairs is consistent with the
          # ordering of the keys.
          output = Array.new

          # First, transform the original input hash into something that has canonicalized keys
          # (which should make them sortable, too). Also canonicalize the values while we are
          # at it...
          sortable_input = {}
          input.each { |k,v| sortable_input[canonicalize(k)] = canonicalize(v) }

          # Sort the keys; guard this operation so we can raise an intelligent error if
          # something is still not sortable even after canonicalization.
          begin
            ordered_keys = sortable_input.keys.sort
          rescue Exception => e
            msg = "SignedHash requires sortable hash keys; cannot sort #{sortable_input.keys.inspect} " +
                  "due to #{e.class.name}: #{e.message}"
            e2 = ArgumentError.new(msg)
            e2.set_backtrace(e.backtrace)
            raise e2
          end

          ordered_keys.each do |key|
            output << [ key, sortable_input[key] ]
          end
        when Array
          output = input.collect { |x| canonicalize(x) }
        when Time
          output = input.to_i
        when Symbol
          output = input.to_s
        else
          output = input
      end

      output
    end
  end
end