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

#include "WidevineProvisioner.h"

#include <cppbor.h>
#include <cppbor_parse.h>
#include <keymaster/cppcose/cppcose.h>
#include <openssl/bn.h>
#include <openssl/curve25519.h>
#include <openssl/ec.h>
#include <openssl/evp.h>
#include <openssl/hkdf.h>
#include <openssl/rand.h>
#include <string.h>

#include <memory>
#include <string>
#include <utility>

#include "log.h"
#include "properties.h"

namespace widevine {

namespace {
const std::vector<std::vector<uint8_t>> kAuthorizedEekRoots = {
    {0x99, 0xB9, 0xEE, 0xDD, 0x5E, 0xE4, 0x52, 0xF6, 0x85, 0xC6, 0x4C,
     0x62, 0xDC, 0x3E, 0x61, 0xAB, 0x57, 0x48, 0x7D, 0x75, 0x37, 0x29,
     0xAD, 0x76, 0x80, 0x32, 0xD2, 0xB3, 0xCB, 0x63, 0x58, 0xD9},
};
}  // namespace

WidevineProvisioner::WidevineProvisioner() {
  InitializeCryptoInterface();
  assert(crypto_interface_ != nullptr);
}

std::vector<uint8_t> WidevineProvisioner::GetBcc() {
  std::vector<uint8_t> bcc;
  OEMCryptoResult result = crypto_interface_->GetBcc(bcc);
  if (result != OEMCrypto_SUCCESS) {
    LOGE("Failed to get BCC.");
  }
  return bcc;
}

bool WidevineProvisioner::GenerateCertificateRequest(
    bool testMode, const std::vector<uint8_t>& endpointEncCertChain,
    std::vector<uint8_t>& deviceInfo, std::vector<uint8_t>& protectedData) {
  if (!GetDeviceInfo(deviceInfo)) {
    LOGE("Failed to get device_info.");
    return false;
  }

  std::vector<uint8_t> bcc;
  OEMCryptoResult result = crypto_interface_->GetBcc(bcc);
  if (result != OEMCrypto_SUCCESS) {
    LOGE("Failed to get BCC.");
    return false;
  }

  if (!GenerateProtectedData(testMode, endpointEncCertChain, bcc,
                             protectedData)) {
    LOGE("Failed to generate protected data.");
    return false;
  }

  return true;
}

bool WidevineProvisioner::GetDeviceInfo(std::vector<uint8_t>& device_info) {
  auto device_info_map = cppbor::Map();
  device_info_map.add(cppbor::Tstr("type"), cppbor::Tstr("widevine"));

  std::string company_name;
  if (!wvcdm::Properties::GetCompanyName(&company_name) ||
      company_name.empty()) {
    LOGE("Failed to get company name.");
    return false;
  }
  device_info_map.add(cppbor::Tstr("manufacturer"), cppbor::Tstr(company_name));

  std::string model_name;
  if (!wvcdm::Properties::GetModelName(&model_name) || model_name.empty()) {
    LOGE("Failed to get model name.");
    return false;
  }
  device_info_map.add(cppbor::Tstr("model"), cppbor::Tstr(model_name));

  std::string arch_name;
  if (!wvcdm::Properties::GetArchitectureName(&arch_name) ||
      arch_name.empty()) {
    LOGE("Failed to get architecture name.");
    return false;
  }
  device_info_map.add(cppbor::Tstr("architecture"), cppbor::Tstr(arch_name));

  std::string device_name;
  if (!wvcdm::Properties::GetDeviceName(&device_name) || device_name.empty()) {
    LOGE("Failed to get device name.");
    return false;
  }
  device_info_map.add(cppbor::Tstr("device"), cppbor::Tstr(device_name));

  std::string product_name;
  if (!wvcdm::Properties::GetProductName(&product_name) ||
      product_name.empty()) {
    LOGE("Failed to get product name.");
    return false;
  }
  device_info_map.add(cppbor::Tstr("product"), cppbor::Tstr(product_name));

  std::string build_info;
  if (!wvcdm::Properties::GetBuildInfo(&build_info) || build_info.empty()) {
    LOGE("Failed to get build info.");
    return false;
  }
  device_info_map.add(cppbor::Tstr("fingerprint"), cppbor::Tstr(build_info));

  std::string oemcrypto_build_info;
  OEMCryptoResult result =
      crypto_interface_->GetOEMCryptoBuildInfo(oemcrypto_build_info);
  if (result != OEMCrypto_SUCCESS) {
    LOGE("Failed to get oemcrypto build info.");
    return false;
  }
  device_info_map.add(cppbor::Tstr("oemcrypto_build_info"),
                      cppbor::Tstr(oemcrypto_build_info));

  device_info = device_info_map.canonicalize().encode();
  return true;
}

bool WidevineProvisioner::GenerateProtectedData(
    bool test_mode, const std::vector<uint8_t>& endpoint_encryption_cert_chain,
    std::vector<uint8_t> bcc, std::vector<uint8_t>& protected_data) const {
  // Encrypt |signedMac| and |bcc_| with GEEK.
  std::vector<uint8_t> eek_pub;
  std::vector<uint8_t> eek_id;
  if (!ValidateAndExtractEekPubAndId(test_mode, endpoint_encryption_cert_chain,
                                     &eek_pub, &eek_id)) {
    LOGE("Failed to validate and extract the endpoint encryption key.");
    return false;
  }

  std::vector<uint8_t> ephemeralPrivKey(X25519_PRIVATE_KEY_LEN);
  std::vector<uint8_t> ephemeralPubKey(X25519_PUBLIC_VALUE_LEN);
  X25519_keypair(ephemeralPubKey.data(), ephemeralPrivKey.data());
  auto sessionKey = cppcose::x25519_HKDF_DeriveKey(
      ephemeralPubKey, ephemeralPrivKey, eek_pub, true /* senderIsA */);
  if (!sessionKey) {
    LOGE("Failed to derive the session key.");
    return false;
  }
  // Generate 4 bytes of random data as IV
  std::vector<uint8_t> iv(cppcose::kAesGcmNonceLength);
  if (RAND_bytes(iv.data(), iv.size()) != 1) {
    LOGE("Failed to generate a random nonce.");
    return false;
  }

  auto bcc_parse_result = cppbor::parse(bcc.data(), bcc.size());
  cppbor::Item* bcc_item = std::get<0>(bcc_parse_result).get();
  auto coseEncrypted = cppcose::constructCoseEncrypt(
      *sessionKey, iv,
      cppbor::Array()            // payload
          .add(cppbor::Array())  // Empty signedMac
          .add(std::move(bcc_item->clone()))
          .encode(),
      {},  // aad
      BuildCertReqRecipients(ephemeralPubKey, eek_id));
  if (!coseEncrypted) {
    LOGE("Failed to construct a COSE_Encrypt ProtectedData structure");
    return false;
  }

  protected_data = coseEncrypted->encode();
  return true;
}

bool WidevineProvisioner::ValidateAndExtractEekPubAndId(
    bool test_mode, const std::vector<uint8_t>& endpoint_encryption_cert_chain,
    std::vector<uint8_t>* eek_pub, std::vector<uint8_t>* eek_id) const {
  auto parse_result = cppbor::parse(endpoint_encryption_cert_chain);
  auto item = std::move(std::get<0>(parse_result));
  if (!item || !item->asArray()) {
    LOGE("Error parsing EEK chain: %s", std::get<2>(parse_result).c_str());
    return false;
  }

  const cppbor::Array* certArr = item->asArray();
  std::vector<uint8_t> lastPubKey;
  for (size_t i = 0; i < certArr->size(); ++i) {
    auto cosePubKey = cppcose::verifyAndParseCoseSign1(
        certArr->get(i)->asArray(), std::move(lastPubKey), {} /* AAD */);
    if (!cosePubKey) {
      LOGE("Failed to validate EEK chain: %s",
           cosePubKey.moveMessage().c_str());
      return false;
    }
    lastPubKey = *std::move(cosePubKey);

    // In prod mode the first pubkey should match a well-known Google public
    // key.
    if (!test_mode && i == 0) {
      auto parsedPubKey = cppcose::CoseKey::parse(lastPubKey);
      if (!parsedPubKey) {
        LOGE("%s", parsedPubKey.moveMessage().c_str());
        return false;
      }
      auto rawPubKey = parsedPubKey->getBstrValue(cppcose::CoseKey::PUBKEY_X);
      if (!rawPubKey) {
        LOGE("Key is missing required label 'PUBKEY_X'");
        return false;
      }
      std::vector<uint8_t> matcher(rawPubKey->data(),
                                   rawPubKey->data() + rawPubKey->size());

      if (std::find(std::begin(kAuthorizedEekRoots),
                    std::end(kAuthorizedEekRoots),
                    matcher) == std::end(kAuthorizedEekRoots)) {
        LOGE("Unrecognized root of EEK chain");
        return false;
      }
    }
  }

  auto eek = cppcose::CoseKey::parseX25519(lastPubKey, true /* requireKid */);
  if (!eek) {
    LOGE("Failed to get EEK: %s", eek.moveMessage().c_str());
    return false;
  }

  *eek_pub = eek->getBstrValue(cppcose::CoseKey::PUBKEY_X).value();
  *eek_id = eek->getBstrValue(cppcose::CoseKey::KEY_ID).value();
  return true;
}

cppbor::Array WidevineProvisioner::BuildCertReqRecipients(
    const std::vector<uint8_t>& pubkey, const std::vector<uint8_t>& kid) const {
  return cppbor::Array()           // Array of recipients
      .add(cppbor::Array()         // Recipient
               .add(cppbor::Map()  // Protected
                        .add(cppcose::ALGORITHM, cppcose::ECDH_ES_HKDF_256)
                        .canonicalize()
                        .encode())
               .add(cppbor::Map()  // Unprotected
                        .add(cppcose::COSE_KEY,
                             cppbor::Map()
                                 .add(cppcose::CoseKey::KEY_TYPE,
                                      cppcose::OCTET_KEY_PAIR)
                                 .add(cppcose::CoseKey::CURVE, cppcose::X25519)
                                 .add(cppcose::CoseKey::PUBKEY_X, pubkey)
                                 .canonicalize())
                        .add(cppcose::KEY_ID, kid)
                        .canonicalize())
               .add(cppbor::Null()));  // No ciphertext
}

void WidevineProvisioner::InitializeCryptoInterface() {
  std::string oemcrypto_path;
  if (!wvcdm::Properties::GetOEMCryptoPath(&oemcrypto_path)) {
    LOGE("Failed to get OEMCrypto path.");
  }
  LOGI("OEMCrypto path is %s", oemcrypto_path.c_str());

  crypto_interface_ = std::make_unique<OEMCryptoInterface>();
  if (!crypto_interface_->Init(oemcrypto_path)) {
    LOGE("Failed to initialize OEMCrypto interface.");
    crypto_interface_.reset(nullptr);
  }
}

}  // namespace widevine