// Copyright 2020 Google LLC. All Rights Reserved.

#include "crypto_utils/aes_ctr_encryptor.h"

#include <cstddef>
#include <cstdint>
#include <vector>

#include "base/logging.h"

namespace widevine {
namespace {

constexpr size_t kAesBlockSize = 16;

// Increment an 8-byte counter by 1. Return true if overflowed.
bool Increment64(uint8_t* counter) {
  DCHECK(counter);
  for (int i = 7; i >= 0; --i) {
    if (++counter[i] != 0)
      return false;
  }
  return true;
}

}  // namespace

bool AesCtrEncryptor::SetKey(const uint8_t* key, size_t key_size) {
  DCHECK(key);

  if (key_size != kAesBlockSize && key_size != kAesBlockSize * 2) {
    LOG(WARNING) << "Incorrect key size " << key_size;
    return false;
  }
  if (AES_set_encrypt_key(key, key_size * 8, &aes_key_) != 0) {
    LOG(WARNING) << "Invalid AES key.";
    return false;
  }
  aes_key_size_ = key_size;
  return true;
}

bool AesCtrEncryptor::Encrypt(const uint8_t* iv,
                              size_t iv_size,
                              const uint8_t* input_data,
                              size_t input_data_size,
                              uint8_t* output_data) {
  DCHECK(iv);
  DCHECK(input_data);
  DCHECK(output_data);

  if (aes_key_size_ == 0) {
    LOG(WARNING) << "This class has not been initialized.";
    return false;
  }
  // IV is allowed to be either AES BLOCK size or half of it.
  if (iv_size != kAesBlockSize && iv_size != kAesBlockSize / 2) {
    LOG(WARNING) << "Invalid IV size " << iv_size;
    return false;
  }
  std::vector<uint8_t> counter(iv, iv + iv_size);
  counter.resize(kAesBlockSize);
  std::vector<uint8_t> encrypted_counter(kAesBlockSize);
  for (size_t i = 0; i < input_data_size; i += 16) {
    AES_encrypt(counter.data(), encrypted_counter.data(), &aes_key_);
    // As mentioned in ISO/IEC 23001-7:2016 CENC spec, of the 16 byte counter
    // block, bytes 8 to 15 (i.e. the least significant bytes) are used as a
    // simple 64 bit unsigned integer that is incremented by one for each
    // subsequent block of sample data processed and is kept in network byte
    // order.
    Increment64(counter.data() + 8);
    for (size_t j = 0; j < kAesBlockSize && i + j < input_data_size; j++)
      output_data[i + j] = input_data[i + j] ^ encrypted_counter[j];
  }
  return true;
}

}  // namespace widevine