// Copyright 2024 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
// This tool extracts BCC by calling OEMCrypto APIs and generates a CSR file in
// JSON format, which can be handled by CE CDM wv_upload_tool.py.
#include <fstream>
#include <iostream>
#include <map>
#include <regex>
#include <string>
#include <vector>

#include "OEMCryptoCENC.h"
#include "string_conversions.h"

namespace {
// Make and Model for system ID resolution.
const std::string kDeviceMake = "widevine_test";
const std::string kDeviceModel = "prov4";

// Informative fields.
const std::string kDeviceArchitecture = "x86_64";
const std::string kDeviceName = "prov40 test client";
const std::string kDeviceProduct = "prov40 test";
const std::string kDeviceBuildInfo = "prov40 test build";

// == Utils ==

std::string StringMapToJson(
    const std::map<std::string, std::string>& string_map) {
  std::string json = "{";
  for (const auto& value_pair : string_map) {
    std::string escaped_value =
        std::regex_replace(value_pair.second, std::regex("\""), "\\\"");
    json.append("\"" + value_pair.first + "\": " + "\"" + escaped_value +
                "\",");
  }
  json.resize(json.size() - 1);  // Remove the last comma.
  json.append("}");
  return json;
}

// == Primary ==

bool GetBccAndBuildInfo(std::vector<uint8_t>* bcc,
                        std::string* oemcrypto_build_info) {
  // Step 1:  Initialize.
  OEMCryptoResult result = OEMCrypto_Initialize();
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to initialize: result = " << result << std::endl;
    return false;
  }

  // Step 2: Get BCC.
  const OEMCrypto_ProvisioningMethod method = OEMCrypto_GetProvisioningMethod();
  if (method != OEMCrypto_BootCertificateChain) {
    std::cerr << "ProvisioningMethod is not BCC type: method = ";
    std::cerr << method << std::endl;
    OEMCrypto_Terminate();
    return false;
  }

  bcc->resize(0);
  size_t bcc_size = 0;
  std::vector<uint8_t> additional_signature;  // It should be empty.
  size_t additional_signature_size = 0;
  result = OEMCrypto_GetBootCertificateChain(bcc->data(), &bcc_size,
                                             additional_signature.data(),
                                             &additional_signature_size);
  if (additional_signature_size != 0) {
    std::cerr << "The additional_signature_size required by OEMCrypto is "
              << additional_signature_size
              << ", while it is expected to be zero." << std::endl;
    OEMCrypto_Terminate();
    return false;
  }
  if (result == OEMCrypto_ERROR_SHORT_BUFFER) {
    bcc->resize(bcc_size);
    additional_signature.resize(additional_signature_size);
    result = OEMCrypto_GetBootCertificateChain(bcc->data(), &bcc_size,
                                               additional_signature.data(),
                                               &additional_signature_size);
  }
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to get BCC: result = " << result << std::endl;
    OEMCrypto_Terminate();
    return false;
  }
  bcc->resize(bcc_size);

  // Step 3: Get oemcrypto build info.
  oemcrypto_build_info->resize(0);
  size_t oemcrypto_build_info_size = 0;
  result = OEMCrypto_BuildInformation(oemcrypto_build_info->data(),
                                      &oemcrypto_build_info_size);
  if (result == OEMCrypto_ERROR_SHORT_BUFFER) {
    oemcrypto_build_info->resize(oemcrypto_build_info_size);
    result = OEMCrypto_BuildInformation(oemcrypto_build_info->data(),
                                        &oemcrypto_build_info_size);
  }
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to get build information: result = " << result
              << std::endl;
    OEMCrypto_Terminate();
    return false;
  }
  oemcrypto_build_info->resize(oemcrypto_build_info_size);

  // Step 4: Cleanup.
  result = OEMCrypto_Terminate();
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to terminate: result = " << result << std::endl;
    return false;
  }
  return true;
}

bool GenerateBccRecord(const std::vector<uint8_t>& bcc,
                       const std::string& oemcrypto_build_info,
                       std::string* bcc_record) {
  std::map<std::string, std::string> record;
  record["company"] = kDeviceMake;
  record["model"] = kDeviceModel;

  record["architecture"] = kDeviceArchitecture;
  record["name"] = kDeviceName;
  record["product"] = kDeviceProduct;
  record["build_info"] = kDeviceBuildInfo;
  record["bcc"] = wvutil::Base64Encode(bcc);
  record["oemcrypto_build_info"] = oemcrypto_build_info;

  const std::string record_json = StringMapToJson(record);
  bcc_record->assign(record_json.begin(), record_json.end());
  return true;
}

bool OutputBccRecord(const std::string& path, const std::string& record) {
  std::cout << "Writing BCC record to file " << path << std::endl;
  std::cout << record << std::endl;
  std::ofstream out(path);
  if (out) out << record;
  if (out.bad()) {
    std::cerr << "Failed to write BCC record to file " << path << std::endl;
    return false;
  }
  return true;
}
}  // namespace

int main(int argc, char** argv) {
  if (argc != 2) {
    std::cerr << "Usage: " << argv[0] << " <output JSON filename>" << std::endl;
    return 1;
  }
  const std::string bcc_path = argv[1];

  std::vector<uint8_t> bcc;
  std::string oemcrypto_build_info;
  if (!GetBccAndBuildInfo(&bcc, &oemcrypto_build_info)) {
    std::cerr << "Failed to get BCC or OEMCrypto build info" << std::endl;
    return 1;
  }
  std::string bcc_record;
  if (!GenerateBccRecord(bcc, oemcrypto_build_info, &bcc_record)) {
    std::cerr << "Failed to generate BCC record" << std::endl;
    return 1;
  }
  if (!OutputBccRecord(bcc_path, bcc_record)) {
    std::cerr << "Failed to output BCC record" << std::endl;
    return 1;
  }
  return 0;
}