////////////////////////////////////////////////////////////////////////////////
// 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 "common/wvm_token_handler.h"

#include "testing/gmock.h"
#include "testing/gunit.h"
#include "absl/strings/escaping.h"
#include "absl/strings/string_view.h"
#include "common/wvm_test_keys.h"

using widevine::wvm_test_keys::kTestSystemId;
using widevine::wvm_test_keys::kTestSystemId3Des;
using widevine::wvm_test_keys::kTestPreprovKeyHex;
using widevine::wvm_test_keys::kTestDeviceKey1Hex;
using widevine::wvm_test_keys::kTestDeviceKey2Hex;
using widevine::wvm_test_keys::kTestDeviceKey3DesHex;
using widevine::wvm_test_keys::kTestToken1Hex;
using widevine::wvm_test_keys::kTestToken2Hex;
using widevine::wvm_test_keys::kTestToken3DesHex;
using widevine::wvm_test_keys::GetPreprovKeyVector;

namespace widevine {

using absl::BytesToHexString;
using absl::HexStringToBytes;

// TODO(user): Add EXPECT_OK macro to testing/gmock.h.
// (b/37545268).
#define EXPECT_OK(expression) \
  EXPECT_EQ(util::error::OK, expression.error_code())


TEST(WvmTokenHandlerTest, GetSystemId) {
  EXPECT_EQ(kTestSystemId,
            WvmTokenHandler::GetSystemId(HexStringToBytes(kTestToken1Hex)));
  EXPECT_EQ(kTestSystemId,
            WvmTokenHandler::GetSystemId(HexStringToBytes(kTestToken2Hex)));
}

TEST(WvmTokenHandlerTest, GetEncryptedUniqueId) {
  EXPECT_EQ(
      HexStringToBytes("8e1ebfe037828096ca6538b4f6f4bcb5"),
      WvmTokenHandler::GetEncryptedUniqueId(HexStringToBytes(kTestToken1Hex)));
  EXPECT_EQ(
      HexStringToBytes("d906feebe1750c5886ff77c2dfa31bb4"),
      WvmTokenHandler::GetEncryptedUniqueId(HexStringToBytes(kTestToken2Hex)));
}

TEST(WvmTokenHandlerTest, DecryptDeviceKeyWithPreprovKey) {
  util::Status status;
  std::string device_key;

  status = WvmTokenHandler::DecryptDeviceKeyWithPreprovKey(
      HexStringToBytes(kTestPreprovKeyHex), HexStringToBytes(kTestToken1Hex),
      &device_key);
  EXPECT_OK(status) << status;
  EXPECT_EQ(kTestDeviceKey1Hex, BytesToHexString(device_key));

  status = WvmTokenHandler::DecryptDeviceKeyWithPreprovKey(
      HexStringToBytes(kTestPreprovKeyHex), HexStringToBytes(kTestToken2Hex),
      &device_key);
  EXPECT_OK(status);
  EXPECT_EQ(kTestDeviceKey2Hex, BytesToHexString(device_key));

  // Test with invalid token. Hash failure should produce PERMISSION_DENIED.
  device_key.clear();
  std::string token = HexStringToBytes(
      "00000002000001129e1ebfe037828096ca6538b4f6f4bcb51c2b7191cf037e98"
      "beaa24924907e128f9ff49b54a165cd9c33e6547537eb4d29fb7e8df3c2c1cd9"
      "2517a12f4922953e");
  status = WvmTokenHandler::DecryptDeviceKeyWithPreprovKey(
      HexStringToBytes(kTestPreprovKeyHex), token, &device_key);
  EXPECT_FALSE(status.ok());
  EXPECT_EQ(util::error::PERMISSION_DENIED, status.error_code());
  EXPECT_TRUE(device_key.empty());
}

// See b/68798704 for background of this test.
// It is important to keep this test *before* all the DecryptDeviceKey* tests
// below, in which WvmTokenHandler::SetPreprovKeys() would be called.
TEST(WvmTokenHandlerTest, DecryptDeviceKey_PreprovKeysNullPtr) {
  // Not calling WvmTokenHandler::SetPreprovKeys()
  // So preprov_keys_ would be nullptr.
  util::Status status;
  std::string device_key;
  status = WvmTokenHandler::DecryptDeviceKey(HexStringToBytes(kTestToken1Hex),
                                             &device_key, nullptr, nullptr);
  EXPECT_FALSE(status.ok());
  EXPECT_EQ(util::error::INVALID_ARGUMENT, status.error_code());
}

// Same tests as DecryptDeviceKeyWithPreprovKey(), but we use the handler's
// table of preprov keys instead of providing our own.
TEST(WvmTokenHandlerTest, DecryptDeviceKey) {
  util::Status status;
  std::string device_key;
  WvmTokenHandler::SetPreprovKeys(GetPreprovKeyVector());

  status = WvmTokenHandler::DecryptDeviceKey(HexStringToBytes(kTestToken1Hex),
                                             &device_key, nullptr, nullptr);
  EXPECT_OK(status);
  EXPECT_EQ(HexStringToBytes(kTestDeviceKey1Hex), device_key);

  status = WvmTokenHandler::DecryptDeviceKey(HexStringToBytes(kTestToken2Hex),
                                             &device_key, nullptr, nullptr);
  EXPECT_OK(status);
  EXPECT_EQ(HexStringToBytes(kTestDeviceKey2Hex), device_key);

  // Test with invalid token. Hash failure should produce PERMISSION_DENIED.
  device_key.clear();
  std::string token = HexStringToBytes(
      "00000002000001129e1ebfe037828096ca6538b4f6f4bcb51c2b7191cf037e98"
      "beaa24924907e128f9ff49b54a165cd9c33e6547537eb4d29fb7e8df3c2c1cd9"
      "2517a12f4922953e");
  status =
      WvmTokenHandler::DecryptDeviceKey(token, &device_key, nullptr, nullptr);
  EXPECT_FALSE(status.ok());
  EXPECT_EQ(util::error::PERMISSION_DENIED, status.error_code());
  EXPECT_TRUE(device_key.empty());

  // Test with nonexistent system id. Should produce NOT_FOUND.
  device_key.clear();
  token = HexStringToBytes(
      "00000002555555559e1ebfe037828096ca6538b4f6f4bcb51c2b7191cf037e98"
      "beaa24924907e128f9ff49b54a165cd9c33e6547537eb4d29fb7e8df3c2c1cd9"
      "2517a12f4922953e");
  status =
      WvmTokenHandler::DecryptDeviceKey(token, &device_key, nullptr, nullptr);
  EXPECT_FALSE(status.ok());
  EXPECT_EQ(util::error::NOT_FOUND, status.error_code());
  EXPECT_TRUE(device_key.empty());
}

TEST(WvmTokenHandlerTest, GetEncryptedAssetKey) {
  WvmTokenHandler::SetPreprovKeys(GetPreprovKeyVector());

  std::string raw_asset_key = "asset-key-000000";
  std::string asset_key;
  std::string make_model;
  util::Status status = WvmTokenHandler::GetEncryptedAssetKey(
      HexStringToBytes(kTestToken1Hex), raw_asset_key, make_model, &asset_key);
  EXPECT_OK(status);
  EXPECT_EQ("305d5f979074b1c4f932be70d3cc850c", BytesToHexString(asset_key));

  status = WvmTokenHandler::GetEncryptedAssetKey(
      HexStringToBytes(kTestToken2Hex), raw_asset_key, make_model, &asset_key);
  EXPECT_OK(status);
  EXPECT_EQ("091802159bf8da12aecfcdfb092075c8", BytesToHexString(asset_key));

  // Check 3DES encryption of asset keys
  status = WvmTokenHandler::EncryptAssetKey(
      HexStringToBytes(kTestDeviceKey3DesHex), raw_asset_key,
      WvmTokenHandler::DES3, &asset_key);
  EXPECT_OK(status);
  EXPECT_EQ("3693a68bdeba192be0ea279e6c165197", BytesToHexString(asset_key));

  asset_key.clear();
  status = WvmTokenHandler::GetEncryptedAssetKey(
      HexStringToBytes(kTestToken3DesHex), raw_asset_key, make_model,
      &asset_key);
  EXPECT_OK(status);
  EXPECT_EQ("3693a68bdeba192be0ea279e6c165197", BytesToHexString(asset_key));

  // Test with pass-thru (Cipher=PASS_THRU).
  asset_key.clear();
  status = WvmTokenHandler::EncryptAssetKey(
      HexStringToBytes(kTestDeviceKey1Hex), raw_asset_key,
      WvmTokenHandler::PASS_THRU, &asset_key);

  EXPECT_OK(status);
  EXPECT_EQ(BytesToHexString(raw_asset_key), BytesToHexString(asset_key));
}

TEST(WvmTokenHandlerTest, FilterOnMakeModel) {
  // Use all good keys, but only the key for LG:BD572 should work. It works
  // by setting only that one to DES3.  All other AES ppks will encrypt the
  // asset key incorrectly.
  std::vector<WvmTokenHandler::PreprovKey> ppks;
  ppks.push_back(WvmTokenHandler::PreprovKey(
      kTestSystemId3Des, HexStringToBytes(kTestPreprovKeyHex),
      WvmTokenHandler::AES, ""));
  ppks.push_back(WvmTokenHandler::PreprovKey(
      kTestSystemId3Des, HexStringToBytes(kTestPreprovKeyHex),
      WvmTokenHandler::DES3, "LG:BD572"));
  ppks.push_back(WvmTokenHandler::PreprovKey(
      kTestSystemId3Des, HexStringToBytes(kTestPreprovKeyHex),
      WvmTokenHandler::AES, ""));
  WvmTokenHandler::SetPreprovKeys(ppks);
  std::string raw_asset_key = "asset-key-000000";
  std::string asset_key;
  // Check 3DES encryption of asset keys
  util::Status status = WvmTokenHandler::EncryptAssetKey(
      HexStringToBytes(kTestDeviceKey3DesHex), raw_asset_key,
      WvmTokenHandler::DES3, &asset_key);
  EXPECT_OK(status);
  EXPECT_EQ("3693a68bdeba192be0ea279e6c165197", BytesToHexString(asset_key));

  asset_key.clear();
  std::string make_model;
  status = WvmTokenHandler::GetEncryptedAssetKey(
      HexStringToBytes(kTestToken3DesHex), raw_asset_key, make_model,
      &asset_key);
  EXPECT_OK(status);
  // Should fail because the asset key was encrypted with AES instead of DES3.
  EXPECT_NE("3693a68bdeba192be0ea279e6c165197", BytesToHexString(asset_key));

  // Set the make/model so we find and use the correct ppk.
  make_model = "LG:BD572";
  status = WvmTokenHandler::GetEncryptedAssetKey(
      HexStringToBytes(kTestToken3DesHex), raw_asset_key, make_model,
      &asset_key);
  EXPECT_OK(status);
  // Should work because the asset key was encrypted with DES3.
  EXPECT_EQ("3693a68bdeba192be0ea279e6c165197", BytesToHexString(asset_key));
}

TEST(WvmTokenHandlerTest, AncientKeybox) {
  util::Status status;
  std::string device_key;

  std::string v1_token(
      std::string(kTestPreprovKeyHex).replace(0, 16, "0000000100000001"));
  status = WvmTokenHandler::DecryptDeviceKeyWithPreprovKey(
      HexStringToBytes(v1_token), HexStringToBytes(kTestToken1Hex),
      &device_key);
  EXPECT_EQ(util::error::PERMISSION_DENIED, status.error_code());
  EXPECT_TRUE(device_key.empty());
}

}  // namespace widevine