// Copyright 2021 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
//
//  Reference implementation utilities of OEMCrypto APIs
//
#include "oemcrypto_key_deriver.h"

#include "log.h"

namespace wvoec {
namespace util {
namespace {
bool Derive128KeyAppend(Cmac* cmac, uint8_t counter, const uint8_t* ctx,
                        size_t ctx_size, std::vector<uint8_t>* derived_key) {
  cmac->Reset();
  if (!cmac->Update(counter)) {
    return false;
  }
  if (!cmac->Update(ctx, ctx_size)) {
    return false;
  }
  if (!cmac->FinalizeAppend(derived_key)) {
    return false;
  }
  return true;
}

bool Derive128Key(Cmac* cmac, uint8_t counter, const uint8_t* ctx,
                  size_t ctx_size, std::vector<uint8_t>* derived_key) {
  derived_key->clear();
  return Derive128KeyAppend(cmac, counter, ctx, ctx_size, derived_key);
}

bool Derive256Key(Cmac* cmac, uint8_t counter_base, const uint8_t* ctx,
                  size_t ctx_size, std::vector<uint8_t>* derived_key) {
  derived_key->clear();
  if (!Derive128KeyAppend(cmac, counter_base, ctx, ctx_size, derived_key)) {
    return false;
  }
  return Derive128KeyAppend(cmac, counter_base + 1, ctx, ctx_size, derived_key);
}

bool NistKdf(Cmac* cmac, const std::vector<uint8_t>& label,
             const std::vector<uint8_t>& context, size_t bits,
             std::vector<uint8_t>* renewed_device_key) {
  const std::vector<uint8_t> size_bits_big_endian = {
      static_cast<uint8_t>(bits >> 24), static_cast<uint8_t>(bits >> 16),
      static_cast<uint8_t>(bits >> 8), static_cast<uint8_t>(bits)};
  const size_t kAesBlockSizeBits = 16 * 8;

  if (bits % kAesBlockSizeBits != 0) return false;
  if (renewed_device_key == nullptr) {
    return false;
  }
  renewed_device_key->clear();
  bool res = false;
  for (size_t counter = 0; counter < bits / kAesBlockSizeBits; counter++) {
    cmac->Reset();
    res = cmac->Update(counter + 1) && cmac->Update(label) &&
          cmac->Update(0x00) && cmac->Update(context) &&
          cmac->Update(size_bits_big_endian) &&
          cmac->FinalizeAppend(renewed_device_key);
    if (!res) break;
  }
  return res;
}
}  // namespace

// static
std::unique_ptr<KeyDeriver> KeyDeriver::Create(const uint8_t* key,
                                               size_t key_size) {
  if (key == nullptr) {
    LOGE("Key deriver key is null");
    return std::unique_ptr<KeyDeriver>();
  }
  std::unique_ptr<KeyDeriver> key_deriver(new KeyDeriver());
  if (!key_deriver->Init(key, key_size)) {
    key_deriver.reset();
  }
  return key_deriver;
}

// static
std::unique_ptr<KeyDeriver> KeyDeriver::Create(
    const std::vector<uint8_t>& key) {
  if (key.empty()) {
    LOGE("Key deriver key is empty");
    return std::unique_ptr<KeyDeriver>();
  }
  return Create(key.data(), key.size());
}

bool KeyDeriver::Init(const uint8_t* key, size_t key_size) {
  cmac_ = Cmac::Create(key, key_size);
  if (!cmac_) {
    LOGE("Failed to create CMAC for key deriver");
    return false;
  }
  return true;
}

bool KeyDeriver::DeriveServerMacKey(const uint8_t* mac_key_context,
                                    size_t mac_key_context_size,
                                    std::vector<uint8_t>* mac_key_server) {
  if (mac_key_context == nullptr) {
    LOGE("Server MAC key context is null");
    return false;
  }
  if (mac_key_server == nullptr) {
    LOGE("Output server MAC key buffer is null");
    return false;
  }
  return Derive256Key(cmac_.get(), 0x01, mac_key_context, mac_key_context_size,
                      mac_key_server);
}

bool KeyDeriver::DeriveServerMacKey(const std::vector<uint8_t>& mac_key_context,
                                    std::vector<uint8_t>* mac_key_server) {
  if (mac_key_context.empty()) {
    LOGE("Server MAC key context is empty");
    return false;
  }
  return DeriveServerMacKey(mac_key_context.data(), mac_key_context.size(),
                            mac_key_server);
}

bool KeyDeriver::DeriveClientMacKey(const uint8_t* mac_key_context,
                                    size_t mac_key_context_size,
                                    std::vector<uint8_t>* mac_key_client) {
  if (mac_key_context == nullptr) {
    LOGE("Client MAC key context is null");
    return false;
  }
  if (mac_key_client == nullptr) {
    LOGE("Output client MAC key buffer is null");
    return false;
  }
  return Derive256Key(cmac_.get(), 0x03, mac_key_context, mac_key_context_size,
                      mac_key_client);
}

bool KeyDeriver::DeriveClientMacKey(const std::vector<uint8_t>& mac_key_context,
                                    std::vector<uint8_t>* mac_key_client) {
  if (mac_key_context.empty()) {
    LOGE("Client MAC key context is empty");
    return false;
  }
  return DeriveClientMacKey(mac_key_context.data(), mac_key_context.size(),
                            mac_key_client);
}

bool KeyDeriver::DeriveEncryptionKey(const uint8_t* enc_key_context,
                                     size_t enc_key_context_size,
                                     std::vector<uint8_t>* enc_key) {
  if (enc_key_context == nullptr) {
    LOGE("Encryption key context is null");
    return false;
  }
  if (enc_key == nullptr) {
    LOGE("Output encryption key buffer is null");
    return false;
  }
  return Derive128Key(cmac_.get(), 0x01, enc_key_context, enc_key_context_size,
                      enc_key);
}

bool KeyDeriver::DeriveEncryptionKey(
    const std::vector<uint8_t>& enc_key_context,
    std::vector<uint8_t>* enc_key) {
  if (enc_key_context.empty()) {
    LOGE("Encryption key context is empty");
    return false;
  }
  return DeriveEncryptionKey(enc_key_context.data(), enc_key_context.size(),
                             enc_key);
}

bool KeyDeriver::DeriveRenewedDeviceKey(
    const std::vector<uint8_t>& context,
    std::vector<uint8_t>* renewed_device_key) {
  if (renewed_device_key == nullptr) {
    LOGE("Output key buffer is null");
    return false;
  }
  const std::string kKeyboxRenewalLabel = "Keyboxv4";
  const std::vector<uint8_t> kKeyboxRenewalLabelVec(kKeyboxRenewalLabel.begin(),
                                                    kKeyboxRenewalLabel.end());

  return NistKdf(cmac_.get(), kKeyboxRenewalLabelVec, context, 0x80,
                 renewed_device_key);
}
}  // namespace util
}  // namespace wvoec