// 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_DRM_KEY_H_
#define WVOEC_UTIL_DRM_KEY_H_

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

#include "OEMCryptoCENCCommon.h"
#include "oemcrypto_ecc_key.h"
#include "oemcrypto_rsa_key.h"

namespace wvoec {
namespace util {
// DRM private key performs all of the operations required by an
// OEMCrypto session's RSA/ECC private key.
class DrmPrivateKey {
 public:
  // Create an RSA-based DRM key.
  static std::unique_ptr<DrmPrivateKey> Create(
      std::shared_ptr<RsaPrivateKey>&& rsa_key);
  static std::unique_ptr<DrmPrivateKey> Create(
      std::unique_ptr<RsaPrivateKey>&& rsa_key);
  // Create an ECC-based DRM key.
  static std::unique_ptr<DrmPrivateKey> Create(
      std::shared_ptr<EccPrivateKey>&& ecc_key);
  static std::unique_ptr<DrmPrivateKey> Create(
      std::unique_ptr<EccPrivateKey>&& ecc_key);

  bool IsRsaKey() const { return static_cast<bool>(rsa_key_); }
  bool IsEccKey() const { return static_cast<bool>(ecc_key_); }

  // Generates a session key from the key source.
  // For RSA keys, |key_source| is an encrypted session key.
  // For ECC keys, |key_source| is a ephemeral public key to be
  // used in ECDH.
  OEMCryptoResult GetSessionKey(const uint8_t* key_source,
                                size_t key_source_size,
                                std::vector<uint8_t>* session_key) const;
  std::vector<uint8_t> GetSessionKey(
      const std::vector<uint8_t>& key_source) const;

  // Generates a encryption key from the key source.
  // For RSA keys, |key_source| is an encrypted encryption key.
  // For ECC keys, this method is not supported.
  std::vector<uint8_t> GetEncryptionKey(
      const std::vector<uint8_t>& key_source) const;

  // Generates a signature for the provided message.
  // For RSA keys, the signature is RSASSA-PSS.
  // For ECC keys, the signature is ECDSA.
  OEMCryptoResult GenerateSignature(const uint8_t* message,
                                    size_t message_length, uint8_t* signature,
                                    size_t* signature_length) const;
  std::vector<uint8_t> GenerateSignature(
      const std::vector<uint8_t>& message) const;
  size_t SignatureSize() const;

  // Generates a signature for the provided message.
  // For RSA keys, the signature is RSASSA-PKCS1.
  // For ECC keys, this is not supported.
  OEMCryptoResult GenerateRsaSignature(const uint8_t* message,
                                       size_t message_length,
                                       uint8_t* signature,
                                       size_t* signature_length) const;
  std::vector<uint8_t> GenerateRsaSignature(
      const std::vector<uint8_t>& message) const;

  ~DrmPrivateKey() {}

 private:
  DrmPrivateKey() {}

  // Only one will be set.
  std::shared_ptr<EccPrivateKey> ecc_key_;
  std::shared_ptr<RsaPrivateKey> rsa_key_;
};
}  // namespace util
}  // namespace wvoec
#endif  // WVOEC_UTIL_DRM_KEY_H_