/*******************************************************************************
 *
 * Copyright 2013 Google Inc. All Rights Reserved.
 *
 * mock implementation of OEMCrypto APIs
 *
 ******************************************************************************/

#ifndef L3CRYPTO_ENGINE_MOCK_H_
#define L3CRYPTO_ENGINE_MOCK_H_

#include <map>
#include <stdint.h>
#include <vector>
#include "lock.h"
#include "l3crypto_key_mock.h"
#include "l3crypto_keybox_mock.h"
#include "wv_cdm_types.h"

namespace wvoec_obfs {

enum SessionState {
  SESSION_STATE_ILLEGAL,
  SESSION_STATE_INITIALIZED,
  SESSION_STATE_HAS_DERIVED_KEYS,
  SESSION_STATE_ACTIVE,
  SESSION_STATE_ERROR
};

enum BufferType {
  BUFFER_TYPE_CLEAR,
  BUFFER_TYPE_SECURE,
  BUFFER_TYPE_DIRECT
};

class SessionContext;
class CryptoEngine;

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);

 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) {
      valid_[i] = false;
    }
  }
  ~NonceTable() {};
  void AddNonce(uint32_t nonce);
  bool CheckNonce(uint32_t nonce);
 private:
  bool valid_[kTableSize];
  uint32_t age_[kTableSize];
  uint32_t nonces_[kTableSize];
};

class SessionContext {
 private:
  SessionContext() {}

 public:
  explicit SessionContext(CryptoEngine* ce, SessionId sid)
      : valid_(true), ce_(ce), id_(sid), current_content_key_(NULL) {}
  ~SessionContext() {}

  void Open();
  void Close();

  bool isValid() { return valid_; }

  bool DeriveKeys(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);

  bool ValidateMessage(const uint8_t* message,
                       size_t message_length,
                       const uint8_t* signature,
                       size_t signature_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);
  bool ParseKeyControl(const std::vector<uint8_t>& key_control_string,
                       KeyControlBlock& key_control_block);
  bool RefreshKey(const KeyId& key_id,
                  const std::vector<uint8_t>& key_control,
                  const std::vector<uint8_t>& key_control_iv);
  bool UpdateMacKey(const std::vector<uint8_t>& mac_key, const std::vector<uint8_t>& iv);
  bool SelectContentKey(const KeyId& key_id);
  const Key* current_content_key(void) {return current_content_key_;}
  void set_mac_key(const std::vector<uint8_t>& mac_key) { mac_key_ = mac_key; }
  const std::vector<uint8_t>& mac_key() { return mac_key_; }

  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_; }

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

 private:

  bool DeriveKey(const std::vector<uint8_t>& key, const std::vector<uint8_t>& context,
                 int counter, std::vector<uint8_t>* out);

  bool valid_;
  CryptoEngine* ce_;
  SessionId id_;
  SessionState state_;
  std::vector<uint8_t> mac_key_;
  std::vector<uint8_t> encryption_key_;
  const Key* current_content_key_;
  SessionKeyTable session_keys_;
  NonceTable nonce_table_;

  CORE_DISALLOW_COPY_AND_ASSIGN(SessionContext);
};

class CryptoEngine {

 private:

  enum CryptoEngineState {
    CE_ILLEGAL,
    CE_INITIALIZED,
    CE_HAS_KEYBOX,
    CE_HAS_SESSIONS,
    CE_ERROR
  };

 public:

  CryptoEngine();
  ~CryptoEngine();

  bool Initialized() { return (ce_state_ != CE_ILLEGAL); }

  void Terminate();

  bool isValid() { return valid_; }

  KeyboxError ValidateKeybox();
  WvKeybox& keybox() { return keybox_; }

  SessionId CreateSession();

  bool DestroySession(SessionId sid);

  SessionContext* FindSession(SessionId sid);

  void set_current_session_(SessionContext* current) {
    current_session_ = current;
  }

  bool DecryptMessage(SessionContext* session,
                      const std::vector<uint8_t>& key,
                      const std::vector<uint8_t>& iv,
                      const std::vector<uint8_t>& message,
                      std::vector<uint8_t>* decrypted);

  bool DecryptCTR(SessionContext* session,
                  const std::vector<uint8_t>& iv,
                  size_t byte_offset,
                  const std::vector<uint8_t>& cipher_data,
                  bool is_encrypted,
                  void* clear_data,
                  BufferType buffer_type);

 private:

  bool valid_;
  CryptoEngineState ce_state_;
  SessionContext* current_session_;
  ActiveSessions sessions_;
  WvKeybox keybox_;
  wvcdm::Lock session_table_lock_;

  CORE_DISALLOW_COPY_AND_ASSIGN(CryptoEngine);
};

};   // namespace wvoec_eng
#endif