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

#include <algorithm>
#include <iostream>
#include <vector>

#include "OEMCryptoCENC.h"
#include "oec_test_data.h"
#include "platform.h"
#include "string_conversions.h"

namespace {

constexpr size_t kAesBlockSize = 16;

/*
This function concatenates the test Prov30 OEM certificate chain and key to the
format below:

+-----------------------+----------------------+--------------------------+
| Cert Chain Length     | Certificate Chain    | Key Length               |
+-----------------------+----------------------+--------------------------+
| (4 bytes, big-endian) | (DER-encoded PKCS#7) | (4 bytes, big-endian)    |
+-----------------------+----------------------+--------------------------+
| Private Key           |
+-----------------------+

|oem_private_key| should be a RSA key in PKCS#8 PrivateKeyInfo format.
|oem_public_cert| should be a DER-encoded PKCS#7 certificate chain.

The output will be consumed by OEMCrypto Prov30 factory functions:
1. It is wrapped by OEMCrypto_WrapKeyboxOrOEMCert(), and
2. The wrapped root of trust will be installed by
OEMCrypto_InstallKeyboxOrOEMCert(). Therefore, the OEMCrypto implementation of
the factory functions and the tool must have an agreement on the format above.
*/
std::vector<uint8_t> PrepareProv30OEMCertAndKey(
    const uint8_t* oem_public_cert, const size_t oem_public_cert_size,
    const uint8_t* oem_private_key, const size_t oem_private_key_size) {
  std::vector<uint8_t> oem_cert_and_key;
  // Calculate total size
  size_t total_size = sizeof(uint32_t) + oem_public_cert_size +
                      sizeof(uint32_t) + oem_private_key_size;
  oem_cert_and_key.resize(total_size);

  // Offset to track where to write in the output vector
  size_t offset = 0;
  // 1. Store public cert size (big-endian)
  uint32_t networkOrderCertSize = htonl((uint32_t)oem_public_cert_size);
  std::copy(reinterpret_cast<const uint8_t*>(&networkOrderCertSize),
            reinterpret_cast<const uint8_t*>(&networkOrderCertSize) +
                sizeof(uint32_t),
            oem_cert_and_key.begin());
  offset += sizeof(uint32_t);

  // 2. Store public cert content
  std::copy(oem_public_cert, oem_public_cert + oem_public_cert_size,
            oem_cert_and_key.begin() + offset);
  offset += oem_public_cert_size;

  // 3. Store private key size (big-endian)
  uint32_t networkOrderKeySize = htonl((uint32_t)oem_private_key_size);
  std::copy(
      reinterpret_cast<const uint8_t*>(&networkOrderKeySize),
      reinterpret_cast<const uint8_t*>(&networkOrderKeySize) + sizeof(uint32_t),
      oem_cert_and_key.begin() + offset);
  offset += sizeof(uint32_t);

  // 4. Store private key content
  std::copy(oem_private_key, oem_private_key + oem_private_key_size,
            oem_cert_and_key.begin() + offset);
  return oem_cert_and_key;
}

OEMCryptoResult InstallTestProv30RootOfTrust(
    const uint8_t* oem_public_cert, const size_t oem_public_cert_size,
    const uint8_t* oem_private_key, const size_t oem_private_key_size) {
  if (oem_public_cert == nullptr || oem_private_key == nullptr ||
      oem_public_cert_size == 0 || oem_private_key_size == 0) {
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }

  // 1. Prepare OEM cert and key.
  std::vector<uint8_t> oem_cert_and_key =
      PrepareProv30OEMCertAndKey(oem_public_cert, oem_public_cert_size,
                                 oem_private_key, oem_private_key_size);
  if (oem_cert_and_key.empty()) {
    std::cerr << "Failed to prepare OEM cert and key" << std::endl;
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }
  // Add padding.
  const uint8_t padding =
      kAesBlockSize - (oem_cert_and_key.size() % kAesBlockSize);
  for (size_t i = 0; i < padding; i++) {
    oem_cert_and_key.push_back(padding);
  }

  // 2:  Initialize OEMCrypto.
  OEMCryptoResult sts = OEMCrypto_Initialize();
  if (sts != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to initialize: result = " << sts << std::endl;
    return sts;
  }

  // 3: Wrap OEM cert and key before calling install function.
  const OEMCrypto_ProvisioningMethod method = OEMCrypto_GetProvisioningMethod();
  if (method != OEMCrypto_OEMCertificate) {
    std::cerr << "OEMCrypto is not OEMCrypto_OEMCertificate: method = ";
    std::cerr << method << std::endl;
    OEMCrypto_Terminate();
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }
  std::vector<uint8_t> wrapped_oem_cert_and_key;
  size_t wrapped_oem_cert_and_key_size = 0;
  sts = OEMCrypto_WrapKeyboxOrOEMCert(
      oem_cert_and_key.data(), oem_cert_and_key.size(),
      wrapped_oem_cert_and_key.data(), &wrapped_oem_cert_and_key_size, nullptr,
      0);
  if (sts != OEMCrypto_ERROR_SHORT_BUFFER) {
    OEMCrypto_Terminate();
    return sts;
  }
  wrapped_oem_cert_and_key.resize(wrapped_oem_cert_and_key_size);
  sts = OEMCrypto_WrapKeyboxOrOEMCert(
      oem_cert_and_key.data(), oem_cert_and_key.size(),
      wrapped_oem_cert_and_key.data(), &wrapped_oem_cert_and_key_size, nullptr,
      0);
  if (sts != OEMCrypto_SUCCESS) {
    OEMCrypto_Terminate();
    return sts;
  }

  // 4: Install the wrapped OEM cert and key.
  sts = OEMCrypto_InstallKeyboxOrOEMCert(wrapped_oem_cert_and_key.data(),
                                         wrapped_oem_cert_and_key_size);
  OEMCrypto_Terminate();
  return sts;
}
}  // namespace

int main() {
  const OEMCryptoResult result = InstallTestProv30RootOfTrust(
      wvoec::kTestOEMPublicCertInfo2, sizeof(wvoec::kTestOEMPublicCertInfo2),
      wvoec::kTestRSAPKCS8PrivateKeyInfo2_2048,
      sizeof(wvoec::kTestRSAPKCS8PrivateKeyInfo2_2048));
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to install OEM cert and key with result: " << result
              << std::endl;
    return 1;
  }
  return 0;
}