////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 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 "provisioning_sdk/internal/provisioning30_session_impl.h"

#include <cstddef>

#include "testing/gmock.h"
#include "testing/gunit.h"
#include "absl/strings/escaping.h"
#include "absl/strings/str_cat.h"
#include "common/aes_cbc_util.h"
#include "common/hash_algorithm.h"
#include "common/hash_algorithm_util.h"
#include "common/mock_rsa_key.h"
#include "common/sha_util.h"
#include "provisioning_sdk/internal/oem_device_cert.h"
#include "provisioning_sdk/internal/provisioning_engine_impl.h"
#include "protos/public/hash_algorithm.pb.h"

using ::testing::_;
using ::testing::ByMove;
using ::testing::DoAll;
using ::testing::IsEmpty;
using ::testing::Return;
using ::testing::SaveArg;
using ::testing::SetArgPointee;
// TODO(user): use a real test key instead of mock_key to testcore message
// code change.
namespace {
const char kCoreMessage[] =
    "00000005000000580000001000000010000000000000004000000000000000000000000000"
    "00000000000000000000000000000000000000000000000000000000000000000000000000"
    "0000000000000000000000000000";
const char kInvalidCoreMessage[] =
    "0000000500000058000000100000001000000000000000400000000000000000000000000";
const char kEmptyCoreMessage[] = "";
const char kEncryptedClientIdIv[] = "sixteen_bytes_iv";
const char kPrivacyKey[] = "privacy_key_16B_";
const char kProviderId[] = "testing_provider";
const char kClientToken[] = "client_id_token";
const char kDevicePublicKey[] = "device_public_key";
const char kEncryptedPrivacyKey[] = "encrypted_privacy_key";
const char kDevicePrivateKey[] = "device_private_key";
const char kWrappingKey[] = "wrapping_key";
const char kDeviceCertificate[] = "device_certificate";
const char kNonce[] = "testing_nonce";
const char kSignature[] = "generated_signature";

// Derives Stable Per-Origin IDentifiers.
std::string DeriveSpoid(const std::string& client_token,
                        const std::string& provider_id,
                        const std::string& secret_sauce) {
  return widevine::Sha256_Hash(client_token + provider_id + secret_sauce)
      .substr(0, 16);
}

}  // namespace

namespace widevine {

class MockProvisioningEngineImpl : public ProvisioningEngineImpl {
 public:
  MOCK_METHOD(ProvisioningStatus, GenerateProviderDeviceDrmCertificate,
              (uint32_t system_id, const std::string& oem_ca_serial_number,
               const std::string& provider_id, const std::string& public_key,
               const std::string& certificate_serial_number,
               std::string* certificate),
              (const, override));
};

class MockOemDeviceCert : public OemDeviceCert {
 public:
  // gmock does not support SetArgPointee on std::unique_ptr, so we have to
  // workaround it with a trick.
  MOCK_METHOD(bool, DoVerifyCertificateChain,
              (const std::string& certificate_chain,
               RsaPublicKey** leaf_public_key, uint32_t* system_id,
               std::string* oem_ca_serial_number),
              (const));
  bool VerifyCertificateChain(
      const std::string& certificate_chain,
      std::unique_ptr<RsaPublicKey>* leaf_public_key, uint32_t* system_id,
      std::string* oem_ca_serial_number) const override {
    RsaPublicKey* raw_leaf_public_key = nullptr;
    if (!DoVerifyCertificateChain(certificate_chain, &raw_leaf_public_key,
                                  system_id, oem_ca_serial_number)) {
      return false;
    }
    *leaf_public_key = std::unique_ptr<RsaPublicKey>(raw_leaf_public_key);
    return true;
  }
};

class Provisioning30SessionImplTest : public ::testing::Test {
 protected:
  Provisioning30SessionImplTest()
      : session_impl_(mock_engine_impl_, mock_oem_device_cert_,
                      mock_service_private_key_) {
    mock_rsa_key_factory_ = new MockRsaKeyFactory;
    session_impl_.set_rsa_key_factory(
        std::unique_ptr<RsaKeyFactory>(mock_rsa_key_factory_));
  }

  Provisioning30SessionImpl session_impl_;
  MockRsaKeyFactory* mock_rsa_key_factory_ = nullptr;
  MockProvisioningEngineImpl mock_engine_impl_;
  MockOemDeviceCert mock_oem_device_cert_;
  MockRsaPrivateKey mock_service_private_key_;
};

class Provisioning30SessionImplProcessTest
    : public Provisioning30SessionImplTest {
 public:
  void SetUp() override {
    MockRsaPublicKey* mock_rsa_public_key = new MockRsaPublicKey;
    EXPECT_CALL(*mock_rsa_key_factory_,
                CreateFromPkcs1PublicKey(kDevicePublicKey))
        .WillOnce(
            Return(ByMove(std::unique_ptr<RsaPublicKey>(mock_rsa_public_key))));
    EXPECT_CALL(*mock_rsa_key_factory_,
                CreateFromPkcs8PrivateKey(kDevicePrivateKey, IsEmpty()))
        .WillOnce(Return(
            ByMove(std::unique_ptr<RsaPrivateKey>(new MockRsaPrivateKey))));
    EXPECT_CALL(*mock_rsa_public_key, MatchesPrivateKey(_))
        .WillOnce(Return(true));
    ASSERT_EQ(OK,
              session_impl_.Initialize(kDevicePublicKey, kDevicePrivateKey));

    // Setup Provisioning Message.
    client_id_.set_type(ClientIdentification::OEM_DEVICE_CERTIFICATE);
    client_id_.set_token(kClientToken);

    EncryptedClientIdentification* encrypted_client_id =
        prov_request_.mutable_encrypted_client_id();
    encrypted_client_id->set_encrypted_client_id(crypto_util::EncryptAesCbc(
        kPrivacyKey, kEncryptedClientIdIv, client_id_.SerializeAsString()));
    encrypted_client_id->set_encrypted_client_id_iv(kEncryptedClientIdIv);
    encrypted_client_id->set_encrypted_privacy_key(kEncryptedPrivacyKey);
    prov_request_.set_provider_id(kProviderId);
    prov_request_.set_nonce(kNonce);

    signed_prov_message_.set_message(prov_request_.SerializeAsString());
    signed_prov_message_.set_signature("testing_signature");
  }

  void ProcessMessage(std::string core_message) {
    signed_prov_message_.set_oemcrypto_core_message(
        absl::HexStringToBytes(core_message));
    const uint32_t kSystemId = 1234;
    EXPECT_CALL(mock_service_private_key_, Decrypt(kEncryptedPrivacyKey, _))
        .WillOnce(DoAll(SetArgPointee<1>(kPrivacyKey), Return(true)));
    MockRsaPublicKey* mock_cert_public_key = new MockRsaPublicKey;
    EXPECT_CALL(mock_oem_device_cert_,
                DoVerifyCertificateChain(kClientToken, _, _, _))
        .WillOnce(DoAll(SetArgPointee<1>(mock_cert_public_key),
                        SetArgPointee<2>(kSystemId), Return(true)));
    const HashAlgorithm hash_algorithm =
        HashAlgorithmProtoToEnum(signed_prov_message_.hash_algorithm());
    EXPECT_CALL(
        *mock_cert_public_key,
        VerifySignature(signed_prov_message_.oemcrypto_core_message() +
                            signed_prov_message_.message(),
                        hash_algorithm, signed_prov_message_.signature()))
        .WillOnce(Return(true));

    EXPECT_CALL(mock_engine_impl_, GenerateProviderDeviceDrmCertificate(
                                       kSystemId, _, _, kDevicePublicKey, _, _))
        .WillOnce(DoAll(SetArgPointee<5>(kDeviceCertificate), Return(OK)));
    EXPECT_CALL(*mock_cert_public_key, Encrypt(_, _))
        .WillOnce(DoAll(SaveArg<0>(&message_key_),
                        SetArgPointee<1>(kWrappingKey), Return(true)));
  }

  void ValidateMessage(std::string core_message) {
    ProvisioningResponse prov_response;
    // Verify the response.
    ASSERT_TRUE(signed_prov_message_.ParseFromString(response_));
    EXPECT_EQ(kSignature, signed_prov_message_.signature());
    ASSERT_TRUE(prov_response.ParseFromString(signed_prov_message_.message()));
    if (!core_message.empty()) {
      EXPECT_TRUE(signed_prov_message_.has_oemcrypto_core_message());
      EXPECT_EQ(signature_input_,
                absl::StrCat(signed_prov_message_.oemcrypto_core_message(),
                             signed_prov_message_.message()));
    }
    EXPECT_EQ(kDevicePrivateKey,
              crypto_util::DecryptAesCbc(message_key_,
                                         prov_response.device_rsa_key_iv(),
                                         prov_response.device_rsa_key()));
    EXPECT_EQ(kDeviceCertificate, prov_response.device_certificate());
    EXPECT_EQ(kNonce, prov_response.nonce());
    EXPECT_EQ(kWrappingKey, prov_response.wrapping_key());
  }

  ClientIdentification client_id_;
  ProvisioningRequest prov_request_;
  SignedProvisioningMessage signed_prov_message_;
  std::string signature_input_;
  std::string message_key_;
  std::string response_;
};

TEST_F(Provisioning30SessionImplProcessTest, InvalidMessage) {
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage("invalid_message", &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, EmptyMessage) {
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage("", &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, MissingMessage) {
  signed_prov_message_.clear_message();
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, MissingSignature) {
  signed_prov_message_.clear_signature();
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, MissingClientId) {
  prov_request_.clear_encrypted_client_id();
  signed_prov_message_.set_message(prov_request_.SerializeAsString());
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, MissingEncryptedClientId) {
  prov_request_.mutable_encrypted_client_id()->clear_encrypted_client_id();
  signed_prov_message_.set_message(prov_request_.SerializeAsString());
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, MissingEncryptedClientIdIv) {
  prov_request_.mutable_encrypted_client_id()->clear_encrypted_client_id_iv();
  signed_prov_message_.set_message(prov_request_.SerializeAsString());
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, MissingEncryptedPrivacyKey) {
  prov_request_.mutable_encrypted_client_id()->clear_encrypted_privacy_key();
  signed_prov_message_.set_message(prov_request_.SerializeAsString());
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, InvalidNonce) {
  // Nonce should be at least 4 buytes.
  const char kNonceWithLessThanFourBytes[] = "xx";
  prov_request_.set_nonce(kNonceWithLessThanFourBytes);
  signed_prov_message_.set_message(prov_request_.SerializeAsString());
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, PrivacyKeyDecryptionFailed) {
  EXPECT_CALL(mock_service_private_key_, Decrypt(kEncryptedPrivacyKey, _))
      .WillOnce(Return(false));
  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, InvalidEncryptedClientId) {
  prov_request_.mutable_encrypted_client_id()->set_encrypted_client_id(
      "invalid_encrypted_client_id");
  signed_prov_message_.set_message(prov_request_.SerializeAsString());
  EXPECT_CALL(mock_service_private_key_, Decrypt(kEncryptedPrivacyKey, _))
      .WillOnce(DoAll(SetArgPointee<1>(kPrivacyKey), Return(true)));

  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, VerifyCertificateChainFailed) {
  EXPECT_CALL(mock_service_private_key_, Decrypt(kEncryptedPrivacyKey, _))
      .WillOnce(DoAll(SetArgPointee<1>(kPrivacyKey), Return(true)));
  EXPECT_CALL(mock_oem_device_cert_,
              DoVerifyCertificateChain(kClientToken, _, _, _))
      .WillOnce(Return(false));

  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest,
       ClearClientIdVerifyCertificateChainFailed) {
  *prov_request_.mutable_client_id() = client_id_;
  prov_request_.clear_encrypted_client_id();
  signed_prov_message_.set_message(prov_request_.SerializeAsString());

  EXPECT_CALL(mock_oem_device_cert_,
              DoVerifyCertificateChain(kClientToken, _, _, _))
      .WillOnce(Return(false));

  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, ClearClientIdInvalidClientIdType) {
  client_id_.set_type(ClientIdentification::KEYBOX);
  *prov_request_.mutable_client_id() = client_id_;
  prov_request_.clear_encrypted_client_id();
  signed_prov_message_.set_message(prov_request_.SerializeAsString());

  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, ClearClientIdMissingToken) {
  client_id_.clear_token();
  *prov_request_.mutable_client_id() = client_id_;
  prov_request_.clear_encrypted_client_id();
  signed_prov_message_.set_message(prov_request_.SerializeAsString());

  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, VerifySignatureFailed) {
  EXPECT_CALL(mock_service_private_key_, Decrypt(kEncryptedPrivacyKey, _))
      .WillOnce(DoAll(SetArgPointee<1>(kPrivacyKey), Return(true)));
  MockRsaPublicKey* mock_cert_public_key = new MockRsaPublicKey;
  EXPECT_CALL(mock_oem_device_cert_,
              DoVerifyCertificateChain(kClientToken, _, _, _))
      .WillOnce(DoAll(SetArgPointee<1>(mock_cert_public_key), Return(true)));
  const HashAlgorithm hash_algorithm =
      HashAlgorithmProtoToEnum(signed_prov_message_.hash_algorithm());
  EXPECT_CALL(*mock_cert_public_key,
              VerifySignature(signed_prov_message_.message(), hash_algorithm,
                              signed_prov_message_.signature()))
      .WillOnce(Return(false));

  std::string response;
  bool done;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, GenerateDeviceCertificateFailed) {
  const uint32_t kSystemId = 1234;
  const char kExpectedOemSerialNumber[] = "test_oem_serial_number";

  EXPECT_CALL(mock_service_private_key_, Decrypt(kEncryptedPrivacyKey, _))
      .WillOnce(DoAll(SetArgPointee<1>(kPrivacyKey), Return(true)));
  MockRsaPublicKey* mock_cert_public_key = new MockRsaPublicKey;
  EXPECT_CALL(mock_oem_device_cert_,
              DoVerifyCertificateChain(kClientToken, _, _, _))
      .WillOnce(DoAll(
          SetArgPointee<1>(mock_cert_public_key), SetArgPointee<2>(kSystemId),
          SetArgPointee<3>(kExpectedOemSerialNumber), Return(true)));
  const HashAlgorithm hash_algorithm =
      HashAlgorithmProtoToEnum(signed_prov_message_.hash_algorithm());
  EXPECT_CALL(*mock_cert_public_key,
              VerifySignature(signed_prov_message_.message(), hash_algorithm,
                              signed_prov_message_.signature()))
      .WillOnce(Return(true));

  EXPECT_CALL(
      mock_engine_impl_,
      GenerateProviderDeviceDrmCertificate(
          kSystemId, kExpectedOemSerialNumber, kProviderId, kDevicePublicKey,
          DeriveSpoid(kClientToken, kProviderId, ""), _))
      .WillOnce(Return(INTERNAL_ERROR));

  std::string response;
  bool done;
  EXPECT_EQ(INTERNAL_ERROR,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response, &done));
}

TEST_F(Provisioning30SessionImplProcessTest, SuccessWithCoreMessage) {
  ProcessMessage(kCoreMessage);
  bool done;
  EXPECT_CALL(mock_service_private_key_, GenerateSignature(_, _, _))
      .WillOnce(DoAll(SaveArg<0>(&signature_input_),
                      SetArgPointee<2>(kSignature), Return(true)));
  ASSERT_EQ(
      OK, session_impl_.ProcessMessage(signed_prov_message_.SerializeAsString(),
                                       &response_, &done));
  ASSERT_TRUE(done);
  ValidateMessage(kCoreMessage);
}

TEST_F(Provisioning30SessionImplProcessTest, SuccessWithEmptyCoreMessage) {
  ProcessMessage(kEmptyCoreMessage);
  bool done;
  EXPECT_CALL(mock_service_private_key_, GenerateSignature(_, _, _))
      .WillOnce(DoAll(SaveArg<0>(&signature_input_),
                      SetArgPointee<2>(kSignature), Return(true)));
  ASSERT_EQ(
      OK, session_impl_.ProcessMessage(signed_prov_message_.SerializeAsString(),
                                       &response_, &done));
  ASSERT_TRUE(done);
  ValidateMessage(kEmptyCoreMessage);
}

TEST_F(Provisioning30SessionImplProcessTest, FailedWithInvalidCoreMessage) {
  ProcessMessage(kInvalidCoreMessage);
  bool done;
  ASSERT_EQ(INTERNAL_ERROR,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response_, &done));
  ASSERT_FALSE(done);
}

TEST_F(Provisioning30SessionImplProcessTest, VerifyHashAlgorithmInResponse) {
  const HashAlgorithm hash_algorithm = HashAlgorithm::kSha256;
  signed_prov_message_.set_hash_algorithm(
      HashAlgorithmEnumToProto(hash_algorithm));
  ProcessMessage(kCoreMessage);
  bool done;
  EXPECT_CALL(mock_service_private_key_, GenerateSignature(_, _, _))
      .WillOnce(DoAll(SaveArg<0>(&signature_input_),
                      SetArgPointee<2>(kSignature), Return(true)));
  ASSERT_EQ(
      OK, session_impl_.ProcessMessage(signed_prov_message_.SerializeAsString(),
                                       &response_, &done));
  ASSERT_TRUE(done);
  SignedProvisioningMessage signed_response;
  signed_response.ParseFromString(response_);
  EXPECT_EQ(signed_response.hash_algorithm(),
            signed_prov_message_.hash_algorithm());
}
}  // namespace widevine