////////////////////////////////////////////////////////////////////////////////
// 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 "google/protobuf/util/message_differencer.h"
#include "testing/gmock.h"
#include "testing/gunit.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;

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());
}

class DrmRootCertificateTest : public testing::Test {
 protected:
  DrmRootCertificateTest() {
    private_keys_.emplace_back(
        RsaPrivateKey::Create(test_keys_.private_test_key_1_3072_bits()));
    private_keys_.emplace_back(
        RsaPrivateKey::Create(test_keys_.private_test_key_2_2048_bits()));
    private_keys_.emplace_back(
        RsaPrivateKey::Create(test_keys_.private_test_key_3_2048_bits()));
  }

  void SetUp() override {
    drm_certificates_[0].set_serial_number("level 0");
    drm_certificates_[0].set_creation_time_seconds(0);
    drm_certificates_[0].set_public_key(
        test_keys_.public_test_key_1_3072_bits());
    drm_certificates_[1].set_serial_number("level 1");
    drm_certificates_[1].set_creation_time_seconds(1);
    drm_certificates_[1].set_public_key(
        test_keys_.public_test_key_2_2048_bits());
    drm_certificates_[2].set_serial_number("level 2");
    drm_certificates_[2].set_creation_time_seconds(2);
    drm_certificates_[2].set_public_key(
        test_keys_.public_test_key_3_2048_bits());

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

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

    current_sc = current_sc->mutable_signer();
    ASSERT_TRUE(drm_certificates_[1].SerializeToString(
        current_sc->mutable_drm_certificate()));
    ASSERT_TRUE(private_keys_[0]->GenerateSignature(
        current_sc->drm_certificate(), current_sc->mutable_signature()));

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

  RsaTestKeys test_keys_;
  std::vector<std::unique_ptr<RsaPrivateKey>> private_keys_;
  SignedDrmCertificate signed_drm_certificate_;
  DrmCertificate drm_certificates_[3];
  std::unique_ptr<DrmRootCertificate> root_cert_;
};

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

TEST_F(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_F(DrmRootCertificateTest, InvalidSignedDrmCertificate) {
  EXPECT_EQ(Status(error_space, INVALID_DRM_CERTIFICATE,
                   "invalid-signed-drm-certificate"),
            root_cert_->VerifyCertificate("pure garbage", nullptr, nullptr));
}

TEST_F(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_F(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_F(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_F(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_F(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_F(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_F(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_F(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_F(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