// Copyright 2020 Google LLC. All Rights Reserved.

#include "crypto_utils/aes_cbc_encryptor.h"

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

#include "base/logging.h"

namespace widevine {
namespace {
constexpr size_t kAesBlockSize = 16;
}  // namespace

bool AesCbcEncryptor::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 AesCbcEncryptor::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;
  }
  if ((input_data_size % kAesBlockSize) != 0) {
    LOG(WARNING) << "Input data size must be multiple of 16: "
                 << input_data_size;
    return false;
  }
  std::vector<uint8_t> local_iv(iv, iv + iv_size);
  local_iv.resize(kAesBlockSize);
  AES_cbc_encrypt(input_data, output_data, input_data_size, &aes_key_,
                  local_iv.data(), AES_ENCRYPT);
  return true;
}

}  // namespace widevine