////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 Google Inc.
//
// 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/provisioning_session_impl.h"

#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "common/aes_cbc_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"

using ::testing::_;
using ::testing::ByMove;
using ::testing::DoAll;
using ::testing::Return;
using ::testing::SaveArg;
using ::testing::SetArgPointee;

namespace {
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_CONST_METHOD6(GenerateProviderDeviceDrmCertificate,
                     ProvisioningStatus(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));
};

class MockOemDeviceCert : public OemDeviceCert {
 public:
  // gmock does not support SetArgPointee on std::unique_ptr, so we have to
  // workaround it with a trick.
  MOCK_CONST_METHOD4(DoVerifyCertificateChain,
                     bool(const std::string& certificate_chain,
                          RsaPublicKey** leaf_public_key, uint32_t* system_id,
                          std::string* oem_ca_serial_number));
  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 ProvisioningSessionImplTest : public ::testing::Test {
 protected:
  ProvisioningSessionImplTest()
      : 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_));
  }

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

TEST_F(ProvisioningSessionImplTest, InitializeWithInvalidPublicKey) {
  EXPECT_CALL(*mock_rsa_key_factory_,
              CreateFromPkcs1PublicKey(kDevicePublicKey))
      .WillOnce(Return(ByMove(nullptr)));
  EXPECT_EQ(
      INVALID_DEVICE_PUBLIC_KEY,
      session_impl_.Initialize(kDevicePublicKey, kDevicePrivateKey));
}

TEST_F(ProvisioningSessionImplTest, InitializeWithInvalidPrivateKey) {
  EXPECT_CALL(*mock_rsa_key_factory_,
              CreateFromPkcs1PublicKey(kDevicePublicKey))
      .WillOnce(
          Return(ByMove(std::unique_ptr<RsaPublicKey>(new MockRsaPublicKey))));
  EXPECT_CALL(*mock_rsa_key_factory_,
              CreateFromPkcs1PrivateKey(kDevicePrivateKey))
      .WillOnce(Return(ByMove(nullptr)));
  EXPECT_EQ(
      INVALID_DEVICE_PRIVATE_KEY,
      session_impl_.Initialize(kDevicePublicKey, kDevicePrivateKey));
}

TEST_F(ProvisioningSessionImplTest, InitializeWithMismatchPublicPrivateKey) {
  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_,
              CreateFromPkcs1PrivateKey(kDevicePrivateKey))
      .WillOnce(Return(
          ByMove(std::unique_ptr<RsaPrivateKey>(new MockRsaPrivateKey))));
  EXPECT_CALL(*mock_rsa_public_key, MatchesPrivateKey(_))
      .WillOnce(Return(false));
  EXPECT_EQ(
      INVALID_DEVICE_PRIVATE_KEY,
      session_impl_.Initialize(kDevicePublicKey, kDevicePrivateKey));
}

class ProvisioningSessionImplProcessTest : public ProvisioningSessionImplTest {
 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_,
                CreateFromPkcs1PrivateKey(kDevicePrivateKey))
        .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");
  }

  ClientIdentification client_id_;
  ProvisioningRequest prov_request_;
  SignedProvisioningMessage signed_prov_message_;
};

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

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

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

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

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

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

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

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

TEST_F(ProvisioningSessionImplProcessTest, 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;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response));
}

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

TEST_F(ProvisioningSessionImplProcessTest, 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;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response));
}

TEST_F(ProvisioningSessionImplProcessTest, 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;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response));
}

TEST_F(ProvisioningSessionImplProcessTest,
       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;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response));
}

TEST_F(ProvisioningSessionImplProcessTest, 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;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response));
}

TEST_F(ProvisioningSessionImplProcessTest, 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;
  EXPECT_EQ(INVALID_REQUEST_MESSAGE,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response));
}

TEST_F(ProvisioningSessionImplProcessTest, 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)));
  EXPECT_CALL(*mock_cert_public_key,
              VerifySignature(signed_prov_message_.message(),
                              signed_prov_message_.signature()))
      .WillOnce(Return(false));

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

TEST_F(ProvisioningSessionImplProcessTest, 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)));
  EXPECT_CALL(*mock_cert_public_key,
              VerifySignature(signed_prov_message_.message(),
                              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;
  EXPECT_EQ(INTERNAL_ERROR,
            session_impl_.ProcessMessage(
                signed_prov_message_.SerializeAsString(), &response));
}

TEST_F(ProvisioningSessionImplProcessTest, Success) {
  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)));
  EXPECT_CALL(*mock_cert_public_key,
              VerifySignature(signed_prov_message_.message(),
                              signed_prov_message_.signature()))
      .WillOnce(Return(true));

  EXPECT_CALL(mock_engine_impl_,
              GenerateProviderDeviceDrmCertificate(kSystemId, _, _,
                                                   kDevicePublicKey, _, _))
      .WillOnce(DoAll(SetArgPointee<5>(kDeviceCertificate), Return(OK)));

  std::string message_key;
  EXPECT_CALL(*mock_cert_public_key, Encrypt(_, _))
      .WillOnce(DoAll(SaveArg<0>(&message_key),
                      SetArgPointee<1>(kWrappingKey), Return(true)));
  std::string message;
  EXPECT_CALL(mock_service_private_key_, GenerateSignature(_, _))
      .WillOnce(DoAll(SaveArg<0>(&message),
                      SetArgPointee<1>(kSignature), Return(true)));

  std::string response;
  ASSERT_EQ(OK, session_impl_.ProcessMessage(
                    signed_prov_message_.SerializeAsString(), &response));

  // Verify the response.
  SignedProvisioningMessage signed_prov_message;
  ASSERT_TRUE(signed_prov_message.ParseFromString(response));
  EXPECT_EQ(message, signed_prov_message.message());
  EXPECT_EQ(kSignature, signed_prov_message.signature());

  ProvisioningResponse prov_response;
  ASSERT_TRUE(prov_response.ParseFromString(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());
}

}  // namespace widevine