#
# 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'
require 'base64'

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.
  #
  # This class supports RSA, DSA and ECDSA signature methods. When used with
  # ECDSA, it _only_ supports JWT envelope format. The intent is to enable
  # consumers of this class to switch to JSON Web Token in the future.
  #
  # @deprecated please use JSON Web Token instead of this class!
  #
  # @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)

    # List of acceptable Hash algorihtms + map of Ruby builtins to OpenSSL
    # counterparts.
    DIGEST_MAP = {
      Digest::MD5  => OpenSSL::Digest::MD5,
      Digest::SHA1 => OpenSSL::Digest::SHA1,
      Digest::SHA2 => OpenSSL::Digest::SHA256,
      OpenSSL::Digest::MD5 => OpenSSL::Digest::MD5,
      OpenSSL::Digest::SHA1 => OpenSSL::Digest::SHA1,
      OpenSSL::Digest::SHA256 => OpenSSL::Digest::SHA256,
      OpenSSL::Digest::SHA384 => OpenSSL::Digest::SHA384,
      OpenSSL::Digest::SHA512 => OpenSSL::Digest::SHA512,
    }.freeze

    # Digest output sizes (in bits)
    HASHSIZE_MAP = {
      Digest::MD5  => 128,
      Digest::SHA1 => 160,
      Digest::SHA2 => 256,
      OpenSSL::Digest::MD5 => 128,
      OpenSSL::Digest::SHA1 => 160,
      OpenSSL::Digest::SHA256 => 256,
      OpenSSL::Digest::SHA384 => 384,
      OpenSSL::Digest::SHA512 => 512,
    }.freeze

    # 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] data the actual data that is to be signed
    # @param [OpenSSL::PKey::PKey] key
    # @param [Digest::Base, OpenSSL::Digest] digest hash function to use
    # @param [#dump] encoding
    # @param [nil,:none,:right_support,:jwt] envelope serialization format and encryption scheme; default depends on key type
    # @param [OpenSSL::PKey::PKey] private_key deprecated -- pass key instead
    # @param [OpenSSL::PKey::PKey] public_key deprecated -- pass key instead
    def initialize(data={}, key=nil, digest:nil, encoding:DefaultEncoding, envelope:nil, private_key:nil, public_key:nil)
      # Cope with legacy parameter
      if key.nil?
        key = private_key || public_key
        warn(':private_key and :public_key are deprecated; please pass key as 2nd parameter') if key
      end

      @data     = data
      @encoding = encoding
      @key      = key

      # Figure out envelope type
      env = case envelope
      when nil then guess_envelope
      when :none, false then :none
      when :right_support, true then :right_support
      when :jwt then :jwt
      else raise ArgumentError.new("Unsupported envelope #{envelope.inspect}")
      end
      @envelope = env

      # Figure out digest algorithm
      @digest = digest || guess_digest

      check_parameters
    end

    # Compute a signature and return a JSON Web Token representation of
    # this hash.
    #
    # @return [String] a signed JWT with the specified expiration timestamp
    def to_jwt(expires_at)
      prefix, sig = sign_with_canonical_representation(expires_at)
      sig = RightSupport::Data::Base64URL.encode(sig)
      "#{prefix}.#{sig}"
    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)
      _, sig = sign_with_canonical_representation(expires_at)
      sig
    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 @key

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

      case @envelope
      when :none
        digest = @digest.new.update(plaintext).digest
        # raw RSA decryption
        actual = @key.public_decrypt(signature)
        raise InvalidSignature, "Signature mismatch: expected #{digest}, got #{actual}" unless actual == digest
      when :jwt
        if @key.respond_to?(:dsa_verify_asn1)
          # DSA signature with JWT-compatible encoding
          digest = @digest.new.update(plaintext).digest
          signature = raw_to_asn1(signature, @key)
          result = @key.dsa_verify_asn1(digest, signature)
          raise InvalidSignature, "Signature mismatch: DSA verify failed" unless result
        elsif @key.respond_to?(:verify)
          digest = @digest.new
          result = @key.verify(digest, signature, plaintext)
          raise InvalidSignature, "Signature mismatch: verify failed" unless result
        else
          raise NotImplementedError, "Cannot verify JWT with #{@key.class.name}"
        end
      when :right_support
        digest = DIGEST_MAP[@digest].new
        if @key.respond_to?(:dsa_verify_asn1)
          # DSA signature with ASN.1 encoding
          @key.dsa_verify_asn1(digest.digest, signature)
        else
          # RSA/DSA signature as specified in PKCS #1 v1.5
          result = @key.verify(digest, signature, plaintext)
        end
        raise InvalidSignature, "Signature verification failed" unless true == result
      end

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

      true
    rescue OpenSSL::PKey::RSAError => e
      raise InvalidSignature, "Signature mismatch: #{e.message}"
    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 ExpiredSignature, InvalidSignature => e
      false
    end

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

    # Free the inner Hash.
    def respond_to?(meth, include_all=false)
      super || @data.respond_to?(meth)
    end

    def respond_to_missing?(meth, include_all=false)
      super || @data.respond_to?(meth, include_all)
    end

    private

    # @return [Array] a pair of strings: the exact message that was signed, and its raw-binary signature
    def sign_with_canonical_representation(expires_at)
      raise ArgumentError, "Cannot sign; missing private_key" unless @key
      raise ArgumentError, "expires_at must be a Time in the future" unless time_check(expires_at)

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

      case @envelope
      when :none
        digest = @digest.new.update(plaintext).digest
        # raw RSA encryption; output is a single bignum
        sig = @key.private_encrypt(digest)
      when :jwt
        if @key.respond_to?(:dsa_sign_asn1)
          # raw ECDSA signature; output are two bignums (r, s) stuck together
          # with no delimiter
          digest = @digest.new.update(plaintext).digest
          asn1 = @key.dsa_sign_asn1(digest)
          sig = asn1_to_raw(asn1, @key)
        elsif @key.respond_to?(:sign)
          # RSA/DSA signature as specified in PKCS #1 v1.5
          digest = @digest.new
          sig = @key.sign(digest, plaintext)
        else
          raise NotImplementedError, "Cannot sign JWT with a #{@key.class.name}"
        end
      when :right_support
        digest = DIGEST_MAP[@digest].new(plaintext)
        if @key.respond_to?(:dsa_sign_asn1)
          # DSA signature with ASN.1 encoding
          sig = @key.dsa_sign_asn1(digest.digest)
        else
          # RSA/DSA signature as specified in PKCS #1 v1.5
          sig = @key.sign(digest, plaintext)
        end
      end

      return plaintext, sig
    end

    # @raise [ArgumentError] if there is anything wrong with the crypto parameters
    def check_parameters
      # raises if key/digest are not compatible with JOSE
      begin
        jose_header if @envelope == :jwt
      rescue NotImplementedError => e
        raise ArgumentError.new(e.message)
      end

      unless DIGEST_MAP.key?(@digest)
        raise ArgumentError, "Digest must be one of Digest::* class or OpenSSL::Digest::*"
      end
      unless @encoding.respond_to?(:dump)
        raise ArgumentError, "Encoding class/module/object must respond to .dump method"
      end

      if @envelope
        if @key && !@key.respond_to?(:verify)
          raise ArgumentError, "Public key must respond to :verify"
        end
        if @key && !@key.respond_to?(:sign)
          raise ArgumentError, "Private key must respond to :sign"
        end
      else
        if @key && !@key.respond_to?(:public_decrypt)
          raise ArgumentError, "Public key must respond to :public_decrypt"
        end
        if @key && !@key.respond_to?(:private_encrypt)
          raise ArgumentError, "Private key must respond to :private_encrypt"
        end
      end
    end

    # Figure out a default envelope if none is provided
    def guess_envelope
      case @key
      when OpenSSL::PKey::EC
        :jwt
      else
        :none
      end
    end

    # Figure out a suitable default digest if none is provided
    def guess_digest
      case @key
      when OpenSSL::PKey::DSA
        OpenSSL::Digest::SHA1
      when OpenSSL::PKey::EC
        case @key.group.degree
        when 256 then OpenSSL::Digest::SHA256
        when 384 then OpenSSL::Digest::SHA384
        when 521 then OpenSSL::Digest::SHA512
        else
          raise ArgumentError, "Cannot guess digest; please pass it as an option"
        end
      when OpenSSL::PKey::RSA
        case @envelope
        when :jwt then OpenSSL::Digest::SHA256
        else OpenSSL::Digest::SHA1
        end
      else
        OpenSSL::Digest::SHA1
      end
    end

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

    # Encode a canonicalized representation of the input.
    def encode(input) # :nodoc:
      case @envelope
      when :none, :right_support
        @encoding.dump(input)
      when :jwt
        bits = input.map do |m|
          RightSupport::Data::Base64URL.encode(JSON.dump(m))
        end
        bits.join '.'
      end
    end


    # Canonicalize a framed message according to the envelope format.
    def canonicalize(input) # :nodoc:
      case @envelope
      when :none, :right_support
        canonrs(input)
      when :jwt
        [input.first, canonjwt(input.last)]
      end
    end

    # Canonicalize the input by rearranging all Hash keys in lexical order and
    # converting data types to JSON-friendly versions.
    def canonjwt(input)
      case input
      when Hash
        output = {}
        input.keys.sort.each { |k| output[k] = input[k] }
      when Array
        output = input.map { |e| canonjwt(input) }
      when Time
        output = input.to_i
      when Symbol
        output = input.to_s
      else
        output = input
      end

      output
    end

    # canonicalize the input 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 canonrs(input)
      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[canonrs(k)] = canonrs(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| canonrs(x) }
        when Time
          output = input.to_i
        when Symbol
          output = input.to_s
        else
          output = input
      end

      output
    end

    # Incorporate the hash and its signature metadata into a Hash or Array
    # that represents the logical layout of the thing to be signed or verified.
    def frame(data, metadata) # :nodoc:
      case @envelope
      when :none, :right_support
        # Proprietary framing
        {:data => data, :metadata => metadata}
      when :jwt
        # JOSE framing; see http://jose.readthedocs.io
        payload = data.dup
        payload['exp'] = metadata[:expires_at].to_i
        [jose_header, payload]
      end
    end

    # @return [Hash] JOSE header defining document type (JWT) and encryption algorithm
    def jose_header
      key = @key || @key

      prefix = case key
      when OpenSSL::PKey::EC
        'ES'
      when OpenSSL::PKey::RSA
        'RS'
      else
        raise NotImplementedError, "Cannot use a #{key.class.name} with JWT envelope"
      end

      size = HASHSIZE_MAP[@digest]
      unless size
        raise NotImplementedError, "Cannot use #{@digest.name} with JWT envelope and this kind of key"
      end

      {'typ' => 'JWT', 'alg' => "#{prefix}#{size}"}
    end

    # Convert ASN1-encoded pair of integers into raw pair of concatenated bignums
    # This only works for OpenSSL::PKey::EC.
    # See https://github.com/jwt/ruby-jwt/blob/master/lib/jwt.rb#L166
    def asn1_to_raw(signature, private_key) # :nodoc:
      byte_size = (private_key.group.degree + 7) / 8
      OpenSSL::ASN1.decode(signature).value.map { |value| value.value.to_s(2).rjust(byte_size, "\x00") }.join
    end

    # Convert raw pair of concatenated bignums into ASN1-encoded pair of integers.
    # This only works for OpenSSL::PKey::EC.
    # https://github.com/jwt/ruby-jwt/blob/master/lib/jwt.rb#L159
    def raw_to_asn1(signature, public_key) # :nodoc:
      byte_size = (public_key.group.degree + 7) / 8
      r = signature[0..(byte_size - 1)]
      s = signature[byte_size..-1] || ''
      OpenSSL::ASN1::Sequence.new([r, s].map { |int| OpenSSL::ASN1::Integer.new(OpenSSL::BN.new(int, 2)) }).to_der
    end
  end
end