////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 Google Inc.
//
// 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_util RSA utilties library.

#include "gtest/gtest.h"
#include "testing/base/public/test_utils.h"
#include "common/rsa_test_keys.h"
#include "common/rsa_util.h"

namespace widevine {
namespace rsa_util {

class RsaUtilTest : public ::testing::Test {
 protected:
  RsaTestKeys test_keys_;
};

TEST_F(RsaUtilTest, SerializeDeserializePrivateKey) {
  RSA* private_key;
  std::string serialized_private_key;
  // Key 1
  EXPECT_TRUE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_1_3072_bits(), &private_key));
  ASSERT_TRUE(private_key != NULL);
  EXPECT_TRUE(SerializeRsaPrivateKey(private_key, &serialized_private_key));
  EXPECT_EQ(serialized_private_key, test_keys_.private_test_key_1_3072_bits());
  EXPECT_EQ(RSA_check_key(private_key), 1);
  RSA_free(private_key);
  // Key 2
  EXPECT_TRUE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_2_2048_bits(), &private_key));
  ASSERT_TRUE(private_key != NULL);
  EXPECT_TRUE(SerializeRsaPrivateKey(private_key, &serialized_private_key));
  EXPECT_EQ(serialized_private_key, test_keys_.private_test_key_2_2048_bits());
  EXPECT_EQ(RSA_check_key(private_key), 1);
  RSA_free(private_key);
  // Key 3
  EXPECT_TRUE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_3_2048_bits(), &private_key));
  ASSERT_TRUE(private_key != NULL);
  EXPECT_TRUE(SerializeRsaPrivateKey(private_key, &serialized_private_key));
  EXPECT_EQ(serialized_private_key, test_keys_.private_test_key_3_2048_bits());
  EXPECT_EQ(RSA_check_key(private_key), 1);
  RSA_free(private_key);
  // Invalid key
  EXPECT_FALSE(DeserializeRsaPrivateKey("this is a bad key", &private_key));
}

TEST_F(RsaUtilTest, SerializeDeserializePublicKey) {
  RSA* public_key;
  std::string serialized_public_key;
  // Key 1
  EXPECT_TRUE(DeserializeRsaPublicKey(
      test_keys_.public_test_key_1_3072_bits(), &public_key));
  ASSERT_TRUE(public_key != NULL);
  EXPECT_TRUE(SerializeRsaPublicKey(public_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_1_3072_bits());
  RSA_free(public_key);
  // Key 2
  EXPECT_TRUE(DeserializeRsaPublicKey(
      test_keys_.public_test_key_2_2048_bits(), &public_key));
  ASSERT_TRUE(public_key != NULL);
  EXPECT_TRUE(SerializeRsaPublicKey(public_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_2_2048_bits());
  RSA_free(public_key);
  // Key 3
  EXPECT_TRUE(DeserializeRsaPublicKey(
      test_keys_.public_test_key_3_2048_bits(), &public_key));
  ASSERT_TRUE(public_key != NULL);
  EXPECT_TRUE(SerializeRsaPublicKey(public_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_3_2048_bits());
  RSA_free(public_key);
  // Invalid key
  EXPECT_FALSE(DeserializeRsaPublicKey("this is a bad key", &public_key));
}

TEST_F(RsaUtilTest, PublicKeyExtraction) {
  RSA* private_key;
  std::string serialized_public_key;
  // Key 1
  EXPECT_TRUE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_1_3072_bits(), &private_key));
  ASSERT_TRUE(private_key != NULL);
  EXPECT_TRUE(SerializeRsaPublicKey(private_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_1_3072_bits());
  RSA_free(private_key);
  // Key 2
  EXPECT_TRUE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_2_2048_bits(), &private_key));
  ASSERT_TRUE(private_key != NULL);
  EXPECT_TRUE(SerializeRsaPublicKey(private_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_2_2048_bits());
  RSA_free(private_key);
  // Key 3
  EXPECT_TRUE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_3_2048_bits(), &private_key));
  ASSERT_TRUE(private_key != NULL);
  EXPECT_TRUE(SerializeRsaPublicKey(private_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_3_2048_bits());
  RSA_free(private_key);
}

TEST_F(RsaUtilTest, Pkcs8PrivateKeyInfo) {
  // The PKCS#1 <-> PKCS#8 conversion routines exercise all the PKCS#8
  // serialization/deserialization functionality , so we test those.
  std::string serialized_pkcs8;
  std::string serialized_pkcs1;
  // Key 1
  EXPECT_TRUE(RsaPrivateKeyToPrivateKeyInfo(
      test_keys_.private_test_key_1_3072_bits(), &serialized_pkcs8));
  EXPECT_TRUE(PrivateKeyInfoToRsaPrivateKey(serialized_pkcs8,
                                            &serialized_pkcs1));
  EXPECT_NE(serialized_pkcs1, serialized_pkcs8);
  EXPECT_EQ(test_keys_.private_test_key_1_3072_bits(), serialized_pkcs1);
  // Key 2
  EXPECT_TRUE(RsaPrivateKeyToPrivateKeyInfo(
      test_keys_.private_test_key_2_2048_bits(), &serialized_pkcs8));
  EXPECT_TRUE(PrivateKeyInfoToRsaPrivateKey(serialized_pkcs8,
                                            &serialized_pkcs1));
  EXPECT_NE(serialized_pkcs1, serialized_pkcs8);
  EXPECT_EQ(test_keys_.private_test_key_2_2048_bits(), serialized_pkcs1);
  // Key 3
  EXPECT_TRUE(RsaPrivateKeyToPrivateKeyInfo(
      test_keys_.private_test_key_3_2048_bits(), &serialized_pkcs8));
  EXPECT_TRUE(PrivateKeyInfoToRsaPrivateKey(serialized_pkcs8,
                                            &serialized_pkcs1));
  EXPECT_NE(serialized_pkcs1, serialized_pkcs8);
  EXPECT_EQ(test_keys_.private_test_key_3_2048_bits(), serialized_pkcs1);
}

TEST_F(RsaUtilTest, Pkcs8EncryptedPrivateKeyInfo) {
  // The PKCS#1 <-> PKCS#8 conversion routines exercise all the PKCS#8
  // serialization/deserialization functionality , so we test those.
  std::string serialized_pkcs8;
  std::string serialized_pkcs1;
  std::string passphrase("passphrase");
  // Key 1
  EXPECT_TRUE(RsaPrivateKeyToEncryptedPrivateKeyInfo(
      test_keys_.private_test_key_1_3072_bits(), passphrase,
      &serialized_pkcs8));
  EXPECT_TRUE(EncryptedPrivateKeyInfoToRsaPrivateKey(serialized_pkcs8,
                                                     passphrase,
                                                     &serialized_pkcs1));
  EXPECT_NE(serialized_pkcs1, serialized_pkcs8);
  EXPECT_EQ(test_keys_.private_test_key_1_3072_bits(), serialized_pkcs1);
  // Key 2
  EXPECT_TRUE(RsaPrivateKeyToEncryptedPrivateKeyInfo(
      test_keys_.private_test_key_2_2048_bits(), passphrase,
      &serialized_pkcs8));
  EXPECT_TRUE(EncryptedPrivateKeyInfoToRsaPrivateKey(serialized_pkcs8,
                                                     passphrase,
                                                     &serialized_pkcs1));
  EXPECT_NE(serialized_pkcs1, serialized_pkcs8);
  EXPECT_EQ(test_keys_.private_test_key_2_2048_bits(), serialized_pkcs1);
  // Key 3
  EXPECT_TRUE(RsaPrivateKeyToEncryptedPrivateKeyInfo(
      test_keys_.private_test_key_3_2048_bits(), passphrase,
      &serialized_pkcs8));
  EXPECT_TRUE(EncryptedPrivateKeyInfoToRsaPrivateKey(serialized_pkcs8,
                                                     passphrase,
                                                     &serialized_pkcs1));
  EXPECT_NE(serialized_pkcs1, serialized_pkcs8);
  EXPECT_EQ(test_keys_.private_test_key_3_2048_bits(), serialized_pkcs1);
}

TEST_F(RsaUtilTest, FailOnInvalidParams) {
  RSA* test_input_key = NULL;
  RSA* test_output_key = NULL;
  std::string test_string;
  std::string pass("password");
  ASSERT_TRUE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_2_2048_bits(), &test_input_key));
  ASSERT_TRUE(test_input_key != NULL);
  EXPECT_FALSE(SerializeRsaPrivateKey(NULL, &test_string));
  EXPECT_FALSE(SerializeRsaPrivateKey(test_input_key, NULL));
  EXPECT_FALSE(SerializeRsaPublicKey(NULL, &test_string));
  EXPECT_FALSE(SerializeRsaPublicKey(test_input_key, NULL));
  EXPECT_FALSE(SerializePrivateKeyInfo(NULL, &test_string));
  EXPECT_FALSE(SerializePrivateKeyInfo(test_input_key, NULL));
  EXPECT_FALSE(SerializeEncryptedPrivateKeyInfo(NULL, pass, &test_string));
  EXPECT_FALSE(SerializeEncryptedPrivateKeyInfo(test_input_key, pass, NULL));
  EXPECT_FALSE(DeserializeRsaPrivateKey("", &test_output_key));
  EXPECT_FALSE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_2_2048_bits(), NULL));
  EXPECT_FALSE(DeserializeRsaPublicKey("", &test_output_key));
  EXPECT_FALSE(DeserializeRsaPublicKey(
      test_keys_.public_test_key_2_2048_bits(), NULL));
  EXPECT_FALSE(DeserializePrivateKeyInfo("", &test_output_key));
  EXPECT_FALSE(DeserializePrivateKeyInfo(
      test_keys_.private_test_key_2_2048_bits(), NULL));
  EXPECT_FALSE(DeserializeEncryptedPrivateKeyInfo("", pass, &test_output_key));
  EXPECT_FALSE(DeserializeEncryptedPrivateKeyInfo(
      test_keys_.private_test_key_2_2048_bits(), pass, NULL));
  RSA_free(test_input_key);
}

TEST_F(RsaUtilTest, Pkcs8FailOnInvalidPassword) {
  RSA* test_input_key = NULL;
  RSA* test_output_key = NULL;
  std::string serialized_pkcs8;
  std::string pass("password");
  ASSERT_TRUE(DeserializeRsaPrivateKey(
      test_keys_.private_test_key_2_2048_bits(), &test_input_key));
  EXPECT_FALSE(SerializeEncryptedPrivateKeyInfo(test_input_key, "",
                                                &serialized_pkcs8));
  ASSERT_TRUE(SerializeEncryptedPrivateKeyInfo(test_input_key, pass,
                                               &serialized_pkcs8));
  EXPECT_FALSE(DeserializeEncryptedPrivateKeyInfo(serialized_pkcs8, pass + "a",
                                                  &test_output_key));
  EXPECT_TRUE(DeserializeEncryptedPrivateKeyInfo(serialized_pkcs8, pass,
                                                 &test_output_key));
  RSA_free(test_input_key);
  RSA_free(test_output_key);
}

}  // namespace rsa_util
}  // namespace widevine