// 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
//
#include "oemcrypto_drm_key.h"

#include <utility>

#include "OEMCryptoCENC.h"
#include "log.h"

namespace wvoec {
namespace util {
// static
std::unique_ptr<DrmPrivateKey> DrmPrivateKey::Create(
    std::shared_ptr<RsaPrivateKey>&& rsa_key) {
  if (!rsa_key) {
    LOGE("No RSA key provided");
    return std::unique_ptr<DrmPrivateKey>();
  }
  std::unique_ptr<DrmPrivateKey> drm_key(new DrmPrivateKey());
  drm_key->rsa_key_ = std::move(rsa_key);
  return drm_key;
}

// static
std::unique_ptr<DrmPrivateKey> DrmPrivateKey::Create(
    std::unique_ptr<RsaPrivateKey>&& rsa_key) {
  if (!rsa_key) {
    LOGE("No RSA key provided");
    return std::unique_ptr<DrmPrivateKey>();
  }
  std::unique_ptr<DrmPrivateKey> drm_key(new DrmPrivateKey());
  drm_key->rsa_key_ = std::move(rsa_key);
  return drm_key;
}

// static
std::unique_ptr<DrmPrivateKey> DrmPrivateKey::Create(
    std::shared_ptr<EccPrivateKey>&& ecc_key) {
  if (!ecc_key) {
    LOGE("No ECC key provided");
    return std::unique_ptr<DrmPrivateKey>();
  }
  std::unique_ptr<DrmPrivateKey> drm_key(new DrmPrivateKey());
  drm_key->ecc_key_ = std::move(ecc_key);
  return drm_key;
}

// static
std::unique_ptr<DrmPrivateKey> DrmPrivateKey::Create(
    std::unique_ptr<EccPrivateKey>&& ecc_key) {
  if (!ecc_key) {
    LOGE("No ECC key provided");
    return std::unique_ptr<DrmPrivateKey>();
  }
  std::unique_ptr<DrmPrivateKey> drm_key(new DrmPrivateKey());
  drm_key->ecc_key_ = std::move(ecc_key);
  return drm_key;
}

OEMCryptoResult DrmPrivateKey::GetSessionKey(
    const uint8_t* key_source, size_t key_source_size,
    std::vector<uint8_t>* session_key) const {
  if (session_key == nullptr) {
    LOGE("Output session key is null");
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  // RSA -> Decrypt session key.
  if (rsa_key_) {
    if (!(rsa_key_->allowed_schemes() & kSign_RSASSA_PSS)) {
      LOGE("RSA key cannot be used for session key decryption");
      return OEMCrypto_ERROR_INVALID_KEY;
    }
    size_t session_key_size = rsa_key_->SessionKeyLength();
    session_key->resize(session_key_size);
    const OEMCryptoResult res = rsa_key_->DecryptSessionKey(
        key_source, key_source_size, session_key->data(), &session_key_size);
    if (res != OEMCrypto_SUCCESS) {
      session_key->clear();
      return res;
    }
    session_key->resize(session_key_size);
    return OEMCrypto_SUCCESS;
  }
  // ECC -> ECDH.
  // Step 1: Parse |key_source| as ECC key.
  std::unique_ptr<EccPublicKey> ephemeral_ecc_key =
      EccPublicKey::Load(key_source, key_source_size);
  if (!ephemeral_ecc_key) {
    LOGE("Failed to load server's ephemeral ECC key");
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  // Step 2: Derive session key.
  size_t session_key_size = ecc_key_->SessionKeyLength();
  session_key->resize(session_key_size);
  const OEMCryptoResult res = ecc_key_->DeriveSessionKey(
      *ephemeral_ecc_key, session_key->data(), &session_key_size);
  if (res != OEMCrypto_SUCCESS) {
    session_key->clear();
    return res;
  }
  session_key->resize(session_key_size);
  return OEMCrypto_SUCCESS;
}

std::vector<uint8_t> DrmPrivateKey::GetSessionKey(
    const std::vector<uint8_t>& key_source) const {
  // RSA -> Decrypt session key.
  if (rsa_key_) {
    if (!(rsa_key_->allowed_schemes() & kSign_RSASSA_PSS)) {
      LOGE("RSA key cannot be used for session key decryption");
      return std::vector<uint8_t>();
    }
    return rsa_key_->DecryptSessionKey(key_source);
  }
  // ECC -> ECDH.
  // Step 1: Parse |key_source| as ECC key.
  std::unique_ptr<EccPublicKey> ephemeral_ecc_key =
      EccPublicKey::Load(key_source);
  if (!ephemeral_ecc_key) {
    LOGE("Failed to load server's ephemeral ECC key");
    return std::vector<uint8_t>();
  }
  // Step 2: Derive session key.
  return ecc_key_->DeriveSessionKey(*ephemeral_ecc_key);
}

std::vector<uint8_t> DrmPrivateKey::GetEncryptionKey(
    const std::vector<uint8_t>& key_source) const {
  if (!rsa_key_) {
    LOGE("Only RSA DRM keys can derive an encryption key");
    return std::vector<uint8_t>();
  }
  return rsa_key_->DecryptEncryptionKey(key_source);
}

OEMCryptoResult DrmPrivateKey::GenerateSignature(
    const uint8_t* message, size_t message_length, uint8_t* signature,
    size_t* signature_length) const {
  if (rsa_key_) {
    return rsa_key_->GenerateSignature(message, message_length, kRsaPssDefault,
                                       signature, signature_length);
  }
  return ecc_key_->GenerateSignature(message, message_length, signature,
                                     signature_length);
}

std::vector<uint8_t> DrmPrivateKey::GenerateSignature(
    const std::vector<uint8_t>& message) const {
  if (rsa_key_) {
    return rsa_key_->GenerateSignature(message, kRsaPssDefault);
  }
  return ecc_key_->GenerateSignature(message);
}

size_t DrmPrivateKey::SignatureSize() const {
  if (rsa_key_) {
    return rsa_key_->SignatureSize();
  }
  return ecc_key_->SignatureSize();
}

OEMCryptoResult DrmPrivateKey::GenerateRsaSignature(
    const uint8_t* message, size_t message_length, uint8_t* signature,
    size_t* signature_length) const {
  if (!rsa_key_) {
    LOGE("Only RSA DRM keys can generate PKCS1 signatures");
    return OEMCrypto_ERROR_INVALID_KEY;
  }
  return rsa_key_->GenerateSignature(message, message_length, kRsaPkcs1Cast,
                                     signature, signature_length);
}

std::vector<uint8_t> DrmPrivateKey::GenerateRsaSignature(
    const std::vector<uint8_t>& message) const {
  if (!rsa_key_) {
    LOGE("Only RSA DRM keys can generate PKCS1 signatures");
    return std::vector<uint8_t>();
  }
  return rsa_key_->GenerateSignature(message, kRsaPkcs1Cast);
}
}  // namespace util
}  // namespace wvoec