////////////////////////////////////////////////////////////////////////////////
// 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.
////////////////////////////////////////////////////////////////////////////////
//
// Description:
//   Unit tests for drm_root_certificate.cc

#include "common/drm_root_certificate.h"

#include <memory>

#include "glog/logging.h"
#include "google/protobuf/util/message_differencer.h"
#include "testing/gmock.h"
#include "testing/gunit.h"
#include "absl/memory/memory.h"
#include "common/ec_key.h"
#include "common/ec_test_keys.h"
#include "common/error_space.h"
#include "common/rsa_key.h"
#include "common/rsa_test_keys.h"
#include "common/test_drm_certificates.h"
#include "protos/public/drm_certificate.pb.h"
#include "protos/public/errors.pb.h"
#include "protos/public/signed_drm_certificate.pb.h"

using google::protobuf::util::MessageDifferencer;
using ::testing::Values;

namespace widevine {

TEST(DrmRootCertificateCreateTest, TestCertificate) {
  const std::string kTestCertificateHash(
      "49f917b1bdfed78002a58e799a58e940"
      "1fffaaed9d8d80752782b066757e2c8c");
  std::unique_ptr<DrmRootCertificate> root_cert;
  ASSERT_EQ(OkStatus(), DrmRootCertificate::CreateByType(
                            kCertificateTypeTesting, &root_cert));
  ASSERT_TRUE(root_cert != nullptr);
  EXPECT_EQ(kTestCertificateHash, root_cert->GetDigest());
}

TEST(DrmRootCertificateCreateTest, DevCertificate) {
  const std::string kDevelopmentCertificateHash(
      "0e25ee95476a770f30b98ac5ef778b3f"
      "137b66c29385b84f547a361b4724b17d");
  std::unique_ptr<DrmRootCertificate> root_cert;
  ASSERT_EQ(OkStatus(), DrmRootCertificate::CreateByType(
                            kCertificateTypeDevelopment, &root_cert));
  ASSERT_TRUE(root_cert != nullptr);
  EXPECT_EQ(kDevelopmentCertificateHash, root_cert->GetDigest());
}

TEST(DrmRootCertificateCreateTest, ProdCertificate) {
  const std::string kProductionCertificateHash(
      "d62fdabc9286648a81f7d3bedaf2f5a5"
      "27bbad39bc38da034ba98a21569adb9b");
  std::unique_ptr<DrmRootCertificate> root_cert;
  ASSERT_EQ(OkStatus(), DrmRootCertificate::CreateByType(
                            kCertificateTypeProduction, &root_cert));
  ASSERT_TRUE(root_cert != nullptr);
  EXPECT_EQ(kProductionCertificateHash, root_cert->GetDigest());
}

TEST(DrmRootCertificateTestCertificatesTest, Success) {
  TestDrmCertificates test_certs;
  std::unique_ptr<DrmRootCertificate> root_cert;
  ASSERT_TRUE(
      DrmRootCertificate::CreateByType(kCertificateTypeTesting, &root_cert)
          .ok());
  EXPECT_TRUE(root_cert
                  ->VerifyCertificate(test_certs.test_root_certificate(),
                                      nullptr, nullptr)
                  .ok());
  EXPECT_TRUE(
      root_cert
          ->VerifyCertificate(test_certs.test_intermediate_certificate(),
                              nullptr, nullptr)
          .ok());
  EXPECT_TRUE(root_cert
                  ->VerifyCertificate(test_certs.test_user_device_certificate(),
                                      nullptr, nullptr)
                  .ok());
  EXPECT_TRUE(
      root_cert
          ->VerifyCertificate(test_certs.test_service_certificate_no_type(),
                              nullptr, nullptr)
          .ok());
}

TEST(DrmRootCertificateTestCertificatesTest, NonWidevineRootSigner) {
  // TODO(b/138929855): Add test to verify certificate chain is signed by
  // Widevine.
}

class SignerPrivateKey {
 public:
  virtual ~SignerPrivateKey() {}
  virtual bool GenerateSignature(const std::string& message,
                                 std::string* signature) const = 0;
  virtual DrmCertificate::Algorithm algorithm() const = 0;
  static std::unique_ptr<SignerPrivateKey> Create(
      const std::string& private_key,
      widevine::DrmCertificate::Algorithm algorithm);

 protected:
  SignerPrivateKey() {}
};

template <typename T>
class SignerPrivateKeyImpl : public SignerPrivateKey {
 public:
  SignerPrivateKeyImpl(std::unique_ptr<T> private_key,
                       DrmCertificate::Algorithm algorithm)
      : private_key_(std::move(private_key)), algorithm_(algorithm) {}
  ~SignerPrivateKeyImpl() override {}
  bool GenerateSignature(const std::string& message,
                         std::string* signature) const override {
    return private_key_->GenerateSignature(message, signature);
  }
  DrmCertificate::Algorithm algorithm() const override { return algorithm_; }

 private:
  std::unique_ptr<T> private_key_;
  DrmCertificate::Algorithm algorithm_;
};

std::unique_ptr<SignerPrivateKey> SignerPrivateKey::Create(
    const std::string& private_key,
    widevine::DrmCertificate::Algorithm algorithm) {
  DCHECK(algorithm != DrmCertificate::UNKNOWN_ALGORITHM);
  switch (algorithm) {
    case DrmCertificate::RSA: {
      auto rsa_key =
          std::unique_ptr<RsaPrivateKey>(RsaPrivateKey::Create(private_key));
      CHECK(rsa_key);
      std::unique_ptr<SignerPrivateKey> new_rsa_signer_private_key =
          absl::make_unique<SignerPrivateKeyImpl<RsaPrivateKey>>(
              std::move(rsa_key), algorithm);
      CHECK(new_rsa_signer_private_key);
      return new_rsa_signer_private_key;
    }
    case DrmCertificate::ECC_SECP256R1:
    case DrmCertificate::ECC_SECP384R1:
    case DrmCertificate::ECC_SECP521R1: {
      auto ec_key = ECPrivateKey::Create(private_key);
      CHECK(ec_key);
      std::unique_ptr<SignerPrivateKey> new_ec_signer_private_key =
          absl::make_unique<SignerPrivateKeyImpl<ECPrivateKey>>(
              std::move(ec_key), algorithm);
      CHECK(new_ec_signer_private_key);
      return new_ec_signer_private_key;
    }
    default:
      break;
  }
  return nullptr;
}

static const int kDrmRootKey = 0;
static const int kInterMediateKey = 1;
static const int kClientKey = 2;

class DrmRootCertificateTest : public testing::TestWithParam<const char*> {
 protected:
  DrmRootCertificateTest() {}

  void SetUp() override {
    bool algorithm_status = false;
    std::string algorithm(GetParam());
    if (algorithm == "RSA") {
      RsaTestSetup();
      algorithm_status = true;
    }
    if (algorithm == "ECC") {
      EcTestSetup();
      algorithm_status = true;
    }

    CHECK(algorithm_status);

    ASSERT_EQ(OkStatus(), DrmRootCertificate::CreateByType(
                              kCertificateTypeTesting, &root_cert_));
  }

  void RsaTestSetup() {
    private_keys_.resize(3);
    private_keys_[kDrmRootKey] =
        SignerPrivateKey::Create(rsa_test_keys_.private_test_key_1_3072_bits(),
                                 widevine::DrmCertificate::RSA);
    drm_certificates_[kDrmRootKey].set_serial_number("level 0");
    drm_certificates_[kDrmRootKey].set_creation_time_seconds(0);
    drm_certificates_[kDrmRootKey].set_public_key(
        rsa_test_keys_.public_test_key_1_3072_bits());

    private_keys_[kInterMediateKey] =
        SignerPrivateKey::Create(rsa_test_keys_.private_test_key_2_2048_bits(),
                                 widevine::DrmCertificate::RSA);
    drm_certificates_[kInterMediateKey].set_serial_number("level 1");
    drm_certificates_[kInterMediateKey].set_creation_time_seconds(1);
    drm_certificates_[kInterMediateKey].set_public_key(
        rsa_test_keys_.public_test_key_2_2048_bits());

    private_keys_[kClientKey] =
        SignerPrivateKey::Create(rsa_test_keys_.private_test_key_1_3072_bits(),
                                 widevine::DrmCertificate::RSA);
    drm_certificates_[kClientKey].set_serial_number("level 2");
    drm_certificates_[kClientKey].set_creation_time_seconds(2);
    drm_certificates_[kClientKey].set_public_key(
        rsa_test_keys_.public_test_key_3_2048_bits());
  }

  void EcTestSetup() {
    private_keys_.resize(3);
    private_keys_[kDrmRootKey] =
        SignerPrivateKey::Create(rsa_test_keys_.private_test_key_1_3072_bits(),
                                 widevine::DrmCertificate::RSA);
    drm_certificates_[kDrmRootKey].set_serial_number("level 0");
    drm_certificates_[kDrmRootKey].set_creation_time_seconds(0);
    drm_certificates_[kDrmRootKey].set_public_key(
        rsa_test_keys_.public_test_key_1_3072_bits());

    private_keys_[kInterMediateKey] =
        SignerPrivateKey::Create(ec_test_keys_.private_test_key_1_secp521r1(),
                                 DrmCertificate::ECC_SECP521R1);
    drm_certificates_[kInterMediateKey].set_serial_number("level 1");
    drm_certificates_[kInterMediateKey].set_creation_time_seconds(1);
    drm_certificates_[kInterMediateKey].set_public_key(
        ec_test_keys_.public_test_key_1_secp521r1());
    drm_certificates_[kInterMediateKey].set_algorithm(
        DrmCertificate::ECC_SECP521R1);

    private_keys_[kClientKey] =
        SignerPrivateKey::Create(ec_test_keys_.private_test_key_1_secp256r1(),
                                 DrmCertificate::ECC_SECP256R1);
    drm_certificates_[kClientKey].set_serial_number("level 2");
    drm_certificates_[kClientKey].set_creation_time_seconds(2);
    drm_certificates_[kClientKey].set_public_key(
        ec_test_keys_.public_test_key_1_secp256r1());
    drm_certificates_[kClientKey].set_algorithm(DrmCertificate::ECC_SECP256R1);

    // Client certificate.

    // Intermediate certificate.
  }

  void GenerateSignedDrmCertificate() {
    SignedDrmCertificate* current_sc(&signed_drm_certificate_);
    drm_certificates_[kClientKey].set_algorithm(
        private_keys_[kClientKey]->algorithm());
    ASSERT_TRUE(drm_certificates_[kClientKey].SerializeToString(
        current_sc->mutable_drm_certificate()));
    ASSERT_TRUE(private_keys_[kInterMediateKey]->GenerateSignature(
        current_sc->drm_certificate(), current_sc->mutable_signature()));

    current_sc = current_sc->mutable_signer();
    drm_certificates_[kInterMediateKey].set_algorithm(
        private_keys_[kInterMediateKey]->algorithm());
    ASSERT_TRUE(drm_certificates_[kInterMediateKey].SerializeToString(
        current_sc->mutable_drm_certificate()));
    ASSERT_TRUE(private_keys_[kDrmRootKey]->GenerateSignature(
        current_sc->drm_certificate(), current_sc->mutable_signature()));

    current_sc = current_sc->mutable_signer();
    drm_certificates_[kDrmRootKey].set_algorithm(
        private_keys_[kDrmRootKey]->algorithm());
    ASSERT_TRUE(drm_certificates_[kDrmRootKey].SerializeToString(
        current_sc->mutable_drm_certificate()));
    ASSERT_TRUE(private_keys_[kDrmRootKey]->GenerateSignature(
        current_sc->drm_certificate(), current_sc->mutable_signature()));
  }

  RsaTestKeys rsa_test_keys_;
  ECTestKeys ec_test_keys_;
  std::vector<std::unique_ptr<SignerPrivateKey>> private_keys_;
  SignedDrmCertificate signed_drm_certificate_;
  DrmCertificate drm_certificates_[3];
  std::unique_ptr<DrmRootCertificate> root_cert_;
};

INSTANTIATE_TEST_SUITE_P(SuccessNoOutput, DrmRootCertificateTest,
                         Values("RSA", "ECC"));

TEST_P(DrmRootCertificateTest, SuccessNoOutput) {
  GenerateSignedDrmCertificate();
  ASSERT_EQ(OkStatus(),
            root_cert_->VerifyCertificate(
                signed_drm_certificate_.SerializeAsString(), nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, SuccessWithOutput) {
  GenerateSignedDrmCertificate();
  SignedDrmCertificate out_signed_cert;
  DrmCertificate out_cert;
  ASSERT_EQ(OkStatus(), root_cert_->VerifyCertificate(
                            signed_drm_certificate_.SerializeAsString(),
                            &out_signed_cert, &out_cert));
  EXPECT_TRUE(
      MessageDifferencer::Equals(out_signed_cert, signed_drm_certificate_));
  EXPECT_TRUE(MessageDifferencer::Equals(out_cert, drm_certificates_[2]));
}

TEST_P(DrmRootCertificateTest, InvalidSignedDrmCertificate) {
  EXPECT_EQ(Status(error_space, INVALID_DRM_CERTIFICATE,
                   "invalid-signed-drm-certificate"),
            root_cert_->VerifyCertificate("pure garbage", nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, InvalidSignerCertificate) {
  GenerateSignedDrmCertificate();
  signed_drm_certificate_.mutable_signer()->set_drm_certificate("more garbage");
  EXPECT_EQ(Status(error_space, INVALID_DRM_CERTIFICATE,
                   "invalid-signer-certificate"),
            root_cert_->VerifyCertificate(
                signed_drm_certificate_.SerializeAsString(), nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, MissingDrmCertificate) {
  GenerateSignedDrmCertificate();
  signed_drm_certificate_.clear_drm_certificate();
  EXPECT_EQ(
      Status(error_space, INVALID_DRM_CERTIFICATE, "invalid-drm-certificate"),
      root_cert_->VerifyCertificate(signed_drm_certificate_.SerializeAsString(),
                                    nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, InvalidDrmCertificate) {
  GenerateSignedDrmCertificate();
  signed_drm_certificate_.set_drm_certificate("junk");
  EXPECT_EQ(
      Status(error_space, INVALID_DRM_CERTIFICATE, "invalid-drm-certificate"),
      root_cert_->VerifyCertificate(signed_drm_certificate_.SerializeAsString(),
                                    nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, InvalidPublicKey) {
  drm_certificates_[0].set_public_key("rubbish");
  GenerateSignedDrmCertificate();
  EXPECT_EQ(
      Status(error_space, INVALID_DRM_CERTIFICATE, "invalid-signer-public-key"),
      root_cert_->VerifyCertificate(signed_drm_certificate_.SerializeAsString(),
                                    nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, MissingPublicKey) {
  drm_certificates_[2].clear_public_key();
  GenerateSignedDrmCertificate();
  EXPECT_EQ(Status(error_space, INVALID_DRM_CERTIFICATE, "missing-public-key"),
            root_cert_->VerifyCertificate(
                signed_drm_certificate_.SerializeAsString(), nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, MissingCreationTime) {
  drm_certificates_[2].clear_creation_time_seconds();
  GenerateSignedDrmCertificate();
  EXPECT_EQ(
      Status(error_space, INVALID_DRM_CERTIFICATE, "missing-creation-time"),
      root_cert_->VerifyCertificate(signed_drm_certificate_.SerializeAsString(),
                                    nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, MissingSerialNumber) {
  drm_certificates_[2].set_serial_number("");
  GenerateSignedDrmCertificate();
  EXPECT_EQ(
      Status(error_space, INVALID_DRM_CERTIFICATE, "missing-serial-number"),
      root_cert_->VerifyCertificate(signed_drm_certificate_.SerializeAsString(),
                                    nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, InvalidSignatureWithNoCache) {
  GenerateSignedDrmCertificate();
  signed_drm_certificate_.mutable_signer()->set_signature(
      "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx");
  EXPECT_EQ(
      Status(error_space, INVALID_SIGNATURE, "cache-miss-invalid-signature"),
      root_cert_->VerifyCertificate(signed_drm_certificate_.SerializeAsString(),
                                    nullptr, nullptr));
}

TEST_P(DrmRootCertificateTest, InvalidSignatureWithCache) {
  GenerateSignedDrmCertificate();
  // Verify and cache.
  ASSERT_EQ(OkStatus(),
            root_cert_->VerifyCertificate(
                signed_drm_certificate_.SerializeAsString(), nullptr, nullptr));

  // Verify success using cache.
  ASSERT_EQ(OkStatus(),
            root_cert_->VerifyCertificate(
                signed_drm_certificate_.SerializeAsString(), nullptr, nullptr));

  // Verify failure using cache.
  signed_drm_certificate_.mutable_signer()->set_signature(
      "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx");
  EXPECT_EQ(Status(error_space, INVALID_SIGNATURE, "cached-signature-mismatch"),
            root_cert_->VerifyCertificate(
                signed_drm_certificate_.SerializeAsString(), nullptr, nullptr));
}

}  // namespace widevine