require 'digest/sha1'
require 'rsa'

module RSA
  module OAEP
    extend self

    # Represents an error that occurs during decoding when using
    # RSA::OAEP.decode or RSA::OAEP.eme_decode. There is one argument which is
    # a brief message detailing the error
    class DecodeError < StandardError; end

    # The algorithms below need the HLEN variable. This is the length of the
    # hashes generated by the hashing function. For now, this only supports SHA1
    # as the hashing function, and this has a hash length of 20
    HLEN = 20

    # Performs the rsa-oaep-mgf1 decrypt algorithm. This is specified in section
    # 7.1.2 of http://www.ietf.org/rfc/rfc2437.txt.
    #
    # This implementation assumes that the sha1 hashing algorithm was used.
    #
    # @param [RSA::Key] k the private key whose public key was used to
    #        encrypt the data
    # @param [String] c a string of raw bytes representing the text to be
    #        decoded
    # @param [String] p the options which were used in the original encoding of
    #        the string. By default this is the empty string.
    #
    # @return [String] the decoded string of bytes
    # @raise [DecodeError] If decoding cannot occur, an error is raised
    def decode k, c, p = ''
      # First, generate how many bytes the key's modulus is
      n = k.modulus
      bytes = 0
      while n > 0
        bytes += 1
        n /= 2
      end
      bytes /= 8

      raise DecodeError, 'input is wrong length!' unless c.length == bytes

      enc = RSA::PKCS1.os2ip c
      m   = RSA::PKCS1.rsadp k, enc
      em  = RSA::PKCS1.i2osp m, bytes - 1

      eme_decode em, p
    end

    # Decodes the encrypted message as specified by the algorithm listed on
    # http://www.ietf.org/rfc/rfc2437.txt in section 9.1.1.2
    #
    # @param [String] em the encoded message that needs to be decoded
    # @param [String] p the flags used in the original encoding scheme.
    #
    # @return [String] the decoded byte string of the supplied message
    # @raise [DecodeError] if decoding goes awry or the message does not pass
    #        sanity checks during decoding
    def eme_decode em, p = ''
      raise DecodeError, 'message is too short!' if em.length < HLEN * 2 + 1

      maskedSeed = em[0...HLEN]
      maskedDB   = em[HLEN..-1]
      seedMask   = mgf1 maskedDB, HLEN
      seed       = xor maskedSeed, seedMask
      dbMask     = mgf1 seed, em.size - HLEN
      db         = xor maskedDB, dbMask
      pHash      = Digest::SHA1.digest p

      ind = db.index("\x01", HLEN)
      raise DecodeError, 'message is invalid!' if ind.nil?

      pHash2 = db[0...HLEN]
      ps     = db[HLEN...ind]
      m      = db[(ind + 1)..-1]

      raise DecodeError, 'message is invalid!' unless ps.bytes.all?(&:zero?)
      raise DecodeError, "specified p = #{p.inspect} is wrong!" unless pHash2 == pHash

      m
    end

    # Defined in seciton 10.2.1 of http://www.ietf.org/rfc/rfc2437.txt, this
    # is the mask generation function used in the eme_decode function
    #
    # @param [String] z this is the seed which the mask function runs off of
    # @param [Integer] l the desired length of the resultant hash
    #
    # @return [String] the mask generated
    def mgf1 z, l
      t = ''

      (0..(l / HLEN)).each{ |i|
        t += Digest::SHA1.digest(z + RSA::PKCS1.i2osp(i, 4))
      }

      t[0...l]
    end

    private

    def xor s1, s2
      b1 = s1.unpack('c*')
      b2 = s2.unpack('c*')

      if b1.length != b2.length
        raise DecodeError, 'cannot xor strings of different lengths!'
      end

      b1.zip(b2).map{ |a, b| a ^ b }.pack('c*')
    end

  end
end