// Copyright 2019 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
//
// OEMCrypto unit tests
//

#include "oec_key_deriver.h"

#include <openssl/aes.h>
#include <openssl/bio.h>
#include <openssl/cmac.h>
#include <openssl/err.h>
#include <openssl/hmac.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/x509.h>
#include <openssl/x509_vfy.h>
#include <stdint.h>

#include <gtest/gtest.h>
#include <iostream>
#include <memory>
#include <string>
#include <vector>

#include "OEMCryptoCENC.h"
#include "disallow_copy_and_assign.h"
#include "log.h"
#include "oec_device_features.h"
#include "oec_test_data.h"
#include "oemcrypto_types.h"
#include "platform.h"
#include "string_conversions.h"

using namespace std;

namespace wvoec {

void Encryptor::set_enc_key(const std::vector<uint8_t>& enc_key) {
  enc_key_ = enc_key;
}

void Encryptor::CBCEncrypt(const uint8_t* data, uint8_t* encrypted_data,
                           size_t data_length,
                           const uint8_t (&iv)[KEY_IV_SIZE]) const {
  ASSERT_EQ(enc_key_.size(), KEY_SIZE);
  ASSERT_NE(data, nullptr);
  ASSERT_NE(encrypted_data, nullptr);
  AES_KEY aes_key;
  static const int key_size = KEY_SIZE * 8;  // in bits.
  AES_set_encrypt_key(enc_key_.data(), key_size, &aes_key);
  uint8_t iv_buffer[KEY_IV_SIZE];
  memcpy(iv_buffer, iv, KEY_IV_SIZE);
  AES_cbc_encrypt(data, encrypted_data, data_length, &aes_key, iv_buffer,
                  AES_ENCRYPT);
}

void Encryptor::PadAndEncryptProvisioningMessage(
    RSAPrivateKeyMessage* data, RSAPrivateKeyMessage* encrypted) const {
  EXPECT_EQ(1, RAND_bytes(data->rsa_key_iv, KEY_IV_SIZE));
  ASSERT_EQ(enc_key_.size(), KEY_SIZE);
  *encrypted = *data;
  if (data->rsa_key_length > sizeof(data->rsa_key)) {
    // OEMCrypto Fuzzing: fuzzed |rsa_key_length| overflows the allocated
    // buffer. Skip encryption in that case.
    return;
  }

  size_t padding = AES_BLOCK_SIZE - (data->rsa_key_length % AES_BLOCK_SIZE);
  memset(data->rsa_key + data->rsa_key_length, static_cast<uint8_t>(padding),
         padding);
  encrypted->rsa_key_length = data->rsa_key_length + padding;
  AES_KEY aes_key;
  static const int key_size = KEY_SIZE * 8;  // in bits.
  AES_set_encrypt_key(enc_key_.data(), key_size, &aes_key);
  uint8_t iv_buffer[KEY_IV_SIZE];
  memcpy(iv_buffer, &data->rsa_key_iv[0], KEY_IV_SIZE);
  AES_cbc_encrypt(&data->rsa_key[0], &encrypted->rsa_key[0],
                  encrypted->rsa_key_length, &aes_key, iv_buffer, AES_ENCRYPT);
}

// This generates the data for deriving one key. If there are failures in
// this function, then there is something wrong with the test program and its
// dependency on BoringSSL.
void KeyDeriver::DeriveKey(const uint8_t* key, size_t master_key_size,
                           const vector<uint8_t>& context, int counter,
                           vector<uint8_t>* out) {
  ASSERT_NE(key, nullptr);
  ASSERT_FALSE(context.empty());
  ASSERT_GE(4, counter);
  ASSERT_LE(1, counter);
  ASSERT_NE(out, nullptr);
  // For RSA, the master key is expected to be 16 bytes; for EC key, 32 bytes.
  ASSERT_TRUE(master_key_size == KEY_SIZE || master_key_size == 2 * KEY_SIZE);

  const EVP_CIPHER* cipher =
      master_key_size == KEY_SIZE ? EVP_aes_128_cbc() : EVP_aes_256_cbc();
  CMAC_CTX* cmac_ctx = CMAC_CTX_new();
  ASSERT_NE(nullptr, cmac_ctx);

  ASSERT_TRUE(CMAC_Init(cmac_ctx, key, master_key_size, cipher, nullptr));

  std::vector<uint8_t> message;
  message.push_back(static_cast<uint8_t>(counter));
  message.insert(message.end(), context.begin(), context.end());

  ASSERT_TRUE(CMAC_Update(cmac_ctx, message.data(), message.size()));

  size_t reslen;
  uint8_t res[128];
  ASSERT_TRUE(CMAC_Final(cmac_ctx, res, &reslen));

  out->assign(res, res + reslen);
  CMAC_CTX_free(cmac_ctx);
}

// This generates the data for deriving a set of keys. If there are failures in
// this function, then there is something wrong with the test program and its
// dependency on BoringSSL.
void KeyDeriver::DeriveKeys(const uint8_t* master_key, size_t master_key_size,
                            const vector<uint8_t>& mac_key_context,
                            const vector<uint8_t>& enc_key_context) {
  // Generate derived key for mac key
  std::vector<uint8_t> mac_key_part2;
  DeriveKey(master_key, master_key_size, mac_key_context, 1, &mac_key_server_);
  DeriveKey(master_key, master_key_size, mac_key_context, 2, &mac_key_part2);
  mac_key_server_.insert(mac_key_server_.end(), mac_key_part2.begin(),
                         mac_key_part2.end());

  DeriveKey(master_key, master_key_size, mac_key_context, 3, &mac_key_client_);
  DeriveKey(master_key, master_key_size, mac_key_context, 4, &mac_key_part2);
  mac_key_client_.insert(mac_key_client_.end(), mac_key_part2.begin(),
                         mac_key_part2.end());

  // Generate derived key for encryption key
  std::vector<uint8_t> enc_key;
  DeriveKey(master_key, master_key_size, enc_key_context, 1, &enc_key);
  set_enc_key(enc_key);
}

void KeyDeriver::set_mac_keys(const uint8_t* mac_keys) {
  ASSERT_EQ(mac_key_server_.size(), MAC_KEY_SIZE);
  ASSERT_EQ(mac_key_client_.size(), MAC_KEY_SIZE);
  memcpy(mac_key_server_.data(), mac_keys, MAC_KEY_SIZE);
  memcpy(mac_key_client_.data(), mac_keys + MAC_KEY_SIZE, MAC_KEY_SIZE);
}

void KeyDeriver::ServerSignBuffer(const uint8_t* data, size_t data_length,
                                  std::vector<uint8_t>* signature) const {
  ASSERT_EQ(mac_key_server_.size(), MAC_KEY_SIZE);
  signature->assign(SHA256_DIGEST_LENGTH, 0);
  unsigned int sig_len = SHA256_DIGEST_LENGTH;
  ASSERT_TRUE(HMAC(EVP_sha256(), mac_key_server_.data(),
                   static_cast<int>(mac_key_server_.size()), data, data_length,
                   signature->data(), &sig_len));
}

void KeyDeriver::ClientSignBuffer(const vector<uint8_t>& buffer,
                                  std::vector<uint8_t>* signature) const {
  ASSERT_EQ(mac_key_client_.size(), MAC_KEY_SIZE);
  signature->assign(SHA256_DIGEST_LENGTH, 0);
  unsigned int sig_len = SHA256_DIGEST_LENGTH;
  ASSERT_TRUE(HMAC(EVP_sha256(), mac_key_client_.data(),
                   static_cast<int>(mac_key_client_.size()), buffer.data(),
                   buffer.size(), signature->data(), &sig_len));
}

void KeyDeriver::ClientSignPstReport(const vector<uint8_t>& pst_report_buffer,
                                     std::vector<uint8_t>* signature) const {
  ASSERT_EQ(mac_key_client_.size(), MAC_KEY_SIZE);
  signature->assign(SHA_DIGEST_LENGTH, 0);
  unsigned int sig_len = SHA_DIGEST_LENGTH;
  ASSERT_TRUE(HMAC(EVP_sha1(), mac_key_client_.data(),
                   static_cast<int>(mac_key_client_.size()),
                   &pst_report_buffer[SHA_DIGEST_LENGTH],
                   pst_report_buffer.size() - SHA_DIGEST_LENGTH,
                   signature->data(), &sig_len));
}

}  // namespace wvoec