module AWS::S3::EncryptionUtils

def check_encryption_materials mode, key
  rsa = OpenSSL::PKey::RSA
  case key
  when rsa
    unless key.private? or mode == :encrypt
      msg = "invalid key, #{rsa} requires a private key"
      raise ArgumentError, msg
    end
  when String # no problem
  else
    msg = "invalid key, must be an #{rsa} or a cipher key string"
    raise ArgumentError, msg
  end
end

def decrypt data, key
  rsa = OpenSSL::PKey::RSA
  begin
    case key
    when rsa # Asymmetric Decryption
        key.private_decrypt(data)
    when String             # Symmetric Decryption
        cipher = get_aes_cipher(:decrypt, :ECB, key)
        cipher.update(data) + cipher.final
    end
  rescue OpenSSL::Cipher::CipherError
    raise RuntimeError, "decryption failed, incorrect key?"
  end
end

def encrypt data, key
  rsa = OpenSSL::PKey::RSA
  data_cipher_size = get_cipher_size(data.length)
  # Encrypting data key
  case key
  when rsa # Asymmetric encryption
    warn UNSAFE_MSG if key.public_key.n.num_bits < data_cipher_size
    key.public_encrypt(data)
  when String             # Symmetric encryption
    warn UNSAFE_MSG if get_cipher_size(key.length) < data_cipher_size
    cipher = get_aes_cipher(:encrypt, :ECB, key)
    cipher.update(data) + cipher.final
  end
end

def generate_aes_key cipher, &block
  key_iv_pair = [cipher.random_key, cipher.random_iv]
  yield(key_iv_pair) if block_given?
end

def get_aes_cipher mode, block_mode, key = nil, iv = nil
  # If no key given, default to 256 bit
  cipher_size = (key) ? get_cipher_size(key.length) : 256
  cipher = OpenSSL::Cipher.new("AES-#{cipher_size}-#{block_mode}")
  (mode == :encrypt) ? cipher.encrypt : cipher.decrypt
  cipher.key = key if key
  cipher.iv  = iv  if iv
  cipher
end

def get_cipher_size(key_length)
  case key_length
  when 32 then 256
  when 24 then 192
  when 16 then 128
  else
    msg = "invalid key, symmetric key required to be 16, 24, or 32 bytes "
    msg << "in length, saw length #{key_length}"
    raise ArgumentError, msg
  end
end

def get_encrypted_size size
  # The next multiple of 16
  ((size / 16) + 1) * 16
end