////////////////////////////////////////////////////////////////////////////////
// Copyright 2019 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/certificate_client_cert.h"

#include "glog/logging.h"
#include "absl/memory/memory.h"
#include "common/crypto_util.h"
#include "common/ec_key.h"
#include "common/ec_util.h"
#include "common/error_space.h"
#include "common/openssl_util.h"
#include "common/random_util.h"
#include "common/rsa_key.h"
#include "common/sha_util.h"
#include "common/signing_key_util.h"
#include "protos/public/drm_certificate.pb.h"
#include "protos/public/errors.pb.h"
#include "protos/public/signed_drm_certificate.pb.h"

namespace widevine {

using EllipticCurve = ECPrivateKey::EllipticCurve;

ECPrivateKey::EllipticCurve CertificateAlgorithmToCurve(
    DrmCertificate::Algorithm algorithm) {
  switch (algorithm) {
    case DrmCertificate::ECC_SECP256R1:
      return ECPrivateKey::SECP256R1;
    case DrmCertificate::ECC_SECP384R1:
      return ECPrivateKey::SECP384R1;
    case DrmCertificate::ECC_SECP521R1:
      return ECPrivateKey::SECP521R1;
    default:
      return ECPrivateKey::UNDEFINED_CURVE;
  }
}

class ClientCertAlgorithmRSA : public ClientCertAlgorithm {
 public:
  ClientCertAlgorithmRSA() {}
  ~ClientCertAlgorithmRSA() override {}
  ClientCertAlgorithmRSA(const ClientCertAlgorithmRSA&) = delete;
  ClientCertAlgorithmRSA& operator=(const ClientCertAlgorithmRSA&) = delete;

  Status Initialize(const std::string& public_key,
                    DrmCertificate::Algorithm /*not_used*/) override {
    rsa_public_key_ =
        std::unique_ptr<RsaPublicKey>(RsaPublicKey::Create(public_key));
    if (!rsa_public_key_) {
      return Status(error_space, INVALID_DRM_CERTIFICATE,
                    "drm-certificate-public-key-failed");
    }
    session_key_ = Random16Bytes();
    if (!rsa_public_key_->Encrypt(session_key_, &wrapped_session_key_)) {
      return Status(error_space, ENCRYPT_ERROR,
                    "drm-certificate-failed-encrypt-session-key");
    }
    return OkStatus();
  }

  Status VerifySignature(const std::string& message,
                         const std::string& signature) const override {
    CHECK(rsa_public_key_);

    if (!rsa_public_key_->VerifySignature(message, signature)) {
      return Status(error_space, INVALID_SIGNATURE, "");
    }
    return OkStatus();
  }

  const std::string& session_key() const override { return session_key_; }

  const std::string& wrapped_session_key() const override {
    return wrapped_session_key_;
  }

 private:
  std::unique_ptr<RsaPublicKey> rsa_public_key_;
  std::string session_key_;
  std::string wrapped_session_key_;
};

// ClientCertAlgorithmECC implements the Widevine protocol using ECC. It
// verifies an ECC based request and generates keys for use in building a
// license. The curve type value is contained in |algorithm|.
class ClientCertAlgorithmECC : public ClientCertAlgorithm {
 public:
  ClientCertAlgorithmECC() = default;
  ~ClientCertAlgorithmECC() override = default;
  ClientCertAlgorithmECC(const ClientCertAlgorithmECC&) = delete;
  ClientCertAlgorithmECC& operator=(const ClientCertAlgorithmECC&) = delete;

  Status Initialize(const std::string& public_key,
                    DrmCertificate::Algorithm algorithm) override {
    ECPrivateKey::EllipticCurve curve_id =
        CertificateAlgorithmToCurve(algorithm);
    if (curve_id == ECPrivateKey::UNDEFINED_CURVE) {
      return Status(error_space, INVALID_DRM_CERTIFICATE,
                    "drm-certificate-unknown-curve");
    }
    // Parse the certifcate ECC public key.
    client_ecc_public_key_ = ECPublicKey::Create(public_key);
    if (client_ecc_public_key_ == nullptr) {
      return Status(error_space, INVALID_DRM_CERTIFICATE,
                    "drm-certificate-public-key-failed");
    }
    // Generate an ephemeral ecc key pair with the same curve as used by the
    // certificate public key.
    ScopedECKEY key = ec_util::GenerateKeyWithCurve(curve_id);
    auto new_private_key = absl::make_unique<ECPrivateKey>(std::move(key));
    if (new_private_key == nullptr) {
      return Status(error_space, DRM_DEVICE_CERTIFICATE_ECC_KEYGEN_FAILED,
                    "drm-certificate-ephemeral-private-key-failed");
    }

    // Serialize the ephemeral public key for inclusion in a license response.
    std::unique_ptr<ECPublicKey> new_public_key = new_private_key->PublicKey();
    if (new_public_key == nullptr ||
        !new_public_key->SerializedKey(&ephemeral_public_key_)) {
      return Status(error_space, DRM_DEVICE_CERTIFICATE_ECC_KEYGEN_FAILED,
                    "drm-certificate-ephemeral-public-key-failed");
    }

    // Generate the session key from the ephemeral private key and the
    // certificate public key.
    if (!new_private_key->DeriveSharedSessionKey(*client_ecc_public_key_,
                                                 &derived_session_key_)) {
      return Status(error_space, DRM_DEVICE_CERTIFICATE_ECC_KEYGEN_FAILED,
                    "drm-certificate-shared-key-gen-failed");
    }
    return OkStatus();
  }

  Status VerifySignature(const std::string& message,
                         const std::string& signature) const override {
    CHECK(client_ecc_public_key_);

    if (!client_ecc_public_key_->VerifySignature(message, signature)) {
      return Status(error_space, INVALID_SIGNATURE, "");
    }
    return OkStatus();
  }
  // Returns an aes key generated from the sha256 hash of the shared ecc secret.
  // This key is used for key derivation.
  const std::string& session_key() const override {
    return derived_session_key_;
  }

  // Returns an ephemeral serialized ecc public key. This key is added to a
  // license response in the SignedMessage::session_key field. The client will
  // use this key to generate the shared secret and derived session key.
  const std::string& wrapped_session_key() const override {
    return ephemeral_public_key_;
  }

 private:
  std::unique_ptr<ECPublicKey> client_ecc_public_key_;
  std::string ephemeral_public_key_;
  std::string derived_session_key_;
};

Status CertificateClientCert::Initialize(
    const DrmRootCertificate* root_certificate,
    const std::string& serialized_certificate) {
  CHECK(root_certificate);

  if (is_initialized_) {
    return Status(error_space, INVALID_PARAMETER,
                  "certificate-is-already-initialized");
  }

  SignedDrmCertificate signed_device_cert;
  Status status = root_certificate->VerifyCertificate(
      serialized_certificate, &signed_device_cert, &device_cert_);
  if (!status.ok()) {
    return status;
  }
  if (device_cert_.type() != DrmCertificate::DEVICE ||
      device_cert_.public_key().empty()) {
    return Status(error_space, INVALID_DRM_CERTIFICATE,
                  "expected-device-certificate-type");
  }

  const SignedDrmCertificate& device_cert_signer = signed_device_cert.signer();

  if (!model_certificate_.ParseFromString(
          device_cert_signer.drm_certificate())) {
    return Status(error_space, INVALID_DRM_CERTIFICATE,
                  "drm-certificate-invalid-signer");
  }
  if (model_certificate_.type() != DrmCertificate::DEVICE_MODEL) {
    return Status(error_space, INVALID_DRM_CERTIFICATE,
                  "expected-device-model-certificate-type");
  }
  if (!model_certificate_.has_serial_number()) {
    return Status(error_space, INVALID_DRM_CERTIFICATE,
                  "missing-signer-serial-number");
  }
  // Check to see if this model certificate is signed by a
  // provisioner (entity using Widevine Provisioning Server SDK).
  if (device_cert_signer.has_signer()) {
    if (!provisioner_certificate_.ParseFromString(
            device_cert_signer.signer().drm_certificate())) {
      return Status(error_space, INVALID_DRM_CERTIFICATE,
                    "model-certificate-invalid-signer");
    }
    if (provisioner_certificate_.type() != DrmCertificate::PROVISIONER) {
      return Status(error_space, INVALID_DRM_CERTIFICATE,
                    "expected-provisioning-provider-certificate-type");
    }
    if (!provisioner_certificate_.has_provider_id() ||
        provisioner_certificate_.provider_id().empty()) {
      return Status(error_space, INVALID_DRM_CERTIFICATE,
                    "missing-provisioning-service-id");
    }
    signed_by_provisioner_ = true;
  }
  if (!model_certificate_.has_system_id()) {
    return Status(error_space, INVALID_DRM_CERTIFICATE,
                  "model-certificate-missing-system-id");
  }

  switch (device_cert_.algorithm()) {
    case DrmCertificate::RSA:
      algorithm_ = absl::make_unique<ClientCertAlgorithmRSA>();
      break;
    case DrmCertificate::ECC_SECP256R1:
    case DrmCertificate::ECC_SECP384R1:
    case DrmCertificate::ECC_SECP521R1:
      algorithm_ = absl::make_unique<ClientCertAlgorithmECC>();
      break;
    default:
      return Status(error_space, INVALID_DRM_CERTIFICATE,
                    "unsupported-certificate-algorithm");
  }

  status = algorithm_->Initialize(device_cert_.public_key(),
                                  device_cert_.algorithm());
  if (!status.ok()) {
    return status;
  }
  is_initialized_ = true;
  return OkStatus();
}

Status CertificateClientCert::VerifySignature(
    const std::string& message, const std::string& signature,
    ProtocolVersion protocol_version) const {
  return algorithm_->VerifySignature(
      protocol_version < VERSION_2_2 ? message : Sha512_Hash(message),
      signature);
}

void CertificateClientCert::GenerateSigningKey(
    const std::string& message, ProtocolVersion protocol_version) {
  signing_key_ = crypto_util::DeriveKey(
      key(), crypto_util::kSigningKeyLabel,
      protocol_version < VERSION_2_2 ? message : Sha512_Hash(message),
      SigningKeyMaterialSizeBits(protocol_version));
}

}  // namespace widevine