// Copyright 2021 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
//
//  Reference implementation utilities of OEMCrypto APIs
//
#ifndef WVOEC_UTIL_RSA_KEY_H_
#define WVOEC_UTIL_RSA_KEY_H_

#include <stddef.h>
#include <stdint.h>

#include <memory>
#include <string>
#include <vector>

#include <openssl/rsa.h>

#include "OEMCryptoCENC.h"

namespace wvoec {
namespace util {
enum RsaFieldSize {
  kRsaFieldUnknown = 0,
  kRsa2048Bit = 2048,
  kRsa3072Bit = 3084
};

// Identifies the RSA signature algorithm to be used when signing
// messages or verifying message signatures.
// The two standard signing algorithms specified by PKCS1 RSA V2.1
// are RSASSA-PKCS1 and RSASSA-PSS.  Each require agreement on a
// set of options.  For OEMCrypto, only one set of options are agreed
// upon for each RSA signature scheme.  CAST receivers specify a
// special implementation of PKCS1 where the message is already
// digested and encoded when provided.
enum RsaSignatureAlgorithm {
  // RSASSA-PSS with default options:
  //  Hash algorithm: SHA-1
  //  MGF: MGF1 with SHA-1
  //  Salt length: 20 bytes
  //  Trailer field: 0xbc
  kRsaPssDefault = 0,
  // RSASSA-PKCS1 for CAST receivers.
  // Assumes message is already digested & encoded.  Max message length
  // is 83 bytes.
  kRsaPkcs1Cast = 1
};

// Returns the string representation of the provided RSA field size.
// Intended for logging purposes.
std::string RsaFieldSizeToString(RsaFieldSize field_size);

// Compares two OpenSSL/BoringSSL RSA keys to see if their public RSA
// components are matching.
// This function assumes both keys are valid.
// Returns true if they are matching, false otherwise.
bool RsaKeysAreMatchingPair(const RSA* public_key, const RSA* private_key);

class RsaPrivateKey;

class RsaPublicKey {
 public:
  // Creates a new public key equivalent of the provided private key.
  static std::unique_ptr<RsaPublicKey> New(const RsaPrivateKey& private_key);

  // Creates an RSA public key from a native OpenSSL/BoringSSL RSA key handle.
  // Ownership of the handle is NOT transferred.
  static std::unique_ptr<RsaPublicKey> FromSslHandle(
      const RSA* rsa_handle, uint32_t allowed_schemes = kSign_RSASSA_PSS);

  // Loads a serialized RSA public key.
  // The provided |buffer| must contain a valid ASN.1 DER encoded
  // SubjectPublicKey.  This API will reject any RSA key that is not
  // approximately to 2048bits or 3072bits.
  //
  //  buffer: SubjectPublicKeyInfo = {
  //    algorithm: AlgorithmIdentifier = {
  //      algorithm: OID = rsaEncryption,
  //      parameters: NULL = null
  //    },
  //    subjectPublicKey: BIT STRING = ... -- ASN.1 DER encoded RSAPublicKey
  //  }
  //
  // Failure will occur if the provided |buffer| does not contain a
  // valid SubjectPublicKey, or if the specified curve is not
  // supported.
  static std::unique_ptr<RsaPublicKey> Load(const uint8_t* buffer,
                                            size_t length);
  static std::unique_ptr<RsaPublicKey> Load(const std::string& buffer);
  static std::unique_ptr<RsaPublicKey> Load(const std::vector<uint8_t>& buffer);

  // Loads a serialized RSA private key, but only converting the public key.
  static std::unique_ptr<RsaPublicKey> LoadPrivateKeyInfo(const uint8_t* buffer,
                                                          size_t length);
  static std::unique_ptr<RsaPublicKey> LoadPrivateKeyInfo(
      const std::string& buffer);
  static std::unique_ptr<RsaPublicKey> LoadPrivateKeyInfo(
      const std::vector<uint8_t>& buffer);

  RsaFieldSize field_size() const { return field_size_; }
  uint32_t allowed_schemes() const { return allowed_schemes_; }
  const RSA* GetRsaKey() const { return key_; }

  // Checks if the provided |private_key| is the RSA private key of this
  // public key.
  bool IsMatchingPrivateKey(const RsaPrivateKey& private_key) const;

  // Serializes the public key into an ASN.1 DER encoded SubjectPublicKey
  // representation.
  // On success, |buffer_size| is populated with the number of bytes
  // written to |buffer|, and OEMCrypto_SUCCESS is returned.
  // If the provided |buffer_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |buffer_size| is set
  // to the required buffer size.
  OEMCryptoResult Serialize(uint8_t* buffer, size_t* buffer_size) const;
  // Same as above, except directly returns the serialized key.
  // Returns an empty vector on error.
  std::vector<uint8_t> Serialize() const;

  // Verifies the |signature| matches the provided |message| by the
  // private equivalent of this public key.
  // The signature algorithm can be specified via the |algorithm| field.
  // See RsaSignatureAlgorithm for details on each algorithm.
  //
  // Returns:
  //    OEMCrypto_SUCCESS if signature is valid
  //    OEMCrypto_ERROR_SIGNATURE_FAILURE if the signature is invalid
  //    OEMCrypto_ERROR_UNKNOWN_FAILURE if any error occurs
  OEMCryptoResult VerifySignature(
      const uint8_t* message, size_t message_length, const uint8_t* signature,
      size_t signature_length,
      RsaSignatureAlgorithm algorithm = kRsaPssDefault) const;
  OEMCryptoResult VerifySignature(
      const std::string& message, const std::string& signature,
      RsaSignatureAlgorithm algorithm = kRsaPssDefault) const;
  OEMCryptoResult VerifySignature(
      const std::vector<uint8_t>& message,
      const std::vector<uint8_t>& signature,
      RsaSignatureAlgorithm algorithm = kRsaPssDefault) const;

  // Encrypts the OEMCrypto session key used for deriving other keys.
  // On success, |enc_session_key_size| is populated with the number
  // of bytes written to |enc_session_key|, and OEMCrypto_SUCCESS is
  // returned.  If the provided |enc_session_key_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and
  // |enc_session_key_size| is set to the required buffer size.
  OEMCryptoResult EncryptSessionKey(const uint8_t* session_key,
                                    size_t session_key_size,
                                    uint8_t* enc_session_key,
                                    size_t* enc_session_key_size) const;
  // Same as above, except directly returns the encrypted key.
  std::vector<uint8_t> EncryptSessionKey(
      const std::vector<uint8_t>& session_key) const;
  std::vector<uint8_t> EncryptSessionKey(const std::string& session_key) const;

  // Encrypts the OEMCrypto encryption key used for encrypting the
  // DRM private key.
  // On success, |enc_encryption_key_size| is populated with the
  // number of bytes written to |enc_encryption_key|, and
  // OEMCrypto_SUCCESS is returned.
  // If the provided |enc_encryption_key_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and
  // |enc_encryption_key_size| is set to the required buffer size.
  OEMCryptoResult EncryptEncryptionKey(const uint8_t* encryption_key,
                                       size_t encryption_key_size,
                                       uint8_t* enc_encryption_key,
                                       size_t* enc_encryption_key_size) const;
  // Same as above, except directly returns the encrypted key.
  std::vector<uint8_t> EncryptEncryptionKey(
      const std::vector<uint8_t>& encryption_key) const;
  std::vector<uint8_t> EncryptEncryptionKey(
      const std::string& encryption_key) const;

  ~RsaPublicKey();

  RsaPublicKey(const RsaPublicKey&) = delete;
  RsaPublicKey(RsaPublicKey&&) = delete;
  const RsaPublicKey& operator=(const RsaPublicKey&) = delete;
  RsaPublicKey& operator=(RsaPublicKey&&) = delete;

 private:
  RsaPublicKey() {}

  // Initializes the public key object using the provided |buffer|.
  // In case of any failure, false is return and the key should be
  // discarded.
  bool InitFromSubjectPublicKeyInfo(const uint8_t* buffer, size_t length);
  bool InitFromPrivateKeyInfo(const uint8_t* buffer, size_t length);
  // Initializes the public key object from a private.
  bool InitFromPrivateKey(const RsaPrivateKey& private_key);
  // Initializes the public key object from an existing
  // OpenSSL/BoringSSL RSA key handle.  The RSA key must be
  // initialized and |allowed_schemes| must be a valid value.
  bool InitFromSslHandle(const RSA* rsa_handle, uint32_t allowed_schemes);

  // Signature specialization functions.
  OEMCryptoResult VerifySignaturePss(const uint8_t* message,
                                     size_t message_length,
                                     const uint8_t* signature,
                                     size_t signature_length) const;
  OEMCryptoResult VerifySignaturePkcs1Cast(const uint8_t* message,
                                           size_t message_length,
                                           const uint8_t* signature,
                                           size_t signature_length) const;

  // RSAES-OAEP encrypt.
  OEMCryptoResult EncryptOaep(const uint8_t* message, size_t message_size,
                              uint8_t* enc_message,
                              size_t* enc_message_length) const;

  // OpenSSL/BoringSSL implementation of an RSA key.
  // Will only include components of an RSA public key.
  RSA* key_ = nullptr;
  uint32_t allowed_schemes_ = 0;
  RsaFieldSize field_size_ = kRsaFieldUnknown;
};  // class RsaPublicKey

class RsaPrivateKey {
 public:
  // Creates a new, pseudorandom RSA private key.
  static std::unique_ptr<RsaPrivateKey> New(RsaFieldSize field_size);

  // Loads a serialized RSA private key.
  // The provided |buffer| must contain a valid ASN.1 DER encoded
  // PrivateKeyInfo (RFC 5208).
  //
  //  buffer: PrivateKeyInfo = {
  //    version: INTEGER = v1(0),
  //    privateKeyAlgorithm: OID = rsaEncryption,
  //    privateKey: OCTET STRING = ...,
  //      -- BER encoding of RSAPrivateKey (RFC 3447)
  //    attributes: Attributes = ...  -- Optional, not used by OEMCrypto
  //  }
  // Note: If the public key is not included, then it is computed from
  // the private.
  //
  // Failure will occur if the provided |buffer| does not contain a
  // valid RSAPrivateKey, or if the specified curve is not supported.
  static std::unique_ptr<RsaPrivateKey> Load(const uint8_t* buffer,
                                             size_t length);
  static std::unique_ptr<RsaPrivateKey> Load(const std::string& buffer);
  static std::unique_ptr<RsaPrivateKey> Load(
      const std::vector<uint8_t>& buffer);

  // Creates a new RSA public key of this private key.
  // Equivalent to calling RsaPublicKey::New with this private
  // key.
  std::unique_ptr<RsaPublicKey> MakePublicKey() const;

  RsaFieldSize field_size() const { return field_size_; }
  uint32_t allowed_schemes() const { return allowed_schemes_; }
  const RSA* GetRsaKey() const { return key_; }

  // Checks if the provided |public_key| is the RSA public key of this
  // private key.
  bool IsMatchingPublicKey(const RsaPublicKey& public_key) const;

  // Serializes the private key into an ASN.1 DER encoded X
  // representation.
  // On success, |buffer_size| is populated with the number of bytes
  // written to |buffer|, and OEMCrypto_SUCCESS is returned.
  // If the provided |buffer_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |buffer_size| is
  // set to the required buffer size.
  OEMCryptoResult Serialize(uint8_t* buffer, size_t* buffer_size) const;
  // Same as above, except directly returns the serialized key.
  // Returns an empty vector on error.
  std::vector<uint8_t> Serialize() const;

  // Signs the provided |message| using the RSA signing algorithm
  // specified by |algorithm|.  See RsaSignatureAlgorithm for
  // details on each algorithm.
  //
  // On success, |signature_length| is populated with the number of
  // bytes written to |signature|, and OEMCrypto_SUCCESS is returned.
  // If the provided |signature_length| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |signature_length|
  // is set to the required signature size.
  OEMCryptoResult GenerateSignature(const uint8_t* message,
                                    size_t message_length,
                                    RsaSignatureAlgorithm algorithm,
                                    uint8_t* signature,
                                    size_t* signature_length) const;
  // Same as above, except directly returns the serialized signature.
  // Returns an empty vector on error.
  std::vector<uint8_t> GenerateSignature(
      const std::vector<uint8_t>& message,
      RsaSignatureAlgorithm algorithm = kRsaPssDefault) const;
  std::vector<uint8_t> GenerateSignature(
      const std::string& message,
      RsaSignatureAlgorithm algorithm = kRsaPssDefault) const;
  // Returns an upper bound for the signature size.  May be larger than
  // the actual signature generated by GenerateSignature().
  size_t SignatureSize() const;

  // Decrypts the OEMCrypto session key used for deriving other keys.
  // On success, |session_key_size| is populated with the number of
  // bytes written to |session_key|, and OEMCrypto_SUCCESS is returned.
  // If the provided |session_key_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |session_key_size|
  // is set to the required buffer size.
  OEMCryptoResult DecryptSessionKey(const uint8_t* enc_session_key,
                                    size_t enc_session_key_size,
                                    uint8_t* session_key,
                                    size_t* session_key_size) const;
  // Same as above, except directly returns the decrypted key.
  std::vector<uint8_t> DecryptSessionKey(
      const std::vector<uint8_t>& enc_session_key) const;
  std::vector<uint8_t> DecryptSessionKey(
      const std::string& enc_session_key) const;
  // Returns the byte length of the symmetric key that would be derived
  // by DecryptSessionKey().
  size_t SessionKeyLength() const;

  // Decrypts the OEMCrypto encryption key used for decrypting DRM
  // private key.
  // On success, |encryption_key_size| is populated with the number of
  // bytes written to |encryption_key|, and OEMCrypto_SUCCESS is
  // returned.
  // If the provided |encryption_key_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |encryption_key_size|
  // is set to the required buffer size.
  OEMCryptoResult DecryptEncryptionKey(const uint8_t* enc_encryption_key,
                                       size_t enc_encryption_key_size,
                                       uint8_t* encryption_key,
                                       size_t* encryption_key_size) const;
  // Same as above, except directly returns the decrypted key.
  std::vector<uint8_t> DecryptEncryptionKey(
      const std::vector<uint8_t>& enc_encryption_key) const;
  std::vector<uint8_t> DecryptEncryptionKey(
      const std::string& enc_encryption_key) const;

  ~RsaPrivateKey();

  RsaPrivateKey(const RsaPrivateKey&) = delete;
  RsaPrivateKey(RsaPrivateKey&&) = delete;
  const RsaPrivateKey& operator=(const RsaPrivateKey&) = delete;
  RsaPrivateKey& operator=(RsaPrivateKey&&) = delete;

 private:
  RsaPrivateKey() {}

  // Initializes the public key object using the provided |buffer|.
  // In case of any failure, false is return and the key should be
  // discarded.
  bool InitFromPrivateKeyInfo(const uint8_t* buffer, size_t length);
  // Generates a new key based on the provided field size.
  bool InitFromFieldSize(RsaFieldSize field_size);

  // Signature specialization functions.
  OEMCryptoResult GenerateSignaturePss(const uint8_t* message,
                                       size_t message_length,
                                       uint8_t* signature,
                                       size_t* signature_length) const;
  OEMCryptoResult GenerateSignaturePkcs1Cast(const uint8_t* message,
                                             size_t message_length,
                                             uint8_t* signature,
                                             size_t* signature_length) const;

  // RSAES-OAEP decrypt.
  OEMCryptoResult DecryptOaep(const uint8_t* enc_message,
                              size_t enc_message_size, uint8_t* message,
                              size_t expected_message_length) const;

  // OpenSSL/BoringSSL implementation of an RSA key.
  // Will include all components of an RSA private key.
  RSA* key_ = nullptr;
  uint32_t allowed_schemes_ = 0;
  // Set true if the deserialized key contained an allowed schemes.
  bool explicit_schemes_ = false;
  RsaFieldSize field_size_ = kRsaFieldUnknown;
};  // class RsaPrivateKey
}  // namespace util
}  // namespace wvoec
#endif  // WVOEC_UTIL_RSA_KEY_H_