////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 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 test for rsa_key RSA encryption and signing.

#include <memory>

#include "testing/gunit.h"
#include "common/rsa_key.h"
#include "common/rsa_test_keys.h"
#include "common/rsa_util.h"

namespace widevine {

static const char kTestMessage[] =
    "A fool thinks himself to be wise, but a wise man knows himself to be a "
    "fool.";

class RsaKeyTest : public ::testing::Test {
 protected:
  void TestEncryption(std::unique_ptr<RsaPrivateKey> private_key,
                      std::unique_ptr<RsaPublicKey> public_key);

  void TestSigning(std::unique_ptr<RsaPrivateKey> private_key,
                   std::unique_ptr<RsaPublicKey> public_key);

  void TestSigningSha256Pkcs7(std::unique_ptr<RsaPrivateKey> private_key,
                              std::unique_ptr<RsaPublicKey> public_key);

  RsaTestKeys test_keys_;
  RsaKeyFactory factory_;
};

TEST_F(RsaKeyTest, CopyConstructor) {
  std::unique_ptr<RsaPrivateKey> private_key(
      RsaPrivateKey::Create(test_keys_.private_test_key_2_2048_bits()));
  std::unique_ptr<RsaPublicKey> public_key(
      RsaPublicKey::Create(test_keys_.public_test_key_2_2048_bits()));

  std::unique_ptr<RsaPrivateKey> private_key_copy(
      new RsaPrivateKey(*private_key));
  std::unique_ptr<RsaPublicKey> public_key_copy(
      new RsaPublicKey(*public_key));

  EXPECT_TRUE(public_key_copy->MatchesPublicKey(*public_key));
  EXPECT_TRUE(public_key_copy->MatchesPrivateKey(*private_key));

  EXPECT_TRUE(private_key_copy->MatchesPublicKey(*public_key));
  EXPECT_TRUE(private_key_copy->MatchesPrivateKey(*private_key));
}

void RsaKeyTest::TestEncryption(std::unique_ptr<RsaPrivateKey> private_key,
                                std::unique_ptr<RsaPublicKey> public_key) {
  ASSERT_TRUE(private_key);
  ASSERT_TRUE(public_key);
  std::string encrypted_message;
  EXPECT_TRUE(public_key->Encrypt(kTestMessage, &encrypted_message));
  std::string decrypted_message;
  EXPECT_TRUE(private_key->Decrypt(encrypted_message, &decrypted_message));
  EXPECT_EQ(kTestMessage, decrypted_message);
  // Add a byte to the encrypted message.
  std::string bad_enc_message(encrypted_message);
  bad_enc_message += '\0';
  EXPECT_FALSE(private_key->Decrypt(bad_enc_message, &decrypted_message));
  // Remove a byte from the encrypted message.
  bad_enc_message = encrypted_message.substr(0, encrypted_message.size() - 1);
  EXPECT_FALSE(private_key->Decrypt(bad_enc_message, &decrypted_message));
  // Change a byte in the encrypted message.
  bad_enc_message = encrypted_message;
  bad_enc_message[128] ^= 0x55;
  EXPECT_FALSE(private_key->Decrypt(bad_enc_message, &decrypted_message));
}

void RsaKeyTest::TestSigning(std::unique_ptr<RsaPrivateKey> private_key,
                             std::unique_ptr<RsaPublicKey> public_key) {
  ASSERT_TRUE(private_key);
  ASSERT_TRUE(public_key);
  std::string signature;
  EXPECT_TRUE(private_key->GenerateSignature(kTestMessage, &signature));
  EXPECT_TRUE(public_key->VerifySignature(kTestMessage, signature));
  // Add a byte to the signature.
  std::string bad_signature(signature);
  bad_signature += '\0';
  EXPECT_FALSE(public_key->VerifySignature(kTestMessage, bad_signature));
  // Remove a byte from the signature.
  bad_signature = signature.substr(0, signature.size() - 1);
  EXPECT_FALSE(public_key->VerifySignature(kTestMessage, bad_signature));
  // Change a byte in the signature.
  bad_signature = signature;
  bad_signature[32] ^= 0x55;
  EXPECT_FALSE(public_key->VerifySignature(kTestMessage, bad_signature));
}

void RsaKeyTest::TestSigningSha256Pkcs7(
    std::unique_ptr<RsaPrivateKey> private_key,
    std::unique_ptr<RsaPublicKey> public_key) {
  ASSERT_TRUE(private_key);
  ASSERT_TRUE(public_key);
  std::string signature;
  EXPECT_TRUE(
      private_key->GenerateSignatureSha256Pkcs7(kTestMessage, &signature));
  EXPECT_TRUE(public_key->VerifySignatureSha256Pkcs7(kTestMessage, signature));
  // Add a byte to the signature.
  std::string bad_signature(signature);
  bad_signature += '\0';
  EXPECT_FALSE(
      public_key->VerifySignatureSha256Pkcs7(kTestMessage, bad_signature));
  // Remove a byte from the signature.
  bad_signature = signature.substr(0, signature.size() - 1);
  EXPECT_FALSE(
      public_key->VerifySignatureSha256Pkcs7(kTestMessage, bad_signature));
  // Change a byte in the signature.
  bad_signature = signature;
  bad_signature[32] ^= 0x55;
  EXPECT_FALSE(
      public_key->VerifySignatureSha256Pkcs7(kTestMessage, bad_signature));
}

TEST_F(RsaKeyTest, BadKey) {
  std::unique_ptr<RsaPrivateKey> private_key(
      RsaPrivateKey::Create("bad_private_key"));
  EXPECT_TRUE(!private_key);
  std::unique_ptr<RsaPublicKey> public_key(
      RsaPublicKey::Create("bad_public_key"));
  EXPECT_TRUE(!public_key);
}

TEST_F(RsaKeyTest, EncryptAndDecrypt_3072) {
  const std::string& private_key = test_keys_.private_test_key_1_3072_bits();
  const std::string& public_key = test_keys_.public_test_key_1_3072_bits();
  TestEncryption(
      std::unique_ptr<RsaPrivateKey>(RsaPrivateKey::Create(private_key)),
      std::unique_ptr<RsaPublicKey>(RsaPublicKey::Create(public_key)));
  TestEncryption(factory_.CreateFromPkcs1PrivateKey(private_key),
                 factory_.CreateFromPkcs1PublicKey(public_key));
}

TEST_F(RsaKeyTest, EncryptAndDecrypt_2048) {
  const std::string& private_key = test_keys_.private_test_key_2_2048_bits();
  const std::string& public_key = test_keys_.public_test_key_2_2048_bits();
  TestEncryption(
      std::unique_ptr<RsaPrivateKey>(RsaPrivateKey::Create(private_key)),
      std::unique_ptr<RsaPublicKey>(RsaPublicKey::Create(public_key)));

  std::string pkcs8_key;
  std::string passphrase("passphrase");
  ASSERT_TRUE(rsa_util::RsaPrivateKeyToEncryptedPrivateKeyInfo(
      private_key, passphrase, &pkcs8_key));
  TestEncryption(factory_.CreateFromPkcs8PrivateKey(pkcs8_key, passphrase),
                 factory_.CreateFromPkcs1PublicKey(public_key));

  ASSERT_TRUE(rsa_util::RsaPrivateKeyToPrivateKeyInfo(private_key, &pkcs8_key));
  TestEncryption(factory_.CreateFromPkcs8PrivateKey(pkcs8_key, std::string()),
                 factory_.CreateFromPkcs1PublicKey(public_key));
}

TEST_F(RsaKeyTest, SignAndVerify_3072) {
  const std::string& private_key = test_keys_.private_test_key_1_3072_bits();
  const std::string& public_key = test_keys_.public_test_key_1_3072_bits();
  TestSigning(
      std::unique_ptr<RsaPrivateKey>(RsaPrivateKey::Create(private_key)),
      std::unique_ptr<RsaPublicKey>(RsaPublicKey::Create(public_key)));
  TestSigning(factory_.CreateFromPkcs1PrivateKey(private_key),
              factory_.CreateFromPkcs1PublicKey(public_key));
}

TEST_F(RsaKeyTest, SignAndVerify_2048) {
  const std::string& private_key = test_keys_.private_test_key_2_2048_bits();
  const std::string& public_key = test_keys_.public_test_key_2_2048_bits();
  TestSigning(
      std::unique_ptr<RsaPrivateKey>(RsaPrivateKey::Create(private_key)),
      std::unique_ptr<RsaPublicKey>(RsaPublicKey::Create(public_key)));

  std::string pkcs8_key;
  std::string passphrase("passphrase");
  ASSERT_TRUE(rsa_util::RsaPrivateKeyToEncryptedPrivateKeyInfo(
      private_key, passphrase, &pkcs8_key));
  TestSigning(factory_.CreateFromPkcs8PrivateKey(pkcs8_key, passphrase),
              factory_.CreateFromPkcs1PublicKey(public_key));

  ASSERT_TRUE(rsa_util::RsaPrivateKeyToPrivateKeyInfo(private_key, &pkcs8_key));
  TestSigning(factory_.CreateFromPkcs8PrivateKey(pkcs8_key, std::string()),
              factory_.CreateFromPkcs1PublicKey(public_key));
}

TEST_F(RsaKeyTest, SignAndVerifySha256Pkcs7_3072) {
  const std::string& private_key = test_keys_.private_test_key_1_3072_bits();
  const std::string& public_key = test_keys_.public_test_key_1_3072_bits();
  TestSigningSha256Pkcs7(
      std::unique_ptr<RsaPrivateKey>(RsaPrivateKey::Create(private_key)),
      std::unique_ptr<RsaPublicKey>(RsaPublicKey::Create(public_key)));
  TestSigningSha256Pkcs7(factory_.CreateFromPkcs1PrivateKey(private_key),
                         factory_.CreateFromPkcs1PublicKey(public_key));
}

TEST_F(RsaKeyTest, SignAndVerifySha256Pkcs7_2048) {
  const std::string& private_key = test_keys_.private_test_key_2_2048_bits();
  const std::string& public_key = test_keys_.public_test_key_2_2048_bits();
  TestSigningSha256Pkcs7(
      std::unique_ptr<RsaPrivateKey>(RsaPrivateKey::Create(private_key)),
      std::unique_ptr<RsaPublicKey>(RsaPublicKey::Create(public_key)));

  std::string pkcs8_key;
  std::string passphrase("passphrase");
  ASSERT_TRUE(rsa_util::RsaPrivateKeyToEncryptedPrivateKeyInfo(
      private_key, passphrase, &pkcs8_key));
  TestSigningSha256Pkcs7(
      factory_.CreateFromPkcs8PrivateKey(pkcs8_key, passphrase),
      factory_.CreateFromPkcs1PublicKey(public_key));

  ASSERT_TRUE(rsa_util::RsaPrivateKeyToPrivateKeyInfo(private_key, &pkcs8_key));
  TestSigningSha256Pkcs7(
      factory_.CreateFromPkcs8PrivateKey(pkcs8_key, std::string()),
      factory_.CreateFromPkcs1PublicKey(public_key));
}

TEST_F(RsaKeyTest, KeySize) {
  std::unique_ptr<RsaPrivateKey> private_key(
      RsaPrivateKey::Create(test_keys_.private_test_key_2_2048_bits()));
  std::unique_ptr<RsaPublicKey> public_key(
      RsaPublicKey::Create(test_keys_.public_test_key_2_2048_bits()));

  EXPECT_EQ(256, private_key->KeySize());
  EXPECT_EQ(256, public_key->KeySize());
}

TEST_F(RsaKeyTest, RsaKeyMatch) {
  std::unique_ptr<RsaPrivateKey> private_key2(
      RsaPrivateKey::Create(test_keys_.private_test_key_2_2048_bits()));
  std::unique_ptr<RsaPrivateKey> private_key3(
      RsaPrivateKey::Create(test_keys_.private_test_key_3_2048_bits()));
  std::unique_ptr<RsaPublicKey> public_key2(
      RsaPublicKey::Create(test_keys_.public_test_key_2_2048_bits()));
  std::unique_ptr<RsaPublicKey> public_key3(
      RsaPublicKey::Create(test_keys_.public_test_key_3_2048_bits()));

  EXPECT_TRUE(public_key2->MatchesPublicKey(*public_key2));
  EXPECT_FALSE(public_key2->MatchesPublicKey(*public_key3));
  EXPECT_TRUE(public_key2->MatchesPrivateKey(*private_key2));
  EXPECT_FALSE(public_key2->MatchesPrivateKey(*private_key3));

  EXPECT_TRUE(private_key2->MatchesPublicKey(*public_key2));
  EXPECT_FALSE(private_key2->MatchesPublicKey(*public_key3));
  EXPECT_TRUE(private_key2->MatchesPrivateKey(*private_key2));
  EXPECT_FALSE(private_key2->MatchesPrivateKey(*private_key3));
}

}  // namespace widevine