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

#include "certificate_provisioning.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <string>

#include "crypto_session.h"
#include "license_protocol.pb.h"
#include "metrics_collections.h"
#include "string_conversions.h"
#include "string_format.h"
#include "test_base.h"
#include "wv_cdm_types.h"

namespace {

using video_widevine::DrmCertificate;
using video_widevine::ProvisioningResponse;
using video_widevine::SignedDrmCertificate;
using video_widevine::SignedProvisioningMessage;

const std::string kSignedDeviceCertificate = wvutil::a2bs_hex(
    "0A3D0802121B7769646576696E655F746573745F73657269616C5F6E756D62657218C09A0C"
    "28D2093A11746573742E7769646576696E652E636F6D6080B51812097369676E617475726"
    "5");
const std::string kSignedDeviceCertificateCreationTimeUnlimited =
    wvutil::a2bs_hex(
        "0A370802121B7769646576696E655F746573745F73657269616C5F6E756D6265721800"
        "28D2093A11746573742E7769646576696E652E636F6D12097369676E6174757265");
const std::string kSignedDeviceCertificateExpirationTimeInvalid =
    wvutil::a2bs_hex(
        "0A390802121B7769646576696E655F746573745F73657269616C5F6E756D62657218C0"
        "9A0C28D2093A11746573742E7769646576696E652E636F6D12097369676E617475726"
        "5");
const std::string kSignedDeviceCertificateExpirationTimeUnlimited =
    wvutil::a2bs_hex(
        "0A3B0802121B7769646576696E655F746573745F73657269616C5F6E756D62657218C0"
        "9A0C28D2093A11746573742E7769646576696E652E636F6D600012097369676E617475"
        "7265");
const std::string kSignedDeviceCertificateInvalid = wvutil::a2bs_hex(
    "76340802121B7769646576696E655F746573745F73657269616C5F6E756D62657228D2093A"
    "11746573742E7769646576696E652E636F6D12097369676E6174757265");
const std::string kSignedDeviceCertificateInvalidCertificateType =
    wvutil::a2bs_hex(
        "0A350801121B7769646576696E655F746573745F73657269616C5F6E756D62657228D2"
        "093A11746573742E7769646576696E652E636F6D12097369676E6174757265");
const std::string kSignedDeviceCertificateNoDrmCertificate =
    wvutil::a2bs_hex("12097369676E6174757265");
const std::string kSignedDeviceCertificateTimesInvalid = wvutil::a2bs_hex(
    "0A350802121B7769646576696E655F746573745F73657269616C5F6E756D62657228D2093A"
    "11746573742E7769646576696E652E636F6D12097369676E6174757265");
const int64_t kCreationTime = 200000;
const int64_t kExpirationTime = 400000;
const std::string kSerialNumber = "widevine_test_serial_number";
const uint32_t kSystemId = 1234;
const uint32_t kNonce = 0x49e81305;
const std::string kNonceString = "\x49\xe8\x13\x05";
const std::string kEmptyString;
const std::string kFakeCoreMessage = wvutil::a2bs_hex("DEADBEEF");
const std::string kFakeProvisioningRequestSignature =
    "a very real provisioning request signature";
const std::string kFakeBuildInfo = "Mock Crypto Session - License Test";
const std::string kFakeAuthority = "a fake authority";
const std::string kOrigin = "an origin for device id";
const std::string kSpoid = "a spoid for identifying a device";
const std::string kFakeCertificateToken = "a testing certificate token";
const std::string kFakePublicKey = "a very real public key for a certificate";
const std::string kFakeSignature = "a very real certificate signature";
const std::string kWrappedPrivateKey = "a wrapped private key";

bool MakeSignedDrmCertificate(const std::string& public_key,
                              const std::string& serial_number,
                              uint32_t system_id, const std::string& signature,
                              std::string* certificate) {
  DrmCertificate drm_certificate;
  drm_certificate.set_public_key(public_key);
  drm_certificate.set_serial_number(serial_number);
  drm_certificate.set_system_id(system_id);
  drm_certificate.set_type(DrmCertificate::DEVICE);
  drm_certificate.set_creation_time_seconds(kCreationTime);
  std::string serialized_drm_certificate;
  drm_certificate.SerializeToString(&serialized_drm_certificate);

  SignedDrmCertificate signed_drm_certificate;
  signed_drm_certificate.set_drm_certificate(serialized_drm_certificate);
  signed_drm_certificate.set_signature(signature);
  return signed_drm_certificate.SerializeToString(certificate);
}

bool MakeProvisioningResponseJson(
    const std::string& certificate, const std::string& nonce,
    ProvisioningResponse::ProvisioningStatus status,
    const std::string& signature, const std::string& core_message,
    std::string* response) {
  ProvisioningResponse provisioning_response;
  provisioning_response.set_device_certificate(certificate);
  provisioning_response.set_nonce(nonce);
  provisioning_response.set_status(status);
  std::string provisioning_response_string;
  if (!provisioning_response.SerializeToString(&provisioning_response_string)) {
    return false;
  }

  SignedProvisioningMessage signed_response;
  signed_response.set_message(provisioning_response_string);
  signed_response.set_signature(signature);
  signed_response.set_oemcrypto_core_message(core_message);
  signed_response.set_protocol_version(SignedProvisioningMessage::VERSION_1_1);
  signed_response.set_provisioning_type(
      SignedProvisioningMessage::DRM_REPROVISIONING);
  std::string signed_response_string;
  if (!signed_response.SerializeToString(&signed_response_string)) {
    return false;
  }

  std::string response_base64 =
      wvutil::Base64SafeEncode(signed_response_string);
  return wvutil::FormatString(response, "{ \"signedResponse\": \"%s\" }",
                              response_base64.c_str());
}

}  // unnamed namespace

namespace wvcdm {

using ::testing::_;
using ::testing::ByMove;
using ::testing::DoAll;
using ::testing::HasSubstr;
using ::testing::NiceMock;
using ::testing::NotNull;
using ::testing::Return;
using ::testing::ReturnArg;
using ::testing::SaveArg;
using ::testing::SetArgPointee;
using ::testing::SetArgReferee;
using ::testing::StrEq;

using wvutil::File;

class MockFile : public File {
 public:
  MockFile() {}
  ~MockFile() override {}

  MOCK_METHOD(ssize_t, Read, (char*, size_t), (override));
  MOCK_METHOD(ssize_t, Write, (const char*, size_t), (override));
};

class MockFileSystem : public wvutil::FileSystem {
 public:
  MockFileSystem() {}
  ~MockFileSystem() override {}

  MOCK_METHOD(std::unique_ptr<File>, Open, (const std::string&, int),
              (override));
};

class MockCryptoSession : public TestCryptoSession {
 public:
  MockCryptoSession(metrics::CryptoMetrics* metrics)
      : TestCryptoSession(metrics) {}
  MOCK_METHOD(CdmResponseType, Open, (wvcdm::RequestedSecurityLevel),
              (override));
  // Usage Table Header.
  MOCK_METHOD(CdmResponseType, CreateUsageTableHeader,
              (wvcdm::RequestedSecurityLevel, UsageTableHeader*), (override));
  MOCK_METHOD(CdmResponseType, LoadUsageTableHeader,
              (wvcdm::RequestedSecurityLevel, const UsageTableHeader&),
              (override));
  MOCK_METHOD(CdmResponseType, ShrinkUsageTableHeader,
              (wvcdm::RequestedSecurityLevel, uint32_t, UsageTableHeader*),
              (override));
  // Usage Entry.
  MOCK_METHOD(CdmResponseType, CreateUsageEntry, (UsageEntryIndex*),
              (override));
  MOCK_METHOD(CdmResponseType, LoadUsageEntry,
              (UsageEntryIndex, const UsageEntry&), (override));
  MOCK_METHOD(CdmResponseType, UpdateUsageEntry,
              (UsageTableHeader*, UsageEntry*), (override));
  MOCK_METHOD(CdmResponseType, MoveUsageEntry, (UsageEntryIndex), (override));

  MOCK_METHOD(bool, IsOpen, (), (override));
  MOCK_METHOD(CdmClientTokenType, GetPreProvisionTokenType, (), (override));
  MOCK_METHOD(CdmResponseType, GetProvisioningToken,
              (std::string*, std::string*), (override));
  MOCK_METHOD(CdmResponseType, GenerateNonce, (uint32_t*), (override));
  MOCK_METHOD(CdmResponseType, PrepareAndSignProvisioningRequest,
              (const std::string&, std::string*, std::string*, bool&,
               OEMCrypto_SignatureHashAlgorithm&),
              (override));
  MOCK_METHOD(bool, GetSupportedCertificateTypes, (SupportedCertificateTypes*),
              (override));
  MOCK_METHOD(bool, GetApiVersion, (uint32_t*), (override));
  MOCK_METHOD(bool, GetResourceRatingTier, (uint32_t*), (override));
  MOCK_METHOD(bool, GetBuildInformation, (std::string*), (override));
  MOCK_METHOD(CdmSecurityLevel, GetSecurityLevel, (), (override));
  MOCK_METHOD(CdmResponseType, LoadProvisioning,
              (const std::string&, const std::string&, const std::string&,
               const std::string&, std::string*),
              (override));
};

class TestStubCryptoSessionFactory : public CryptoSessionFactory {
  CryptoSession* MakeCryptoSession(
      metrics::CryptoMetrics* crypto_metrics) override {
    return new NiceMock<MockCryptoSession>(crypto_metrics);
  }
};

class CertificateProvisioningTest
    : public WvCdmTestBase,
      public testing::WithParamInterface<CdmClientTokenType> {
 public:
  void SetUp() override {
    WvCdmTestBase::SetUp();
    CryptoSession::SetCryptoSessionFactory(new TestStubCryptoSessionFactory());

    metrics_.reset(new metrics::CryptoMetrics());
    certificate_provisioning_.reset(
        new CertificateProvisioning(metrics_.get()));
    GetCryptoSession()->pre_provision_token_type_ = GetParam();
  }

  void TearDown() override {}

 protected:
  MockCryptoSession* GetCryptoSession() {
    return static_cast<MockCryptoSession*>(
        certificate_provisioning_->crypto_session_.get());
  }

  std::unique_ptr<metrics::CryptoMetrics> metrics_;
  std::unique_ptr<CertificateProvisioning> certificate_provisioning_;
};

void MockValidCryptoSession(MockCryptoSession* crypto_session,
                              CdmClientTokenType token_type) {
  const uint32_t crypto_session_api_version = 18;
  const uint32_t resource_rating_tier = RESOURCE_RATING_TIER_LOW;

  ON_CALL(*crypto_session, IsOpen()).WillByDefault(Return(true));
  ON_CALL(*crypto_session, Open(_))
      .WillByDefault(Return(CdmResponseType(NO_ERROR)));
  ON_CALL(*crypto_session, GenerateNonce(NotNull()))
      .WillByDefault(
          DoAll(SetArgPointee<0>(kNonce), Return(CdmResponseType(NO_ERROR))));
  ON_CALL(*crypto_session, GetApiVersion(NotNull()))
      .WillByDefault(
          DoAll(SetArgPointee<0>(crypto_session_api_version), Return(true)));
  ON_CALL(*crypto_session, GetResourceRatingTier(NotNull()))
      .WillByDefault(
          DoAll(SetArgPointee<0>(resource_rating_tier), Return(true)));
  ON_CALL(*crypto_session, GetBuildInformation(NotNull()))
      .WillByDefault(DoAll(SetArgPointee<0>(kFakeBuildInfo), Return(true)));
  ON_CALL(*crypto_session, GetSupportedCertificateTypes(NotNull()))
      .WillByDefault(Return(false));
  ON_CALL(*crypto_session, GetPreProvisionTokenType())
      .WillByDefault(Return(token_type));
}

// Tests ExtractDeviceInfo failure scenarios
// * invalid output parmeters
// * invalid signed drm device certificate
// * signed drm device certificate contains no drm certificate
// * drm certificate has an invalid certificate type
TEST_P(CertificateProvisioningTest, ExtractDeviceInfo_InvalidInput) {
  std::string serial_number;
  uint32_t system_id;

  EXPECT_FALSE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificate, nullptr, nullptr, nullptr, nullptr));

  int64_t creation_time_seconds, expiration_time_seconds;
  EXPECT_FALSE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateInvalid, &serial_number, &system_id,
      &creation_time_seconds, &expiration_time_seconds));

  EXPECT_FALSE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateNoDrmCertificate, &serial_number, &system_id,
      &creation_time_seconds, &expiration_time_seconds));

  EXPECT_FALSE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateInvalidCertificateType, &serial_number,
      &system_id, &creation_time_seconds, &expiration_time_seconds));
}

// Tests ExtractDeviceInfo success scenarios
// * able to extract both |serial_number| and |system_id|
// * able to extract only |serial_number|
// * able to extract only |system_id|
TEST_P(CertificateProvisioningTest, ExtractDeviceInfo) {
  std::string serial_number;
  uint32_t system_id;
  int64_t creation_time_seconds, expiration_time_seconds;

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateTimesInvalid, &serial_number, &system_id,
      &creation_time_seconds, &expiration_time_seconds));
  EXPECT_EQ(kSerialNumber, serial_number);
  EXPECT_EQ(kSystemId, system_id);
  EXPECT_EQ(INVALID_TIME, creation_time_seconds);
  EXPECT_EQ(INVALID_TIME, expiration_time_seconds);

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateTimesInvalid, nullptr, &system_id, nullptr,
      nullptr));
  EXPECT_EQ(kSystemId, system_id);

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateTimesInvalid, &serial_number, nullptr, nullptr,
      nullptr));
  EXPECT_EQ(kSerialNumber, serial_number);

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateCreationTimeUnlimited, &serial_number, &system_id,
      &creation_time_seconds, &expiration_time_seconds));
  EXPECT_EQ(kSerialNumber, serial_number);
  EXPECT_EQ(kSystemId, system_id);
  EXPECT_EQ(UNLIMITED_DURATION, creation_time_seconds);
  EXPECT_EQ(INVALID_TIME, expiration_time_seconds);

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateExpirationTimeInvalid, &serial_number, &system_id,
      &creation_time_seconds, &expiration_time_seconds));
  EXPECT_EQ(kSerialNumber, serial_number);
  EXPECT_EQ(kSystemId, system_id);
  EXPECT_EQ(kCreationTime, creation_time_seconds);
  EXPECT_EQ(INVALID_TIME, expiration_time_seconds);

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateExpirationTimeUnlimited, &serial_number,
      &system_id, &creation_time_seconds, &expiration_time_seconds));
  EXPECT_EQ(kSerialNumber, serial_number);
  EXPECT_EQ(kSystemId, system_id);
  EXPECT_EQ(kCreationTime, creation_time_seconds);
  EXPECT_EQ(UNLIMITED_DURATION, expiration_time_seconds);

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificate, &serial_number, &system_id,
      &creation_time_seconds, &expiration_time_seconds));
  EXPECT_EQ(kSerialNumber, serial_number);
  EXPECT_EQ(kSystemId, system_id);
  EXPECT_EQ(kCreationTime, creation_time_seconds);
  EXPECT_EQ(kExpirationTime, expiration_time_seconds);

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificate, &serial_number, &system_id, nullptr,
      &expiration_time_seconds));
  EXPECT_EQ(kSerialNumber, serial_number);
  EXPECT_EQ(kSystemId, system_id);
  EXPECT_EQ(kExpirationTime, expiration_time_seconds);

  EXPECT_TRUE(certificate_provisioning_->ExtractDeviceInfo(
      kSignedDeviceCertificateExpirationTimeUnlimited, &serial_number,
      &system_id, &creation_time_seconds, nullptr));
  EXPECT_EQ(kSerialNumber, serial_number);
  EXPECT_EQ(kSystemId, system_id);
  EXPECT_EQ(kCreationTime, creation_time_seconds);
  EXPECT_EQ(kExpirationTime, expiration_time_seconds);
}

TEST_P(CertificateProvisioningTest, ProvisioningRequestIsValid) {
  certificate_provisioning_->Init("");
  MockFileSystem file_system;
  MockCryptoSession* crypto_session = GetCryptoSession();
  MockValidCryptoSession(crypto_session, GetParam());
  ON_CALL(*crypto_session, GetProvisioningToken(NotNull(), _))
      .WillByDefault(DoAll(SetArgPointee<0>(kFakeCertificateToken),
                           Return(CdmResponseType(NO_ERROR))));
  EXPECT_CALL(*crypto_session,
              PrepareAndSignProvisioningRequest(_, NotNull(), NotNull(), _, _))
      .WillOnce(DoAll(SetArgPointee<1>(kFakeCoreMessage),
                      SetArgPointee<2>(kFakeProvisioningRequestSignature),
                      SetArgReferee<3>(false),
                      Return(CdmResponseType(NO_ERROR))));

  CdmProvisioningRequest request;
  std::string default_url;
  EXPECT_EQ(NO_ERROR,
            certificate_provisioning_->GetProvisioningRequest(
                &file_system, kLevelDefault, kCertificateWidevine,
                kFakeAuthority, kOrigin, kSpoid, &request, &default_url));
}

TEST_P(CertificateProvisioningTest, ProvisioningRequestFailsEmptySignature) {
  certificate_provisioning_->Init("");
  MockFileSystem file_system;
  MockCryptoSession* crypto_session = GetCryptoSession();
  MockValidCryptoSession(crypto_session, GetParam());
  ON_CALL(*crypto_session, GetProvisioningToken(NotNull(), _))
      .WillByDefault(DoAll(SetArgPointee<0>(kFakeCertificateToken),
                           Return(CdmResponseType(NO_ERROR))));
  EXPECT_CALL(*crypto_session,
              PrepareAndSignProvisioningRequest(_, NotNull(), NotNull(), _, _))
      .WillOnce(DoAll(SetArgPointee<1>(kFakeCoreMessage),
                      SetArgPointee<2>(kEmptyString), SetArgReferee<3>(false),
                      Return(CdmResponseType(NO_ERROR))));

  CdmProvisioningRequest request;
  std::string default_url;
  EXPECT_EQ(CERT_PROVISIONING_REQUEST_ERROR_4,
            certificate_provisioning_->GetProvisioningRequest(
                &file_system, kLevelDefault, kCertificateWidevine,
                kFakeAuthority, kOrigin, kSpoid, &request, &default_url));
}

TEST_P(CertificateProvisioningTest,
       ProvisioningResponseFailsWithEmptyResponse) {
  certificate_provisioning_->Init("");
  // Must set state if not generating request.
  certificate_provisioning_->SetStateForTesting(
      CertificateProvisioning::kDrmRequestSent);

  MockFileSystem file_system;
  std::string certificate;
  std::string wrapped_key;
  EXPECT_EQ(CERT_PROVISIONING_RESPONSE_ERROR_1,
            certificate_provisioning_->HandleProvisioningResponse(
                &file_system, /*response=*/"", &certificate, &wrapped_key));
}

TEST_P(CertificateProvisioningTest,
       ProvisioningResponseFailsIfDeviceIsRevoked) {
  certificate_provisioning_->Init("");
  // Must set state if not generating request.
  certificate_provisioning_->SetStateForTesting(
      CertificateProvisioning::kDrmRequestSent);

  MockFileSystem file_system;
  std::string response_certificate;
  std::string response;
  ASSERT_TRUE(MakeSignedDrmCertificate(kFakePublicKey, kSerialNumber, kSystemId,
                                       kFakeSignature, &response_certificate));
  ASSERT_TRUE(MakeProvisioningResponseJson(
      response_certificate, kNonceString,
      ProvisioningResponse::REVOKED_DEVICE_CREDENTIALS, kFakeSignature,
      kFakeCoreMessage, &response));

  std::string certificate;
  std::string wrapped_key;
  EXPECT_EQ(DEVICE_REVOKED,
            certificate_provisioning_->HandleProvisioningResponse(
                &file_system, response, &certificate, &wrapped_key));
}

TEST_P(CertificateProvisioningTest, ProvisioningResponseSuccess) {
  certificate_provisioning_->Init("");
  // Must set state if not generating request.
  certificate_provisioning_->SetStateForTesting(
      CertificateProvisioning::kDrmRequestSent);

  std::string expected_certificate;
  std::string response;
  ASSERT_TRUE(MakeSignedDrmCertificate(kFakePublicKey, kSerialNumber, kSystemId,
                                       kFakeSignature, &expected_certificate));
  ASSERT_TRUE(MakeProvisioningResponseJson(
      expected_certificate, kNonceString, ProvisioningResponse::NO_ERROR,
      kFakeSignature, kFakeCoreMessage, &response));

  MockCryptoSession* crypto_session = GetCryptoSession();
  ON_CALL(*crypto_session, IsOpen()).WillByDefault(Return(true));
  ON_CALL(*crypto_session, GetSecurityLevel())
      .WillByDefault(Return(kSecurityLevelL3));
  EXPECT_CALL(*crypto_session, LoadProvisioning)
      .Times(1)
      .WillOnce(DoAll(SetArgPointee<4>(kWrappedPrivateKey),
                      Return(CdmResponseType(NO_ERROR))));

  MockFile* file = new MockFile();
  std::string stored_certificate;
  EXPECT_CALL(*file, Write(_, _))
      .Times(1)
      .WillOnce(DoAll(SaveArg<0>(&stored_certificate), ReturnArg<1>()));

  MockFileSystem file_system;
  EXPECT_CALL(file_system, Open(HasSubstr(wvutil::kLegacyCertificateFileName), _))
      .Times(1)
      .WillOnce(Return(ByMove(std::unique_ptr<File>(file))));

  std::string certificate;
  std::string wrapped_key;
  EXPECT_EQ(NO_ERROR, certificate_provisioning_->HandleProvisioningResponse(
                          &file_system, response, &certificate, &wrapped_key));
  EXPECT_NE(std::string::npos, stored_certificate.find(expected_certificate));
}

INSTANTIATE_TEST_SUITE_P(
    CertificateProvisioningTests, CertificateProvisioningTest,
    testing::Values(kClientTokenKeybox, kClientTokenOemCert,
                    kClientTokenDrmCertificateReprovisioning),
    [](const testing::TestParamInfo<CertificateProvisioningTest::ParamType>&
           param_type) {
      return CdmClientTokenTypeToString(param_type.param);
    });

}  // namespace wvcdm