// Copyright 2018 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
// These tests are for the generic crypto operations.  They call on the
// CdmEngine class and exercise the classes below it as well.  In
// particular, we assume that the OEMCrypo layer works, and has a valid keybox.
// This is because we need a valid RSA certificate, and will attempt to connect
// to the provisioning server to request one if we don't.

#include <gtest/gtest.h>
#include <memory>
#include <string>

#include "cdm_engine.h"
#include "config_test_env.h"
#include "license_holder.h"
#include "log.h"
#include "oec_session_util.h"
#include "oemcrypto_session_tests_helper.h"
#include "oemcrypto_types.h"
#include "platform.h"
#include "properties.h"
#include "string_conversions.h"
#include "test_base.h"
#include "test_printers.h"
#include "url_request.h"
#include "wv_cdm_constants.h"
#include "wv_cdm_types.h"

using wvutil::a2b_hex;

namespace wvcdm {

class WvGenericCryptoTest : public WvCdmTestBaseWithEngine {
 public:
  WvGenericCryptoTest() : holder_("CDM_GenericCrypto", &cdm_engine_, config_) {}

  void SetUp() override {
    WvCdmTestBase::SetUp();
    if (!wvoec::global_features.generic_crypto) {
      GTEST_SKIP() << "Test for devices with generic crypto API only";
    }
    // TODO: b/305093063 - Remove when Drm Reprovisioning server is implemented.
    if (wvoec::global_features.provisioning_method ==
        OEMCrypto_DrmReprovisioning) {
      GTEST_SKIP()
          << "Skipping until Drm Reprovisioning server support is implemented.";
    }
    EnsureProvisioned();
    ASSERT_NO_FATAL_FAILURE(holder_.OpenSession());
    ASSERT_NO_FATAL_FAILURE(holder_.FetchLicense());
    ASSERT_NO_FATAL_FAILURE(holder_.LoadLicense());

    ency_id_ = "encrypt-key-----";
    ency_key_ = a2b_hex("0102030405060708090a0b0c0d0e0f10");
    dency_id_ = "decrypt-key-----";
    dency_key_ = a2b_hex("AA02030405060708090a0b0c0d0e0f10");
    siggy_id_ = "sign-key--------";
    siggy_key_ = a2b_hex(
        "BB02030405060708090a0b0c0d0e0f100102030405060708090a0b0c0d0e0f10");
    vou_id_ = "verify-key------";
    vou_key_ = a2b_hex(
        "CC02030405060708090a0b0c0d0e0f100102030405060708090a0b0c0d0e0f10");
    both_id_ = "enc-and-dec-key-";
    both_key_ = a2b_hex("DD02030405060708090a0b0c0d0e0f10");
    sign_and_verify_id_ = "sign-and-verify-";
    sign_and_verify_key_ = a2b_hex(
        "EE02030405060708090a0b0c0d0e0f100102030405060708090a0b0c0d0e0f10");

    StripeBuffer(&in_vector_, CONTENT_KEY_SIZE * 15, '1');
    in_buffer_ = std::string(in_vector_.begin(), in_vector_.end());

    StripeBuffer(&iv_vector_, KEY_IV_SIZE, 'a');
    iv_ = std::string(iv_vector_.begin(), iv_vector_.end());
  }

  void TearDown() override {
    // TODO: b/305093063 - Remove when Drm Reprovisioning server is implemented.
    if (IsSkipped()) return;
    holder_.CloseSession();
  }

 protected:
  LicenseHolder holder_;

  KeyId ency_id_;
  KeyId dency_id_;
  KeyId siggy_id_;
  KeyId vou_id_;
  KeyId both_id_;
  KeyId sign_and_verify_id_;

  std::vector<uint8_t> ency_key_;
  std::vector<uint8_t> dency_key_;
  std::vector<uint8_t> siggy_key_;
  std::vector<uint8_t> vou_key_;
  std::vector<uint8_t> both_key_;
  std::vector<uint8_t> sign_and_verify_key_;

  std::vector<uint8_t> in_vector_;
  std::vector<uint8_t> iv_vector_;
  std::string in_buffer_;
  std::string iv_;
};

TEST_F(WvGenericCryptoTest, GenericEncryptGood) {
  std::string encrypted = Aes128CbcEncrypt(ency_key_, in_vector_, iv_vector_);
  std::string out_buffer;
  EXPECT_EQ(NO_ERROR, cdm_engine_.GenericEncrypt(
                          holder_.session_id(), in_buffer_, ency_id_, iv_,
                          wvcdm::kEncryptionAlgorithmAesCbc128, &out_buffer));
  EXPECT_EQ(encrypted, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericEncryptNoKey) {
  std::string encrypted = Aes128CbcEncrypt(ency_key_, in_vector_, iv_vector_);
  std::string out_buffer;
  KeyId key_id("no_key");
  EXPECT_NE(NO_ERROR, cdm_engine_.GenericEncrypt(
                          holder_.session_id(), in_buffer_, key_id, iv_,
                          wvcdm::kEncryptionAlgorithmAesCbc128, &out_buffer));
  EXPECT_NE(encrypted, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericEncryptKeyNotAllowed) {
  // Trying to use Decrypt key to encrypt, which is not allowed.
  KeyId key_id = dency_id_;
  std::string encrypted = Aes128CbcEncrypt(dency_key_, in_vector_, iv_vector_);
  std::string out_buffer;
  EXPECT_EQ(UNKNOWN_ERROR,
            cdm_engine_.GenericEncrypt(
                holder_.session_id(), in_buffer_, key_id, iv_,
                wvcdm::kEncryptionAlgorithmAesCbc128, &out_buffer));
  EXPECT_NE(encrypted, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericDecryptGood) {
  std::string decrypted = Aes128CbcDecrypt(dency_key_, in_vector_, iv_vector_);
  std::string out_buffer;
  EXPECT_EQ(NO_ERROR, cdm_engine_.GenericDecrypt(
                          holder_.session_id(), in_buffer_, dency_id_, iv_,
                          wvcdm::kEncryptionAlgorithmAesCbc128, &out_buffer));
  EXPECT_EQ(decrypted, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericDecryptNoKey) {
  std::string decrypted = Aes128CbcDecrypt(dency_key_, in_vector_, iv_vector_);
  std::string out_buffer;
  KeyId key_id = "no_key";
  EXPECT_NE(NO_ERROR, cdm_engine_.GenericDecrypt(
                          holder_.session_id(), in_buffer_, key_id, iv_,
                          wvcdm::kEncryptionAlgorithmAesCbc128, &out_buffer));
  EXPECT_NE(decrypted, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericDecryptKeyNotAllowed) {
  // Trying to use Encrypt key to decrypt, which is not allowed.
  KeyId key_id = ency_id_;
  std::string decrypted = Aes128CbcDecrypt(ency_key_, in_vector_, iv_vector_);
  std::string out_buffer;
  EXPECT_EQ(UNKNOWN_ERROR,
            cdm_engine_.GenericDecrypt(
                holder_.session_id(), in_buffer_, key_id, iv_,
                wvcdm::kEncryptionAlgorithmAesCbc128, &out_buffer));
  EXPECT_NE(decrypted, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericSignGood) {
  std::string out_buffer;
  std::string signature = SignHMAC(in_buffer_, siggy_key_);
  EXPECT_EQ(NO_ERROR, cdm_engine_.GenericSign(
                          holder_.session_id(), in_buffer_, siggy_id_,
                          wvcdm::kSigningAlgorithmHmacSha256, &out_buffer));
  EXPECT_EQ(signature, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericSignKeyNotAllowed) {
  // Wrong key
  std::string key_id = vou_id_;
  std::string out_buffer;
  std::string signature = SignHMAC(in_buffer_, siggy_key_);
  EXPECT_EQ(
      UNKNOWN_ERROR,
      cdm_engine_.GenericSign(holder_.session_id(), in_buffer_, key_id,
                              wvcdm::kSigningAlgorithmHmacSha256, &out_buffer));
  EXPECT_NE(signature, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericVerifyGood) {
  std::string signature = SignHMAC(in_buffer_, vou_key_);
  EXPECT_EQ(NO_ERROR, cdm_engine_.GenericVerify(
                          holder_.session_id(), in_buffer_, vou_id_,
                          wvcdm::kSigningAlgorithmHmacSha256, signature));
}

TEST_F(WvGenericCryptoTest, GenericVerifyKeyNotAllowed) {
  // Wrong key
  std::string key_id = siggy_id_;
  std::string signature = SignHMAC(in_buffer_, siggy_key_);
  EXPECT_EQ(UNKNOWN_ERROR, cdm_engine_.GenericVerify(
                               holder_.session_id(), in_buffer_, key_id,
                               wvcdm::kSigningAlgorithmHmacSha256, signature));
}

TEST_F(WvGenericCryptoTest, GenericVerifyBadSignature) {
  std::string signature(MAC_KEY_SIZE, 's');
  // OEMCrypto error is OEMCrypto_ERROR_SIGNATURE_FAILURE
  EXPECT_EQ(UNKNOWN_ERROR, cdm_engine_.GenericVerify(
                               holder_.session_id(), in_buffer_, vou_id_,
                               wvcdm::kSigningAlgorithmHmacSha256, signature));
}

TEST_F(WvGenericCryptoTest, GenericEncryptDecrypt) {
  std::string encrypted = Aes128CbcEncrypt(both_key_, in_vector_, iv_vector_);
  std::string out_buffer;
  EXPECT_EQ(NO_ERROR, cdm_engine_.GenericEncrypt(
                          holder_.session_id(), in_buffer_, both_id_, iv_,
                          wvcdm::kEncryptionAlgorithmAesCbc128, &out_buffer));
  EXPECT_EQ(encrypted, out_buffer);

  std::string decrypted = Aes128CbcDecrypt(dency_key_, in_vector_, iv_vector_);
  EXPECT_EQ(NO_ERROR, cdm_engine_.GenericDecrypt(
                          holder_.session_id(), encrypted, both_id_, iv_,
                          wvcdm::kEncryptionAlgorithmAesCbc128, &out_buffer));
  EXPECT_EQ(in_buffer_, out_buffer);
}

TEST_F(WvGenericCryptoTest, GenericSignVerify) {
  std::string out_buffer;
  std::string signature = SignHMAC(in_buffer_, sign_and_verify_key_);
  EXPECT_EQ(NO_ERROR, cdm_engine_.GenericSign(
                          holder_.session_id(), in_buffer_, sign_and_verify_id_,
                          wvcdm::kSigningAlgorithmHmacSha256, &out_buffer));
  EXPECT_EQ(signature, out_buffer);
  EXPECT_EQ(NO_ERROR, cdm_engine_.GenericVerify(
                          holder_.session_id(), in_buffer_, sign_and_verify_id_,
                          wvcdm::kSigningAlgorithmHmacSha256, signature));
}
}  // namespace wvcdm