////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 Google Inc.
//
// 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 <cstdint>
#include <limits>
#include <memory>
#include <string>
#include <utility>

#include "glog/logging.h"
#include "common/random_util.h"
#include "common/rsa_key.h"
#include "provisioning_sdk/internal/certificates/root_certificates.h"
#include "provisioning_sdk/public/provisioning_status.h"

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

namespace widevine {
namespace {

// Verify that |certificate| is signed by |public_key|. If |public_key| is null,
// |certificate| should be self signed.
bool VerifyAndExtractCertificate(const RsaPublicKey* public_key,
                                 const std::string& certificate,
                                 SignedDrmDeviceCertificate* signed_drm_cert,
                                 DrmDeviceCertificate* drm_cert) {
  DCHECK(signed_drm_cert);
  DCHECK(drm_cert);
  if (!signed_drm_cert->ParseFromString(certificate)) {
    LOG(WARNING) << "Failed to parse SignedDrmDeviceCertificate.";
    return false;
  }
  if (signed_drm_cert->drm_certificate().empty()) {
    LOG_WITH_PROTO("Missing drm_certificate", *signed_drm_cert);
    return false;
  }
  if (signed_drm_cert->signature().empty()) {
    LOG_WITH_PROTO("Missing signature", *signed_drm_cert);
    return false;
  }

  if (!drm_cert->ParseFromString(signed_drm_cert->drm_certificate())) {
    LOG_WITH_PROTO("Failed to parse DrmDeviceCertificate", *signed_drm_cert);
    return false;
  }
  if (drm_cert->public_key().empty()) {
    LOG_WITH_PROTO("Missing public_key", *drm_cert);
    return false;
  }

  std::unique_ptr<RsaPublicKey> local_public_key;
  if (!public_key) {
    local_public_key.reset(RsaPublicKey::Create(drm_cert->public_key()));
    if (!local_public_key) {
      LOG_WITH_PROTO("Invalid root public key", *drm_cert);
      return false;
    }
    public_key = local_public_key.get();
  }

  if (!public_key->VerifySignature(signed_drm_cert->drm_certificate(),
                                   signed_drm_cert->signature())) {
    LOG_WITH_PROTO("Signature verification failed", *signed_drm_cert);
    return false;
  }
  return true;
}

bool GenerateCertificate(DrmDeviceCertificate::CertificateType 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 SignedDrmDeviceCertificate& signer,
                         std::string* certificate) {
  DCHECK(type == DrmDeviceCertificate::DRM_INTERMEDIATE ||
         type == DrmDeviceCertificate::DRM_USER_DEVICE);
  if (serial_number.empty()) {
    LOG(WARNING) << "Require an non-empty serial number.";
    return false;
  }

  DrmDeviceCertificate 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);

  SignedDrmDeviceCertificate signed_cert;
  if (!drm_cert.SerializeToString(
          signed_cert.mutable_drm_certificate())) {
    LOG(WARNING) << "Error serializing DrmDeviceCertificate.";
    return false;
  }
  if (!signing_key.GenerateSignature(signed_cert.drm_certificate(),
                                     signed_cert.mutable_signature())) {
    LOG(WARNING) << "Failed to generate signature for DrmDeviceCertificate.";
    return false;
  }
  *signed_cert.mutable_signer() = signer;
  if (!signed_cert.SerializeToString(certificate)) {
    LOG(WARNING) << "Failed to serialize SignedDrmDeviceCertificate 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

ProvisioningEngineImpl::ProvisioningEngineImpl()
    : rsa_key_factory_(new RsaKeyFactory) {}

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_phassphrase,
    const std::string& provisioning_drm_certificate,
    const std::string& provisioning_private_key,
    const std::string& provisioning_private_key_phassphrase,
    const std::string& secret_spoid_sauce) {
  if (!LoadDrmRootPublicKey(certificate_type)) return INVALID_CERTIFICATE_TYPE;

  SignedDrmDeviceCertificate signed_drm_cert;
  DrmDeviceCertificate drm_cert;
  if (!VerifyAndExtractCertificate(root_public_key_.get(),
                                   drm_service_certificate, &signed_drm_cert,
                                   &drm_cert)) {
    return INVALID_SERVICE_DRM_CERTIFICATE;
  }
  if (drm_cert.type() != DrmDeviceCertificate::SERVICE) {
    LOG(WARNING) << "Expecting SERVICE certificate.";
    return INVALID_SERVICE_DRM_CERTIFICATE;
  }
  service_public_key_ =
      rsa_key_factory_->CreateFromPkcs1PublicKey(drm_cert.public_key());
  if (!service_public_key_) return INVALID_SERVICE_DRM_CERTIFICATE;
  service_private_key_ = rsa_key_factory_->CreateFromPkcs8PrivateKey(
      service_private_key, service_private_key_phassphrase);
  if (!service_private_key_) return INVALID_SERVICE_PRIVATE_KEY;
  if (!service_public_key_->MatchesPrivateKey(*service_private_key_)) {
    LOG(WARNING) << "Services public key and private key do not match.";
    return INVALID_SERVICE_PRIVATE_KEY;
  }

  if (!VerifyAndExtractCertificate(root_public_key_.get(),
                                   provisioning_drm_certificate,
                                   &signed_provisioning_cert_, &drm_cert)) {
    return INVALID_PROVISIONER_DRM_CERTIFICATE;
  }
  if (drm_cert.type() != DrmDeviceCertificate::ROOT &&
      drm_cert.type() != DrmDeviceCertificate::PROVISIONER) {
    LOG(WARNING) << "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_phassphrase);
  if (!provisioning_private_key_) return INVALID_PROVISIONER_PRIVATE_KEY;
  if (!provisioning_public_key_->MatchesPrivateKey(
          *provisioning_private_key_)) {
    LOG(WARNING) << "Provisioning public key and private key do not match.";
    return INVALID_PROVISIONER_PRIVATE_KEY;
  }

  if (secret_spoid_sauce.empty()) {
    LOG(WARNING) << "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;
  }

  SignedCertificateStatusList signed_cert_status_list;
  if (!signed_cert_status_list.ParseFromString(certificate_status_list)) {
    LOG(WARNING) << "Error parsing SignedCertificateStatusList.";
    return INVALID_STATUS_LIST;
  }
  if (!root_public_key_->VerifySignature(
          signed_cert_status_list.certificate_status_list(),
          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;
  }

  WriterMutexLock writer_lock(&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 kCertificateSerialNumberSize = 16;
  std::string serial_number;
  if (!RandomBytes(kCertificateSerialNumberSize, &serial_number)) {
    LOG(WARNING) << "Failed to generate serial_number.";
    return INTERNAL_ERROR;
  }
  if (!GenerateCertificate(DrmDeviceCertificate::DRM_INTERMEDIATE,
                           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) {
  SignedDrmDeviceCertificate intermediate_signed_cert;
  DrmDeviceCertificate intermediate_drm_cert;
  if (!VerifyAndExtractCertificate(provisioning_public_key_.get(),
                                   intermediate_cert, &intermediate_signed_cert,
                                   &intermediate_drm_cert)) {
    return INVALID_INTERMEDIATE_DRM_CERTIFICATE;
  }
  if (intermediate_drm_cert.type() !=
      DrmDeviceCertificate::DRM_INTERMEDIATE) {
    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 status = CheckDeviceStatus(
      intermediate_drm_cert.system_id(), empty_oem_ca_serial_number);
  if (status != OK) return 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;
  }

  WriterMutexLock writer_lock(&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 SignedDrmDeviceCertificate* intermediate_cert = nullptr;
  {
    ReaderMutexLock reader_lock(&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_DRM_INTERMEDIATE_CERT;
    }
    intermediate_private_key = info_it->second.private_key;
    intermediate_cert = &info_it->second.signed_drm_certificate;
  }

  if (!GenerateCertificate(DrmDeviceCertificate::DRM_USER_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 {
  ReaderMutexLock reader_lock(&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;
}

bool ProvisioningEngineImpl::LoadDrmRootPublicKey(
    CertificateType certificate_type) {
  const std::string* root_cert_string = nullptr;
  RootCertificates root_certificates;
  switch (certificate_type) {
    case kCertTesting:
      root_cert_string = &root_certificates.drm_root_test_certificate();
      break;
    case kCertDevelopment:
      root_cert_string = &root_certificates.drm_root_dev_certificate();
      break;
    case kCertProduction:
      root_cert_string = &root_certificates.drm_root_prod_certificate();
      break;
    default:
      LOG(WARNING) << "Invalid certificate type " << certificate_type;
      return false;
  }

  SignedDrmDeviceCertificate signed_root_cert;
  DrmDeviceCertificate root_cert;
  if (!VerifyAndExtractCertificate(nullptr /* self signed */, *root_cert_string,
                                   &signed_root_cert, &root_cert)) {
    LOG(WARNING) << "Failed to extract root certificate.";
    return false;
  }
  if (root_cert.type() != DrmDeviceCertificate::ROOT) {
    LOG(WARNING) << "Expecting ROOT certificate.";
    return false;
  }
  root_public_key_ =
      rsa_key_factory_->CreateFromPkcs1PublicKey(root_cert.public_key());
  CHECK(root_public_key_);
  return true;
}

ProvisioningStatus ProvisioningEngineImpl::CheckDeviceStatus(
    uint32_t system_id, const std::string& oem_ca_serial_number) const {
  ReaderMutexLock reader_lock(&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::VALID) {
    LOG(WARNING) << "Device revoked " << system_id;
    return DEVICE_REVOKED;
  }
  return OK;
}

}  // namespace widevine