// Copyright 2012 Google Inc. All Rights Reserved.
//
// Crypto - wrapper classes for OEMCrypto interface
//

#include "crypto_session.h"

#include <iostream>

#include "crypto_engine.h"
#include "log.h"
// TODO(gmorgan,jtinker): decide if OEMCryptoCENC is needed here.
#include "OEMCryptoCENC.h"
#include "string_conversions.h"
#include "wv_cdm_constants.h"

namespace {
// Encode unsigned integer into a big endian formatted string
std::string EncodeUint32(unsigned int u) {
  std::string s;
  s.append(1, (u >> 24) & 0xFF);
  s.append(1, (u >> 16) & 0xFF);
  s.append(1, (u >>  8) & 0xFF);
  s.append(1, (u >>  0) & 0xFF);
  return s;
}
}

namespace wvcdm {

// wrapper classes for OEMCrypto interface
// CryptoEngine -- top-level interface
// CryptoSession -- session-specific interface
// CryptoKey -- key interface

// CryptoSession methods

CryptoSession::CryptoSession() : valid_(false), open_(false) {}

CryptoSession::CryptoSession(const std::string& sname) : valid_(true),
    open_(false), cdm_session_id_(sname) {}

CryptoSession::~CryptoSession() {
  if (open_) {
    Close();
  }
  LOGV("CryptoSession::dtor: SLock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  crypto_engine->sessions_.erase(cdm_session_id_);
  if (0 == crypto_engine->sessions_.size()) {
    crypto_engine->DeleteInstance();
  }

  CryptoKeyMap::iterator i(keys_.begin());
  for (; i != keys_.end(); ++i)
    delete i->second;
  keys_.clear();
}

bool CryptoSession::Open() {
  LOGV("CryptoSession::Open: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);
  OEMCrypto_SESSION sid;
  OEMCryptoResult sts;
  if (open_)
    return false;
  sts = OEMCrypto_OpenSession(&sid);
  if (OEMCrypto_SUCCESS != sts) {
    open_ = false;
  } else {
    oec_session_id_ = static_cast<CryptoSessionId>(sid);
    LOGV("OpenSession: id= %ld", (uint32_t) oec_session_id_);
    open_ = true;
  }
  return open_;
}

void CryptoSession::Close() {
  LOGV("CryptoSession::Close: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);
  LOGV("CloseSession: id=%ld open=%s", (uint32_t) oec_session_id_, open_? "true" : "false") ;
  if (open_) {
    OEMCryptoResult sts = OEMCrypto_CloseSession(oec_session_id_);
    if (OEMCrypto_SUCCESS == sts) {
      open_ = false;
    }
  }
}

void CryptoSession::GenerateRequestId(std::string& req_id_str) {
  LOGV("CryptoSession::GenerateRequestId: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);
  // TODO(gmorgan): Get unique ID from OEMCrypto
  req_id_str.assign("987654321");
}

bool CryptoSession::PrepareRequest(const std::string& message,
                                   std::string* signature) {
  LOGV("CryptoSession::PrepareRequest: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);

  if (!signature) {
    LOGE("CryptoSession::PrepareRequest : No output destination provided.");
    return false;
  }

  uint8_t signature_buf[32];
  size_t length = 32;
  OEMCryptoResult sts;

  std::string mac_deriv_message;
  std::string enc_deriv_message;
  GenerateMacContext(message, &mac_deriv_message);
  GenerateEncryptContext(message, &enc_deriv_message);

  LOGV("GenerateDerivedKeys: id=%ld", (uint32_t) oec_session_id_);
  sts = OEMCrypto_GenerateDerivedKeys(
      oec_session_id_,
      reinterpret_cast<const uint8_t*>(mac_deriv_message.data()),
      mac_deriv_message.size(),
      reinterpret_cast<const uint8_t*>(enc_deriv_message.data()),
      enc_deriv_message.size());
  if (OEMCrypto_SUCCESS != sts) {
    return false;
  }

  LOGV("GenerateSignature: id=%ld", (uint32_t) oec_session_id_);
  sts = OEMCrypto_GenerateSignature(
      oec_session_id_,
      reinterpret_cast<const uint8_t*>(message.data()),
      message.size(),
      signature_buf,
      &length);
  if (OEMCrypto_SUCCESS != sts) {
    return false;
  }

  signature->assign(reinterpret_cast<const char*>(signature_buf), length);

  return true;
}

bool CryptoSession::PrepareRenewalRequest(const std::string& message,
                                          std::string* signature) {
  LOGV("CryptoSession::PrepareRenewalRequest: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);
  uint8_t signature_buf[32];
  size_t length = 32;

  OEMCryptoResult sts = OEMCrypto_GenerateSignature(
      oec_session_id_, reinterpret_cast<const uint8_t*>(message.data()),
      message.size(), signature_buf, &length);
  if (OEMCrypto_SUCCESS != sts) {
    return false;
  }

  signature->assign(reinterpret_cast<const char*>(signature_buf), length);

  return true;
}

void CryptoSession::GenerateMacContext(const std::string& input_context,
                                       std::string* deriv_context) {
  if (!deriv_context) {
    LOGE("CryptoSession::GenerateMacContext : No output destination provided.");
    return;
  }

  const std::string kSigningKeyLabel = "AUTHENTICATION";
  const size_t kSigningKeySizeBits = MAC_KEY_SIZE * 8;

  deriv_context->assign(kSigningKeyLabel);
  deriv_context->append(1, '\0');
  deriv_context->append(input_context);
  deriv_context->append(EncodeUint32(kSigningKeySizeBits));
}

void CryptoSession::GenerateEncryptContext(const std::string& input_context,
                                           std::string* deriv_context) {
  if (!deriv_context) {
    LOGE("CryptoSession::GenerateEncryptContext : No output destination provided.");
    return;
  }

  const std::string kEncryptionKeyLabel = "ENCRYPTION";
  const size_t kEncryptionKeySizeBits = KEY_SIZE * 8;

  deriv_context->assign(kEncryptionKeyLabel);
  deriv_context->append(1, '\0');
  deriv_context->append(input_context);
  deriv_context->append(EncodeUint32(kEncryptionKeySizeBits));
}

size_t CryptoSession::GetOffset(std::string message, std::string field) {
  size_t pos = message.find(field);
  if (pos == std::string::npos) {
    LOGE("CryptoSession::GetOffset : Cannot find offset for %s", field.c_str());
    pos = 0;
  }
  return pos;
}

bool CryptoSession::LoadKeys(const std::string& message,
                             const std::string& signature,
                             const std::string& mac_key_iv,
                             const std::string& mac_key,
                             int num_keys,
                             const CryptoKey* key_array) {
  LOGV("CryptoSession::LoadKeys: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);

  const uint8_t* msg = reinterpret_cast<const uint8_t*>(message.data());
  const uint8_t* enc_mac_key = NULL;
  const uint8_t* enc_mac_key_iv = NULL;
  if (mac_key.size() >= MAC_KEY_SIZE && mac_key_iv.size() >= KEY_IV_SIZE) {
    enc_mac_key = msg + GetOffset(message, mac_key);
    enc_mac_key_iv = msg + GetOffset(message, mac_key_iv);
  }
  OEMCrypto_KeyObject load_key_array[num_keys];
  for (int i=0; i<num_keys; ++i) {
    const CryptoKey* ki = &key_array[i];
    OEMCrypto_KeyObject* ko = &load_key_array[i];
    ko->key_id = msg + GetOffset(message, ki->key_id());
    ko->key_id_length = ki->key_id().length();
    ko->key_data_iv = msg + GetOffset(message, ki->key_data_iv());
    ko->key_data = msg + GetOffset(message, ki->key_data());
    if (ki->HasKeyControl()) {
      ko->key_control_iv = msg + GetOffset(message, ki->key_control_iv());
      ko->key_control = msg + GetOffset(message, ki->key_control());
    }
    else {
      LOGE("For key %d: XXX key has no control block. size=%d", i, ki->key_control().size());
      ko->key_control_iv = NULL;
      ko->key_control = NULL;
    }
  }
  LOGV("LoadKeys: id=%ld", (uint32_t) oec_session_id_);
  return (OEMCrypto_SUCCESS == OEMCrypto_LoadKeys(
      oec_session_id_, msg, message.size(),
      reinterpret_cast<const uint8_t*>(signature.data()),
      signature.size(), enc_mac_key_iv, enc_mac_key,
      num_keys, load_key_array));
}

bool CryptoSession::RefreshKeys(const std::string& message,
                                const std::string& signature,
                                int num_keys,
                                const CryptoKey* key_array) {
  LOGV("CryptoSession::RefreshKeys: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);

  const uint8_t* msg = reinterpret_cast<const uint8_t*>(message.data());
  OEMCrypto_KeyRefreshObject load_key_array[num_keys];
  for (int i=0; i<num_keys; ++i) {
    const CryptoKey* ki = &key_array[i];
    OEMCrypto_KeyRefreshObject* ko = &load_key_array[i];
    if (ki->key_id().empty()) {
      ko->key_id = NULL;
    } else {
      ko->key_id = msg + GetOffset(message, ki->key_id());
    }
    if (ki->HasKeyControl()) {
      if (ki->key_control_iv().empty()) {
        ko->key_control_iv = NULL;
      } else {
        ko->key_control_iv = msg + GetOffset(message, ki->key_control_iv());
      }
      ko->key_control = msg + GetOffset(message, ki->key_control());
    }
    else {
      ko->key_control_iv = NULL;
      ko->key_control = NULL;
    }
  }
  LOGV("RefreshKeys: id=%ld", static_cast<uint32_t>(oec_session_id_));
  return (OEMCrypto_SUCCESS == OEMCrypto_RefreshKeys(
      oec_session_id_, msg, message.size(),
      reinterpret_cast<const uint8_t*>(signature.data()),
      signature.size(),
      num_keys, load_key_array));
}

bool CryptoSession::SelectKey(const std::string& key_id) {
  LOGV("CryptoSession::SelectKey: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);
  const uint8_t* key_id_string =
      reinterpret_cast<const uint8_t*>(key_id.data());

  LOGV("SelectKey: id=%ld", static_cast<uint32_t>(oec_session_id_));
  OEMCryptoResult sts = OEMCrypto_SelectKey(oec_session_id_, key_id_string,
                                            key_id.size());
  if (OEMCrypto_SUCCESS != sts) {
    return false;
  }
  return true;
}

bool CryptoSession::Decrypt(const InputDescriptor input,
                            const OutputDescriptor output) {
  LOGV("CryptoSession::Decrypt: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);
  // TODO(gmorgan): handle inputs and outputs to decrypt call
  const uint8_t* data_addr = NULL;
  uint32_t data_length = 0;
  bool is_encrypted = false;
  uint8_t* iv = NULL;
  uint32_t offset = 0;
  const OEMCrypto_DestBufferDesc* out_buffer = NULL;
  OEMCryptoResult sts = OEMCrypto_DecryptCTR(oec_session_id_, data_addr,
                                             data_length, is_encrypted, iv,
                                             offset, out_buffer);
  if (OEMCrypto_SUCCESS != sts) {
    return false;
  }
  return true;
}

bool CryptoSession::GenerateNonce(uint32_t* nonce) {
  if (!nonce) {
    LOGE("input parameter is null");
    return false;
  }

  LOGV("CryptoSession::GenerateNonce: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);
  return(OEMCrypto_SUCCESS == OEMCrypto_GenerateNonce(oec_session_id_, nonce));
}
};  // namespace wvcdm