// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

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

#include "api/license_whitebox_test_base.h"
#include "api/test_data.h"
#include "api/test_license_builder.h"
#include "crypto_utils/crypto_util.h"
#include "crypto_utils/rsa_key.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace widevine {

class LicenseWhiteboxSignRenewalRequestTest : public LicenseWhiteboxTestBase {
 protected:
  void SetUp() override {
    LicenseWhiteboxTestBase::SetUp();

    // We don't know the actual signature size, so make the buffer large enough
    // that it should accomidate any change in signature size.
    signature_size_ = 256;
    signature_.resize(signature_size_);
  }

  void LoadLicense(const std::vector<uint8_t>& padding) {
    TestLicenseBuilder builder;
    builder.AddSigningKey(signing_key_, padding);
    // Add a throw away key. We just need a key in the license since a license
    // should always have a content key.
    builder.AddStubbedContentKey();

    License license;
    builder.Build(*public_key_, &license);

    ASSERT_EQ(
        WB_License_ProcessLicenseResponse(
            whitebox_, license.core_message.data(), license.core_message.size(),
            license.message.data(), license.message.size(),
            license.signature.data(), license.signature.size(),
            license.session_key.data(), license.session_key.size(),
            license.request.data(), license.request.size()),
        WB_RESULT_OK);
  }

  // Get the expected signature for |message|. By returning the message, this
  // will allow us to assert in the test, giving us better output in the case of
  // a failure.
  std::vector<uint8_t> GetSignature(const std::vector<uint8_t>& message) {
    // The client key is the second half of the signing key.
    std::string key_str(signing_key_.begin(), signing_key_.end());
    key_str.erase(0, crypto_util::kSigningKeySizeBytes);

    const std::string signature = crypto_util::CreateSignatureHmacSha256(
        key_str, std::string(message.begin(), message.end()));

    return std::vector<uint8_t>(signature.begin(), signature.end());
  }

  const std::string session_key_ = "0123456789ABCDEF";

  const std::vector<uint8_t> signing_key_ =
      TestLicenseBuilder::DefaultSigningKey();

  size_t signature_size_;
  std::vector<uint8_t> signature_;

  const std::vector<uint8_t> garbage_request_ = {
      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,
  };
};

TEST_F(LicenseWhiteboxSignRenewalRequestTest, SuccessWithInvalidRequest) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_SignRenewalRequest(whitebox_, garbage_request_.data(),
                                          garbage_request_.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_OK);

  signature_.resize(signature_size_);
  ASSERT_EQ(signature_, GetSignature(garbage_request_));
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest,
       SuccessWithSigningKeyPKSC8Padding) {
  LoadLicense(TestLicenseBuilder::PKSC8Padding());

  ASSERT_EQ(WB_License_SignRenewalRequest(whitebox_, garbage_request_.data(),
                                          garbage_request_.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_OK);

  signature_.resize(signature_size_);
  ASSERT_EQ(signature_, GetSignature(garbage_request_));
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest, InvalidParameterForNullWhitebox) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_SignRenewalRequest(nullptr, garbage_request_.data(),
                                          garbage_request_.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest, InvalidParameterForNullMessage) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_SignRenewalRequest(whitebox_, nullptr, garbage_request_.size(),
                                    signature_.data(), &signature_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest,
       InvalidParameterForZeroMessageSize) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_SignRenewalRequest(whitebox_, garbage_request_.data(), 0,
                                          signature_.data(), &signature_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest,
       InvalidParameterForNullSignature) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_SignRenewalRequest(whitebox_, garbage_request_.data(),
                                          garbage_request_.size(), nullptr,
                                          &signature_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest,
       InvalidParameterForNullSignatureSize) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_SignRenewalRequest(whitebox_, garbage_request_.data(),
                                          garbage_request_.size(),
                                          signature_.data(), nullptr),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest, BufferTooSmall) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  // We need the signature to be too small. While it would be possible to use
  // zero, using a non-zero value ensures that we are not combining "empty" and
  // "too small".
  signature_size_ = 1;

  ASSERT_EQ(WB_License_SignRenewalRequest(whitebox_, garbage_request_.data(),
                                          garbage_request_.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_BUFFER_TOO_SMALL);

  // Since the API does not limit the signature size, we can't specify the
  // actual expected size, however, it should at least be greater than our "too
  // small" size.
  ASSERT_GT(signature_size_, 1u);
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest, InvalidStateForNoLicense) {
  // Unlike the other tests, we do not call LoadLicense() because we need to
  // have no license loaded in order to have no renewal key, which is the
  // criteria WB_RESULT_INVALID_STATE.

  ASSERT_EQ(WB_License_SignRenewalRequest(whitebox_, garbage_request_.data(),
                                          garbage_request_.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_INVALID_STATE);
}

TEST_F(LicenseWhiteboxSignRenewalRequestTest, InvalidStateForNoSigningKey) {
  // Make a license with no signing key but has a content key. Every license
  // must have a content key.
  TestLicenseBuilder builder;
  builder.AddStubbedContentKey();

  License license;
  builder.Build(*public_key_, &license);

  ASSERT_EQ(
      WB_License_ProcessLicenseResponse(
          whitebox_, license.core_message.data(), license.core_message.size(),
          license.message.data(), license.message.size(),
          license.signature.data(), license.signature.size(),
          license.session_key.data(), license.session_key.size(),
          license.request.data(), license.request.size()),
      WB_RESULT_OK);

  ASSERT_EQ(WB_License_SignRenewalRequest(whitebox_, garbage_request_.data(),
                                          garbage_request_.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_INVALID_STATE);
}

}  // namespace widevine