// 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) {
  ASSERT_NO_FATAL_FAILURE(GenerateRequest(cert_type, binary_provisioning))
      << "Failed to generate request";
  ASSERT_NO_FATAL_FAILURE(FetchResponse()) << "Failed to fetch response";
  ASSERT_NO_FATAL_FAILURE(LoadResponse(binary_provisioning))
      << "Failed to load response";
}

void ProvisioningHolder::GenerateRequest(CdmCertificateType cert_type,
                                         bool binary_provisioning) {
  // Preconditions.
  ASSERT_NE(cdm_engine_, nullptr) << "CdmEngine instance not set";

  // Set cert authority if using X.509.
  std::string cert_authority;
  if (cert_type == kCertificateX509) {
    cert_authority = "cast.google.com";
  }
  LOGV("Provision with type %s.", CdmCertificateTypeToString(cert_type));
  LOGV("cert_authority = %s", cert_authority.c_str());

  CdmResponseType status(CERT_PROVISIONING_NONCE_GENERATION_ERROR);
  CdmProvisioningRequest cdm_prov_request;
  std::string provisioning_server_url_unused;
  // Get a provisioning request. We might need one retry if there is a nonce
  // flood failure.
  for (int i = 0; i < 2 && status == CERT_PROVISIONING_NONCE_GENERATION_ERROR;
       i++) {
    status = cdm_engine_->GetProvisioningRequest(
        cert_type, cert_authority, provisioning_service_certificate_,
        kLevelDefault, &cdm_prov_request, &provisioning_server_url_unused);
    if (status == CERT_PROVISIONING_NONCE_GENERATION_ERROR) {
      wvutil::TestSleep::Sleep(2);
    }
  }
  ASSERT_EQ(status, NO_ERROR) << "Failed to generate provisioning request";

  // |binary_provisioning| implies the CdmEngine is using binary
  // provisioning messages; however, ProvisioningHolder can only
  // operate using Base64 requests.
  if (binary_provisioning) {
    request_ = wvutil::Base64SafeEncodeNoPad(cdm_prov_request);
  } else {
    request_ = std::move(cdm_prov_request);
  }
  cert_type_ = cert_type;

  if (config_.dump_golden_data()) {
    const std::vector<uint8_t> binary_request =
        wvutil::Base64SafeDecode(request_);
    const CdmProvisioningRequest binary_request_string(binary_request.begin(),
                                                       binary_request.end());
    MessageDumper::DumpProvisioningRequest(binary_request_string);
  }

  if (log_core_message_) {
    const std::vector<uint8_t> binary_request =
        wvutil::Base64SafeDecode(request_);
    const CdmProvisioningRequest binary_request_string(binary_request.begin(),
                                                       binary_request.end());
    MessageDumper::PrintProvisioningRequest(binary_request_string);
  }
}

void ProvisioningHolder::FetchResponse() {
  // Preconditions.
  ASSERT_NE(cdm_engine_, nullptr) << "CdmEngine instance not set";
  ASSERT_FALSE(request_.empty()) << "No request was set";
  ASSERT_FALSE(provisioning_server_url_.empty())
      << "Test config is missing provisioning URL";

  UrlRequest url_request(provisioning_server_url_);
  ASSERT_TRUE(url_request.is_connected())
      << "Failed to connect to provisoining server: "
      << provisioning_server_url_;
  ASSERT_TRUE(url_request.PostCertRequestInQueryString(request_));
  std::string response;
  ASSERT_NO_FATAL_FAILURE(url_request.AssertOkResponseWithRetry(&response))
      << "Failed to fetch provisioning response. "
      << DumpProvAttempt(request_, response, cert_type_);
  ASSERT_FALSE(response.empty()) << "Missing response";
  response_ = std::move(response);
}

void ProvisioningHolder::LoadResponse(bool binary_provisioning) {
  // Preconditions.
  ASSERT_FALSE(response_.empty()) << "No response was fetched";

  std::string cdm_prov_response;
  if (binary_provisioning) {
    // CDM is expecting the response to be in binary form, response
    // must be extracted and decoded.
    std::string base_64_response;
    ASSERT_TRUE(ExtractSignedMessage(response_, &base_64_response))
        << "Failed to extract signed serialized response from JSON response";
    ASSERT_FALSE(base_64_response.empty())
        << "Base64 encoded provisioning response is unexpectedly empty";
    LOGV("Extracted response message: \n%s\n", base_64_response.c_str());

    const std::vector<uint8_t> response_vec =
        wvutil::Base64SafeDecode(base_64_response);
    ASSERT_FALSE(response_vec.empty())
        << "Failed to decode base64 response: " << base_64_response;
    cdm_prov_response.assign(response_vec.begin(), response_vec.end());
  } else {
    cdm_prov_response = response_;
  }

  const CdmResponseType status = cdm_engine_->HandleProvisioningResponse(
      cdm_prov_response, kLevelDefault, &certificate_, &wrapped_key_);
  ASSERT_EQ(status, NO_ERROR)
      << (binary_provisioning ? "Binary provisioning failed. "
                              : "Non-binary provisioning failed. ")
      << DumpProvAttempt(request_, response_, cert_type_);
  if (config_.dump_golden_data()) {
    MessageDumper::DumpProvisioning(response_);
  }
  if (log_core_message_) MessageDumper::PrintProvisioningResponse(response_);
}

CdmResponseType ProvisioningHolder::LoadResponseReturnStatus(
    bool binary_provisioning) {
  // Preconditions.
  if (response_.empty()) {
    ADD_FAILURE() << "No response was fetched";
    return CdmResponseType(UNKNOWN_ERROR);
  }

  std::string cdm_prov_response;
  if (binary_provisioning) {
    // CDM is expecting the response to be in binary form, response
    // must be extracted and decoded.
    std::string base_64_response;
    if (!ExtractSignedMessage(response_, &base_64_response)) {
      ADD_FAILURE()
          << "Failed to extract signed serialized response from JSON response";
      return CdmResponseType(UNKNOWN_ERROR);
    }
    if (base_64_response.empty()) {
      ADD_FAILURE()
          << "Base64 encoded provisioning response is unexpectedly empty";
      return CdmResponseType(UNKNOWN_ERROR);
    }
    LOGV("Extracted response message: \n%s\n", base_64_response.c_str());

    const std::vector<uint8_t> response_vec =
        wvutil::Base64SafeDecode(base_64_response);
    if (response_vec.empty()) {
      ADD_FAILURE() << "Failed to decode base64 response: " << base_64_response;
      return CdmResponseType(UNKNOWN_ERROR);
    }
    cdm_prov_response.assign(response_vec.begin(), response_vec.end());
  } else {
    cdm_prov_response = response_;
  }
  // HandleProvisioningResponse() may or may not succeed,
  // left to caller to determine if this is considered a
  // test failure.
  const CdmResponseType status = cdm_engine_->HandleProvisioningResponse(
      cdm_prov_response, kLevelDefault, &certificate_, &wrapped_key_);
  if (status == NO_ERROR) {
    // Only dump data if successful.
    if (config_.dump_golden_data()) MessageDumper::DumpProvisioning(response_);
    if (log_core_message_) MessageDumper::PrintProvisioningResponse(response_);
  }
  return status;
}

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) const {
  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