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

#ifndef COMMON_WVM_TOKEN_HANDLER_H_
#define COMMON_WVM_TOKEN_HANDLER_H_

#include <map>
#include <vector>

#include "base/macros.h"
#include "absl/strings/string_view.h"
#include "common/status.h"

namespace widevine {

// Class for decoding the encrypted token that comes from a WVM classic keybox.
//
// Internally, this class keeps a multimap from system ID to the preprov key(s)
// for that system, so construction is relatively expensive; don't create an
// instance of this class per-request.
//
// Errors in this file are returned in the canonical space, but are chosen so
// that it's possible to map different failure modes to the appropriate codes
// in the WVM or server SDK error spaces:
//   OK - success.
//   NOT_FOUND - system id from token wasn't in the preprov key table.
//   PERMISSION_DENIED - hash of device key didn't match.
//   INVALID_ARGUMENT - token or a key is wrong size or otherwise invalid.
//   INTERNAL_ERROR - something went wrong that shouldn't have been able to.
class WvmTokenHandler {
 public:
  // Cipher type to use for encrypting asset keys. This matches the enum in
  // video/widevine/lockbox/public/key.proto.
  enum Cipher {
    DES3 = 0,
    AES = 1,
    PASS_THRU = 2,
  };

  struct PreprovKey {
    // Utility constructor.
    PreprovKey(uint32_t system_id, const std::string& key_bytes, Cipher cipher,
               const std::string& model_filter);
    PreprovKey(uint32_t system_id, const std::string& key_bytes, Cipher cipher);
    // Constructor if the cipher isn't needed (i.e. modular DRM).
    PreprovKey(uint32_t system_id, const std::string& key_bytes);
    uint32_t system_id;
    std::string key_bytes;
    Cipher cipher;
    // If set, the make/model in the license request must match this value.
    std::string model_filter;
  };

  // Set pre-provisioning keys from the given vector. This may be called
  // concurrently with other methods.
  static void SetPreprovKeys(const std::vector<PreprovKey>& preprov_keys);

  // Returns true if system_id is in the preprov key table.
  static bool IsSystemIdKnown(uint32_t system_id);

  // Decrypt a token using the preprov key for its system ID, and
  // return the decrypted device key in result.
  // On failure, returns one of the errors listed above.
  // cipher_out may be null; if not, *cipher_out will be set to the cipher type
  // to use with the device key.
  // insecure_out may be null; if not, *insecure_out will be set to the
  // decrypted value of the 'insecure keybox' flag.
  static util::Status DecryptDeviceKey(absl::string_view token,
                                       std::string* device_key_out,
                                       Cipher* cipher_out, bool* insecure_out);
  // Same as above, except takes in the make/model from the license request.
  // For legacy WVM license, we have some special cases where we need to inspect
  // the make/model as we apply alternate keys.
  static util::Status DecryptDeviceKey(absl::string_view token,
                                       const std::string& make_model,
                                       std::string* device_key_out,
                                       Cipher* cipher_out, bool* insecure_out);

  // 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.
  static util::Status GetEncryptedAssetKey(absl::string_view token,
                                           absl::string_view raw_asset_key,
                                           const std::string& make_model,
                                           std::string* result);

  // Extract the system ID component of a token (bytes 4-8).
  static uint32_t GetSystemId(absl::string_view token);

  // Extract the encrypted unique ID component of a token (bytes 8-24).
  static std::string GetEncryptedUniqueId(absl::string_view token);

  // Try to decrypt a token using the provided preprov key, and return the
  // decrypted device key in result.
  //
  // Note that the if the input std::string lengths are correct (16 and 72 bytes),
  // the only possible cause of failure is the decrypted device key hash
  // being incorrect.
  static util::Status DecryptDeviceKeyWithPreprovKey(
      absl::string_view preprov_key_bytes, absl::string_view token,
      std::string* device_key_out);

  // Same as above, but allows extracting the 'insecure keybox' flag and keybox
  // version.
  static util::Status DecryptDeviceKeyWithPreprovKey(
      absl::string_view preprov_key_bytes, absl::string_view token,
      std::string* device_key_out, bool* insecure_out, uint32_t* version);

  // Given a decrypted device key as returned by DecryptToken(), use it to
  // encrypt an asset key with the given cipher.
  static util::Status EncryptAssetKey(absl::string_view device_key,
                                      absl::string_view raw_asset_key,
                                      Cipher cipher, std::string* result);

 private:
  DISALLOW_IMPLICIT_CONSTRUCTORS(WvmTokenHandler);
};

}  // namespace widevine

#endif  // COMMON_WVM_TOKEN_HANDLER_H_