# From https://gist.github.com/Zapotek/981959/raw/d7353edad1bd88110dc1f03dcc47257219d1624e/rsa_aes_cbc.rb
# To generate a public / private key pair:
# require 'openssl'
# new_key = OpenSSL::PKey::RSA.generate( 15360 )  # 4096 probably OK
# File.open("./new_public.pem", "w") { |f| f.puts new_key.public_key }
# File.open("./new_private.pem", "w") { |f| f.puts new_key.to_pem }

require 'openssl'
require 'yaml'
require 'base64'

# SECURE: TVB Mon 14 Oct 2013 13:33:06 BST
# All the encryption keys are generated by standard libraries.
# The data is encrypted with AES symmetric key. The symmetric key
# is encrypted with RSA.
#

# KH: 20160914 - changed:
#
#   Base64.encode64 to Base64.strict_encode64 in encrypt()
#
# to remove newlines, to be consistent with the rest of the ndr_pseudonymise code.
#   Base64.decode64 is kept in decrypt(), as this (correctly) ignores newlines produced
#   by existing encrypted data using encrypt().
#
# http://ruby-doc.org/stdlib-2.3.1/libdoc/base64/rdoc/Base64.html
# https://tools.ietf.org/html/rfc4648#section-3.1



# Simple hybrid crypto class using RSA for public key encryption and AES with CBC
# for bulk data encryption/decryption.
#
# RSA is used to encrypt the AES primitives which are used to encrypt the plaintext.
#
# @author: Tasos "Zapotek" Laskos
#                                      <tasos.laskos@gmail.com>
#                                      <zapotek@segfault.gr>
# @version: 0.1
class RSA_AES_CBC
  #
  # If only encryption is required the private key parameter can be omitted.
  #
  # @param  [String]  public_pem   location of the Public key in PEM format
  # @param  [String]  private_pem  location of the Private key in PEM format
  #
  def initialize(public_pem, private_pem = nil)
    @public_pem = public_pem
    @private_pem = private_pem
  end

  #
  # Encrypts data and returns a Base64 representation of the ciphertext
  # and AES CBC primitives encrypted using the public key.
  #
  # @param  [String]  data
  #
  # @return  [String]   Base64 representation of the ciphertext
  #                       and AES CBC primitives encrypted using the public key.
  #
  def encrypt(data)
    rsa = OpenSSL::PKey::RSA.new(@public_pem)

    # encrypt with 256 bit AES with CBC
    aes = OpenSSL::Cipher.new('aes-256-cbc')
    aes.encrypt

    # use random key and IV
    aes.key = key = aes.random_key
    aes.iv  = iv  = aes.random_iv

    # this will hold all primitives and ciphertext
    primitives = {}

    primitives['ciphertext'] = aes.update(data)
    primitives['ciphertext'] << aes.final

    primitives['key'] = rsa.public_encrypt(key)
    primitives['iv']  = rsa.public_encrypt(iv)

    # serialize everything and base64 encode it
    Base64.strict_encode64(primitives.to_yaml)
  end

  #
  # Decrypts data.
  #
  # @param  [String]  data
  #
  # @return  [String]   plaintext
  #
  def decrypt(data)
    rsa = OpenSSL::PKey::RSA.new(@private_pem)

    # decrypt with 256 bit AES with CBC
    aes = OpenSSL::Cipher.new('aes-256-cbc')
    aes.decrypt

    # unencode and unserialize to get the primitives and ciphertext
    primitives = YAML.load(Base64.decode64(data))

    aes.key = rsa.private_decrypt(primitives['key'])
    aes.iv  = rsa.private_decrypt(primitives['iv'])

    plaintext = aes.update(primitives['ciphertext'])
    plaintext << aes.final

    plaintext
  end
end

# crypto = RSA_AES_CBC.new( 'public.pem', 'private.pem' )
# ciphered = crypto.encrypt( 'Foo' )
# puts crypto.decrypt( ciphered )