// 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 "WidevineOemcryptoInterface.h"
#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);
}

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, result = %d", result);
    return false;
  }

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

  return true;
}

bool WidevineProvisioner::TryAddVerifiedDeviceInfo(
    cppbor::Map& device_info_map) {
  std::vector<uint8_t> verified_device_info;
  OEMCryptoResult result =
      crypto_interface_->GetVerifiedDeviceInformation(verified_device_info);
  if (result == OEMCrypto_ERROR_NOT_IMPLEMENTED) {
    // OEMCrypto v17 and earlier doesn't support GetDeviceInformation()
    LOGI("OEMCrypto_GetDeviceInformation is not implemented.");
    return true;
  }
  if (result != OEMCrypto_SUCCESS) {
    LOGE("Failed to get verified device information, result = %d", result);
    return false;
  }
  auto [parsed, _, err] = cppbor::parse(
      reinterpret_cast<const uint8_t*>(verified_device_info.data()),
      verified_device_info.size());
  if (!parsed || !parsed->asMap()) {
    LOGE("Failed to parse the verified device info cbor: %s", err.c_str());
    return false;
  }
  const cppbor::Map* verified_device_info_map = parsed->asMap();
  for (size_t i = 0; i < verified_device_info_map->size(); i++) {
    auto& [key_item, value_item] = (*verified_device_info_map)[i];
    LOGI("Found device info %s", key_item->asTstr()->value().data());
    if (value_item != nullptr && value_item->asTstr() != nullptr &&
        value_item->asTstr()->value().empty()) {
      LOGI("Value is empty. Skip");
      continue;
    }
    device_info_map.add(key_item->clone(), value_item->clone());
  }
  return true;
}

bool WidevineProvisioner::GetDeviceInfoCommon(cppbor::Map& device_info_map) {
  if (!TryAddVerifiedDeviceInfo(device_info_map)) return false;
  // Add device information from OS properties if the verified device info is
  // not present
  if (device_info_map.get("brand") == nullptr ||
      device_info_map.get("brand")->asTstr()->value().empty()) {
    std::string brand_name;
    if (!wvcdm::Properties::GetBrandName(&brand_name) || brand_name.empty()) {
      LOGE("Failed to get brand name.");
      return false;
    }
    device_info_map.add(cppbor::Tstr("brand"), cppbor::Tstr(brand_name));
    LOGI("use OS property brand: %s", brand_name.data());
  } else {
    LOGI("use verified brand: %s",
         device_info_map.get("brand")->asTstr()->value().data());
  }

  if (device_info_map.get("manufacturer") == nullptr ||
      device_info_map.get("manufacturer")->asTstr()->value().empty()) {
    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));
    LOGI("use OS property manufacturer: %s", company_name.data());
  } else {
    LOGI("use verified manufacture: %s",
         device_info_map.get("manufacturer")->asTstr()->value().data());
  }

  if (device_info_map.get("model") == nullptr ||
      device_info_map.get("model")->asTstr()->value().empty()) {
    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));
    LOGI("use OS property model: %s", model_name.data());
  } else {
    LOGI("use verified model: %s",
         device_info_map.get("model")->asTstr()->value().data());
  }

  if (device_info_map.get("device") == nullptr ||
      device_info_map.get("device")->asTstr()->value().empty()) {
    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));
    LOGI("use OS property device: %s", device_name.data());
  } else {
    LOGI("use verified device: %s",
         device_info_map.get("device")->asTstr()->value().data());
  }

  if (device_info_map.get("product") == nullptr ||
      device_info_map.get("product")->asTstr()->value().empty()) {
    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));
    LOGI("use OS property product: %s", product_name.data());
  } else {
    LOGI("use verified product: %s",
         device_info_map.get("product")->asTstr()->value().data());
  }

  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 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));
  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"));
  device_info_map.add(cppbor::Tstr("version"), cppbor::Uint(2));
  if (!GetDeviceInfoCommon(device_info_map)) return false;
  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
}

bool WidevineProvisioner::GetDeviceInfoV2(cppbor::Map& device_info_map) {
  if (!GetDeviceInfoCommon(device_info_map)) return false;
  device_info_map.canonicalize();
  return true;
}

bool WidevineProvisioner::GenerateCertificateRequestV2(
    const std::vector<uint8_t>& challenge, std::vector<uint8_t>* csr) {
  if (csr == nullptr) {
    LOGE("CSR is null.");
    return false;
  }
  // Prepare BCC
  std::vector<uint8_t> bcc;
  OEMCryptoResult result = crypto_interface_->GetBcc(bcc);
  if (result != OEMCrypto_SUCCESS) {
    LOGE("Failed to get BCC, result = %d", result);
    return false;
  }
  // Prepare device info
  auto device_info_map = cppbor::Map();
  if (!GetDeviceInfoV2(device_info_map)) {
    LOGE("Failed to get device_info.");
    return false;
  }
  // Prepare signed CSR payload
  auto device_info = device_info_map.encode();
  std::vector<uint8_t> signed_csr_payload;
  result = crypto_interface_->GetSignedCsrPayload(challenge, device_info,
                                                  signed_csr_payload);
  if (result != OEMCrypto_SUCCESS) {
    LOGE("Failed to get the signed CSR payload, result = %d", result);
    return false;
  }
  // https://source.corp.google.com/android-internal/hardware/interfaces/security/rkp/aidl/android/hardware/security/keymint/IRemotelyProvisionedComponent.aidl
  *csr = cppbor::Array()
             .add(1 /* version */)
             .add(cppbor::Map() /* UdsCerts */)
             .add(cppbor::EncodedItem(std::move(bcc)))
             .add(cppbor::EncodedItem(std::move(signed_csr_payload)))
             .encode();
  return true;
}

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