////////////////////////////////////////////////////////////////////////////////
// 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.
////////////////////////////////////////////////////////////////////////////////

#include "common/aes_cbc_util.h"
#include "testing/gmock.h"
#include "testing/gunit.h"

namespace {

const uint8_t kKey[] = {
    0x87, 0x27, 0xa4, 0x0e, 0xbd, 0x82, 0x32, 0x9e,
    0x6b, 0x3b, 0x4e, 0x29, 0xfa, 0x3b, 0x00, 0x4b,
};

const uint8_t kIv[] = {
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
    0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
};

}  // namespace

namespace widevine {
namespace crypto_util {

TEST(CryptoUtilTest, EncryptAndDecryptAesCbc) {
  std::string plain_text("Foo");
  std::string ciphertext = EncryptAesCbc(std::string(kKey, kKey + sizeof(kKey)),
                                    std::string(kIv, kIv + sizeof(kIv)), plain_text);
  std::string expected_ciphertext(
      "\xCF\x1A\x3\x1C\x9C\x8C\xB9Z\xEC\xC0\x17\xDCRxX\xD7");
  ASSERT_EQ(0, ciphertext.size() % 16);
  ASSERT_GT(ciphertext.size(), plain_text.size());
  ASSERT_EQ(expected_ciphertext, ciphertext);

  std::string decrypted = DecryptAesCbc(std::string(kKey, kKey + sizeof(kKey)),
                                   std::string(kIv, kIv + sizeof(kIv)), ciphertext);
  ASSERT_EQ(plain_text, decrypted);
}

TEST(CryptoUtilTest, DecryptAesCbcNoPad) {
  const uint8_t kKey[] = {
      0xdd, 0x71, 0x39, 0xea, 0xfa, 0xce, 0xed, 0x7c,
      0xda, 0x9f, 0x25, 0xda, 0x8a, 0xa9, 0x15, 0xea,
  };
  const uint8_t kIv[] = {
      0x5d, 0x16, 0x44, 0xea, 0xec, 0x11, 0xf9, 0x83,
      0x14, 0x75, 0x41, 0xe4, 0x6e, 0xeb, 0x27, 0x74,
  };
  const uint8_t kCiphertext[] = {
      0x6d, 0xa6, 0xda, 0xe4, 0xee, 0x40, 0x09, 0x17,
      0x54, 0x7b, 0xba, 0xa5, 0x27, 0xb8, 0x82, 0x1b,
  };
  const uint8_t kExpectedPlaintext[] = {
      0xb3, 0x49, 0xd4, 0x80, 0x9e, 0x91, 0x06, 0x87,
      0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x10,
  };
  const uint8_t kExpectedPlaintextEmptyIv[] = {
      0xee, 0x5f, 0x90, 0x6a, 0x72, 0x80, 0xff, 0x04,
      0x14, 0x75, 0x41, 0xa4, 0x6e, 0xeb, 0x27, 0x64,
  };

  std::string decrypted = DecryptAesCbcNoPad(
      std::string(kKey, kKey + sizeof(kKey)), std::string(kIv, kIv + sizeof(kIv)),
      std::string(kCiphertext, kCiphertext + sizeof(kCiphertext)));
  ASSERT_EQ(std::string(kExpectedPlaintext,
                   kExpectedPlaintext + sizeof(kExpectedPlaintext)),
            decrypted);

  std::string dummy_iv;
  decrypted = DecryptAesCbcNoPad(
      std::string(kKey, kKey + sizeof(kKey)), dummy_iv,
      std::string(kCiphertext, kCiphertext + sizeof(kCiphertext)));
  ASSERT_EQ(
      std::string(kExpectedPlaintextEmptyIv,
             kExpectedPlaintextEmptyIv + sizeof(kExpectedPlaintextEmptyIv)),
      decrypted);
}

TEST(CryptoUtilTest, TestFailedEncrypt) {
  // Test with bogus initialization vector.
  std::string plain_text("Foo");
  std::string bogus_iv("bogus");
  std::string ciphertext =
      EncryptAesCbc(std::string(kKey, kKey + sizeof(kKey)), bogus_iv, plain_text);
  ASSERT_EQ(ciphertext.size(), 0);

  // Test with bogus key.
  std::string bogus_key("bogus");
  ciphertext =
      EncryptAesCbc(bogus_key, std::string(kIv, kIv + sizeof(kIv)), plain_text);
  ASSERT_EQ(ciphertext.size(), 0);
}

TEST(CryptoUtilTest, TestFailedEncryptNoPad) {
  std::string plaintext("0123456789abcdef");
  std::string key(kKey, kKey + sizeof(kKey));
  std::string iv(kIv, kIv + sizeof(kIv));

  // Control.
  std::string ciphertext = EncryptAesCbcNoPad(key, iv, plaintext);
  ASSERT_EQ(plaintext.size(), ciphertext.size());

  // Bogus key.
  std::string bogus_key("bogus");
  ciphertext = EncryptAesCbcNoPad(bogus_key, iv, plaintext);
  EXPECT_EQ(ciphertext.size(), 0);

  // Bogus IV.
  std::string bogus_iv("bogus");
  ciphertext = EncryptAesCbcNoPad(key, bogus_iv, plaintext);
  EXPECT_EQ(ciphertext.size(), 0);

  // Incorrectly-sized plaintext.
  std::string bad_plaintext("Foo");
  ciphertext = EncryptAesCbcNoPad(key, iv, bad_plaintext);
  EXPECT_EQ(ciphertext.size(), 0);
}

TEST(CryptoUtilTest, TestFailedDecrypt) {
  // First, encrypt the data.
  std::string plain_text("Foo");
  std::string ciphertext = EncryptAesCbc(std::string(kKey, kKey + sizeof(kKey)),
                                    std::string(kIv, kIv + sizeof(kIv)), plain_text);
  ASSERT_NE(ciphertext.size(), 0);

  // Test Decrypt with bogus iv.
  std::string bogus_iv("bogus");
  plain_text =
      DecryptAesCbc(std::string(kKey, kKey + sizeof(kKey)), bogus_iv, ciphertext);
  ASSERT_EQ(plain_text.size(), 0);

  // Test Decrypt with bogus key.
  std::string bogus_key("bogus");
  plain_text =
      DecryptAesCbc(bogus_key, std::string(kIv, kIv + sizeof(kIv)), ciphertext);
  ASSERT_EQ(plain_text.size(), 0);
}

TEST(CryptoUtilTest, TestEmptyEncrypt) {
  EXPECT_EQ("\xDBx\xD9\x91\xE8\x1D\xD9\x19\x80r\x12\x89\xD7Kp\xEB",
            EncryptAesCbc(std::string(kKey, kKey + sizeof(kKey)),
                          std::string(kIv, kIv + sizeof(kIv)), ""));
}

TEST(CryptoUtilTest, TestEmptyDecryptAesCbc) {
  EXPECT_EQ("", DecryptAesCbc(std::string(kKey, kKey + sizeof(kKey)),
                              std::string(kIv, kIv + sizeof(kIv)), ""));
}

TEST(CryptoUtilTest, TestEmptyDecryptAesCbcNoPad) {
  EXPECT_EQ("", DecryptAesCbcNoPad(std::string(kKey, kKey + sizeof(kKey)),
                                   std::string(kIv, kIv + sizeof(kIv)), ""));
}

}  // namespace crypto_util
}  // namespace widevine