// Copyright 2013 Google Inc. All Rights Reserved.
//
//  Mock implementation of OEMCrypto APIs
//
#ifndef OEMCRYPTO_ENGINE_MOCK_H_
#define OEMCRYPTO_ENGINE_MOCK_H_

#include <stdint.h>
#include <time.h>
#include <map>
#include <vector>

#include <openssl/rsa.h>

#include "OEMCryptoCENC.h"  // Needed for enums only.
#include "file_store.h"
#include "lock.h"
#include "oemcrypto_auth_mock.h"
#include "oemcrypto_key_mock.h"
#include "oemcrypto_rsa_key_shared.h"
#include "wv_cdm_types.h"

namespace wvoec_mock {

class SessionContext;
class CryptoEngine;
class UsageTable;
class UsageTableEntry;

typedef uint32_t SessionId;
typedef std::map<SessionId, SessionContext*> ActiveSessions;

typedef std::vector<uint8_t>  KeyId;
typedef std::map<KeyId, Key*> KeyMap;

// SessionKeyTable holds the keys for the current session
class SessionKeyTable {
 public:
  SessionKeyTable() {}
  ~SessionKeyTable();

  bool Insert(const KeyId key_id, const Key& key_data);
  Key* Find(const KeyId key_id);
  void Remove(const KeyId key_id);
  void UpdateDuration(const KeyControlBlock& control);

 private:
  KeyMap keys_;

  CORE_DISALLOW_COPY_AND_ASSIGN(SessionKeyTable);
};

class NonceTable {
 public:
  static const int kTableSize = 16;
  NonceTable() {
    for (int i = 0; i < kTableSize; ++i) {
      state_[i] = kNTStateInvalid;
    }
  }
  ~NonceTable() {};
  void AddNonce(uint32_t nonce);
  bool CheckNonce(uint32_t nonce);
  void Flush();

 private:
  enum NonceTableState {
    kNTStateInvalid,
    kNTStateValid,
    kNTStateFlushPending
  };
  NonceTableState state_[kTableSize];
  uint32_t age_[kTableSize];
  uint32_t nonces_[kTableSize];
};

class SessionContext {
 private:
  SessionContext() {}

 public:
  SessionContext(CryptoEngine* ce, SessionId sid,
                 const RSA_shared_ptr& rsa_key)
      : valid_(true),
        ce_(ce),
        id_(sid),
        current_content_key_(NULL),
        rsa_key_(rsa_key),
        allowed_schemes_(kSign_RSASSA_PSS),
        usage_entry_(NULL) {}
  ~SessionContext();

  bool isValid() { return valid_; }

  bool DeriveKeys(const std::vector<uint8_t>& master_key,
                  const std::vector<uint8_t>& mac_context,
                  const std::vector<uint8_t>& enc_context);
  bool RSADeriveKeys(const std::vector<uint8_t>& enc_session_key,
                     const std::vector<uint8_t>& mac_context,
                     const std::vector<uint8_t>& enc_context);
  bool GenerateSignature(const uint8_t* message,
                         size_t message_length,
                         uint8_t* signature,
                         size_t* signature_length);
  size_t RSASignatureSize();
  OEMCryptoResult GenerateRSASignature(const uint8_t* message,
                                       size_t message_length,
                                       uint8_t* signature,
                                       size_t* signature_length,
                                       RSA_Padding_Scheme padding_scheme);
  bool ValidateMessage(const uint8_t* message,
                       size_t message_length,
                       const uint8_t* signature,
                       size_t signature_length);
  OEMCryptoResult DecryptCENC(const uint8_t* iv, size_t block_offset,
                              const OEMCrypto_CENCEncryptPatternDesc* pattern,
                              const uint8_t* cipher_data,
                              size_t cipher_data_length, bool is_encrypted,
                              uint8_t* clear_data,
                              OEMCryptoBufferType buffer_type);

  OEMCryptoResult Generic_Encrypt(const uint8_t* in_buffer,
                                  size_t buffer_length, const uint8_t* iv,
                                  OEMCrypto_Algorithm algorithm,
                                  uint8_t* out_buffer);
  OEMCryptoResult Generic_Decrypt(const uint8_t* in_buffer,
                                  size_t buffer_length, const uint8_t* iv,
                                  OEMCrypto_Algorithm algorithm,
                                  uint8_t* out_buffer);
  OEMCryptoResult Generic_Sign(const uint8_t* in_buffer, size_t buffer_length,
                               OEMCrypto_Algorithm algorithm,
                               uint8_t* signature, size_t* signature_length);
  OEMCryptoResult Generic_Verify(const uint8_t* in_buffer, size_t buffer_length,
                                 OEMCrypto_Algorithm algorithm,
                                 const uint8_t* signature,
                                 size_t signature_length);
  void StartTimer();
  uint32_t CurrentTimer();  // (seconds).
  OEMCryptoResult LoadKeys(const uint8_t* message, size_t message_length,
                           const uint8_t* signature, size_t signature_length,
                           const uint8_t* enc_mac_key_iv,
                           const uint8_t* enc_mac_keys, size_t num_keys,
                           const OEMCrypto_KeyObject* key_array,
                           const uint8_t* pst, size_t pst_length);
  bool InstallKey(const KeyId& key_id,
                  const std::vector<uint8_t>& key_data,
                  const std::vector<uint8_t>& key_data_iv,
                  const std::vector<uint8_t>& key_control,
                  const std::vector<uint8_t>& key_control_iv,
                  const std::vector<uint8_t>& pst,
                  bool ctr_mode);
  bool InstallRSAEncryptedKey(const uint8_t *encrypted_message_key,
                              size_t encrypted_message_key_length);
  bool DecryptRSAKey(const uint8_t* enc_rsa_key,
                     size_t enc_rsa_key_length,
                     const uint8_t* wrapped_rsa_key_iv,
                     uint8_t* pkcs8_rsa_key);
  bool EncryptRSAKey(const uint8_t* pkcs8_rsa_key,
                     size_t enc_rsa_key_length,
                     const uint8_t* enc_rsa_key_iv,
                     uint8_t* enc_rsa_key);
  bool LoadRSAKey(const uint8_t* pkcs8_rsa_key,
                  size_t rsa_key_length);
  OEMCryptoResult RefreshKey(const KeyId& key_id,
                             const std::vector<uint8_t>& key_control,
                             const std::vector<uint8_t>& key_control_iv);
  bool UpdateMacKeys(const std::vector<uint8_t>& mac_keys,
                     const std::vector<uint8_t>& iv);
  bool QueryKeyControlBlock(const KeyId& key_id, uint32_t* data);
  OEMCryptoResult SelectContentKey(const KeyId& key_id);
  const Key* current_content_key(void) { return current_content_key_; }
  void set_mac_key_server(const std::vector<uint8_t>& mac_key_server) {
    mac_key_server_ = mac_key_server;
  }
  const std::vector<uint8_t>& mac_key_server() { return mac_key_server_; }
  void set_mac_key_client(const std::vector<uint8_t>& mac_key_client) {
    mac_key_client_ = mac_key_client;
  }
  const std::vector<uint8_t>& mac_key_client() { return mac_key_client_; }

  void set_encryption_key(const std::vector<uint8_t>& enc_key) {
    encryption_key_ = enc_key;
  }
  const std::vector<uint8_t>& encryption_key() { return encryption_key_; }
  uint32_t allowed_schemes() const { return allowed_schemes_; }

  void AddNonce(uint32_t nonce);
  bool CheckNonce(uint32_t nonce);
  void FlushNonces();
  void ReleaseUsageEntry();

 private:
  bool DeriveKey(const std::vector<uint8_t>& key,
                 const std::vector<uint8_t>& context,
                 int counter, std::vector<uint8_t>* out);
  bool DecryptMessage(const std::vector<uint8_t>& key,
                      const std::vector<uint8_t>& iv,
                      const std::vector<uint8_t>& message,
                      std::vector<uint8_t>* decrypted);
  bool CheckNonceOrEntry(const KeyControlBlock& key_control_block,
                         const std::vector<uint8_t>& pst);
  bool IsUsageEntryValid();
  OEMCryptoResult DecryptCBC(const uint8_t* key, const uint8_t* iv,
                             const OEMCrypto_CENCEncryptPatternDesc* pattern,
                             const uint8_t* cipher_data,
                             size_t cipher_data_length, uint8_t* clear_data);
  OEMCryptoResult PatternDecryptCTR(
      const uint8_t* key, const uint8_t* iv, size_t block_offset,
      const OEMCrypto_CENCEncryptPatternDesc* pattern,
      const uint8_t* cipher_data, size_t cipher_data_length,
      uint8_t* clear_data);
  OEMCryptoResult DecryptCTR(const uint8_t* key_u8, const uint8_t* iv,
                             size_t block_offset, const uint8_t* cipher_data,
                             size_t cipher_data_length, uint8_t* clear_data);
  OEMCryptoResult AllowKeyUse(const std::string& log_string, uint32_t use_type,
                              OEMCryptoBufferType buffer_type);
  RSA* rsa_key() { return rsa_key_.get(); }

  bool valid_;
  CryptoEngine* ce_;
  SessionId id_;
  std::vector<uint8_t> mac_key_server_;
  std::vector<uint8_t> mac_key_client_;
  std::vector<uint8_t> encryption_key_;
  std::vector<uint8_t> session_key_;
  const Key* current_content_key_;
  SessionKeyTable session_keys_;
  NonceTable nonce_table_;
  RSA_shared_ptr rsa_key_;
  uint32_t allowed_schemes_;  // for RSA signatures.
  time_t timer_start_;
  UsageTableEntry* usage_entry_;

  CORE_DISALLOW_COPY_AND_ASSIGN(SessionContext);
};

class CryptoEngine {
 public:
  CryptoEngine(wvcdm::FileSystem* file_system);
  ~CryptoEngine();

  bool Initialized() { return true; }

  bool ValidRootOfTrust() { return root_of_trust_.Validate(); }

  bool InstallKeybox(const uint8_t* keybox, size_t keybox_length) {
    return root_of_trust_.InstallKeybox(keybox, keybox_length);
  }

  void UseTestKeybox() { root_of_trust_.UseTestKeybox(); }

  bool LoadTestRsaKey() { return root_of_trust_.LoadTestRsaKey(); }

  KeyboxError ValidateKeybox() { return root_of_trust_.ValidateKeybox(); }

  const std::vector<uint8_t>& DeviceRootKey(bool override_to_real = false) {
    return root_of_trust_.DeviceKey(override_to_real);
  }

  const std::vector<uint8_t>& DeviceRootId() {
    return root_of_trust_.DeviceId();
  }

  size_t DeviceRootTokenLength() {
    return root_of_trust_.DeviceTokenLength();
  }

  const uint8_t* const  DeviceRootToken() {
    return root_of_trust_.DeviceToken();
  }

  void Terminate();

  SessionId CreateSession();

  bool DestroySession(SessionId sid);

  SessionContext* FindSession(SessionId sid);

  size_t GetNumberOfOpenSessions() { return sessions_.size(); }

  size_t GetMaxNumberOfSessions() {
    // An arbitrary limit for mock implementation.
    static const size_t kMaxSupportedOEMCryptoSessions = 64;
    return kMaxSupportedOEMCryptoSessions;
  }

  // Configuration constants - controls behavior of this CryptoEngine
  OEMCrypto_HDCP_Capability config_current_hdcp_capability();
  OEMCrypto_HDCP_Capability config_maximum_hdcp_capability();

  UsageTable* usage_table() { return usage_table_; }
  wvcdm::FileSystem* file_system() { return file_system_; }

  bool config_local_display_only();
  bool config_closed_platform();
  bool config_supports_usage_table();
  bool config_supports_keybox();
  OEMCrypto_ProvisioningMethod config_provisioning_method();
  OEMCryptoResult get_oem_certificate(SessionContext* session,
                                      uint8_t* public_cert,
                                      size_t* public_cert_length);
  bool config_is_anti_rollback_hw_present();
  const char* config_security_level();
  uint8_t config_security_patch_level();

 private:
  ActiveSessions sessions_;
  AuthenticationRoot root_of_trust_;
  wvcdm::Lock session_table_lock_;
  wvcdm::FileSystem* file_system_;
  UsageTable* usage_table_;

  CORE_DISALLOW_COPY_AND_ASSIGN(CryptoEngine);
};

}  // namespace wvoec_mock

#endif  // OEMCRYPTO_ENGINE_MOCK_H_