////////////////////////////////////////////////////////////////////////////////
// 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 <vector>

#include "glog/logging.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "util/endian/endian.h"
#include "util/gtl/map_util.h"
#include "util/status.h"
#include "common/aes_cbc_util.h"
#include "common/ecb_util.h"
#include "common/sha_util.h"

namespace widevine {

namespace {
const int kKeyboxFlagInsecure = 1;
const int kSystemIdSamsungTVFactory = 8;
const int kPreProvisioningKeySizeBytes = 16;
const int kKeyboxSizeBytes = 72;
const std::string kZeroIV(16, '\0');

class PreprovKeysMap {
 public:
  void Update(const std::vector<WvmTokenHandler::PreprovKey>& key_vector);
  std::vector<WvmTokenHandler::PreprovKey> GetPreprovKeys(uint32_t system_id);
  bool IsSystemIdKnown(uint32_t system_id);
  bool IsEmpty();

  static PreprovKeysMap* GetSingleton();

 private:
  absl::Mutex mutex_;
  std::multimap<uint32_t, WvmTokenHandler::PreprovKey> preprov_keys_;
};

void PreprovKeysMap::Update(
    const std::vector<WvmTokenHandler::PreprovKey>& key_vector) {
  absl::WriterMutexLock lock(&mutex_);
  preprov_keys_.clear();
  for (const WvmTokenHandler::PreprovKey& ppk : key_vector) {
    if (ppk.key_bytes.size() != kPreProvisioningKeySizeBytes) {
      LOG(WARNING) << "Invalid preprov key for system id: " << ppk.system_id;
      continue;
    }
    preprov_keys_.insert(std::make_pair(ppk.system_id, ppk));
  }
}

std::vector<WvmTokenHandler::PreprovKey> PreprovKeysMap::GetPreprovKeys(
    uint32_t system_id) {
  absl::ReaderMutexLock lock(&mutex_);
  std::vector<WvmTokenHandler::PreprovKey> key_vector;
  auto range = preprov_keys_.equal_range(system_id);
  for (auto it = range.first; it != range.second; ++it)
    key_vector.push_back(it->second);
  return key_vector;
}

bool PreprovKeysMap::IsSystemIdKnown(uint32_t system_id) {
  absl::ReaderMutexLock lock(&mutex_);
  return gtl::ContainsKey(preprov_keys_, system_id);
}

bool PreprovKeysMap::IsEmpty() {
  absl::ReaderMutexLock lock(&mutex_);
  return preprov_keys_.empty();
}

PreprovKeysMap* PreprovKeysMap::GetSingleton() {
  static auto* const kInstance = new PreprovKeysMap();
  return kInstance;
}

}  // namespace

WvmTokenHandler::PreprovKey::PreprovKey(uint32_t system_id,
                                        const std::string& key_bytes, Cipher cipher,
                                        const std::string& model_filter)
    : system_id(system_id),
      key_bytes(key_bytes),
      cipher(cipher),
      model_filter(model_filter) {}

WvmTokenHandler::PreprovKey::PreprovKey(uint32_t system_id,
                                        const std::string& key_bytes, Cipher cipher)
    : system_id(system_id), key_bytes(key_bytes), cipher(cipher) {}

WvmTokenHandler::PreprovKey::PreprovKey(uint32_t system_id,
                                        const std::string& key_bytes)
    : system_id(system_id), key_bytes(key_bytes), cipher(AES) {}

void WvmTokenHandler::SetPreprovKeys(const std::vector<PreprovKey>& keyvec) {
  PreprovKeysMap::GetSingleton()->Update(keyvec);
}

bool WvmTokenHandler::IsSystemIdKnown(uint32_t system_id) {
  return PreprovKeysMap::GetSingleton()->IsSystemIdKnown(system_id);
}

util::Status WvmTokenHandler::DecryptDeviceKey(absl::string_view token,
                                               std::string* device_key_out,
                                               Cipher* cipher_out,
                                               bool* insecure_out) {
  const std::string default_make_model;
  return DecryptDeviceKey(token, default_make_model, device_key_out, cipher_out,
                          insecure_out);
}

util::Status WvmTokenHandler::DecryptDeviceKey(absl::string_view token,
                                               const std::string& make_model,
                                               std::string* device_key_out,
                                               Cipher* cipher_out,
                                               bool* insecure_out) {
  DCHECK(device_key_out);
  // DCHECK below is commented out because preprov_keys_ being nullptr
  // is a valid test in wvm_token_handler_test.cc. If we have
  // DCHECK(preprov_keys_) here it would thrown an failure in Kokoro
  // presubmit because evidently Kokoro does debug build.
  // DCHECK(preprov_keys_);
  if (token.size() < kKeyboxSizeBytes) {
    return util::Status(util::error::INVALID_ARGUMENT,
                        "Keybox token is too short.");
  }
  if (PreprovKeysMap::GetSingleton()->IsEmpty()) {
    return util::Status(util::error::INVALID_ARGUMENT,
                        "Pre-provisioning key map is nullptr.");
  }

  uint32_t system_id = GetSystemId(token);

  // There may be multiple preprov keys for a system ID; try them all.
  std::vector<PreprovKey> key_vector =
      PreprovKeysMap::GetSingleton()->GetPreprovKeys(system_id);

  util::Status status;
  // First pass through the matching system Ids is an attempt to find an
  // alternate preprov key specific to this make/model.
  const PreprovKey* preferred_ppk = NULL;
  for (const PreprovKey& ppk : key_vector) {
    if (!ppk.model_filter.empty() && ppk.model_filter == make_model) {
      preferred_ppk = &ppk;
      break;
    }
  }
  for (const PreprovKey& ppk : key_vector) {
    if (preferred_ppk && preferred_ppk != &ppk) {
      continue;
    }
    uint32_t version;
    status = DecryptDeviceKeyWithPreprovKey(
        ppk.key_bytes, token, device_key_out, insecure_out, &version);
    if (version != 2) {
      // Only version 2 keyboxes supported.
      return util::Status(util::error::PERMISSION_DENIED,
                          absl::StrCat("invalid-keybox-version ", version));
    }
    if (status.ok()) {
      if (cipher_out) {
        *cipher_out = ppk.cipher;
      }
      return status;
    }
  }
  if (!status.ok()) {
    // Return error from last attempt.
    return status;
  }
  return util::Status(
      util::error::NOT_FOUND,
      absl::StrCat("Unknown system id: ", system_id).c_str());  // NOLINT
}

// Decrypt a token using the preprov key for its system ID, and use the
// decrypted device key to encrypt the given asset key. Returns the encrypted
// asset key in |result|.
// On failure, returns an error from the Widevine Server SDK error space.
util::Status WvmTokenHandler::GetEncryptedAssetKey(
    absl::string_view token, absl::string_view raw_asset_key,
    const std::string& make_model, std::string* result) {
  std::string device_key;
  Cipher cipher = AES;
  util::Status status =
      DecryptDeviceKey(token, make_model, &device_key, &cipher, nullptr);
  if (!status.ok()) {
    return status;
  }
  return EncryptAssetKey(device_key, raw_asset_key, cipher, result);
}

uint32_t WvmTokenHandler::GetSystemId(absl::string_view token) {
  uint32_t system_id = 0;
  if (token.size() >= 8) {
    // Bytes 4-8 contain the little-endian system ID.
    system_id = BigEndian::Load32(token.data() + 4);
  }
  return system_id;
}

std::string WvmTokenHandler::GetEncryptedUniqueId(absl::string_view token) {
  std::string encrypted_unique_id("");
  if (token.size() >= 24) {
    encrypted_unique_id = std::string(token.substr(8, 16));
  }
  return encrypted_unique_id;
}

util::Status WvmTokenHandler::DecryptDeviceKeyWithPreprovKey(
    absl::string_view preprov_key, absl::string_view token,
    std::string* device_key_out) {
  return DecryptDeviceKeyWithPreprovKey(preprov_key, token, device_key_out,
                                        nullptr, nullptr);
}

util::Status WvmTokenHandler::DecryptDeviceKeyWithPreprovKey(
    absl::string_view preprov_key, absl::string_view token,
    std::string* device_key_out, bool* insecure_out, uint32_t* version) {
  CHECK(device_key_out);
  if (token.size() < kKeyboxSizeBytes) {
    return util::Status(util::error::INVALID_ARGUMENT,
                        "Keybox token is too short.");
  }
  if (version) {
    *version = BigEndian::Load32(token.data());
  }
  // This was checked at initialization, so if it fails now something is wrong.
  CHECK_EQ(preprov_key.size(), kPreProvisioningKeySizeBytes);

  absl::string_view encrypted_unique_id = token.substr(8, 16);
  std::string unique_id = crypto_util::DecryptAesCbcNoPad(
      std::string(preprov_key), kZeroIV, std::string(encrypted_unique_id));

  if (unique_id.size() != 16) {
    // Decrypting 16 bytes should result in 16 bytes.
    LOG(WARNING) << "Internal error decrypting unique id from token.";
    return util::Status(util::error::INTERNAL, "Wrong size after decrypt/16.");
  }

  absl::string_view encrypted_bits = token.substr(24, 48);
  std::string decrypted_bits = crypto_util::DecryptAesCbcNoPad(
      unique_id, kZeroIV, std::string(encrypted_bits));
  if (decrypted_bits.size() != 48) {
    // Decrypting 48 bytes should result in 48 bytes.
    LOG(WARNING) << "Internal error decrypting device key from token.";
    return util::Status(util::error::INTERNAL, "Wrong size after decrypt/48.");
  }
  uint8_t keybox_flags = decrypted_bits[36];
  absl::string_view device_key =
      absl::string_view(decrypted_bits).substr(0, 16);
  absl::string_view expected_hash =
      absl::string_view(decrypted_bits).substr(16, 20);
  std::string actual_hash = Sha1_Hash(std::string(device_key));

  if (GetSystemId(token) == kSystemIdSamsungTVFactory) {
    // Keyboxes with this system ID have corrupted bytes starting after the
    // first 16 bytes of the hash, so we use only the uncorrupted part.
    expected_hash = expected_hash.substr(0, 16);
    actual_hash.resize(16);
    keybox_flags = 0;
  }
  if (expected_hash != actual_hash) {
    return util::Status(util::error::PERMISSION_DENIED,
                        "Keybox validation failed.");
  }
  *device_key_out = std::string(device_key);
  if (insecure_out) {
    *insecure_out = (keybox_flags & kKeyboxFlagInsecure) != 0;
  }
  return util::OkStatus();
}

util::Status WvmTokenHandler::EncryptAssetKey(absl::string_view device_key,
                                              absl::string_view raw_asset_key,
                                              Cipher cipher, std::string* result) {
  CHECK(result);
  if (device_key.size() != 16) {
    return util::Status(util::error::INVALID_ARGUMENT,
                        "Invalid device key: size != 16");
  }
  if (raw_asset_key.size() < 16) {
    return util::Status(util::error::INVALID_ARGUMENT,
                        "Invalid asset key: size < 16");
  }
  // Truncate extra characters in the key; wvm always uses 16.
  absl::string_view asset_key = raw_asset_key.substr(0, 16);
  switch (cipher) {
    case DES3:
      if (!crypto_util::Encrypt3DesEcb(device_key, asset_key, result)) {
        return util::Status(util::error::INTERNAL,
                            "Error encrypting asset key with 3DES.");
      }
      return util::OkStatus();
    case AES:
      if (!crypto_util::EncryptAesEcb(device_key, asset_key, result)) {
        return util::Status(util::error::INTERNAL,
                            "Error encrypting asset key with AES.");
      }
      return util::OkStatus();
    case PASS_THRU:
      result->assign(raw_asset_key.data(), raw_asset_key.size());
      return util::OkStatus();
    default:
      return util::Status(util::error::INVALID_ARGUMENT, "Unknown cipher type");
  }
}

}  // namespace widevine