////////////////////////////////////////////////////////////////////////////////
// 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 <cstdint>
#include <vector>

#include <cstdint>
#include "glog/logging.h"
#include "openssl/aes.h"

namespace widevine {
namespace crypto_util {

// Encrypts the provided plantext std::string using AES-CBC encryption.
std::string EncryptAesCbc(const std::string& key, const std::string& iv,
                          const std::string& plaintext) {
  const size_t num_padding_bytes =
      AES_BLOCK_SIZE - (plaintext.size() % AES_BLOCK_SIZE);
  std::string padded_text = plaintext;
  padded_text.append(num_padding_bytes, static_cast<char>(num_padding_bytes));
  return EncryptAesCbcNoPad(key, iv, padded_text);
}

std::string EncryptAesCbcNoPad(const std::string& key, const std::string& iv,
                               const std::string& plaintext) {
  if (iv.size() != AES_BLOCK_SIZE) {
    LOG(WARNING) << "Invalid CBC IV size: " << iv.size();
    return std::string();
  }
  if (plaintext.size() % AES_BLOCK_SIZE) {
    LOG(WARNING) << "Invalid AEC-CBC plaintext size: " << plaintext.size();
    return std::string();
  }

  AES_KEY aes_key;
  if (AES_set_encrypt_key(reinterpret_cast<const uint8_t*>(&key[0]),
                          key.size() * 8, &aes_key) != 0) {
    LOG(WARNING) << "Invalid AES key.";
    return std::string();
  }

  std::string encrypted(plaintext.size(), 0);
  std::vector<uint8_t> local_iv(iv.begin(), iv.end());
  AES_cbc_encrypt(reinterpret_cast<const uint8_t*>(plaintext.data()),
                  reinterpret_cast<uint8_t*>(&encrypted[0]), plaintext.size(),
                  &aes_key, &local_iv[0], AES_ENCRYPT);
  return encrypted;
}

// Decrypts the AES-CBC encrypted text.  Returns an empty std::string on error or
// the plaintext on success.
std::string DecryptAesCbc(const std::string& key, const std::string& iv,
                          const std::string& ciphertext) {
  if (ciphertext.empty()) {
    LOG(WARNING) << "Empty ciphertext.";
    return std::string();
  }
  if (iv.size() != AES_BLOCK_SIZE) {
    LOG(WARNING) << "Invalid CBC IV size: " << iv.size();
    return std::string();
  }
  if ((ciphertext.size() % AES_BLOCK_SIZE) != 0) {
    LOG(WARNING) << "Ciphertext not a multiple of AES block size: "
                 << ciphertext.size();
    return std::string();
  }

  AES_KEY aes_key;
  if (AES_set_decrypt_key(reinterpret_cast<const uint8_t*>(&key[0]),
                          key.size() * 8, &aes_key) != 0) {
    LOG(WARNING) << "Invalid AES key.";
    return std::string();
  }

  std::string cleartext(ciphertext.size(), 0);
  std::vector<uint8_t> local_iv(iv.begin(), iv.end());
  AES_cbc_encrypt(reinterpret_cast<const uint8_t*>(ciphertext.data()),
                  reinterpret_cast<uint8_t*>(&cleartext[0]), ciphertext.size(),
                  &aes_key, &local_iv[0], AES_DECRYPT);

  const uint8_t num_padding_bytes = cleartext[cleartext.size() - 1];
  if (num_padding_bytes > AES_BLOCK_SIZE) {
    LOG(WARNING) << "AES padding too long: " << num_padding_bytes;
    return std::string();
  }
  for (uint8_t i = 0; i < num_padding_bytes; ++i) {
    if (cleartext[cleartext.size() - 1 - i] != num_padding_bytes) {
      LOG(WARNING) << "Padding verification failure.";
      return std::string();
    }
  }
  cleartext.resize(cleartext.size() - num_padding_bytes);
  return cleartext;
}

std::string DecryptAesCbcNoPad(const std::string& key, const std::string& iv,
                               const std::string& ciphertext) {
  std::vector<uint8_t> local_iv(iv.begin(), iv.end());
  if (local_iv.empty()) {
    local_iv.resize(AES_BLOCK_SIZE, '\0');
  } else if (local_iv.size() != AES_BLOCK_SIZE) {
    LOG(WARNING) << "Invalid CBC IV size: " << iv.size();
    return std::string();
  }
  if ((ciphertext.size() % AES_BLOCK_SIZE) != 0) {
    LOG(WARNING) << "Ciphertext not a multiple of AES block size: "
                 << ciphertext.size();
    return std::string();
  }

  AES_KEY aes_key;
  if (AES_set_decrypt_key(reinterpret_cast<const uint8_t*>(&key[0]),
                          key.size() * 8, &aes_key) != 0) {
    LOG(WARNING) << "Invalid AES key.";
    return std::string();
  }

  std::string cleartext(ciphertext.size(), 0);
  AES_cbc_encrypt(reinterpret_cast<const uint8_t*>(ciphertext.data()),
                  reinterpret_cast<uint8_t*>(&cleartext[0]), ciphertext.size(),
                  &aes_key, &local_iv[0], AES_DECRYPT);
  return cleartext;
}

}  // namespace crypto_util
}  // namespace widevine