// 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 "utils/StrongPointer.h"
#include <utils/Log.h>

#include <stdio.h>
#include <stdlib.h>

#include "WVCDMSingleton.h"
#include "WVDrmPlugin.h"
#include "WVGenericCryptoInterface.h"
#include "certificate_provisioning.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "string_conversions.h"
#include "url_request.h"
#include "wv_cdm_constants.h"
#include "wv_cdm_types.h"
#include "wv_content_decryption_module.h"

namespace {
static const std::string kDeviceUniqueId = "deviceUniqueId";
static const std::string kSecurityLevel = "securityLevel";
static const std::string kSystemId = "systemId";
static const std::string kTestAppPackageName = "test_app_package";
static const std::string kTestAppPackageName1 = "test_app_package_1";
static const std::string kTestAppPackageName2 = "test_app_package_2";
static const std::string kTestAppPackageName3 = "test_app_package_3";
static const char* const kUnprovisionResponse = "unprovision";

// HTTP response codes.
const int kHttpOk = 200;
}  // namespace

namespace wvdrm {
namespace hardware {
namespace drm {
namespace widevine {

class WVPluginTest : public testing::Test {

 public :
  virtual void SetUp() {
    mDrmPlugin = ::ndk::SharedRefBase::make<WVDrmPlugin>(
        getCDM(), kTestAppPackageName, &mOemCryptoInterface, true);
  }

  std::shared_ptr<WVDrmPlugin> createDrmPlugin(
      const std::string& app_package_name) {
    return ::ndk::SharedRefBase::make<WVDrmPlugin>(
        getCDM(), app_package_name, &mOemCryptoInterface, true);
  }

  // This will create a provisioning request, send it to the service and
  // process the response. The provisioning response will be passed back
  // if needed for inspection. If first stage (for provisioning 4.0)
  // or OKP is needed, the first roundtrip will be completed
  // silently. The method will return after the second round trip.
  // Verify the return value in case of errors.
  bool ensureProvisioned(
      std::shared_ptr<WVDrmPlugin>& plugin,
      wvcdm::CdmProvisioningResponse* provisioning_response) {
    ::aidl::android::hardware::drm::ProvisionRequest request;

    std::string value;
    auto status =
        mDrmPlugin->getPropertyString(kSystemId, &value);
    EXPECT_TRUE(status.isOk()) << "Failed to get System Id";

    long val = std::strtol(value.c_str(), nullptr, 10);
    EXPECT_NE(val, 0) << "System ID invalid: " << errno;
    EXPECT_NE(val, LONG_MIN) << "System ID invalid: " << errno;
    EXPECT_NE(val, LONG_MAX) << "System ID invalid: " << errno;

    uint32_t system_id = val;

    if (system_id == wvcdm::NULL_SYSTEM_ID) {
      // Either first stage prov 4.0 or OKP
      bool result = provision(plugin, provisioning_response);
      if (!result)
        return false;
    }

    return provision(plugin, provisioning_response);
  }

  void unprovision(const std::shared_ptr<WVDrmPlugin>& plugin) {
    std::vector<uint8_t> unprovisionResponse;
    unprovisionResponse
        .assign(kUnprovisionResponse,
                kUnprovisionResponse + strlen(kUnprovisionResponse));
    ::aidl::android::hardware::drm::ProvideProvisionResponseResult result;
    auto status =
        plugin->provideProvisionResponse(unprovisionResponse, &result);
    ASSERT_TRUE(status.isOk()) << "Failed to unprovision";
  }

  bool isL1() {
    std::string security_level;
    auto status =
        mDrmPlugin->getPropertyString(kSecurityLevel, &security_level);
    EXPECT_TRUE(status.isOk()) << "Failed to get Security level";
    return security_level == wvcdm::QUERY_VALUE_SECURITY_LEVEL_L1;
  }

  std::string getDeviceUniqueId(std::shared_ptr<WVDrmPlugin>& plugin) {
    std::vector<uint8_t> device_unique_id;
    auto status =
        plugin->getPropertyByteArray(kDeviceUniqueId, &device_unique_id);
    EXPECT_TRUE(status.isOk()) << "Failed to get Device Unique ID";
    return std::string(device_unique_id.begin(), device_unique_id.end());
  }

  bool getSerialNumber(const wvcdm::CdmProvisioningResponse& response,
                       std::string* serial_number) {
    if (serial_number == nullptr) {
      ADD_FAILURE();
      return false;
    }

    wvcdm::CdmProvisioningResponse provisioning_response;
    if (!wvcdm::CertificateProvisioning::ExtractAndDecodeSignedMessage(
            response, &provisioning_response)) {
      ADD_FAILURE();
      return false;
    }

    video_widevine::SignedProvisioningMessage signed_response;
    if (!signed_response.ParseFromString(provisioning_response) ||
        !signed_response.has_message()) {
      ADD_FAILURE();
      return false;
    }

    video_widevine::ProvisioningResponse prov_response;
    if (!prov_response.ParseFromString(signed_response.message()) ||
        !prov_response.has_device_certificate()) {
      ADD_FAILURE();
      return false;
    }

    if (!wvcdm::CertificateProvisioning::ExtractDeviceInfo(
            prov_response.device_certificate(), serial_number, nullptr, nullptr,
            nullptr)) {
      ADD_FAILURE();
      return false;
    }

    return true;
  }

 private:
  // Create a Provisioning request, send it to the service and process the
  // response. Pass back the provisioning response if needed for inspection
  bool provision(
      std::shared_ptr<WVDrmPlugin>& plugin,
      wvcdm::CdmProvisioningResponse* provisioning_response) {

    ::aidl::android::hardware::drm::ProvisionRequest request;
    auto status =
        plugin->getProvisionRequest(std::string(""), std::string(""), &request);
    EXPECT_TRUE(status.isOk()) << "Failed to get Provisioning request";
    if (!status.isOk()) return false;

    wvcdm::CdmProvisioningRequest
        provisioning_request(request.request.begin(), request.request.end());
    bool sts = GetProvisioningResponse(provisioning_request,
                                       request.defaultUrl,
                                       provisioning_response);
    if (!sts) {
      EXPECT_TRUE(sts) << "Failed to get provisioning response";
      return false;
    }

    wvcdm::CdmProvisioningResponse prov_response;

    std::vector<uint8_t> response(provisioning_response->begin(),
                                  provisioning_response->end());
    ::aidl::android::hardware::drm::ProvideProvisionResponseResult result;
    status = plugin->provideProvisionResponse(response, &result);
    EXPECT_TRUE(status.isOk()) << "Failed to process provisioning response";
    if (!status.isOk()) return false;
    return true;
  }

  // Post a request and extract the signed provisioning message from
  // the HTTP response.
  // Setting a non-zero value to |duration_seconds| will request a
  // limited duration DRM certificate.
  bool GetProvisioningResponse(const wvcdm::CdmProvisioningRequest& request,
                               const std::string& url,
                               wvcdm::CdmProvisioningResponse* response) {
    EXPECT_NE(response, nullptr) << "No |response| parameter provided";
    if (response == nullptr)
      return false;

    // Use secure connection and chunk transfer coding.
    wvcdm::UrlRequest url_request(url);
    EXPECT_TRUE(url_request.is_connected()) << "Fail to connect to " << url;
    url_request.PostCertRequestInQueryString(request);
    EXPECT_TRUE(url_request.GetResponse(response));

    int http_status_code = url_request.GetStatusCode(*response);
    EXPECT_EQ(kHttpOk, http_status_code) << "provisioning_response = "
                                         << *response;
    if (kHttpOk != http_status_code) {
      LogResponseError(*response, http_status_code);
      return false;
    }
    return true;
  }

  void LogResponseError(const std::string& message, int http_status_code) {
    LOGD("HTTP Status code = %d", http_status_code);
    LOGD("HTTP response(%zu): %s", message.size(),
         wvutil::b2a_hex(message).c_str());
  }

 private:
  std::shared_ptr<WVDrmPlugin> mDrmPlugin;
  wvdrm::WVGenericCryptoInterface mOemCryptoInterface;
};

// Verify that SPOIDs and DRM certificate serial number are stable between
// factory resets/provisioning attempts for the same app and different between
// different apps. Test using three different apps.
TEST_F(WVPluginTest, ProvisioningStableSpoidTestL1) {
  std::shared_ptr<WVDrmPlugin> drm_plugin1 =
      createDrmPlugin(kTestAppPackageName1);

  if (!isL1())  // declare success if device does not natively support L1
    return;

  wvcdm::CdmProvisioningResponse provisioning_response;

  unprovision(drm_plugin1);
  std::string spoid1_1 = getDeviceUniqueId(drm_plugin1);
  ASSERT_TRUE(ensureProvisioned(drm_plugin1, &provisioning_response));
  std::string drm_cert_serial_number1_1;
  ASSERT_TRUE(getSerialNumber(provisioning_response,
                              &drm_cert_serial_number1_1));

  unprovision(drm_plugin1);
  std::string spoid1_2 = getDeviceUniqueId(drm_plugin1);
  ASSERT_TRUE(ensureProvisioned(drm_plugin1, &provisioning_response));
  std::string drm_cert_serial_number1_2;
  ASSERT_TRUE(getSerialNumber(provisioning_response,
                              &drm_cert_serial_number1_2));

  std::shared_ptr<WVDrmPlugin> drm_plugin2 =
      createDrmPlugin(kTestAppPackageName2);

  unprovision(drm_plugin2);
  std::string spoid2_1 = getDeviceUniqueId(drm_plugin2);
  ASSERT_TRUE(ensureProvisioned(drm_plugin2, &provisioning_response));
  std::string drm_cert_serial_number2_1;
  ASSERT_TRUE(getSerialNumber(provisioning_response,
                              &drm_cert_serial_number2_1));

  unprovision(drm_plugin2);
  std::string spoid2_2 = getDeviceUniqueId(drm_plugin2);
  ASSERT_TRUE(ensureProvisioned(drm_plugin2, &provisioning_response));
  std::string drm_cert_serial_number2_2;
  ASSERT_TRUE(getSerialNumber(provisioning_response,
                              &drm_cert_serial_number2_2));

  std::shared_ptr<WVDrmPlugin> drm_plugin3 =
      createDrmPlugin(kTestAppPackageName3);

  unprovision(drm_plugin3);
  std::string spoid3_1 = getDeviceUniqueId(drm_plugin3);
  ASSERT_TRUE(ensureProvisioned(drm_plugin3, &provisioning_response));
  std::string drm_cert_serial_number3_1;
  ASSERT_TRUE(getSerialNumber(provisioning_response,
                              &drm_cert_serial_number3_1));

  unprovision(drm_plugin3);
  std::string spoid3_2 = getDeviceUniqueId(drm_plugin3);
  ASSERT_TRUE(ensureProvisioned(drm_plugin3, &provisioning_response));
  std::string drm_cert_serial_number3_2;
  ASSERT_TRUE(getSerialNumber(provisioning_response,
                              &drm_cert_serial_number3_2));

  // Verify that the spoids are stable after a reprovision
  EXPECT_EQ(spoid1_1, spoid1_2);
  EXPECT_EQ(spoid2_1, spoid2_2);
  EXPECT_EQ(spoid3_1, spoid3_2);

  // Verify that the spoids are different for different apps
  EXPECT_NE(spoid1_1, spoid2_1);
  EXPECT_NE(spoid1_1, spoid3_1);
  EXPECT_NE(spoid2_1, spoid3_1);

  // Verify that the DRM certificate serial numbers are stable after a
  // reprovision
  EXPECT_EQ(drm_cert_serial_number1_1, drm_cert_serial_number1_2);
  EXPECT_EQ(drm_cert_serial_number2_1, drm_cert_serial_number2_2);
  EXPECT_EQ(drm_cert_serial_number3_1, drm_cert_serial_number3_2);

  // Verify that the DRM certificate serial numbers are different for different
  // apps
  EXPECT_NE(drm_cert_serial_number1_1, drm_cert_serial_number2_1);
  EXPECT_NE(drm_cert_serial_number1_1, drm_cert_serial_number3_1);
  EXPECT_NE(drm_cert_serial_number2_1, drm_cert_serial_number3_1);
}

} // namespace widevine
} // namespace drm
} // namespace hardware
} // namespace wvdrm