////////////////////////////////////////////////////////////////////////////////
// Copyright 2017 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_CLIENT_CERT_H__
#define COMMON_CLIENT_CERT_H__

#include <map>
#include <memory>
#include <string>

#include "common/rsa_key.h"
#include "common/status.h"
#include "protos/public/client_identification.pb.h"
#include "protos/public/license_protocol.pb.h"

namespace widevine {

class DrmRootCertificate;
class SignedDrmCertificate;

// Handler class for LicenseRequests; validates requests and encrypts licenses.
// TODO(user): Remove extra accessors after Keybox parsing is moved
// to a separate class in KeyboxClientCert class.
class ClientCert {
 public:
  virtual ~ClientCert() {}
  static Status Create(
      const DrmRootCertificate* root_certificate,
      widevine::ClientIdentification::TokenType token_type,
      const std::string& token, ClientCert** client_cert);
  // Creates a Keybox based ClientCert.
  static Status CreateWithKeybox(const std::string& keybox_token,
                                 ClientCert** client_cert);
  // Creates a Device Certificate based ClientCert.
  static Status CreateWithDrmCertificate(
      const DrmRootCertificate* root_certificate, const std::string& drm_certificate,
      ClientCert** client_cert);
  // Creates a HMAC SHA256 signature based on the message and the key().
  // signature is owned by the caller and can not be NULL.
  virtual void CreateSignature(const std::string& message, std::string* signature);
  // Checks the passed in signature against a signature created used the
  // classes information and the passed in message.  Returns OK if signature
  // is valid.
  virtual Status VerifySignature(const std::string& message, const std::string& signature,
                                 ProtocolVersion protocol_version) = 0;
  // Creates a signing_key that is accessible using signing_key().  Signing_key
  // is constructed by doing a key derivation using the key() and message.
  virtual void GenerateSigningKey(const std::string& message,
                                  ProtocolVersion protocol_version);
  // Used to create signing keys. For Keybox token types this is the device key.
  // For Device Certificate token types this the session key.
  virtual const std::string& key() const = 0;
  virtual void set_key(const std::string& key) = 0;
  virtual const std::string& encrypted_key() const = 0;
  virtual uint32_t system_id() const { return system_id_; }
  virtual const std::string& signing_key() const { return signing_key_; }
  virtual const std::string& public_key() const { return public_key_; }
  virtual const std::string& serial_number() const { return serial_number_; }
  virtual void set_serial_number(const std::string& serial_number) {
    serial_number_ = serial_number;
  }
  virtual const std::string& signer_serial_number() const {
    return signer_serial_number_;
  }
  virtual uint32_t signer_creation_time_seconds() const {
    return signer_creation_time_seconds_;
  }
  virtual widevine::ClientIdentification::TokenType type() const = 0;
  virtual std::string service_id() const { return service_id_; }
  virtual bool signed_by_provisioner() const { return signed_by_provisioner_; }

 protected:
  ClientCert() {}

  virtual void set_system_id(uint32_t system_id) { system_id_ = system_id; }
  virtual void set_signing_key(const std::string& signing_key) {
    signing_key_ = signing_key;
  }
  virtual void set_service_id(const std::string& service_id) {
    service_id_ = service_id;
  }
  virtual void set_signed_by_provisioner(bool provisioner_signed_flag) {
    signed_by_provisioner_ = provisioner_signed_flag;
  }

  std::string public_key_;
  std::string serial_number_;
  std::string signer_serial_number_;
  uint32_t signer_creation_time_seconds_ = 0;
  bool signed_by_provisioner_ = false;

 private:
  uint32_t system_id_ = 0;
  std::string signing_key_;
  std::string service_id_;

  DISALLOW_COPY_AND_ASSIGN(ClientCert);
};

// This class implements the crypto operations based on the Widevine keybox.
// It will unpack token and perform all the crypto operations for securing
// the key material in the license response.
class KeyboxClientCert : public ClientCert {
 public:
  ~KeyboxClientCert() override;

  // Set the system-wide pre-provisioning keys; argument must be human-readable
  // hex digits.
  // Must be called before any other method of this class is called, unless
  // created by ClientCert::CreateWithPreProvisioningKey(...).
  static void SetPreProvisioningKeys(const std::multimap<uint32_t, std::string>& keys);
  static bool IsSystemIdKnown(const uint32_t system_id);
  static uint32_t GetSystemId(const std::string& keybox_bytes);

  Status Initialize(const std::string& keybox_bytes);

  Status VerifySignature(const std::string& message, const std::string& signature,
                         ProtocolVersion protocol_version) override;
  const std::string& key() const override { return device_key_; }
  void set_key(const std::string& key) override { device_key_ = key; }
  const std::string& encrypted_key() const override { return encrypted_device_key_; }
  widevine::ClientIdentification::TokenType type() const override {
    return widevine::ClientIdentification::KEYBOX;
  }

 private:
  KeyboxClientCert();

  friend class ClientCert;
  friend class MockKeyboxClientCert;

  std::string device_key_;
  std::string encrypted_device_key_;

  DISALLOW_COPY_AND_ASSIGN(KeyboxClientCert);
};
// This class implements the device certificate operations based on RSA keys.
// It will unpack token and perform all the crypto operations for securing
// the key material in the license response.
using widevine::RsaPublicKey;
class CertificateClientCert : public ClientCert {
 public:
  ~CertificateClientCert() override;

  Status VerifySignature(const std::string& message, const std::string& signature,
                         ProtocolVersion protocol_version) override;
  const std::string& key() const override { return session_key_; }
  void set_key(const std::string& key) override { session_key_ = key; }
  const std::string& encrypted_key() const override {
    return encrypted_session_key_;
  }
  widevine::ClientIdentification::TokenType type() const override {
    return widevine::ClientIdentification::DRM_DEVICE_CERTIFICATE;
  }

 protected:
  friend class ClientCert;
  friend class MockCertificateClientCert;
  Status Initialize(const DrmRootCertificate* drm_root_certificate,
                    const std::string& serialized_certificate);
  virtual void set_public_key(const std::string& public_key) {
    public_key_ = public_key;
  }
  virtual void set_signer_serial_number(const std::string& signer_serial_number) {
    signer_serial_number_ = signer_serial_number;
  }
  virtual void set_signer_creation_time_seconds(uint32_t creation_time_seconds) {
    signer_creation_time_seconds_ = creation_time_seconds;
  }

  std::string session_key_;
  std::string encrypted_session_key_;
  std::unique_ptr<RsaPublicKey> rsa_public_key_;

 private:
  CertificateClientCert();

  DISALLOW_COPY_AND_ASSIGN(CertificateClientCert);
};

}  // namespace widevine
#endif  // COMMON_CLIENT_CERT_H__