// 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 "properties.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),
    is_destination_buffer_type_valid_(false) {}

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

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

  OEMCryptoResult sts;

  if (!Properties::use_certificates_as_identification()) {
    if (!GenerateDerivedKeys(message)) {
      return false;
    }
  }

  if (!GenerateSignature(message, signature)) {
    return false;
  }

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

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

  if (!GenerateSignature(message, signature)) {
    return false;
  }

  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);
  }
  std::vector<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());
    ko->key_data_length = ki->key_data().length();
    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[0]));
}

bool CryptoSession::LoadCertificatePrivateKey(std::string& wrapped_key) {
  LOGV("CryptoSession::LoadKeys: Lock");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);

  LOGV("LoadDeviceRSAKey: id=%ld", (uint32_t) oec_session_id_);
  OEMCryptoResult sts = OEMCrypto_LoadDeviceRSAKey(
      oec_session_id_,
      reinterpret_cast<const uint8_t*>(wrapped_key.data()),
      wrapped_key.size());

  if (OEMCrypto_SUCCESS != sts) {
    LOGD("LoadCertificatePrivateKey: OEMCrypto_LoadDeviceRSAKey error=%d", sts);
    return false;
  }

  return true;
}

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());
  std::vector<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[0]));
}

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::GenerateDerivedKeys(const std::string& message) {
  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_);
  OEMCryptoResult 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) {
    LOGD("GenerateDerivedKeys: OEMCrypto_GenerateDerivedKeys error=%d", sts);
    return false;
  }

  return true;
}

bool CryptoSession::GenerateDerivedKeys(const std::string& message,
                                        const std::string& session_key) {
  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_);
  OEMCryptoResult sts = OEMCrypto_DeriveKeysFromSessionKey(
      oec_session_id_,
      reinterpret_cast<const uint8_t*>(session_key.data()),
      session_key.size(),
      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) {
    LOGD("GenerateDerivedKeys: OEMCrypto_DeriveKeysFromSessionKey err=%d", sts);
    return false;
  }

  return true;
}

bool CryptoSession::GenerateSignature(const std::string& message,
                                      std::string* signature) {
  LOGV("GenerateSignature: id=%ld", (uint32_t) oec_session_id_);
  uint8_t signature_buf[32];
  uint32_t length = 32;
  OEMCryptoResult sts;
  if (Properties::use_certificates_as_identification()) {
    sts = OEMCrypto_GenerateRSASignature(
        oec_session_id_,
        reinterpret_cast<const uint8_t*>(message.data()),
        message.size(),
        signature_buf,
        &length);
  }
  else {
    sts = OEMCrypto_GenerateSignature(
        oec_session_id_,
        reinterpret_cast<const uint8_t*>(message.data()),
        message.size(),
        signature_buf,
        &length);
  }

  if (OEMCrypto_SUCCESS != sts) {
    LOGD("GenerateSignature: OEMCrypto_GenerateSignature err=%d", sts);
    return false;
  }

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

CdmResponseType CryptoSession::Decrypt(bool is_encrypted,
                                       const uint8_t* encrypt_buffer,
                                       size_t encrypt_length,
                                       const std::vector<uint8_t>& iv,
                                       size_t block_offset,
                                       void* decrypt_buffer,
                                       bool is_video) {
  if (!is_destination_buffer_type_valid_) {
    if (!SetDestinationBufferType())
      return UNKNOWN_ERROR;
  }

  OEMCrypto_DestBufferDesc buffer_descriptor;
  buffer_descriptor.type = destination_buffer_type_;

  switch (buffer_descriptor.type) {
    case OEMCrypto_BufferType_Clear:
      buffer_descriptor.buffer.clear.address =
          static_cast<uint8_t*>(decrypt_buffer);
      buffer_descriptor.buffer.clear.max_length = encrypt_length;
      break;
    case OEMCrypto_BufferType_Secure:
      buffer_descriptor.buffer.secure.handle = decrypt_buffer;
      buffer_descriptor.buffer.secure.max_length = encrypt_length;
      break;
    case OEMCrypto_BufferType_Direct:
      buffer_descriptor.type = OEMCrypto_BufferType_Direct;
      buffer_descriptor.buffer.direct.is_video = is_video;
      break;
  }

  OEMCryptoResult sts = OEMCrypto_DecryptCTR(oec_session_id_, encrypt_buffer,
                                             encrypt_length, is_encrypted,
                                             &iv[0], block_offset,
                                             &buffer_descriptor);

  if (OEMCrypto_SUCCESS != sts) {
    return UNKNOWN_ERROR;
  }
  return NO_ERROR;
}

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

bool CryptoSession::SetDestinationBufferType() {
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();

  if (Properties::oem_crypto_use_secure_buffers()) {
    if (crypto_engine->GetSecurityLevel() == CryptoEngine::kSecurityLevelL1) {
      destination_buffer_type_ = OEMCrypto_BufferType_Secure;
    }
    else {
      destination_buffer_type_ = OEMCrypto_BufferType_Clear;
    }
  }
  else if (Properties::oem_crypto_use_fifo()) {
    destination_buffer_type_ = OEMCrypto_BufferType_Direct;
  }
  else if (Properties::oem_crypto_use_userspace_buffers()) {
    destination_buffer_type_ = OEMCrypto_BufferType_Clear;
  }
  else {
    return false;
  }

  is_destination_buffer_type_valid_ = true;
  return true;
}

bool CryptoSession::RewrapDeviceRSAKey(const std::string& message,
                                       const uint32_t* nonce,
                                       const uint8_t* enc_rsa_key,
                                       size_t enc_rsa_key_length,
                                       const uint8_t* enc_rsa_key_iv,
                                       uint8_t* wrapped_rsa_key,
                                       size_t* wrapped_rsa_key_length) {
  LOGV("CryptoSession::RewrapDeviceRSAKey: Lock+++");
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  AutoLock auto_lock(crypto_engine->crypto_lock_);

  LOGV("crypto session id=%ld", static_cast<uint32_t>(oec_session_id_));

  // HMAC-SHA256 signature
  uint8_t signature[kSignatureSize];
  size_t signature_length = kSignatureSize;
  OEMCryptoResult status = OEMCrypto_GenerateSignature(
      oec_session_id_,
      reinterpret_cast<const uint8_t*>(message.data()),
      message.size(),
      signature,
      &signature_length);
  if (OEMCrypto_SUCCESS != status) {
    LOGE("CryptoSession::RewrapDeviceRSAKey: GenerateSiganture failed");
    return false;
  }

  status = OEMCrypto_RewrapDeviceRSAKey(
      oec_session_id_,
      reinterpret_cast<const uint8_t*>(message.data()), message.length(),
      signature, signature_length,
      nonce,
      enc_rsa_key, enc_rsa_key_length,
      enc_rsa_key_iv,
      wrapped_rsa_key,
      wrapped_rsa_key_length);

  if (OEMCrypto_SUCCESS != status) {
    return false;
  }
  return true;
}

};  // namespace wvcdm