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

#include <cstddef>
#include <memory>
#include <tuple>

#include "testing/gmock.h"
#include "testing/gunit.h"
#include "absl/strings/escaping.h"
#include "common/ec_key.h"
#include "common/ec_key_source.h"
#include "common/ec_test_keys.h"
#include "common/ec_util.h"
#include "common/fake_ec_key_source.h"

using ::testing::_;
using ::testing::Return;

namespace {
const size_t kMacSizeBytes = 32;
}  // anonymous namespace

namespace widevine {

class EciesCryptoTest
    : public ::testing::Test,
      public ::testing::WithParamInterface<
          std::tuple<std::string, std::string, std::string, std::string> > {
 public:
  static std::vector<
      std::tuple<std::string, std::string, std::string, std::string> >
  GetTestKeyList() {
    ECTestKeys test_keys;

    std::vector<std::tuple<std::string, std::string, std::string, std::string> >
        keys({std::make_tuple(test_keys.private_test_key_1_secp521r1(),
                              test_keys.public_test_key_1_secp521r1(),
                              test_keys.private_test_key_2_secp521r1(),
                              test_keys.public_test_key_2_secp521r1()),
              std::make_tuple(test_keys.private_test_key_1_secp384r1(),
                              test_keys.public_test_key_1_secp384r1(),
                              test_keys.private_test_key_2_secp384r1(),
                              test_keys.public_test_key_2_secp384r1()),
              std::make_tuple(test_keys.private_test_key_1_secp256r1(),
                              test_keys.public_test_key_1_secp256r1(),
                              test_keys.private_test_key_2_secp256r1(),
                              test_keys.public_test_key_2_secp256r1())});
    return keys;
  }

 protected:
  EciesCryptoTest() {
    test_private_key_ = std::get<0>(GetParam());
    test_public_key_ = std::get<1>(GetParam());
    test_ephemeral_private_key_ = std::get<2>(GetParam());
    test_ephemeral_public_key_ = std::get<3>(GetParam());
    private_key_ = ECPrivateKey::Create(test_private_key_);
    public_key_ = ECPublicKey::Create(test_public_key_);
    ephemeral_private_key_ = ECPrivateKey::Create(test_ephemeral_private_key_);
    ephemeral_public_key_ = ECPublicKey::Create(test_ephemeral_public_key_);
  }

  std::string test_private_key_;
  std::string test_public_key_;
  std::string test_ephemeral_private_key_;
  std::string test_ephemeral_public_key_;
  std::unique_ptr<ECPrivateKey> private_key_;
  std::unique_ptr<ECPublicKey> public_key_;
  std::unique_ptr<ECPrivateKey> ephemeral_private_key_;
  std::unique_ptr<ECPublicKey> ephemeral_public_key_;
};

TEST_P(EciesCryptoTest, EciesEncryptSuccess) {
  std::string serialized_public_key;

  const std::string context = "test context";
  const std::string plaintext = "test plaintext";

  std::string ecies_message;

  // Use the test ephemeral key in the key source.
  FakeECKeySource fake_key_source;
  ASSERT_TRUE(fake_key_source.SetKey(ephemeral_public_key_->Curve(),
                                     test_ephemeral_private_key_,
                                     test_ephemeral_public_key_));
  std::unique_ptr<EciesEncryptor> encryptor =
      EciesEncryptor::Create(test_public_key_, &fake_key_source);

  ASSERT_TRUE(encryptor->Encrypt(plaintext, context, &ecies_message));
  EXPECT_TRUE(ecies_message.find(plaintext) == std::string::npos);

  // Verify the decrypted_message.
  std::string decrypted_message;
  std::unique_ptr<EciesDecryptor> decryptor =
      EciesDecryptor::Create(test_private_key_);
  ASSERT_TRUE(decryptor != nullptr);
  ASSERT_TRUE(decryptor->Decrypt(ecies_message, context, &decrypted_message));
  EXPECT_EQ(plaintext, decrypted_message);
}

TEST_P(EciesCryptoTest, EciesDecryptShortEciesMessage) {
  std::unique_ptr<EciesDecryptor> decryptor =
      EciesDecryptor::Create(test_private_key_);
  ASSERT_NE(nullptr, decryptor.get());
  std::string decrypted_message;
  std::string short_message(
      ec_util::GetPublicKeyPointSize(private_key_->Curve()) + kMacSizeBytes - 1,
      'a');

  ASSERT_FALSE(
      decryptor->Decrypt(short_message, "test context", &decrypted_message));
}

TEST_P(EciesCryptoTest, EciesDecryptBadSignature) {
  std::string serialized_public_key;

  const std::string context = "test context";
  const std::string plaintext = "test plaintext";

  std::string ecies_message;

  // Use the test ephemeral key in the key source.
  FakeECKeySource fake_key_source;
  ASSERT_TRUE(fake_key_source.SetKey(ephemeral_public_key_->Curve(),
                                     test_ephemeral_private_key_,
                                     test_ephemeral_public_key_));
  std::unique_ptr<EciesEncryptor> encryptor =
      EciesEncryptor::Create(test_public_key_, &fake_key_source);

  ASSERT_TRUE(encryptor->Encrypt(plaintext, context, &ecies_message));

  // Corrupt the signature (at end of message) and verify that decryption fails.
  ecies_message[ecies_message.size() - 1]++;
  std::string decrypted_message;
  std::unique_ptr<EciesDecryptor> decryptor =
      EciesDecryptor::Create(test_private_key_);
  ASSERT_TRUE(decryptor != nullptr);
  ASSERT_FALSE(decryptor->Decrypt(ecies_message, context, &decrypted_message));
  EXPECT_EQ("", decrypted_message);
}

TEST_P(EciesCryptoTest, EciesEncryptMismatchContext) {
  std::string serialized_public_key;

  const std::string context = "test context";
  const std::string bogus_context = "bogus_context";
  const std::string plaintext = "test plaintext";

  std::string ecies_message;

  // Use the test ephemeral key in the key source.
  FakeECKeySource fake_key_source;
  ASSERT_TRUE(fake_key_source.SetKey(ephemeral_public_key_->Curve(),
                                     test_ephemeral_private_key_,
                                     test_ephemeral_public_key_));
  std::unique_ptr<EciesEncryptor> encryptor =
      EciesEncryptor::Create(test_public_key_, &fake_key_source);
  ASSERT_TRUE(encryptor != nullptr);
  ASSERT_TRUE(encryptor->Encrypt(plaintext, context, &ecies_message));
  EXPECT_GT(
      ecies_message.size(),
      kMacSizeBytes + ec_util::GetPublicKeyPointSize(public_key_->Curve()));

  // Verify the decrypted_message using the invalid context.
  std::string decrypted_message;
  std::unique_ptr<EciesDecryptor> decryptor =
      EciesDecryptor::Create(test_private_key_);
  ASSERT_TRUE(decryptor != nullptr);
  ASSERT_FALSE(
      decryptor->Decrypt(ecies_message, bogus_context, &decrypted_message));
  EXPECT_TRUE(decrypted_message.empty());
}

INSTANTIATE_TEST_SUITE_P(
    EciesCryptoTest, EciesCryptoTest,
    ::testing::ValuesIn(EciesCryptoTest::GetTestKeyList()));

TEST(EciesEncryptorTest, EciesEncryptBadPublicKey) {
  // Use the test ephemeral key in the key source.
  FakeECKeySource fake_key_source;
  std::unique_ptr<EciesEncryptor> encryptor =
      EciesEncryptor::Create("bad public key.", &fake_key_source);
  ASSERT_TRUE(encryptor == nullptr);
}

TEST(EciesEncryptorTest, EciesEncryptNullKeySource) {
  std::unique_ptr<EciesEncryptor> encryptor =
      EciesEncryptor::Create("bad public key.", nullptr);
  ASSERT_TRUE(encryptor == nullptr);
}

class MockEcKeySource : public ECKeySource {
 public:
  MockEcKeySource() = default;
  MOCK_METHOD3(GetECKey,
               bool(ECPrivateKey::EllipticCurve curve, std::string* private_key,
                    std::string* public_key));
};

TEST(EciesEncryptorTest, EciesEncryptKeysourceFail) {
  MockEcKeySource mock_ec_key_source;
  EXPECT_CALL(mock_ec_key_source, GetECKey(_, _, _)).WillOnce(Return(false));

  ECTestKeys test_keys;
  std::unique_ptr<EciesEncryptor> encryptor = EciesEncryptor::Create(
      test_keys.public_test_key_1_secp256r1(), &mock_ec_key_source);
  ASSERT_TRUE(encryptor != nullptr);
  std::string ecies_message;
  ASSERT_FALSE(
      encryptor->Encrypt("test plaintext", "test context", &ecies_message));
}

TEST(EciesEncryptorTest, EciesEncryptNullEciesMessage) {
  FakeECKeySource fake_key_source;
  ECTestKeys test_keys;
  std::unique_ptr<EciesEncryptor> encryptor = EciesEncryptor::Create(
      test_keys.public_test_key_1_secp256r1(), &fake_key_source);
  ASSERT_TRUE(encryptor != nullptr);
  ASSERT_FALSE(encryptor->Encrypt("test plaintext", "test context", nullptr));
}

TEST(EciesDecryptorTest, EciesDecryptBadPrivateKey) {
  ECTestKeys test_keys;
  std::string invalid_private_key(test_keys.private_test_key_1_secp521r1());
  invalid_private_key[0]++;

  std::unique_ptr<EciesDecryptor> decryptor =
      EciesDecryptor::Create(invalid_private_key);
  ASSERT_TRUE(decryptor == nullptr);
}

TEST(EciesDecryptorTest, EciesDecryptEmptyEciesMessage) {
  ECTestKeys test_keys;
  std::unique_ptr<EciesDecryptor> decryptor =
      EciesDecryptor::Create(test_keys.private_test_key_1_secp521r1());
  ASSERT_NE(nullptr, decryptor);
  std::string decrypted_message;
  ASSERT_FALSE(decryptor->Decrypt("", "test context", &decrypted_message));
}

TEST(EciesDecryptorTest, EciesDecryptNullPlaintext) {
  ECTestKeys test_keys;
  std::unique_ptr<EciesDecryptor> decryptor =
      EciesDecryptor::Create(test_keys.private_test_key_1_secp521r1());
  ASSERT_NE(nullptr, decryptor);
  std::string decrypted_message;
  ASSERT_FALSE(decryptor->Decrypt("foo", "test context", nullptr));
}

}  // namespace widevine