////////////////////////////////////////////////////////////////////////////////
// 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/signer_public_key.h"

#include <memory>

#include "testing/gunit.h"
#include "common/ec_key.h"
#include "common/ec_test_keys.h"
#include "common/hash_algorithm.h"
#include "common/rsa_key.h"
#include "common/rsa_test_keys.h"
#include "protos/public/drm_certificate.pb.h"

namespace widevine {

static const char kMessage[] = "The rain in Spain falls mainly in the blank?";
const HashAlgorithm kHashAlgorithm = HashAlgorithm::kSha256;

class SignerPublicKeyTest : public ::testing::Test {
 public:
  RsaTestKeys rsa_test_keys_;
  ECTestKeys ec_test_keys_;
};

TEST_F(SignerPublicKeyTest, RSA) {
  std::unique_ptr<RsaPrivateKey> private_key(
      RsaPrivateKey::Create(rsa_test_keys_.private_test_key_1_3072_bits()));

  std::string signature;
  ASSERT_TRUE(
      private_key->GenerateSignature(kMessage, kHashAlgorithm, &signature));

  std::unique_ptr<SignerPublicKey> public_key = SignerPublicKey::Create(
      rsa_test_keys_.public_test_key_1_3072_bits(), DrmCertificate::RSA);
  ASSERT_NE(public_key, nullptr);
  EXPECT_TRUE(public_key->VerifySignature(kMessage, kHashAlgorithm, signature));
}

TEST_F(SignerPublicKeyTest, ECC) {
  std::unique_ptr<ECPrivateKey> private_key =
      ECPrivateKey::Create(ec_test_keys_.private_test_key_1_secp521r1());

  std::string signature;
  ASSERT_TRUE(
      private_key->GenerateSignature(kMessage, kHashAlgorithm, &signature));

  std::unique_ptr<SignerPublicKey> public_key =
      SignerPublicKey::Create(ec_test_keys_.public_test_key_1_secp521r1(),
                              DrmCertificate::ECC_SECP521R1);
  ASSERT_NE(public_key, nullptr);
  EXPECT_TRUE(public_key->VerifySignature(kMessage, kHashAlgorithm, signature));
}

TEST_F(SignerPublicKeyTest, IncorrectAlgorithm) {
  std::unique_ptr<SignerPublicKey> rsa_public_key =
      SignerPublicKey::Create(rsa_test_keys_.public_test_key_1_3072_bits(),
                              DrmCertificate::ECC_SECP521R1);
  ASSERT_EQ(rsa_public_key, nullptr);

  std::unique_ptr<SignerPublicKey> ec_public_key = SignerPublicKey::Create(
      ec_test_keys_.public_test_key_1_secp521r1(), DrmCertificate::RSA);
  ASSERT_EQ(ec_public_key, nullptr);
}

TEST_F(SignerPublicKeyTest, BadKey) {
  std::unique_ptr<SignerPublicKey> rsa_public_key =
      SignerPublicKey::Create("GobbletyGook", DrmCertificate::RSA);
  ASSERT_EQ(rsa_public_key, nullptr);

  std::unique_ptr<SignerPublicKey> ec_public_key = SignerPublicKey::Create(
      "MoreGobbletyGook", DrmCertificate::ECC_SECP521R1);
  ASSERT_EQ(ec_public_key, nullptr);
}

}  // namespace widevine