#include "crypto_utils/aes_ctr_encryptor.h"

#include <cstdint>
#include <vector>

#include "testing/include/gmock/gmock.h"
#include "testing/include/gtest/gtest.h"

using ::testing::ElementsAreArray;

namespace widevine {
namespace {

// From NIST SP 800-38a test case: - F.5.1 CTR-AES128.Encrypt
// http://csrc.nist.gov/publications/nistpubs/800-38a/sp800-38a.pdf
const uint8_t kAesKey[] = {0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
                           0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c};
const uint8_t kAesIv[] = {0xf0, 0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7,
                          0xf8, 0xf9, 0xfa, 0xfb, 0xfc, 0xfd, 0xfe, 0xff};
const uint8_t kAesCtrPlaintext[] = {
    // Block #1
    0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96, 0xe9, 0x3d, 0x7e, 0x11,
    0x73, 0x93, 0x17, 0x2a,
    // Block #2
    0xae, 0x2d, 0x8a, 0x57, 0x1e, 0x03, 0xac, 0x9c, 0x9e, 0xb7, 0x6f, 0xac,
    0x45, 0xaf, 0x8e, 0x51,
    // Block #3
    0x30, 0xc8, 0x1c, 0x46, 0xa3, 0x5c, 0xe4, 0x11, 0xe5, 0xfb, 0xc1, 0x19,
    0x1a, 0x0a, 0x52, 0xef,
    // Block #4
    0xf6, 0x9f, 0x24, 0x45, 0xdf, 0x4f, 0x9b, 0x17, 0xad, 0x2b, 0x41, 0x7b,
    0xe6, 0x6c, 0x37, 0x10};
const uint8_t kAesCtrCiphertext[] = {
    // Block #1
    0x87, 0x4d, 0x61, 0x91, 0xb6, 0x20, 0xe3, 0x26, 0x1b, 0xef, 0x68, 0x64,
    0x99, 0x0d, 0xb6, 0xce,
    // Block #2
    0x98, 0x06, 0xf6, 0x6b, 0x79, 0x70, 0xfd, 0xff, 0x86, 0x17, 0x18, 0x7b,
    0xb9, 0xff, 0xfd, 0xff,
    // Block #3
    0x5a, 0xe4, 0xdf, 0x3e, 0xdb, 0xd5, 0xd3, 0x5e, 0x5b, 0x4f, 0x09, 0x02,
    0x0d, 0xb0, 0x3e, 0xab,
    // Block #4
    0x1e, 0x03, 0x1d, 0xda, 0x2f, 0xbe, 0x03, 0xd1, 0x79, 0x21, 0x70, 0xa0,
    0xf3, 0x00, 0x9c, 0xee};

const uint8_t kIv128Zero[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
const uint8_t kIv128Max64[] = {0,    0,    0,    0,    0,    0,    0,    0,
                               0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};

}  // namespace

TEST(AesCtrEncryptorTest, NistTestCase) {
  AesCtrEncryptor encryptor;
  ASSERT_TRUE(encryptor.SetKey(kAesKey, sizeof(kAesKey)));

  std::vector<uint8_t> output(sizeof(kAesCtrPlaintext));
  ASSERT_TRUE(encryptor.Encrypt(kAesIv, sizeof(kAesIv), kAesCtrPlaintext,
                                sizeof(kAesCtrPlaintext), output.data()));
  EXPECT_THAT(output,
              ElementsAreArray(kAesCtrCiphertext, sizeof(kAesCtrCiphertext)));
}

TEST(AesCtrEncryptorTest, 128BitIVMax64BoundaryCaseEncryption) {
  AesCtrEncryptor encryptor;
  ASSERT_TRUE(encryptor.SetKey(kAesKey, sizeof(kAesKey)));

  std::vector<uint8_t> output(sizeof(kAesCtrPlaintext));
  ASSERT_TRUE(encryptor.Encrypt(kIv128Max64, sizeof(kIv128Max64),
                                kAesCtrPlaintext, sizeof(kAesCtrPlaintext),
                                output.data()));

  // There are four blocks in plain text. The first block should be encrypted
  // with IV = kIv128Max64, the subsequent blocks should be encrypted with iv 0
  // to 3.
  std::vector<uint8_t> output2(sizeof(kAesCtrPlaintext));
  ASSERT_TRUE(encryptor.Encrypt(kIv128Max64, sizeof(kIv128Max64),
                                kAesCtrPlaintext, 16, &output2[0]));
  ASSERT_TRUE(encryptor.Encrypt(kIv128Zero, sizeof(kIv128Zero),
                                &kAesCtrPlaintext[16],
                                sizeof(kAesCtrPlaintext) - 16, &output2[16]));

  EXPECT_EQ(output, output2);
}

}  // namespace widevine