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

#include "provisioning_holder.h"

#include <gtest/gtest.h>

#include "cdm_engine.h"
#include "config_test_env.h"

#include "log.h"
#include "message_dumper.h"
#include "oec_device_features.h"
#include "test_printers.h"
#include "test_sleep.h"
#include "url_request.h"

namespace wvcdm {

void ProvisioningHolder::Provision(CdmCertificateType cert_type,
                                   bool binary_provisioning) {
  CdmProvisioningRequest request;
  std::string provisioning_server_url;
  std::string cert_authority;

  CdmSessionId session_id;

  if (cert_type == kCertificateX509) {
    cert_authority = "cast.google.com";
  }
  LOGV("Provision with type %s.", CdmCertificateTypeToString(cert_type));
  CdmResponseType result(CERT_PROVISIONING_NONCE_GENERATION_ERROR);
  // Get a provisioning request. We might need one retry if there is a nonce
  // flood failure.
  for (int i = 0; i < 2 && result == CERT_PROVISIONING_NONCE_GENERATION_ERROR;
       i++) {
    result = cdm_engine_->GetProvisioningRequest(
        cert_type, cert_authority, provisioning_service_certificate_,
        kLevelDefault, &request, &provisioning_server_url);
    if (result == CERT_PROVISIONING_NONCE_GENERATION_ERROR) {
      wvutil::TestSleep::Sleep(2);
    }
  }
  ASSERT_EQ(NO_ERROR, result);
  LOGV("cert_authority = %s", cert_authority.c_str());

  if (binary_provisioning) {
    request = wvutil::Base64SafeEncodeNoPad(request);
  }
  if (config_.dump_golden_data()) {
    std::vector<uint8_t> binary_request = wvutil::Base64SafeDecode(request);
    CdmProvisioningRequest binary_request_string(binary_request.begin(),
                                                 binary_request.end());
    MessageDumper::DumpProvisioningRequest(binary_request_string);
  }
  LOGV("Provisioning request: req = %s", request.c_str());

  // Ignore URL provided by CdmEngine.  Use ours, as configured
  // for test vs. production server.
  provisioning_server_url.assign(provisioning_server_url_);

  // Make request.
  UrlRequest url_request(provisioning_server_url);
  if (!url_request.is_connected()) {
    LOGE("Failed to connect to provisioning server: url = %s",
         provisioning_server_url.c_str());
  }
  url_request.PostCertRequestInQueryString(request);

  // Receive and parse response.
  ASSERT_NO_FATAL_FAILURE(url_request.AssertOkResponse(&response_))
      << "Failed to fetch provisioning response. "
      << DumpProvAttempt(request, response_, cert_type);

  if (binary_provisioning) {
    // extract provisioning response from received message
    // Extracts signed response from JSON string, result is serialized
    // protobuf.
    std::string protobuf_response;
    const bool extract_ok = ExtractSignedMessage(response_, &protobuf_response);
    ASSERT_TRUE(extract_ok) << "Failed to extract signed serialized "
                               "response from JSON response";
    LOGV("Extracted response message: \n%s\n", protobuf_response.c_str());

    ASSERT_FALSE(protobuf_response.empty())
        << "Protobuf response is unexpectedly empty";

    // base64 decode response to yield binary protobuf
    const std::vector<uint8_t> response_vec(
        wvutil::Base64SafeDecode(protobuf_response));
    ASSERT_FALSE(response_vec.empty())
        << "Failed to decode base64 of response: response = "
        << protobuf_response;

    response_.assign(response_vec.begin(), response_vec.end());
  }

  ASSERT_EQ(NO_ERROR,
            cdm_engine_->HandleProvisioningResponse(
                response_, kLevelDefault, &certificate_, &wrapped_key_))
      << (binary_provisioning ? "Binary provisioning failed. "
                              : "Non-binary provisioning failed. ")
      << DumpProvAttempt(request, response_, cert_type);
  if (config_.dump_golden_data()) {
    MessageDumper::DumpProvisioning(response_);
  }
}

bool ProvisioningHolder::ExtractSignedMessage(const std::string& response,
                                              std::string* result) {
  static const std::string kMessageStart = "\"signedResponse\": \"";
  static const std::string kMessageEnd = "\"";
  std::string response_string;
  size_t start = response.find(kMessageStart);

  if (start == response.npos) {
    // Assume serialized protobuf message.
    result->assign(response);
  } else {
    // Assume JSON-wrapped protobuf.
    size_t end = response.find(kMessageEnd, start + kMessageStart.length());
    if (end == response.npos) {
      LOGE("ExtractSignedMessage cannot locate end substring");
      result->clear();
      return false;
    }
    size_t result_string_size = end - start - kMessageStart.length();
    result->assign(response, start + kMessageStart.length(),
                   result_string_size);
  }

  if (result->empty()) {
    LOGE("ExtractSignedMessage: Response message is empty");
    return false;
  }
  return true;
}

std::string ProvisioningHolder::DumpProvAttempt(const std::string& request,
                                                const std::string& response,
                                                CdmCertificateType cert_type) {
  std::stringstream info;
  info << "Cert Type: ";
  PrintTo(cert_type, &info);
  info << "\n";
  info << "Provisioning url: " << provisioning_server_url_ << "\n";
  info << "Provisioning Request: " << request << "\n\n";
  info << "Provisioning Response: " << response << "\n\n";
  std::string system_id;
  cdm_engine_->QueryStatus(kLevelDefault, QUERY_KEY_SYSTEM_ID, &system_id);
  info << "system id: " << system_id << "\n";
  if (wvoec::global_features.derive_key_method ==
      wvoec::DeviceFeatures::TEST_PROVISION_30) {
    std::vector<uint8_t> cert;
    size_t cert_length = 0;
    OEMCryptoResult result = OEMCrypto_GetOEMPublicCertificate(
        cert.data(), &cert_length, kLevelDefault);
    if (result == OEMCrypto_ERROR_SHORT_BUFFER) {
      cert.resize(cert_length);
      result = OEMCrypto_GetOEMPublicCertificate(cert.data(), &cert_length,
                                                 kLevelDefault);
    }
    if (result != OEMCrypto_SUCCESS) {
      info << "--- ERROR GETTING CERT. result=" << result;
    } else {
      info << "OEM Cert = (len=" << cert_length << ") "
           << wvutil::unlimited_b2a_hex(cert);
    }
  }
  if (wvoec::global_features.derive_key_method ==
      wvoec::DeviceFeatures::TEST_PROVISION_40) {
    std::vector<uint8_t> bcc;
    size_t bcc_length = 0;
    std::vector<uint8_t> signature;
    size_t signature_length = 0;
    OEMCryptoResult result = OEMCrypto_GetBootCertificateChain(
        bcc.data(), &bcc_length, signature.data(), &signature_length);
    if (result == OEMCrypto_ERROR_SHORT_BUFFER) {
      bcc.resize(bcc_length);
      signature.resize(signature_length);
      result = OEMCrypto_GetBootCertificateChain(
          bcc.data(), &bcc_length, signature.data(), &signature_length);
    }
    if (result != OEMCrypto_SUCCESS) {
      info << "--- ERROR GETTING BCC. result=" << result;
    } else {
      bcc.resize(bcc_length);
      info << "BCC = (len=" << bcc_length << ") "
           << wvutil::unlimited_b2a_hex(bcc) << "\n"
           << "Additional Sig = (len=" << signature_length << ") "
           << wvutil::unlimited_b2a_hex(signature) << "\n";
    }
  }
  return info.str();
}
}  // namespace wvcdm