// 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 "message_dumper.h"

#include "license_request.h"
#include "odk.h"
#include "odk_message.h"
#include "odk_serialize.h"
#include "odk_structs.h"
#include "odk_structs_priv.h"
#include "oec_device_features.h"
#include "test_base.h"

using video_widevine::License;
using video_widevine::LicenseRequest;
using video_widevine::SignedMessage;
using video_widevine::SignedProvisioningMessage;

namespace wvcdm {
namespace {
void DumpHeader(std::ofstream* out, const std::string& type) {
  const ::testing::TestInfo* const test_info =
      ::testing::UnitTest::GetInstance()->current_test_info();
  const std::string suite = test_info->test_case_name();
  const std::string name = test_info->name();
  std::string new_test_name = suite + "_" + name;
  // Replace the slashes with underscores so we can use it as a name again.
  std::replace(new_test_name.begin(), new_test_name.end(), '/', '_');
  *out << "\nTEST_F(ODKGolden" << type << "V" << ODK_MAJOR_VERSION << ", "
       << new_test_name << ") {\n";
}

void DumpHex(std::ofstream* out, const std::string& name,
             const std::string& value) {
  *out << "const uint8_t " << name << "_raw[] = {\n";
  *out << "  ";
  for (unsigned int i = 0; i < value.length(); i++) {
    if ((i > 0) && (i % 10 == 0)) *out << "\n     ";
    uint8_t c = value[i];
    *out << "0x" << std::hex << std::setw(2) << std::setfill('0') << int(c)
         << ", ";
  }
  *out << "\n};\n";
  *out << name << "_ = std::string (\n"
       << "  reinterpret_cast<const char *>(" << name << "_raw), \n"
       << "      sizeof(" << name << "_raw));\n";
  *out << std::dec;  // Turn off hex when we're done.
}
}  // namespace

std::ofstream MessageDumper::license_file;
std::ofstream MessageDumper::renewal_file;
std::ofstream MessageDumper::provision_file;

void MessageDumper::SetUp() {
  LOGD("Creating golden data files for ODK golden data tests.");
  license_file.open("license_data.cpp");
  if (!license_file) LOGE("Could not open dump file license_data.cpp");
  renewal_file.open("renewal_data.cpp");
  if (!renewal_file) LOGE("Could not open dump file renewal_data.cpp");
  provision_file.open("provision_data.cpp");
  if (!provision_file) LOGE("Could not open dump file provision_data.cpp");
}

void MessageDumper::TearDown() {
  LOGD("Closing golden data files.");
  license_file.close();
  renewal_file.close();
  provision_file.close();
}

void MessageDumper::DumpLicenseRequest(const CdmKeyRequest& request) {
  SignedMessage signed_message;
  DumpHeader(&license_file, "License");
  EXPECT_TRUE(signed_message.ParseFromString(request.message));
  EXPECT_TRUE(signed_message.has_oemcrypto_core_message());
  DumpHex(&license_file, "core_request",
          signed_message.oemcrypto_core_message());
  // Since this is run within a test, we can also verify that the
  // request is valid.
  video_widevine::LicenseRequest license_request;
  EXPECT_TRUE(license_request.ParseFromString(signed_message.msg()));
}

void MessageDumper::DumpLicense(const std::string& response) {
  SignedMessage signed_response;
  EXPECT_TRUE(signed_response.ParseFromString(response));
  EXPECT_TRUE(signed_response.has_oemcrypto_core_message());
  DumpHex(&license_file, "core_response",
          signed_response.oemcrypto_core_message());

  video_widevine::License license;
  EXPECT_TRUE(license.ParseFromString(signed_response.msg()));
  DumpHex(&license_file, "serialized_license", signed_response.msg());

  std::string message =
      signed_response.oemcrypto_core_message() + signed_response.msg();
  ODK_Message odk_msg = ODK_Message_Create(
      reinterpret_cast<uint8_t*>(const_cast<char*>(message.c_str())),
      message.length());
  ODK_Message_SetSize(&odk_msg,
                      signed_response.oemcrypto_core_message().length());
  ODK_ParsedLicense odk_parsed_license = {};
  ODK_LicenseResponse odk_license_response = {};
  odk_license_response.parsed_license = &odk_parsed_license;
  Unpack_ODK_LicenseResponse(&odk_msg, &odk_license_response);
  EXPECT_EQ(ODK_Message_GetStatus(&odk_msg), MESSAGE_STATUS_OK);
  // Valid hash is only needed for v16 messages.
  std::string hash(ODK_SHA256_HASH_SIZE, ' ');
  DumpHex(&license_file, "core_request_sha256", hash);
  license_file << "  nonce_required_ = "
               << (odk_parsed_license.nonce_required ? "true" : "false")
               << ";\n";
  license_file << "  RunTest();\n";
  license_file << "}\n\n";
}

void MessageDumper::DumpRenewalRequest(const CdmKeyRequest& request) {
  DumpHeader(&renewal_file, "Renewal");
  SignedMessage signed_message;
  EXPECT_TRUE(signed_message.ParseFromString(request.message));
  EXPECT_TRUE(signed_message.has_oemcrypto_core_message());
  DumpHex(&renewal_file, "core_request",
          signed_message.oemcrypto_core_message());

  video_widevine::LicenseRequest renewal_request;
  EXPECT_TRUE(renewal_request.ParseFromString(signed_message.msg()));
}

void MessageDumper::DumpRenewal(const std::string& response) {
  SignedMessage signed_response;
  EXPECT_TRUE(signed_response.ParseFromString(response))
      << "Response = " << wvutil::b2a_hex(response);
  EXPECT_TRUE(signed_response.has_oemcrypto_core_message());
  DumpHex(&renewal_file, "core_response",
          signed_response.oemcrypto_core_message());

  video_widevine::License renewal;
  EXPECT_TRUE(renewal.ParseFromString(signed_response.msg()));
  DumpHex(&renewal_file, "renewal", signed_response.msg());

  std::string message =
      signed_response.oemcrypto_core_message() + signed_response.msg();
  ODK_Message odk_msg = ODK_Message_Create(
      reinterpret_cast<uint8_t*>(const_cast<char*>(message.c_str())),
      message.length());
  ODK_Message_SetSize(&odk_msg,
                      signed_response.oemcrypto_core_message().length());
  ODK_RenewalResponse odk_renewal_response = {};
  Unpack_ODK_RenewalResponse(&odk_msg, &odk_renewal_response);
  EXPECT_EQ(ODK_Message_GetStatus(&odk_msg), MESSAGE_STATUS_OK);
  renewal_file << "  renewal_duration_seconds_ = "
               << odk_renewal_response.renewal_duration_seconds << ";\n";
  renewal_file << "  RunTest();\n";
  renewal_file << "}\n\n";
}

void MessageDumper::DumpProvisioningRequest(
    const CdmProvisioningRequest& request) {
  if (wvoec::global_features.derive_key_method ==
      wvoec::DeviceFeatures::TEST_PROVISION_40) {
    LOGD("Provisioning 4.0 does not have a v17 or v18 core message.");
  } else {
    DumpHeader(&provision_file, "Provision");
    SignedProvisioningMessage signed_message;
    EXPECT_TRUE(signed_message.ParseFromString(request))
        << "Request = " << wvutil::b2a_hex(request);
    EXPECT_TRUE(signed_message.has_oemcrypto_core_message());
    DumpHex(&provision_file, "core_request",
            signed_message.oemcrypto_core_message());
  }
}

void MessageDumper::DumpProvisioning(const CdmProvisioningResponse& response) {
  if (wvoec::global_features.derive_key_method ==
      wvoec::DeviceFeatures::TEST_PROVISION_40) {
    LOGD("Provisioning 4.0 does not have a core message.");
  } else {
    SignedProvisioningMessage signed_response;
    if (!signed_response.ParseFromString(response)) {
      // A binary provisioning response is buried within a json structure.
      std::string extracted_message;
      EXPECT_TRUE(CertificateProvisioning::ExtractAndDecodeSignedMessage(
          response, &extracted_message));
      EXPECT_TRUE(signed_response.ParseFromString(extracted_message));
    }
    EXPECT_TRUE(signed_response.has_oemcrypto_core_message());
    DumpHex(&provision_file, "core_response",
            signed_response.oemcrypto_core_message());
    DumpHex(&provision_file, "provisioning_response",
            signed_response.message());
    // The choice of ECC or RSA key is decided at the server, based on
    // information in the DCSL. We can only reproduce this by looking
    // at the current response.
    std::string message =
        signed_response.oemcrypto_core_message() + signed_response.message();
    ODK_Message odk_msg = ODK_Message_Create(
        reinterpret_cast<uint8_t*>(const_cast<char*>(message.c_str())),
        message.length());
    ODK_Message_SetSize(&odk_msg,
                        signed_response.oemcrypto_core_message().length());
    ODK_ParsedProvisioning odk_parsed_response;
    ODK_ProvisioningResponse provisioning_response;
    provisioning_response.parsed_provisioning = &odk_parsed_response;
    Unpack_ODK_ProvisioningResponse(&odk_msg, &provisioning_response);
    EXPECT_EQ(ODK_Message_GetStatus(&odk_msg), MESSAGE_STATUS_OK);
    provision_file << "  device_key_type_ = "
                   << ((odk_parsed_response.key_type ==
                        OEMCrypto_RSA_Private_Key)
                           ? "OEMCrypto_RSA_Private_Key;\n"
                           : "OEMCrypto_ECC_Private_Key;\n");
    provision_file << "  RunTest();\n";
    provision_file << "}\n\n";
  }
}
}  // namespace wvcdm