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

#include "Utils.h"

#include <unordered_map>
#include "WVCDMSingleton.h"
#include "log.h"

#include <json/json.h>
#include <sys/stat.h>

namespace {
// Please keep the order stable.
// If APIs are renamed please keep the old enum and add new enum to bottom.
enum CryptoSessionApi : int32_t {
  API_UNKNOWN = 0,
  GET_PROVISIONING_METHOD = 1,
  GET_TOKEN_FROM_KEYBOX = 2,
  GET_TOKEN_FROM_OEM_CERT = 3,
  GET_PROVISIONING_TOKEN = 4,
  GET_INTERNAL_DEVICE_UNIQUE_ID = 5,
  GET_EXTERNAL_DEVICE_UNIQUE_ID = 6,
  GET_PROVISIONING_ID = 7,
  OPEN = 8,
  PREPARE_AND_SIGN_LICENSE_REQUEST = 9,
  USE_SECONDARY_KEY = 10,
  LOAD_LICENSE = 11,
  PREPARE_AND_SIGN_RENEWAL_REQUEST = 12,
  LOAD_RENEWAL = 13,
  PREPARE_AND_SIGN_PROVISIONING_REQUEST = 14,
  LOAD_ENTITLED_CONTENT_KEYS = 15,
  LOAD_CERTIFICATE_PRIVATE_KEY = 16,
  GET_BOOT_CERTIFICATE_CHAIN = 17,
  GENERATE_CERTIFICATE_KEY_PAIR = 18,
  LOAD_OEM_CERTIFICATE_PRIVATE_KEY = 19,
  SELECT_KEY = 20,
  GENERATE_DERIVED_KEYS = 21,
  GENERATE_RSA_SIGNATURE = 22,
  DECRYPT = 23,
  DEACTIVATE_USAGE_INFORMATION = 24,
  GENERATE_USAGE_REPORT = 25,
  GENERATE_NONCE = 26,
  LOAD_PROVISIONING = 27,
  GET_HDCP_CAPABILITIES = 28,
  GET_RANDOM = 29,
  GET_NUMBER_OF_OPEN_SESSIONS = 30,
  GET_MAX_NUMBER_OF_SESSIONS = 31,
  GET_SRM_VERSION = 32,
  SET_DECRYPT_HASH = 33,
  GET_DECRYPT_HASH_ERROR = 34,
  GENERIC_ENCRYPT = 35,
  GENERIC_DECRYPT = 36,
  GENERIC_SIGN = 37,
  GENERIC_VERIFY = 38,
  CREATE_USAGE_TABLE_HEADER = 39,
  LOAD_USAGE_TABLE_HEADER = 40,
  SHRINK_USAGE_TABLE_HEADER = 41,
  CREATE_USAGE_ENTRY = 42,
  LOAD_USAGE_ENTRY = 43,
  UPDATE_USAGE_ENTRY = 44,
  MOVE_USAGE_ENTRY = 45,
  SET_DEBUG_IGNORE_KEYBOX_COUNT = 46,
  SET_ALLOW_TEST_KEYBOX = 47,
  PREPARE_OTA_PROVISIONING_REQUEST = 48,
  LOAD_OTA_PROVISIONING = 49,
  // insert new or renamed methods' enum below
};

static std::unordered_map<std::string, CryptoSessionApi> const table = {
  {"ApiUnknown", API_UNKNOWN},
  {"GetProvisioningMethod", GET_PROVISIONING_METHOD},
  {"GetTokenFromKeybox", GET_TOKEN_FROM_KEYBOX},
  {"GetTokenFromOemCert", GET_TOKEN_FROM_OEM_CERT},
  {"GetProvisioningToken", GET_PROVISIONING_TOKEN},
  {"GetInternalDeviceUniqueId", GET_INTERNAL_DEVICE_UNIQUE_ID},
  {"GetExternalDeviceUniqueId", GET_EXTERNAL_DEVICE_UNIQUE_ID},
  {"GetProvisioningId", GET_PROVISIONING_ID},
  {"Open", OPEN},
  {"PrepareAndSignLicenseRequest", PREPARE_AND_SIGN_LICENSE_REQUEST},
  {"UseSecondaryKey", USE_SECONDARY_KEY},
  {"LoadLicense", LOAD_LICENSE},
  {"PrepareAndSignLicenseRequest", PREPARE_AND_SIGN_LICENSE_REQUEST},
  {"LoadRenewal", LOAD_RENEWAL},
  {"PrepareAndSignProvisioningRequest", PREPARE_AND_SIGN_PROVISIONING_REQUEST},
  {"LoadEntitledContentKeys", LOAD_ENTITLED_CONTENT_KEYS},
  {"LoadCertificatePrivateKey", LOAD_CERTIFICATE_PRIVATE_KEY},
  {"GetBootCertificateChain", GET_BOOT_CERTIFICATE_CHAIN},
  {"GenerateCertificateKeyPair", GENERATE_CERTIFICATE_KEY_PAIR},
  {"LoadOemCertificatePrivateKey", LOAD_OEM_CERTIFICATE_PRIVATE_KEY},
  {"SelectKey", SELECT_KEY},
  {"GenerateDerivedKeys", GENERATE_DERIVED_KEYS},
  {"GenerateRsaSignature", GENERATE_RSA_SIGNATURE},
  {"Decrypt", DECRYPT},
  {"DeactivateUsageInformation", DEACTIVATE_USAGE_INFORMATION},
  {"GenerateUsageReport", GENERATE_USAGE_REPORT},
  {"GenerateNonce", GENERATE_NONCE},
  {"LoadProvisioning", LOAD_PROVISIONING},
  {"GetHdcpCapabilities", GET_HDCP_CAPABILITIES},
  {"GetRandom", GET_RANDOM},
  {"GetNumberOfOpenSessions", GET_NUMBER_OF_OPEN_SESSIONS},
  {"GetMaxNumberOfSessions", GET_MAX_NUMBER_OF_SESSIONS},
  {"GetSrmVersion", GET_SRM_VERSION},
  {"SetDecryptHash", SET_DECRYPT_HASH},
  {"GetDecryptHashError", GET_DECRYPT_HASH_ERROR},
  {"GenericEncrypt", GENERIC_ENCRYPT},
  {"GenericDecrypt", GENERIC_DECRYPT},
  {"GenericSign", GENERIC_SIGN},
  {"GenericVerify", GENERIC_VERIFY},
  {"CreateUsageTableHeader", CREATE_USAGE_TABLE_HEADER},
  {"LoadUsageTableHeader", LOAD_USAGE_TABLE_HEADER},
  {"ShrinkUsageTableHeader", SHRINK_USAGE_TABLE_HEADER},
  {"CreateUsageEntry", CREATE_USAGE_ENTRY},
  {"LoadUsageEntry", LOAD_USAGE_ENTRY},
  {"UpdateUsageEntry", UPDATE_USAGE_ENTRY},
  {"MoveUsageEntry", MOVE_USAGE_ENTRY},
  {"SetDebugIgnoreKeyboxCount", SET_DEBUG_IGNORE_KEYBOX_COUNT},
  {"SetAllowTestKeybox", SET_ALLOW_TEST_KEYBOX},
  {"PrepareOtaProvisioningRequest", PREPARE_OTA_PROVISIONING_REQUEST},
  {"LoadOtaProvisioning", LOAD_OTA_PROVISIONING},
  // insert new or renamed methods below
};

} // namespace

namespace wvdrm {

CryptoSessionApi getCryptoSessionMethodEnum(const std::string& method) {
  if (!method.empty()) {
    auto it = table.find(method);
    if (it != table.end()) {
      return it->second;
    }
  }
  return API_UNKNOWN;
}

::ndk::ScopedAStatus toNdkScopedAStatus(::wvdrm::WvStatus status,
                                        const char* msg) {
  if (status == ::aidl::android::hardware::drm::Status::OK) {
    return ::ndk::ScopedAStatus::ok();
  }
  const int32_t err = static_cast<int32_t>(status);
  Json::Value jsonMsg(Json::objectValue);
  const auto& cdmErr = status.getCdmErr();

  if (!cdmErr.IsOk()) {
    jsonMsg["cdmError"] = cdmErr.code();
  }
  if (cdmErr.HasOemcResult()) {
    jsonMsg["oemError"] = cdmErr.oemc_result();
  }
  if (cdmErr.HasCryptoSessionMethod()) {
    jsonMsg["context"] =
        getCryptoSessionMethodEnum(cdmErr.crypto_session_method());
  }
  if (msg != nullptr) {
    jsonMsg["errorMessage"] = msg;
  }
  Json::FastWriter writer;
  writer.omitEndingLineFeed();
  return ::ndk::ScopedAStatus::fromServiceSpecificErrorWithMessage(
      err, writer.write(jsonMsg).c_str());
}

constexpr const char* kJsonKeySocVendor = "soc_vendor";
constexpr const char* kJsonKeySocModel = "soc_model";
constexpr const char* kJsonKeyFormFactor = "form_factor";

struct DeviceInfo {
  std::string soc_vendor;
  std::string soc_model;
  std::string form_factor;
};

static std::vector<DeviceInfo> const kMultiThreadBinderEnabledDevices = {
  {"MediaTek", "mt*_multi_thread", "TV"},
};

bool matchesSocModelPattern(const std::string& socModel) {
  const std::string prefix = "mt";
  const std::string suffix = "_multi_thread";
  const size_t minLength = prefix.length() + suffix.length();

  if (socModel.length() < minLength) {
      return false;
  }

  if (socModel.compare(0, prefix.length(), prefix) != 0) {
      return false;
  }

  if (socModel.compare(socModel.length() - suffix.length(), suffix.length(), suffix) != 0) {
      return false;
  }

  return true;
}

bool checkIfEnableMultiThreadBinder() {
  // On first boot, the /data/vendor/mediadrm path may not have been created
  // by init yet. Calling QueryStatus at this stage would fail because the L3
  // service cannot initialize its filesystem, leading to a persistent error
  // state. This check defers enabling multi-threading until the basic DRM
  // filesystem is ready. b/425974655.
  const char* kMediaDrmPath = "/data/vendor/mediadrm";
  struct stat sb;

  // check if the MediaDRM path exists and is a directory
  if (stat(kMediaDrmPath, &sb) != 0 || !S_ISDIR(sb.st_mode)) {
    LOGW("MediaDRM path does not exist, deferring multi-thread binder check.");
    return false;
  }

  android::sp<wvcdm::WvContentDecryptionModule> cdm = wvdrm::getCDM();
  if (cdm == nullptr) {
    LOGW("Failed to get CDM when checking if multi-thread binder is enabled.");
    return false;
  }
  std::string buildInfoValue;
  cdm->QueryStatus(wvcdm::RequestedSecurityLevel::kLevelDefault,
                   wvcdm::QUERY_KEY_OEMCRYPTO_BUILD_INFORMATION, &buildInfoValue);

  Json::Reader reader;
  Json::Value buildInfoJson;
  if (!reader.parse(buildInfoValue, buildInfoJson)) {
    LOGW("Failed to parse build info to JSON: %s", buildInfoValue.c_str());
    return false;
  }

  if (!buildInfoJson.isMember(kJsonKeySocVendor) ||
      !buildInfoJson.isMember(kJsonKeySocModel) ||
      !buildInfoJson.isMember(kJsonKeyFormFactor)) {
    return false;
  }

  std::string currentSocVendor = buildInfoJson[kJsonKeySocVendor].asString();
  std::string currentSocModel = buildInfoJson[kJsonKeySocModel].asString();
  std::string currentFormFactor = buildInfoJson[kJsonKeyFormFactor].asString();

  for (const auto& allowedDevice : kMultiThreadBinderEnabledDevices) {
    if (allowedDevice.soc_vendor == currentSocVendor &&
      allowedDevice.form_factor == currentFormFactor) {

      if (allowedDevice.soc_vendor == "MediaTek" &&
          allowedDevice.soc_model == "mt*_multi_thread" &&
          allowedDevice.form_factor == "TV")
      {
        if (matchesSocModelPattern(currentSocModel)) {
          LOGI("Multi-thread binder enabled for device via MediaTek TV pattern: "
                "Model=%s matches pattern=%s",
                currentSocModel.c_str(),
                allowedDevice.soc_model.c_str());
          return true;
        }
      } else {
        // For all other rules, perform an exact match on soc_model
        if (allowedDevice.soc_model == currentSocModel) {
          LOGI("Multi-thread binder enabled for device via exact match rule: "
              "Vendor=%s, Model=%s, FormFactor=%s",
              allowedDevice.soc_vendor.c_str(),
              allowedDevice.soc_model.c_str(),
              allowedDevice.form_factor.c_str());
          return true;
        }
      }
    }
  }

  return false;
}

}  // namespace wvdrm