////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 Google LLC.
//
// This software is licensed under the terms defined in the Widevine Master
// License Agreement. For a copy of this agreement, please contact
// widevine-licensing@google.com.
////////////////////////////////////////////////////////////////////////////////

#include "provisioning_sdk/internal/provisioning_engine_impl.h"

#include <stddef.h>
#include <time.h>

#include <limits>
#include <memory>
#include <string>
#include <utility>

#include "glog/logging.h"
#include "absl/synchronization/mutex.h"
#include "common/aes_cbc_util.h"
#include "common/certificate_type.h"
#include "common/crypto_util.h"
#include "common/drm_root_certificate.h"
#include "common/drm_service_certificate.h"
#include "common/hash_algorithm_util.h"
#include "common/random_util.h"
#include "common/rsa_key.h"
#include "provisioning_sdk/internal/certificates/root_oem_certificates.h"
#include "provisioning_sdk/public/provisioning_status.h"

#define LOG_WITH_PROTO(message, proto) \
  LOG(WARNING) << (message) << " [proto: " << (proto).ShortDebugString() << "]"

namespace widevine {
namespace {

const size_t kContextEncryptionKeySize(32);
const size_t kContextEncryptionIvSize(16);
const size_t kContextMacKeySize(32);

bool GenerateCertificate(DrmCertificate::Type type, uint32_t system_id,
                         const std::string& provider_id,
                         const std::string& serial_number,
                         const std::string& public_key,
                         const RsaPrivateKey& signing_key,
                         const SignedDrmCertificate& signer,
                         std::string* certificate) {
  DCHECK(type == DrmCertificate::DEVICE_MODEL ||
         type == DrmCertificate::DEVICE);
  if (serial_number.empty()) {
    LOG(WARNING) << "Require an non-empty serial number.";
    return false;
  }

  DrmCertificate drm_cert;
  drm_cert.set_type(type);
  drm_cert.set_system_id(system_id);
  if (!provider_id.empty()) drm_cert.set_provider_id(provider_id);
  drm_cert.set_serial_number(serial_number);
  drm_cert.set_creation_time_seconds(time(nullptr));
  drm_cert.set_public_key(public_key);

  SignedDrmCertificate signed_cert;
  if (!drm_cert.SerializeToString(signed_cert.mutable_drm_certificate())) {
    LOG(WARNING) << "Error serializing DrmCertificate.";
    return false;
  }
  if (!signing_key.GenerateSignature(
          signed_cert.drm_certificate(),
          HashAlgorithmProtoToEnum(signed_cert.hash_algorithm()),
          signed_cert.mutable_signature())) {
    LOG(WARNING) << "Failed to generate signature for DrmCertificate.";
    return false;
  }
  *signed_cert.mutable_signer() = signer;
  if (!signed_cert.SerializeToString(certificate)) {
    LOG(WARNING) << "Failed to serialize SignedDrmCertificate to string.";
    return false;
  }
  return true;
}

// Compares oem serial number, which is a 16-byte big number.
bool IsSerialNumberEq(const std::string& a, const std::string& b) {
  // An empty serial number indicates that it does not need to be matched.
  if (a.empty() || b.empty()) return true;
  int a_index = a.size() - 1;
  int b_index = b.size() - 1;
  // Matching a and b backwards.
  for (; a_index >= 0 && b_index >= 0; --a_index, --b_index)
    if (a[a_index] != b[b_index]) return false;
  // The remaining characters should be 0.
  for (; a_index >= 0; --a_index)
    if (a[a_index] != 0) return false;
  for (; b_index >= 0; --b_index)
    if (b[b_index] != 0) return false;
  return true;
}

}  // namespace

// NOTE: certificate_expiration_seconds_utc_ not initialized in initialization
// list due to bug in clang and/or CLIF which causes the initialization to be
// skipped.
ProvisioningEngineImpl::ProvisioningEngineImpl()
    : rsa_key_factory_(new RsaKeyFactory),
      certificate_expiration_seconds_utc_(0) {}

ProvisioningEngineImpl::~ProvisioningEngineImpl() {}

ProvisioningStatus ProvisioningEngineImpl::Initialize(
    CertificateType certificate_type,
    const std::string& drm_service_certificate,
    const std::string& service_private_key,
    const std::string& service_private_key_passphrase,
    const std::string& provisioning_drm_certificate,
    const std::string& provisioning_private_key,
    const std::string& provisioning_private_key_passphrase,
    const std::string& secret_spoid_sauce) {
  Status status;

  if (!drm_root_certificate_) {
    status = DrmRootCertificate::CreateByType(certificate_type,
                                              &drm_root_certificate_);
    if (!status.ok()) {
      LOG(ERROR) << status;
      return INVALID_CERTIFICATE_TYPE;
    }
  }
  drm_root_public_key_ = rsa_key_factory_->CreateFromPkcs1PublicKey(
      drm_root_certificate_->public_key());
  if (!drm_root_public_key_) {
    LOG(ERROR) << "Failed to instiate DRM root public key.";
    return INTERNAL_ERROR;
  }

  status = DrmServiceCertificate::AddDrmServiceCertificate(
      drm_root_certificate_.get(), drm_service_certificate, service_private_key,
      service_private_key_passphrase);
  if (!status.ok()) {
    LOG(ERROR) << status;
    return INVALID_SERVICE_DRM_CERTIFICATE;
  }

  DrmCertificate drm_cert;
  status = drm_root_certificate_->VerifyCertificate(
      provisioning_drm_certificate, &signed_provisioning_cert_, &drm_cert);
  if (!status.ok()) {
    LOG(ERROR) << status;
    return INVALID_PROVISIONER_DRM_CERTIFICATE;
  }
  if (drm_cert.type() != DrmCertificate::ROOT &&
      drm_cert.type() != DrmCertificate::PROVISIONER) {
    LOG(ERROR) << "Expecting ROOT or PROVISIONER certificate.";
    return INVALID_PROVISIONER_DRM_CERTIFICATE;
  }

  provisioning_public_key_ =
      rsa_key_factory_->CreateFromPkcs1PublicKey(drm_cert.public_key());
  if (!provisioning_public_key_) return INVALID_PROVISIONER_DRM_CERTIFICATE;
  provisioning_private_key_ = rsa_key_factory_->CreateFromPkcs8PrivateKey(
      provisioning_private_key, provisioning_private_key_passphrase);
  if (!provisioning_private_key_) return INVALID_PROVISIONER_PRIVATE_KEY;
  if (!provisioning_public_key_->MatchesPrivateKey(
          *provisioning_private_key_)) {
    LOG(ERROR) << "Provisioning public key and private key do not match.";
    return INVALID_PROVISIONER_PRIVATE_KEY;
  }

  if (secret_spoid_sauce.empty()) {
    LOG(ERROR) << "SPOID secret sauce is empty!";
    return INVALID_SPOID_SAUCE;
  }
  secret_spoid_sauce_ = secret_spoid_sauce;

  if (!oem_device_cert_.Initialize(certificate_type)) return INTERNAL_ERROR;
  return OK;
}

ProvisioningStatus ProvisioningEngineImpl::SetCertificateStatusList(
    const std::string& certificate_status_list,
    uint32_t expiration_period_seconds) {
  if (certificate_status_list.empty()) {
    LOG(WARNING) << "Empty certificate_status_list.";
    return INVALID_STATUS_LIST;
  }

  SignedDeviceCertificateStatusList signed_cert_status_list;
  if (!signed_cert_status_list.ParseFromString(certificate_status_list)) {
    LOG(WARNING) << "Error parsing SignedDeviceCertificateStatusList.";
    return INVALID_STATUS_LIST;
  }

  if (!drm_root_public_key_->VerifySignature(
          signed_cert_status_list.certificate_status_list(),
          HashAlgorithmProtoToEnum(signed_cert_status_list.hash_algorithm()),
          signed_cert_status_list.signature())) {
    LOG_WITH_PROTO("Signature verification failed", signed_cert_status_list);
    return INVALID_STATUS_LIST;
  }

  DeviceCertificateStatusList cert_status_list;
  if (!cert_status_list.ParseFromString(
          signed_cert_status_list.certificate_status_list())) {
    LOG_WITH_PROTO("Error parsing DeviceCertificateStatusList",
                   signed_cert_status_list);
    return INVALID_STATUS_LIST;
  }

  absl::WriterMutexLock writer_lock(&cert_status_mutex_);

  if (expiration_period_seconds == 0) {
    certificate_expiration_seconds_utc_ = std::numeric_limits<uint32_t>::max();
  } else {
    certificate_expiration_seconds_utc_ =
        cert_status_list.creation_time_seconds() + expiration_period_seconds;
  }

  certificate_status_map_.clear();
  for (DeviceCertificateStatus& cert_status :
       *cert_status_list.mutable_certificate_status()) {
    const uint32_t system_id = cert_status.device_info().system_id();
    certificate_status_map_[system_id].Swap(&cert_status);
  }

  // Remove items that are no longer valid.
  for (auto it = intermediate_certs_info_.begin();
       it != intermediate_certs_info_.end();) {
    auto certificate_status_it = certificate_status_map_.find(it->first);
    if (certificate_status_it == certificate_status_map_.end())
      intermediate_certs_info_.erase(it++);
    else
      ++it;
  }
  // Set / Update device info.
  for (const auto& cert_status : certificate_status_map_) {
    auto device_info = std::make_shared<ProvisionedDeviceInfo>();
    *device_info = cert_status.second.device_info();
    intermediate_certs_info_[cert_status.first].device_info = device_info;
    // intermediate certificate and private key are not changed if exists.
  }
  return OK;
}

ProvisioningStatus ProvisioningEngineImpl::GenerateDrmIntermediateCertificate(
    uint32_t system_id, const std::string& public_key,
    std::string* certificate) const {
  auto intermediate_public_key =
      rsa_key_factory_->CreateFromPkcs1PublicKey(public_key);
  if (!intermediate_public_key) return INVALID_INTERMEDIATE_PUBLIC_KEY;

  const size_t RootCertificateSerialNumberSize = 16;
  std::string serial_number;
  if (!RandomBytes(RootCertificateSerialNumberSize, &serial_number)) {
    LOG(WARNING) << "Failed to generate serial_number.";
    return INTERNAL_ERROR;
  }
  if (!GenerateCertificate(DrmCertificate::DEVICE_MODEL, system_id,
                           std::string(), serial_number, public_key,
                           *provisioning_private_key_,
                           signed_provisioning_cert_, certificate)) {
    return INTERNAL_ERROR;
  }
  return OK;
}

ProvisioningStatus ProvisioningEngineImpl::AddDrmIntermediateCertificate(
    const std::string& intermediate_cert, const std::string& cert_private_key,
    const std::string& cert_private_key_passphrase) {
  SignedDrmCertificate intermediate_signed_cert;
  DrmCertificate intermediate_drm_cert;
  Status status = drm_root_certificate_->VerifyCertificate(
      intermediate_cert, &intermediate_signed_cert, &intermediate_drm_cert);
  if (!status.ok()) {
    LOG(ERROR) << status;
    return INVALID_INTERMEDIATE_DRM_CERTIFICATE;
  }

  if (intermediate_drm_cert.type() != DrmCertificate::DEVICE_MODEL) {
    LOG_WITH_PROTO("Invalid device certificate type", intermediate_drm_cert);
    return INVALID_INTERMEDIATE_DRM_CERTIFICATE;
  }
  if (!intermediate_drm_cert.has_system_id()) {
    LOG_WITH_PROTO("Missing system_id", intermediate_drm_cert);
    return UNKNOWN_SYSTEM_ID;
  }

  const std::string empty_oem_ca_serial_number;
  ProvisioningStatus provisioning_status = CheckDeviceStatus(
      intermediate_drm_cert.system_id(), empty_oem_ca_serial_number);
  if (provisioning_status != OK) return provisioning_status;

  auto intermediate_public_key = rsa_key_factory_->CreateFromPkcs1PublicKey(
      intermediate_drm_cert.public_key());
  if (!intermediate_public_key) return INVALID_INTERMEDIATE_DRM_CERTIFICATE;
  std::unique_ptr<RsaPrivateKey> intermediate_private_key =
      rsa_key_factory_->CreateFromPkcs8PrivateKey(cert_private_key,
                                                  cert_private_key_passphrase);
  if (!intermediate_private_key) return INVALID_INTERMEDIATE_PRIVATE_KEY;
  if (!intermediate_public_key->MatchesPrivateKey(*intermediate_private_key)) {
    LOG(WARNING) << "Intermediate public key and private key do not match.";
    return INVALID_INTERMEDIATE_PRIVATE_KEY;
  }

  absl::WriterMutexLock writer_lock(&cert_status_mutex_);
  auto& certificate_info =
      intermediate_certs_info_[intermediate_drm_cert.system_id()];
  certificate_info.signed_drm_certificate.Swap(&intermediate_signed_cert);
  certificate_info.private_key = std::move(intermediate_private_key);
  return OK;
}

ProvisioningStatus ProvisioningEngineImpl::GenerateDeviceDrmCertificate(
    uint32_t system_id, const std::string& oem_ca_serial_number,
    const std::string& public_key, const std::string& certificate_serial_number,
    std::string* certificate) const {
  // |oem_ca_serial_number| could be empty if it is called directly from
  // ProvisioningEngine::GenerateDeviceDrmCertificate.
  DCHECK(!certificate_serial_number.empty());

  return GenerateProviderDeviceDrmCertificate(
      system_id, oem_ca_serial_number, std::string(), public_key,
      certificate_serial_number, certificate);
}

ProvisioningStatus ProvisioningEngineImpl::GenerateProviderDeviceDrmCertificate(
    uint32_t system_id, const std::string& oem_ca_serial_number,
    const std::string& provider_id, const std::string& public_key,
    const std::string& certificate_serial_number,
    std::string* certificate) const {
  // |oem_ca_serial_number| could be empty if it is called directly from
  // ProvisioningEngine::GenerateDeviceDrmCertificate.
  DCHECK(!certificate_serial_number.empty());

  ProvisioningStatus status =
      CheckDeviceStatus(system_id, oem_ca_serial_number);
  if (status != OK) return status;

  std::shared_ptr<RsaPrivateKey> intermediate_private_key;
  const SignedDrmCertificate* intermediate_cert = nullptr;
  {
    absl::ReaderMutexLock reader_lock(&cert_status_mutex_);

    auto info_it = intermediate_certs_info_.find(system_id);
    if (info_it == intermediate_certs_info_.end() ||
        !info_it->second.private_key) {
      LOG(WARNING) << "Cannot find the intermediate certificate for system: "
                   << system_id;
      return MISSING_DEVICE_MODEL_CERT;
    }
    intermediate_private_key = info_it->second.private_key;
    intermediate_cert = &info_it->second.signed_drm_certificate;
  }

  if (!GenerateCertificate(DrmCertificate::DEVICE, system_id, provider_id,
                           certificate_serial_number, public_key,
                           *intermediate_private_key, *intermediate_cert,
                           certificate)) {
    return INTERNAL_ERROR;
  }
  return OK;
}

std::shared_ptr<ProvisionedDeviceInfo> ProvisioningEngineImpl::GetDeviceInfo(
    uint32_t system_id) const {
  absl::ReaderMutexLock reader_lock(&cert_status_mutex_);
  auto info_it = intermediate_certs_info_.find(system_id);
  if (info_it == intermediate_certs_info_.end()) {
    LOG(WARNING) << "Cannot find the system id in device certificate list: "
                 << system_id;
    return std::shared_ptr<ProvisionedDeviceInfo>();
  }
  return info_it->second.device_info;
}

ProvisioningStatus ProvisioningEngineImpl::StoreContext(
    const std::string& context_data,
    SignedProvisioningContext* signed_context) const {
  DCHECK(signed_context);

  ProvisioningContextKeyData key_data;
  if (!RandomBytes(kContextEncryptionKeySize,
                   key_data.mutable_encryption_key()) ||
      !RandomBytes(kContextEncryptionIvSize,
                   key_data.mutable_encryption_iv())) {
    LOG(ERROR) << "Failed to generate random context key data.";
    return INTERNAL_ERROR;
  }

  const DrmServiceCertificate* service_cert =
      DrmServiceCertificate::GetDefaultDrmServiceCertificateOrDie();

  ProvisioningContext context;
  context.set_context_data(crypto_util::EncryptAesCbc(
      key_data.encryption_key(), key_data.encryption_iv(), context_data));
  if (!service_cert->public_key()->Encrypt(key_data.SerializeAsString(),
                                           context.mutable_key_data())) {
    LOG(WARNING) << "Context key data encryption failed";
    return INTERNAL_ERROR;
  }

  signed_context->set_provisioning_context(context.SerializeAsString());
  if (!service_cert->private_key()->GenerateSignature(
          signed_context->provisioning_context(),
          HashAlgorithmProtoToEnum(signed_context->hash_algorithm()),
          signed_context->mutable_signature())) {
    LOG(WARNING) << "Failed to generate signature for ProvisioningContext.";
    return INTERNAL_ERROR;
  }
  return OK;
}

ProvisioningStatus ProvisioningEngineImpl::RetrieveContext(
    const SignedProvisioningContext& signed_context,
    std::string* context_data) const {
  DCHECK(context_data);

  const DrmServiceCertificate* service_cert =
      DrmServiceCertificate::GetDefaultDrmServiceCertificateOrDie();

  if (!service_cert->public_key()->VerifySignature(
          signed_context.provisioning_context(),
          HashAlgorithmProtoToEnum(signed_context.hash_algorithm()),
          signed_context.signature())) {
    LOG(WARNING) << "ProvisioningContext signature verification failed.";
    return INVALID_CONTEXT;
  }

  ProvisioningContext context;
  if (!context.ParseFromString(signed_context.provisioning_context())) {
    LOG(WARNING) << "Invalid context.";
    return INVALID_CONTEXT;
  }

  std::string serialized_key_data;
  if (!service_cert->private_key()->Decrypt(context.key_data(),
                                            &serialized_key_data)) {
    LOG(WARNING) << "Could not decrypt context key data";
    return INVALID_CONTEXT_KEY_DATA;
  }
  if (serialized_key_data.empty()) {
    LOG(WARNING) << "Context key data is missing.";
    return INVALID_CONTEXT_KEY_DATA;
  }
  ProvisioningContextKeyData key_data;
  if (!key_data.ParseFromString(serialized_key_data)) {
    LOG(WARNING) << "Invalid context key data.";
    return INVALID_CONTEXT_KEY_DATA;
  }
  *context_data = crypto_util::DecryptAesCbc(key_data.encryption_key(),
                                             key_data.encryption_iv(),
                                             context.context_data());
  if (context_data->empty()) {
    LOG(WARNING) << "Provisioning context decryption failed.";
    return INVALID_CONTEXT;
  }
  return OK;
}

ProvisioningStatus ProvisioningEngineImpl::CheckDeviceStatus(
    uint32_t system_id, const std::string& oem_ca_serial_number) const {
  absl::ReaderMutexLock reader_lock(&cert_status_mutex_);

  if (certificate_expiration_seconds_utc_ < time(nullptr))
    return STATUS_LIST_EXPIRED;

  auto certificate_status_it = certificate_status_map_.find(system_id);
  if (certificate_status_it == certificate_status_map_.end()) {
    LOG(WARNING) << "Cannot find the system id in device certificate list: "
                 << system_id;
    return UNKNOWN_SYSTEM_ID;
  }
  if (!IsSerialNumberEq(certificate_status_it->second.oem_serial_number(),
                        oem_ca_serial_number)) {
    LOG(WARNING) << "Provided serial number does not match with stored serial "
                    "number. It may come from a revoked device. System Id: "
                 << system_id;
    return DEVICE_REVOKED;
  }
  if (certificate_status_it->second.status() ==
      DeviceCertificateStatus::STATUS_REVOKED) {
    LOG(WARNING) << "Device revoked " << system_id;
    return DEVICE_REVOKED;
  }
  return OK;
}

}  // namespace widevine