////////////////////////////////////////////////////////////////////////////////
// Copyright 2017 Google LLC.
//
// This software is licensed under the terms defined in the Widevine Master
// License Agreement. For a copy of this agreement, please contact
// widevine-licensing@google.com.
////////////////////////////////////////////////////////////////////////////////

#include "common/device_status_list.h"

#include <stddef.h>
#include <memory>
#include <type_traits>
#include <utility>

#include "glog/logging.h"
#include "testing/gmock.h"
#include "testing/gunit.h"
#include "absl/strings/str_cat.h"
#include "common/client_cert.h"
#include "common/rsa_key.h"
#include "common/rsa_test_keys.h"
#include "protos/public/client_identification.pb.h"
#include "protos/public/errors.pb.h"
#include "protos/public/provisioned_device_info.pb.h"
#include "protos/public/signed_drm_certificate.pb.h"

namespace widevine {

using ::testing::_;
using ::testing::Return;
using ::testing::ReturnRef;
using ::testing::ReturnRefOfCopy;

const uint32_t kValidCertSystemId = 100;
const uint32_t kRevokedCertSystemId = 101;
const uint32_t kValidPpkSystemId = 102;
const uint32_t kTestOnlyCertSystemId = 103;
const uint32_t kRevokedAllowedDeviceCertSystemId = 104;
const uint32_t kUnknownSystemId = 666;
const char kValidSerialNumber[] = "valid-serial-number";
const char kRevokedSerialNumber[] = "revoked-serial-number";
const char kRevokedAllowDeviceSerialNumber[] =
    "revoked-allow-device-serial-number";
const char kTestOnlySerialNumber[] = "test_only-serial-number";
const char kMismatchSerialNumber[] = "mismatch-serial-number";
const char kDeviceModel[] = "device-model-x";
const char kTestPreprovKey[] = "00112233445566778899aabbccddeeff";
const uint32_t kStatusListCreationTime = 17798001;
const uint32_t kDefaultExpirePeriod = 0;

class MockCertificateClientCert : public CertificateClientCert {
 public:
  MockCertificateClientCert() {}
  MOCK_CONST_METHOD0(system_id, uint32_t());
  MOCK_CONST_METHOD0(signer_serial_number, std::string &());
  MOCK_CONST_METHOD0(signer_creation_time_seconds, uint32_t());
  MOCK_CONST_METHOD0(type, ClientIdentification::TokenType());
  MOCK_CONST_METHOD0(signed_by_provisioner, bool());
};

class MockKeyboxClientCert : public KeyboxClientCert {
 public:
  MockKeyboxClientCert() {}
  MOCK_CONST_METHOD0(system_id, uint32_t());
  MOCK_CONST_METHOD0(type, ClientIdentification::TokenType());
};

class DeviceStatusListTest : public ::testing::Test {
 public:
  ~DeviceStatusListTest() override {}

  void SetUp() override {
    DeviceCertificateStatus *cert_status;

    // Device cert with status RELEASED.
    cert_status = cert_status_list_.add_certificate_status();
    cert_status->mutable_device_info()->set_system_id(kValidCertSystemId);
    cert_status->set_drm_serial_number(kValidSerialNumber);
    cert_status->mutable_device_info()->set_model(kDeviceModel);
    cert_status->set_status(DeviceCertificateStatus::STATUS_RELEASED);

    // Device cert with status REVOKED.
    cert_status = cert_status_list_.add_certificate_status();
    cert_status->mutable_device_info()->set_system_id(kRevokedCertSystemId);
    cert_status->set_drm_serial_number(kRevokedSerialNumber);
    cert_status->set_status(DeviceCertificateStatus::STATUS_REVOKED);

    // Device cert with status REVOKED ALLOWED DEVICE.
    cert_status = cert_status_list_.add_certificate_status();
    cert_status->mutable_device_info()->set_system_id(
        kRevokedAllowedDeviceCertSystemId);
    cert_status->set_drm_serial_number(kRevokedAllowDeviceSerialNumber);
    cert_status->set_status(DeviceCertificateStatus::STATUS_REVOKED);
    device_status_list_.AllowRevokedDevices(
        absl::StrCat(kRevokedAllowedDeviceCertSystemId));

    // Device cert with status TEST_ONLY.
    cert_status = cert_status_list_.add_certificate_status();
    cert_status->mutable_device_info()->set_system_id(kTestOnlyCertSystemId);
    cert_status->set_drm_serial_number(kTestOnlySerialNumber);
    cert_status->set_status(DeviceCertificateStatus::STATUS_TEST_ONLY);

    cert_status_list_.set_creation_time_seconds(kStatusListCreationTime);
    cert_status_list_.SerializeToString(
        signed_cert_status_list_.mutable_certificate_status_list());
    std::unique_ptr<RsaPrivateKey> root_key(
        RsaPrivateKey::Create(test_keys_.private_test_key_1_3072_bits()));
    ASSERT_TRUE(root_key);

    ASSERT_TRUE(root_key->GenerateSignature(
        signed_cert_status_list_.certificate_status_list(),
        signed_cert_status_list_.mutable_signature()));
    ASSERT_TRUE(
        signed_cert_status_list_.SerializeToString(&serialized_status_list_));

    ASSERT_EQ(OkStatus(), device_status_list_.UpdateStatusList(
                              test_keys_.public_test_key_1_3072_bits(),
                              serialized_status_list_, kDefaultExpirePeriod));
  }

  DeviceStatusList device_status_list_;
  RsaTestKeys test_keys_;
  DeviceCertificateStatusList cert_status_list_;
  SignedDeviceCertificateStatusList signed_cert_status_list_;
  std::string serialized_status_list_;
};

// Returns the number of DevcieCertificateStatus messages in the list.

TEST_F(DeviceStatusListTest, CheckForValidAndRevokedCert) {
  // Test case where the Certificate status is set to Valid.
  ProvisionedDeviceInfo device_info;
  MockCertificateClientCert valid_client_cert;
  std::string valid_drm_serial_number(kValidSerialNumber);
  EXPECT_CALL(valid_client_cert, type())
      .WillRepeatedly(Return(ClientIdentification::DRM_DEVICE_CERTIFICATE));
  EXPECT_CALL(valid_client_cert, system_id())
      .WillRepeatedly(Return(kValidCertSystemId));
  EXPECT_CALL(valid_client_cert, signer_serial_number())
      .WillRepeatedly(ReturnRef(valid_drm_serial_number));
  EXPECT_EQ(OkStatus(),
            device_status_list_.GetCertStatus(valid_client_cert, &device_info));
  EXPECT_TRUE(device_info.has_model());
  EXPECT_EQ(kDeviceModel, device_info.model());

  // Test case where the Certificate status is Revoked.
  MockCertificateClientCert revoked_client_cert;
  std::string revoked_drm_serial_number(kRevokedSerialNumber);
  EXPECT_CALL(revoked_client_cert, type())
      .WillRepeatedly(Return(ClientIdentification::DRM_DEVICE_CERTIFICATE));
  EXPECT_CALL(revoked_client_cert, system_id())
      .WillRepeatedly(Return(kRevokedCertSystemId));
  EXPECT_CALL(revoked_client_cert, signer_serial_number())
      .WillRepeatedly(ReturnRef(revoked_drm_serial_number));
  EXPECT_EQ(DRM_DEVICE_CERTIFICATE_REVOKED,
            device_status_list_.GetCertStatus(revoked_client_cert, &device_info)
                .error_code());

  // Test case where the revoked cert is allowed.
  device_status_list_.AllowRevokedDevices(absl::StrCat(kRevokedCertSystemId));
  EXPECT_OK(
      device_status_list_.GetCertStatus(revoked_client_cert, &device_info));
}

TEST_F(DeviceStatusListTest, TestOnlyCertAllowed) {
  ProvisionedDeviceInfo device_info;
  MockCertificateClientCert test_only_client_cert;
  std::string test_only_drm_serial_number(kTestOnlySerialNumber);
  EXPECT_CALL(test_only_client_cert, type())
      .WillRepeatedly(Return(ClientIdentification::DRM_DEVICE_CERTIFICATE));
  EXPECT_CALL(test_only_client_cert, system_id())
      .WillRepeatedly(Return(kTestOnlyCertSystemId));
  EXPECT_CALL(test_only_client_cert, signer_serial_number())
      .WillRepeatedly(ReturnRef(test_only_drm_serial_number));
  EXPECT_EQ(
      DEVELOPMENT_CERTIFICATE_NOT_ALLOWED,
      device_status_list_.GetCertStatus(test_only_client_cert, &device_info)
          .error_code());
}

TEST_F(DeviceStatusListTest, TestOnlyCertNotAllowed) {
  ProvisionedDeviceInfo device_info;
  MockCertificateClientCert test_only_client_cert;
  std::string test_only_drm_serial_number(kTestOnlySerialNumber);
  device_status_list_.set_allow_test_only_devices(true);
  EXPECT_CALL(test_only_client_cert, type())
      .WillRepeatedly(Return(ClientIdentification::DRM_DEVICE_CERTIFICATE));
  EXPECT_CALL(test_only_client_cert, system_id())
      .WillRepeatedly(Return(kTestOnlyCertSystemId));
  EXPECT_CALL(test_only_client_cert, signer_serial_number())
      .WillRepeatedly(ReturnRef(test_only_drm_serial_number));
  EXPECT_EQ(OkStatus(), device_status_list_.GetCertStatus(test_only_client_cert,
                                                          &device_info));
}

TEST_F(DeviceStatusListTest, ValidAndUnknownKeybox) {
  std::multimap<uint32_t, std::string> preprov_keys;
  preprov_keys.insert(std::make_pair(kValidCertSystemId, kTestPreprovKey));
  KeyboxClientCert::SetPreProvisioningKeys(preprov_keys);

  // Test case where the Certificate status is set to Valid.
  ProvisionedDeviceInfo device_info;
  MockKeyboxClientCert valid_client_keybox;
  std::string valid_drm_serial_number(kValidSerialNumber);
  EXPECT_CALL(valid_client_keybox, type())
      .WillRepeatedly(Return(ClientIdentification::KEYBOX));
  EXPECT_CALL(valid_client_keybox, system_id())
      .WillRepeatedly(Return(kValidCertSystemId));
  EXPECT_EQ(OkStatus(), device_status_list_.GetCertStatus(valid_client_keybox,
                                                          &device_info));
  EXPECT_TRUE(device_info.has_model());
  EXPECT_EQ(kDeviceModel, device_info.model());

  MockKeyboxClientCert unknown_client_keybox;
  EXPECT_CALL(unknown_client_keybox, type())
      .WillRepeatedly(Return(ClientIdentification::KEYBOX));
  EXPECT_CALL(unknown_client_keybox, system_id())
      .WillRepeatedly(Return(kUnknownSystemId));
  EXPECT_EQ(
      UNSUPPORTED_SYSTEM_ID,
      device_status_list_.GetCertStatus(unknown_client_keybox, &device_info)
          .error_code());
  EXPECT_TRUE(device_info.has_model());
  EXPECT_EQ(kDeviceModel, device_info.model());
}

TEST_F(DeviceStatusListTest, SignerSerialNumberMismatch) {
  device_status_list_.set_allow_unknown_devices(true);

  // Test case where the signer certificate is older than the current status
  // list.
  MockCertificateClientCert older_client_cert;
  ProvisionedDeviceInfo device_info;
  std::string mismatch_drm_serial_number(kMismatchSerialNumber);
  EXPECT_CALL(older_client_cert, type())
      .WillRepeatedly(Return(ClientIdentification::DRM_DEVICE_CERTIFICATE));
  EXPECT_CALL(older_client_cert, system_id())
      .WillRepeatedly(Return(kValidCertSystemId));
  EXPECT_CALL(older_client_cert, signer_serial_number())
      .WillRepeatedly(ReturnRef(mismatch_drm_serial_number));
  EXPECT_CALL(older_client_cert, signer_creation_time_seconds())
      .WillRepeatedly(Return(kStatusListCreationTime - 1));
  EXPECT_EQ(INVALID_DRM_CERTIFICATE,
            device_status_list_.GetCertStatus(older_client_cert, &device_info)
                .error_code());

  // We allow this case only for certs signed by a provisioner cert.
  EXPECT_CALL(older_client_cert, signed_by_provisioner())
      .WillOnce(Return(true));
  EXPECT_EQ(OkStatus(),
            device_status_list_.GetCertStatus(older_client_cert, &device_info));
  EXPECT_TRUE(device_info.has_system_id());
  EXPECT_EQ(kValidCertSystemId, device_info.system_id());

  // Test case where the signer certificate is newer than the current status
  // list, and unknown devices are allowed.
  MockCertificateClientCert newer_client_cert1;
  EXPECT_CALL(newer_client_cert1, type())
      .WillRepeatedly(Return(ClientIdentification::DRM_DEVICE_CERTIFICATE));
  EXPECT_CALL(newer_client_cert1, system_id())
      .WillRepeatedly(Return(kValidCertSystemId));
  EXPECT_CALL(newer_client_cert1, signer_serial_number())
      .WillRepeatedly(ReturnRef(mismatch_drm_serial_number));
  EXPECT_CALL(newer_client_cert1, signer_creation_time_seconds())
      .WillRepeatedly(Return(kStatusListCreationTime));
  EXPECT_EQ(DRM_DEVICE_CERTIFICATE_UNKNOWN,
            device_status_list_.GetCertStatus(newer_client_cert1, &device_info)
                .error_code());

  // Test case where the signer certificate is newer than the current status
  // list, and unknown devices are not allowed.
  device_status_list_.set_allow_unknown_devices(false);
  MockCertificateClientCert newer_client_cert2;
  EXPECT_CALL(newer_client_cert2, type())
      .WillRepeatedly(Return(ClientIdentification::DRM_DEVICE_CERTIFICATE));
  EXPECT_CALL(newer_client_cert2, system_id())
      .WillRepeatedly(Return(kValidCertSystemId));
  EXPECT_CALL(newer_client_cert2, signer_serial_number())
      .WillRepeatedly(ReturnRef(mismatch_drm_serial_number));
  EXPECT_CALL(newer_client_cert2, signer_creation_time_seconds())
      .WillRepeatedly(Return(kStatusListCreationTime + 1));
  EXPECT_EQ(DRM_DEVICE_CERTIFICATE_UNKNOWN,
            device_status_list_.GetCertStatus(newer_client_cert2, &device_info)
                .error_code());
}

TEST_F(DeviceStatusListTest, InvalidStatusList) {
  EXPECT_EQ(INVALID_CERTIFICATE_STATUS_LIST,
            device_status_list_
                .UpdateStatusList(test_keys_.public_test_key_2_2048_bits(),
                                  serialized_status_list_, 0)
                .error_code());

  ++(*signed_cert_status_list_.mutable_certificate_status_list())[4];
  ASSERT_TRUE(
      signed_cert_status_list_.SerializeToString(&serialized_status_list_));
  EXPECT_EQ(INVALID_CERTIFICATE_STATUS_LIST,
            device_status_list_
                .UpdateStatusList(test_keys_.public_test_key_1_3072_bits(),
                                  serialized_status_list_, 0)
                .error_code());
}

class MockDeviceStatusList : public DeviceStatusList {
 public:
  MOCK_CONST_METHOD0(GetCurrentTime, uint32_t());
};

TEST_F(DeviceStatusListTest, ExpiredStatusListOnSet) {
  MockDeviceStatusList mock_device_status_list;
  EXPECT_CALL(mock_device_status_list, GetCurrentTime())
      .Times(2)
      .WillOnce(Return(kStatusListCreationTime + 100))
      .WillOnce(Return(kStatusListCreationTime + 101));
  EXPECT_EQ(OkStatus(), mock_device_status_list.UpdateStatusList(
                            test_keys_.public_test_key_1_3072_bits(),
                            serialized_status_list_, 100));
  EXPECT_EQ(EXPIRED_CERTIFICATE_STATUS_LIST,
            mock_device_status_list
                .UpdateStatusList(test_keys_.public_test_key_1_3072_bits(),
                                  serialized_status_list_, 100)
                .error_code());
}

TEST_F(DeviceStatusListTest, ExpiredStatusListOnCertCheck) {
  MockDeviceStatusList mock_device_status_list;
  EXPECT_CALL(mock_device_status_list, GetCurrentTime())
      .Times(3)
      .WillOnce(Return(kStatusListCreationTime + 100))
      .WillOnce(Return(kStatusListCreationTime + 100))
      .WillOnce(Return(kStatusListCreationTime + 101));
  EXPECT_EQ(OkStatus(), mock_device_status_list.UpdateStatusList(
                            test_keys_.public_test_key_1_3072_bits(),
                            serialized_status_list_, 100));

  ProvisionedDeviceInfo device_info;
  MockCertificateClientCert valid_client_cert;
  std::string valid_drm_serial_number(kValidSerialNumber);
  EXPECT_CALL(valid_client_cert, type())
      .WillRepeatedly(Return(ClientIdentification::DRM_DEVICE_CERTIFICATE));
  EXPECT_CALL(valid_client_cert, system_id())
      .WillRepeatedly(Return(kValidCertSystemId));
  EXPECT_CALL(valid_client_cert, signer_serial_number())
      .WillRepeatedly(ReturnRef(valid_drm_serial_number));
  EXPECT_CALL(valid_client_cert, signer_creation_time_seconds())
      .WillRepeatedly(Return(kStatusListCreationTime - 1));
  EXPECT_EQ(OkStatus(), mock_device_status_list.GetCertStatus(valid_client_cert,
                                                              &device_info));

  EXPECT_EQ(
      EXPIRED_CERTIFICATE_STATUS_LIST,
      mock_device_status_list.GetCertStatus(valid_client_cert, &device_info)
          .error_code());
}

TEST_F(DeviceStatusListTest, IsSystemIdActive) {
  std::multimap<uint32_t, std::string> preprov_keys;
  preprov_keys.insert(
      std::make_pair(kValidPpkSystemId, "00112233445566778899aabbccddeeff"));
  KeyboxClientCert::SetPreProvisioningKeys(preprov_keys);
  device_status_list_.set_allow_unknown_devices(false);
  EXPECT_TRUE(device_status_list_.IsSystemIdActive(kValidCertSystemId));
  EXPECT_TRUE(device_status_list_.IsSystemIdActive(kValidPpkSystemId));
  EXPECT_FALSE(device_status_list_.IsSystemIdActive(kRevokedCertSystemId));
  EXPECT_FALSE(device_status_list_.IsSystemIdActive(kUnknownSystemId));
  device_status_list_.set_allow_unknown_devices(true);
  EXPECT_TRUE(device_status_list_.IsSystemIdActive(kValidCertSystemId));
  EXPECT_TRUE(device_status_list_.IsSystemIdActive(kValidPpkSystemId));
  EXPECT_FALSE(device_status_list_.IsSystemIdActive(kRevokedCertSystemId));
  EXPECT_TRUE(device_status_list_.IsSystemIdActive(kUnknownSystemId));
  EXPECT_TRUE(
      device_status_list_.IsSystemIdActive(kRevokedAllowedDeviceCertSystemId));
}

}  // namespace widevine