// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

#include <ctime>
#include <memory>
#include <string>
#include <vector>

#include "api/license_test_helper.h"
#include "api/test_data.h"
#include "crypto_utils/rsa_key.h"
#include "testing/include/gtest/gtest.h"

using RsaPublicKey = widevine::RsaPublicKey;

namespace {

const uint8_t kDummyMessage[] = {
    0x1e, 0x70, 0xbd, 0xeb, 0x24, 0xf2, 0x9d, 0x05, 0xc5, 0xb5,
    0xf4, 0xca, 0xe6, 0x1d, 0x01, 0x97, 0x29, 0xf4, 0xe0, 0x7c,
    0xfd, 0xcc, 0x97, 0x8d, 0xc2, 0xbb, 0x2d, 0x9b, 0x6b, 0x45,
    0x06, 0xbd, 0x2c, 0x66, 0x10, 0x42, 0x73, 0x8d, 0x88, 0x9b,
    0x18, 0xcc, 0xcb, 0x7e, 0x43, 0x23, 0x06, 0xe9, 0x8f, 0x8f,
};
const size_t kDummyMessageSize = sizeof(kDummyMessage);

// Size of a license signature. This must be big enough to hold the signature
// returned by WB_License_SignLicenseRequest().
constexpr size_t kSignatureSize = 256;

std::string AsString(const uint8_t* buffer, size_t buffer_size) {
  return std::string(reinterpret_cast<const char*>(buffer), buffer_size);
}

// Used to test the rest of the whitebox API as this class creates |whitebox_|
// during setup, so it is available for all tests. It also creates |public_key_|
// which is the matching RSA public key that can be used to verify signing
// operations inside the whitebox.
class LicenseWhiteboxTest : public ::testing::Test {
 protected:
  void SetUp() override {
    const std::vector<uint8_t> init_data = GetLicenseInitData();
    ASSERT_EQ(WB_License_Create(init_data.data(), init_data.size(), &whitebox_),
              WB_RESULT_OK);

    std::vector<uint8_t> public_key = GetMatchingLicensePublicKey();
    public_key_.reset(
        RsaPublicKey::Create(AsString(public_key.data(), public_key.size())));
    ASSERT_TRUE(public_key_);
  }

  void TearDown() override { WB_License_Delete(whitebox_); }

  // Public key as an usable object.
  std::unique_ptr<RsaPublicKey> public_key_;

  WB_License_Whitebox* whitebox_;
};

TEST_F(LicenseWhiteboxTest, SignAndVerifySucceedsWithRandomData) {
  uint8_t signature[kSignatureSize];
  size_t signature_size = kSignatureSize;
  ASSERT_EQ(
      WB_License_SignLicenseRequest(whitebox_, kDummyMessage, kDummyMessageSize,
                                    signature, &signature_size),
      WB_RESULT_OK);

  // Verify the signature.
  ASSERT_TRUE(
      public_key_->VerifySignature(AsString(kDummyMessage, kDummyMessageSize),
                                   AsString(signature, signature_size)));
}

TEST_F(LicenseWhiteboxTest, SignAndVerifySucceedsWithLicenseRequest) {
  const std::string license_request = widevine::CreateLicenseRequest();

  uint8_t signature[kSignatureSize];
  size_t signature_size = kSignatureSize;
  ASSERT_EQ(
      WB_License_SignLicenseRequest(
          whitebox_, reinterpret_cast<const uint8_t*>(license_request.data()),
          license_request.size(), signature, &signature_size),
      WB_RESULT_OK);

  ASSERT_TRUE(public_key_->VerifySignature(
      license_request, AsString(signature, signature_size)));
}

TEST_F(LicenseWhiteboxTest, SignLicenseRequestFailsWithWhiteboxNull) {
  uint8_t signature[kSignatureSize];
  size_t signature_size = kSignatureSize;
  ASSERT_EQ(
      WB_License_SignLicenseRequest(nullptr, kDummyMessage, kDummyMessageSize,
                                    signature, &signature_size),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, SignLicenseRequestFailsWithMessageNull) {
  uint8_t signature[kSignatureSize];
  size_t signature_size = kSignatureSize;
  ASSERT_EQ(WB_License_SignLicenseRequest(whitebox_, nullptr, kDummyMessageSize,
                                          signature, &signature_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, SignLicenseRequestFailsWithMessageSizeZero) {
  uint8_t signature[kSignatureSize];
  size_t signature_size = kSignatureSize;
  ASSERT_EQ(WB_License_SignLicenseRequest(whitebox_, kDummyMessage, 0,
                                          signature, &signature_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, SignLicenseRequestFailsWithSignatureNull) {
  size_t signature_size = kSignatureSize;
  ASSERT_EQ(
      WB_License_SignLicenseRequest(whitebox_, kDummyMessage, kDummyMessageSize,
                                    nullptr, &signature_size),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, SignLicenseRequestFailsWithSignatureSizeNull) {
  uint8_t signature[kSignatureSize];
  ASSERT_EQ(
      WB_License_SignLicenseRequest(whitebox_, kDummyMessage, kDummyMessageSize,
                                    signature, nullptr),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, SignLicenseRequestFailsWithSignatureSizeZero) {
  uint8_t signature[kSignatureSize];
  size_t signature_size = 0;
  ASSERT_EQ(
      WB_License_SignLicenseRequest(whitebox_, kDummyMessage, kDummyMessageSize,
                                    signature, &signature_size),
      WB_RESULT_BUFFER_TOO_SMALL);
}

TEST_F(LicenseWhiteboxTest, SignLicenseRequestFailsWithSignatureTooSmall) {
  uint8_t signature[10];
  size_t signature_size = sizeof(signature);
  ASSERT_EQ(
      WB_License_SignLicenseRequest(whitebox_, kDummyMessage, kDummyMessageSize,
                                    signature, &signature_size),
      WB_RESULT_BUFFER_TOO_SMALL);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithWhiteboxNull) {
  std::vector<uint8_t> message{1, 2, 3};
  std::vector<uint8_t> signature{4, 5, 6, 7};
  std::vector<uint8_t> session_key{1, 2, 3, 4, 5, 6, 7, 8};
  std::vector<uint8_t> license_request{10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                nullptr, message.data(), message.size(), signature.data(),
                signature.size(), session_key.data(), session_key.size(),
                license_request.data(), license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithMessageNull) {
  std::vector<uint8_t> message{1, 2, 3};
  std::vector<uint8_t> signature{4, 5, 6, 7};
  std::vector<uint8_t> session_key{1, 2, 3, 4, 5, 6, 7, 8};
  std::vector<uint8_t> license_request{10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, nullptr, message.size(), signature.data(),
                signature.size(), session_key.data(), session_key.size(),
                license_request.data(), license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithMessageSizeZero) {
  std::vector<uint8_t> message{/* empty */};
  std::vector<uint8_t> signature{4, 5, 6, 7};
  std::vector<uint8_t> session_key{1, 2, 3, 4, 5, 6, 7, 8};
  std::vector<uint8_t> license_request{10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, message.data(), message.size(), signature.data(),
                signature.size(), session_key.data(), session_key.size(),
                license_request.data(), license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithSignatureNull) {
  std::vector<uint8_t> message{1, 2, 3};
  std::vector<uint8_t> signature{4, 5, 6, 7};
  std::vector<uint8_t> session_key{1, 2, 3, 4, 5, 6, 7, 8};
  std::vector<uint8_t> license_request{10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, nullptr, message.size(), nullptr, signature.size(),
                session_key.data(), session_key.size(), license_request.data(),
                license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithSignatureSizeZero) {
  std::vector<uint8_t> message{1, 2, 3};
  std::vector<uint8_t> signature{/* empty */};
  std::vector<uint8_t> session_key{1, 2, 3, 4, 5, 6, 7, 8};
  std::vector<uint8_t> license_request{10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, message.data(), message.size(), signature.data(),
                signature.size(), session_key.data(), session_key.size(),
                license_request.data(), license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithSessionKeyNull) {
  std::vector<uint8_t> message{1, 2, 3};
  std::vector<uint8_t> signature{4, 5, 6, 7};
  std::vector<uint8_t> session_key{1, 2, 3, 4, 5, 6, 7, 8};
  std::vector<uint8_t> license_request{10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, message.data(), message.size(), signature.data(),
                signature.size(), nullptr, session_key.size(),
                license_request.data(), license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithSessionKeySizeZero) {
  std::vector<uint8_t> message{1, 2, 3};
  std::vector<uint8_t> signature{4, 5, 6, 7};
  std::vector<uint8_t> session_key{/* empty */};
  std::vector<uint8_t> license_request{10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, message.data(), message.size(), signature.data(),
                signature.size(), session_key.data(), session_key.size(),
                license_request.data(), license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithRequestNull) {
  std::vector<uint8_t> message{1, 2, 3};
  std::vector<uint8_t> signature{4, 5, 6, 7};
  std::vector<uint8_t> session_key{1, 2, 3, 4, 5, 6, 7, 8};
  std::vector<uint8_t> license_request{10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, message.data(), message.size(), signature.data(),
                signature.size(), session_key.data(), session_key.size(),
                nullptr, license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxTest, ProcessLicenseResponseFailsWithRequestSizeZero) {
  std::vector<uint8_t> message{1, 2, 3};
  std::vector<uint8_t> signature{4, 5, 6, 7};
  std::vector<uint8_t> session_key{1, 2, 3, 4, 5, 6, 7, 8};
  std::vector<uint8_t> license_request{/* empty */};
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, message.data(), message.size(), signature.data(),
                signature.size(), session_key.data(), session_key.size(),
                license_request.data(), license_request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

}  // namespace