////////////////////////////////////////////////////////////////////////////////
// Copyright 2013 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 "common/drm_service_certificate.h"

#include <map>
#include <memory>
#include <utility>

#include "glog/logging.h"
#include "base/thread_annotations.h"
#include "absl/strings/escaping.h"
#include "absl/synchronization/mutex.h"
#include "util/gtl/map_util.h"
#include "common/aes_cbc_util.h"
#include "common/certificate_type.h"
#include "common/drm_root_certificate.h"
#include "common/error_space.h"
#include "common/rsa_util.h"
#include "protos/public/client_identification.pb.h"
#include "protos/public/drm_certificate.pb.h"
#include "protos/public/errors.pb.h"
#include "protos/public/signed_drm_certificate.pb.h"

namespace widevine {

namespace {

// Class used to hold global service certificate map.
class DrmServiceCertificateMap {
 public:
  DrmServiceCertificateMap();
  ~DrmServiceCertificateMap();

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

  void Reset();
  void AddCert(std::unique_ptr<DrmServiceCertificate> new_cert);
  void ClearDefaultDrmServiceCertificate();
  const DrmServiceCertificate* GetDefaultCert();
  const DrmServiceCertificate* GetCert(const std::string& serial_number);

  static DrmServiceCertificateMap* GetInstance();

 private:
  absl::Mutex mutex_;
  // Certificate serial number to certificate map.
  std::map<std::string, std::unique_ptr<DrmServiceCertificate>> map_
      GUARDED_BY(mutex_);
  DrmServiceCertificate* default_cert_ GUARDED_BY(mutex_);
};

DrmServiceCertificateMap::DrmServiceCertificateMap() : default_cert_(nullptr) {}
DrmServiceCertificateMap::~DrmServiceCertificateMap() { Reset(); }

void DrmServiceCertificateMap::Reset() {
  absl::WriterMutexLock lock(&mutex_);
  map_.clear();
  default_cert_ = nullptr;
}

void DrmServiceCertificateMap::AddCert(
    std::unique_ptr<DrmServiceCertificate> new_cert) {
  absl::WriterMutexLock lock(&mutex_);

  std::unique_ptr<DrmServiceCertificate>* previous_cert =
      gtl::FindOrNull(map_, new_cert->serial_number());
  if (previous_cert != nullptr) {
    if (default_cert_ == previous_cert->get()) {
      default_cert_ = nullptr;
    }
  }

  if (default_cert_ == nullptr) {
    default_cert_ = new_cert.get();
  }
  const std::string& serial_number = new_cert->serial_number();
  map_[serial_number] = std::move(new_cert);
}

void DrmServiceCertificateMap::ClearDefaultDrmServiceCertificate() {
  absl::WriterMutexLock lock(&mutex_);
  default_cert_ = nullptr;
}

const DrmServiceCertificate* DrmServiceCertificateMap::GetDefaultCert() {
  absl::ReaderMutexLock lock(&mutex_);
  return default_cert_;
}

const DrmServiceCertificate* DrmServiceCertificateMap::GetCert(
    const std::string& serial_number) {
  absl::ReaderMutexLock lock(&mutex_);
  return map_[serial_number].get();
}

DrmServiceCertificateMap* DrmServiceCertificateMap::GetInstance() {
  static auto* const kInstance = new DrmServiceCertificateMap();
  return kInstance;
}

}  // namespace

util::Status DrmServiceCertificate::AddDrmServiceCertificate(
    const std::string& root_public_key, const std::string& service_certificate,
    const std::string& service_private_key,
    const std::string& service_private_key_passphrase) {
  std::unique_ptr<RsaPublicKey> root_key(RsaPublicKey::Create(root_public_key));
  if (root_key == nullptr) {
    return util::Status(error_space, INVALID_DRM_CERTIFICATE,
                        "root-certificate-rsa-public-key-failed");
  }
  SignedDrmCertificate signed_cert;
  if (!signed_cert.ParseFromString(service_certificate)) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "signed-certificate-parse-failed");
  }
  if (!root_key->VerifySignature(signed_cert.drm_certificate(),
                                 signed_cert.signature())) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "certificate-signature-verification-failed");
  }
  DrmCertificate drm_cert;
  if (!drm_cert.ParseFromString(signed_cert.drm_certificate())) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "certificate-parse-failed");
  }
  if (drm_cert.type() != DrmCertificate::SERVICE) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "not-service-certificate");
  }
  if (drm_cert.serial_number().empty()) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "missing-certificate-serial-number");
  }
  if (drm_cert.provider_id().empty()) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "missing-certificate-service-id");
  }
  if (!drm_cert.has_creation_time_seconds()) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "missing-certificate-creation-time");
  }
  if (drm_cert.public_key().empty()) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "missing-certificate-public-key");
  }
  std::unique_ptr<RsaPublicKey> public_key(
      RsaPublicKey::Create(drm_cert.public_key()));
  if (!public_key) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "invalid-certificate-public-key");
  }
  std::string pkcs1_key;
  if (!rsa_util::EncryptedPrivateKeyInfoToRsaPrivateKey(
          service_private_key, service_private_key_passphrase, &pkcs1_key)) {
    return util::Status(error_space, INVALID_SERVICE_PRIVATE_KEY,
                        "key-decryption-failed");
  }
  std::unique_ptr<RsaPrivateKey> private_key(RsaPrivateKey::Create(pkcs1_key));
  if (private_key == nullptr) {
    return util::Status(error_space, INVALID_SERVICE_PRIVATE_KEY,
                        "invalid-private-key");
  }

  std::unique_ptr<DrmServiceCertificate> new_cert(new DrmServiceCertificate(
      service_certificate, drm_cert.provider_id(), drm_cert.serial_number(),
      drm_cert.creation_time_seconds(), std::move(public_key),
      std::move(private_key)));
  DrmServiceCertificateMap::GetInstance()->AddCert(std::move(new_cert));

  return util::OkStatus();
}

util::Status DrmServiceCertificate::AddDrmServiceCertificate(
    CertificateType root_cert_type, const std::string& service_certificate,
    const std::string& service_private_key,
    const std::string& service_private_key_passphrase) {
  std::unique_ptr<DrmRootCertificate> root_cert;
  util::Status status =
      DrmRootCertificate::CreateByType(root_cert_type, &root_cert);
  if (!status.ok()) {
    return status;
  }
  return AddDrmServiceCertificate(root_cert->public_key(), service_certificate,
                                  service_private_key,
                                  service_private_key_passphrase);
}

const DrmServiceCertificate*
DrmServiceCertificate::GetDefaultDrmServiceCertificate() {
  return DrmServiceCertificateMap::GetInstance()->GetDefaultCert();
}

const DrmServiceCertificate*
DrmServiceCertificate::GetDefaultDrmServiceCertificateOrDie() {
  const DrmServiceCertificate* default_cert =
      DrmServiceCertificateMap::GetInstance()->GetDefaultCert();
  CHECK(default_cert) << "Service Certificate not set!";
  return default_cert;
}

const DrmServiceCertificate* DrmServiceCertificate::GetDrmServiceCertificate(
    const std::string& serial_number) {
  return DrmServiceCertificateMap::GetInstance()->GetCert(serial_number);
}

util::Status DrmServiceCertificate::SetDefaultDrmServiceCertificate(
    const std::string& root_public_key, const std::string& service_certificate,
    const std::string& service_private_key,
    const std::string& service_private_key_passphrase) {
  DrmServiceCertificateMap::GetInstance()->ClearDefaultDrmServiceCertificate();
  return AddDrmServiceCertificate(root_public_key, service_certificate,
                                  service_private_key,
                                  service_private_key_passphrase);
}

util::Status DrmServiceCertificate::SetDefaultDrmServiceCertificate(
    CertificateType root_cert_type, const std::string& service_certificate,
    const std::string& service_private_key,
    const std::string& service_private_key_passphrase) {
  std::unique_ptr<DrmRootCertificate> root_cert;
  util::Status status =
      DrmRootCertificate::CreateByType(root_cert_type, &root_cert);
  if (!status.ok()) {
    return status;
  }
  return SetDefaultDrmServiceCertificate(
      root_cert->public_key(), service_certificate, service_private_key,
      service_private_key_passphrase);
}

util::Status DrmServiceCertificate::DecryptClientIdentification(
    const EncryptedClientIdentification& encrypted_client_id,
    ClientIdentification* client_id) {
  DCHECK(client_id);
  if (encrypted_client_id.service_certificate_serial_number().empty()) {
    return util::Status(error_space, INVALID_ENCRYPTED_CLIENT_IDENTIFICATION,
                        "missing-service-certificate-serial-number");
  }
  if (encrypted_client_id.provider_id().empty()) {
    return util::Status(error_space, INVALID_ENCRYPTED_CLIENT_IDENTIFICATION,
                        "missing-service-id");
  }
  if (encrypted_client_id.encrypted_client_id().empty()) {
    return util::Status(error_space, INVALID_ENCRYPTED_CLIENT_IDENTIFICATION,
                        "missing-encrypted-client-id");
  }
  if (encrypted_client_id.encrypted_client_id_iv().empty()) {
    return util::Status(error_space, INVALID_ENCRYPTED_CLIENT_IDENTIFICATION,
                        "missing-encrypted-client-id-iv");
  }
  if (encrypted_client_id.encrypted_privacy_key().empty()) {
    return util::Status(error_space, INVALID_ENCRYPTED_CLIENT_IDENTIFICATION,
                        "missing-encrypted-privacy-key");
  }
  std::string privacy_key;
  std::string provider_id;
  const DrmServiceCertificate* cert = GetDrmServiceCertificate(
      encrypted_client_id.service_certificate_serial_number());
  if (!cert) {
    return util::Status(
        error_space, SERVICE_CERTIFICATE_NOT_FOUND,
        "service-certificate-not-found (SN " +
            absl::BytesToHexString(
                encrypted_client_id.service_certificate_serial_number()) +
            ")");
  }
  if (!cert->private_key()->Decrypt(encrypted_client_id.encrypted_privacy_key(),
                                    &privacy_key)) {
    return util::Status(error_space, INVALID_ENCRYPTED_CLIENT_IDENTIFICATION,
                        "privacy-key-decryption-failed");
  }
  if (cert->provider_id() != encrypted_client_id.provider_id()) {
    return util::Status(error_space, SERVICE_CERTIFICATE_NOT_FOUND,
                        std::string("provider-id-mismatch (") + cert->provider_id() +
                            " / " + encrypted_client_id.provider_id() + ")");
  }
  std::string serialized_client_id(crypto_util::DecryptAesCbc(
      privacy_key, encrypted_client_id.encrypted_client_id_iv(),
      encrypted_client_id.encrypted_client_id()));
  if (serialized_client_id.empty()) {
    return util::Status(error_space, INVALID_ENCRYPTED_CLIENT_IDENTIFICATION,
                        "client-id-decryption-failed");
  }
  if (!client_id->ParseFromString(serialized_client_id)) {
    return util::Status(error_space, INVALID_ENCRYPTED_CLIENT_IDENTIFICATION,
                        "client-id-parse-failed");
  }
  return util::OkStatus();
}

void DrmServiceCertificate::ResetServiceCertificates() {
  DrmServiceCertificateMap::GetInstance()->Reset();
}

util::Status DrmServiceCertificate::ValidateDrmServiceCertificate() {
  const DrmServiceCertificate* service_certificate =
      GetDefaultDrmServiceCertificate();
  if (!service_certificate) {
    return util::Status(error_space, SERVICE_CERTIFICATE_NOT_FOUND,
                        "drm service certificate is not found.");
  }
  SignedDrmCertificate signed_cert;
  if (!signed_cert.ParseFromString(service_certificate->certificate())) {
    return util::Status(error_space, INVALID_DRM_CERTIFICATE,
                        "signed drm service certificate is failed to parse.");
  }
  DrmCertificate drm_cert;
  if (!drm_cert.ParseFromString(signed_cert.drm_certificate())) {
    return util::Status(error_space, INVALID_DRM_CERTIFICATE,
                        "Drm service certificate is failed to parse.");
  }
  if (!drm_cert.has_creation_time_seconds()) {
    return util::Status(error_space, INVALID_SERVICE_CERTIFICATE,
                        "missing certificate creation time");
  }
  // TODO(user): Check creation_time_seconds field in DrmCertificate and also
  // export the absl/time dependency through moe.
  return util::OkStatus();
}

DrmServiceCertificate::DrmServiceCertificate(
    const std::string& service_certificate, const std::string& provider_id,
    const std::string& serial_number, const uint32_t creation_time_seconds,
    std::unique_ptr<RsaPublicKey> public_key,
    std::unique_ptr<RsaPrivateKey> private_key)
    : certificate_(service_certificate),
      provider_id_(provider_id),
      serial_number_(serial_number),
      creation_time_seconds_(creation_time_seconds),
      public_key_(std::move(public_key)),
      private_key_(std::move(private_key)) {}

}  // namespace widevine