// Copyright 2019 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
//
#ifndef CDM_OEC_KEY_DERIVER_H_
#define CDM_OEC_KEY_DERIVER_H_

#include <openssl/aes.h>
#include <openssl/rsa.h>
#include <time.h>
#include <string>
#include <vector>

#include "oec_device_features.h"
#include "oemcrypto_types.h"
#include "pst_report.h"

namespace wvoec {

constexpr size_t kMaxTestRSAKeyLength = 2000;  // Rough estimate.
constexpr size_t kMaxCoreProvRequest = 150;    // Rough estimate.
constexpr size_t kMaxX509CertLength = 4000;    // Rough estimate.

// This structure will be signed to simulate a provisioning response from the
// server.
struct RSAPrivateKeyMessage {
  uint8_t rsa_key[kMaxTestRSAKeyLength];
  uint8_t rsa_key_iv[KEY_IV_SIZE];
  size_t rsa_key_length;
  uint8_t enc_message_key[kMaxTestRSAKeyLength];
  size_t enc_message_key_length;
  uint32_t nonce;
};

// This structure simulates a provisioning 4.0 response from the server.
// However, OEMCrypto doesn't need to load this response, since it doesn't have
// any secrets to be handled. It is just a dummy struct for the tests to
// compile.
struct Prov40CertMessage {
  uint8_t device_certificate[kMaxX509CertLength];
  uint32_t nonce;
};

// Holds an encryption key and can encrypt a provisioning message.  It also can
// encrypt short buffers using CBC, such as content keys in a license.
class Encryptor {
 public:
  Encryptor() : enc_key_(KEY_SIZE, 0) {}
  Encryptor(const std::vector<uint8_t>& enc_key) { set_enc_key(enc_key); };
  Encryptor& operator=(const Encryptor&) = default;
  void set_enc_key(const std::vector<uint8_t>& enc_key);

  // This encrypts an RSAPrivateKeyMessage with encryption_key so that it may be
  // loaded with OEMCrypto_LoadProvisioningResponse.
  // This modifies the clear data: it adds padding and generates a random iv.
  void PadAndEncryptProvisioningMessage(RSAPrivateKeyMessage* data,
                                        RSAPrivateKeyMessage* encrypted) const;

  void CBCEncrypt(const uint8_t* data, uint8_t* encrypted_data,
                  size_t data_length, const uint8_t (&iv)[KEY_IV_SIZE]) const;

 private:
  std::vector<uint8_t> enc_key_;
};

// Holds encryption and mac keys derived from a master key.
// Can be used to sign a buffer as either a server or client.
class KeyDeriver : public Encryptor {
 public:
  KeyDeriver()
      : mac_key_server_(MAC_KEY_SIZE, 0), mac_key_client_(MAC_KEY_SIZE, 0) {}
  KeyDeriver& operator=(const KeyDeriver&) = default;

  // Generate mac and enc keys give the master key.
  void DeriveKeys(const uint8_t* master_key, size_t master_key_size,
                  const std::vector<uint8_t>& context);
  void DeriveKeys(const uint8_t* master_key, size_t master_key_size,
                  const std::vector<uint8_t>& context, const char* mac_label,
                  const char* enc_label);
  // Sign the buffer with server's mac key.
  void ServerSignBuffer(const uint8_t* data, size_t data_length,
                        std::vector<uint8_t>* signature) const;
  // Sign the buffer with client's known mac key.  Known test keys must be
  // installed first. This uses HMAC with SHA256, so is suitable for a message.
  void ClientSignBuffer(const std::vector<uint8_t>& buffer,
                        std::vector<uint8_t>* signature) const;
  // Sign the pst buffer with client's known mac key.  Known test keys must be
  // installed first. This uses HMAC with SHA128, and skips the beginning of the
  // buffer, so is only suitable for a pst report.
  void ClientSignPstReport(const std::vector<uint8_t>& pst_report_buffer,
                           std::vector<uint8_t>* signature) const;
  void set_mac_keys(const uint8_t* mac_keys);

 private:
  // Internal utility function to derive key using CMAC-128 or CMAC-256 based on
  // master_key_size.
  void DeriveKey(const uint8_t* key, size_t master_key_size,
                 const std::vector<uint8_t>& context, int counter,
                 std::vector<uint8_t>* out);

  std::vector<uint8_t> mac_key_server_;
  std::vector<uint8_t> mac_key_client_;
};

}  // namespace wvoec

#endif  // CDM_OEC_KEY_DERIVER_H_