// Copyright 2013 Google Inc. All Rights Reserved.

#include "cdm_engine.h"

#include <stdlib.h>

#include <iostream>
#include <sstream>

#include "cdm_session.h"
#include "clock.h"
#include "device_files.h"
#include "license_protocol.pb.h"
#include "log.h"
#include "properties.h"
#include "scoped_ptr.h"
#include "string_conversions.h"
#include "wv_cdm_constants.h"
#include "wv_cdm_event_listener.h"

namespace {
  const uint32_t kUpdateUsageInformationPeriod = 60;  // seconds
  const size_t kUsageReportsPerRequest = 1;
}  // unnamed namespace

namespace wvcdm {

bool CdmEngine::seeded_ = false;

CdmEngine::CdmEngine()
    : cert_provisioning_(NULL),
      cert_provisioning_requested_security_level_(kLevelDefault),
      usage_session_(NULL),
      last_usage_information_update_time(0) {
  Properties::Init();
  if (!seeded_) {
    Clock clock;
    srand(clock.GetCurrentTime());
    seeded_ = true;
  }
}

CdmEngine::~CdmEngine() {
  CdmSessionMap::iterator i(sessions_.begin());
  for (; i != sessions_.end(); ++i) {
    delete i->second;
  }
  sessions_.clear();
}

CdmResponseType CdmEngine::OpenSession(
    const CdmKeySystem& key_system,
    const CdmClientPropertySet* property_set,
    CdmSessionId* session_id) {
  LOGI("CdmEngine::OpenSession");

  if (!ValidateKeySystem(key_system)) {
    LOGI("CdmEngine::OpenSession: invalid key_system = %s", key_system.c_str());
    return KEY_ERROR;
  }

  if (!session_id) {
    LOGE("CdmEngine::OpenSession: no session ID destination provided");
    return KEY_ERROR;
  }

  scoped_ptr<CdmSession> new_session(new CdmSession(property_set));
  if (new_session->session_id().empty()) {
    LOGE("CdmEngine::OpenSession: failure to generate session ID");
    return UNKNOWN_ERROR;
  }

  CdmResponseType sts = new_session->Init();
  if (sts != NO_ERROR) {
    if (sts == NEED_PROVISIONING) {
      cert_provisioning_requested_security_level_ =
          new_session->GetRequestedSecurityLevel();
    } else {
      LOGE("CdmEngine::OpenSession: bad session init: %u", sts);
    }
    return sts;
  }
  *session_id = new_session->session_id();
  sessions_[*session_id] = new_session.release();
  return NO_ERROR;
}

CdmResponseType CdmEngine::OpenKeySetSession(const CdmKeySetId& key_set_id) {
  LOGI("CdmEngine::OpenKeySetSession");

  if (key_set_id.empty()) {
    LOGE("CdmEngine::OpenKeySetSession: invalid key set id");
    return KEY_ERROR;
  }

  CdmSessionId session_id;
  CdmResponseType sts = OpenSession(KEY_SYSTEM, NULL, &session_id);

  if (sts != NO_ERROR)
    return sts;

  release_key_sets_[key_set_id] = session_id;
  return NO_ERROR;
}

CdmResponseType CdmEngine::CloseSession(const CdmSessionId& session_id) {
  LOGI("CdmEngine::CloseSession");

  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::CloseSession: session not found = %s", session_id.c_str());
    return KEY_ERROR;
  }

  CdmSession* session = iter->second;
  sessions_.erase(session_id);
  delete session;
  return NO_ERROR;
}

CdmResponseType CdmEngine::CloseKeySetSession(const CdmKeySetId& key_set_id) {
  LOGI("CdmEngine::CloseKeySetSession");

  CdmReleaseKeySetMap::iterator iter = release_key_sets_.find(key_set_id);
  if (iter == release_key_sets_.end()) {
    LOGE("CdmEngine::CloseKeySetSession: key set id not found = %s",
        key_set_id.c_str());
    return KEY_ERROR;
  }

  CdmResponseType sts = CloseSession(iter->second);
  release_key_sets_.erase(iter);
  return sts;
}

CdmResponseType CdmEngine::GenerateKeyRequest(
    const CdmSessionId& session_id,
    const CdmKeySetId& key_set_id,
    const InitializationData& init_data,
    const CdmLicenseType license_type,
    CdmAppParameterMap& app_parameters,
    CdmKeyMessage* key_request,
    std::string* server_url) {
  LOGI("CdmEngine::GenerateKeyRequest");

  CdmSessionId id = session_id;
  CdmResponseType sts;

  if (license_type == kLicenseTypeRelease) {
    if (key_set_id.empty()) {
      LOGE("CdmEngine::GenerateKeyRequest: invalid key set ID");
      return UNKNOWN_ERROR;
    }

    if (!session_id.empty()) {
      LOGE("CdmEngine::GenerateKeyRequest: invalid session ID = %s",
          session_id.c_str());
      return UNKNOWN_ERROR;
    }

    CdmReleaseKeySetMap::iterator iter = release_key_sets_.find(key_set_id);
    if (iter == release_key_sets_.end()) {
      LOGE("CdmEngine::GenerateKeyRequest: key set ID not found = %s",
          key_set_id.c_str());
      return UNKNOWN_ERROR;
    }

    id = iter->second;
  }

  CdmSessionMap::iterator iter = sessions_.find(id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::GenerateKeyRequest: session_id not found = %s",
        id.c_str());
    return KEY_ERROR;
  }

  if (!key_request) {
    LOGE("CdmEngine::GenerateKeyRequest: no key request destination provided");
    return KEY_ERROR;
  }

  key_request->clear();

  if (license_type == kLicenseTypeRelease) {
    sts = iter->second->RestoreOfflineSession(key_set_id, kLicenseTypeRelease);
    if (sts != KEY_ADDED) {
      LOGE("CdmEngine::GenerateKeyRequest: key release restoration failed,"
          "sts = %d", (int)sts);
      return sts;
    }
  }

  sts = iter->second->GenerateKeyRequest(init_data, license_type,
                                         app_parameters, key_request,
                                         server_url);

  if (KEY_MESSAGE != sts) {
    if (sts == NEED_PROVISIONING) {
      cert_provisioning_requested_security_level_ =
          iter->second->GetRequestedSecurityLevel();
    }
    LOGE("CdmEngine::GenerateKeyRequest: key request generation failed, "
        "sts = %d", (int)sts);
    return sts;
  }

  if (license_type == kLicenseTypeRelease) {
    OnKeyReleaseEvent(key_set_id);
  }

  return KEY_MESSAGE;
}

CdmResponseType CdmEngine::AddKey(
    const CdmSessionId& session_id,
    const CdmKeyResponse& key_data,
    CdmKeySetId* key_set_id) {
  LOGI("CdmEngine::AddKey");

  CdmSessionId id = session_id;
  bool license_type_release = session_id.empty();

  if (license_type_release) {
    if (!key_set_id) {
      LOGE("CdmEngine::AddKey: no key set id provided");
      return KEY_ERROR;
    }

    if (key_set_id->empty()) {
      LOGE("CdmEngine::AddKey: invalid key set id");
      return KEY_ERROR;
    }

    CdmReleaseKeySetMap::iterator iter = release_key_sets_.find(*key_set_id);
    if (iter == release_key_sets_.end()) {
      LOGE("CdmEngine::AddKey: key set id not found = %s", key_set_id->c_str());
      return KEY_ERROR;
    }

    id = iter->second;
  }

  CdmSessionMap::iterator iter = sessions_.find(id);

  if (iter == sessions_.end()) {
    LOGE("CdmEngine::AddKey: session id not found = %s", id.c_str());
    return KEY_ERROR;
  }

  if (key_data.empty()) {
    LOGE("CdmEngine::AddKey: no key_data");
    return KEY_ERROR;
  }

  CdmResponseType sts = iter->second->AddKey(key_data, key_set_id);

  if (KEY_ADDED != sts) {
    LOGE("CdmEngine::AddKey: keys not added, result = %d", (int)sts);
    return sts;
  }

  return KEY_ADDED;
}

CdmResponseType CdmEngine::RestoreKey(
    const CdmSessionId& session_id,
    const CdmKeySetId& key_set_id) {
  LOGI("CdmEngine::RestoreKey");

  if (key_set_id.empty()) {
    LOGI("CdmEngine::RestoreKey: invalid key set id");
    return KEY_ERROR;
  }

  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::RestoreKey: session_id not found = %s ",
        session_id.c_str());
    return UNKNOWN_ERROR;
  }

  CdmResponseType sts =
      iter->second->RestoreOfflineSession(key_set_id, kLicenseTypeOffline);
  if (sts == NEED_PROVISIONING) {
    cert_provisioning_requested_security_level_ =
        iter->second->GetRequestedSecurityLevel();
  }
  return sts;
}

CdmResponseType CdmEngine::CancelKeyRequest(const CdmSessionId& session_id) {
  LOGI("CdmEngine::CancelKeyRequest");

  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::CancelKeyRequest: session_id not found = %s",
         session_id.c_str());
    return KEY_ERROR;
  }

  iter->second->CancelKeyRequest();
  return NO_ERROR;
}

CdmResponseType CdmEngine::GenerateRenewalRequest(
    const CdmSessionId& session_id,
    CdmKeyMessage* key_request,
    std::string* server_url) {
  LOGI("CdmEngine::GenerateRenewalRequest");

  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::GenerateRenewalRequest: session_id not found = %s",
         session_id.c_str());
    return KEY_ERROR;
  }

  if (!key_request) {
    LOGE("CdmEngine::GenerateRenewalRequest: no key request destination");
    return KEY_ERROR;
  }

  key_request->clear();

  CdmResponseType sts = iter->second->GenerateRenewalRequest(key_request,
                                                             server_url);

  if (KEY_MESSAGE != sts) {
    LOGE("CdmEngine::GenerateRenewalRequest: key request gen. failed, sts=%d",
         (int)sts);
    return sts;
  }

  return KEY_MESSAGE;
}

CdmResponseType CdmEngine::RenewKey(
    const CdmSessionId& session_id,
    const CdmKeyResponse& key_data) {
  LOGI("CdmEngine::RenewKey");

  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::RenewKey: session_id not found = %s", session_id.c_str());
    return KEY_ERROR;
  }

  if (key_data.empty()) {
    LOGE("CdmEngine::RenewKey: no key_data");
    return KEY_ERROR;
  }

  CdmResponseType sts = iter->second->RenewKey(key_data);
  if (KEY_ADDED != sts) {
    LOGE("CdmEngine::RenewKey: keys not added, sts=%d", (int)sts);
    return sts;
  }

  return KEY_ADDED;
}

CdmResponseType CdmEngine::QueryStatus(CdmQueryMap* key_info) {
  LOGI("CdmEngine::QueryStatus");
  CryptoSession crypto_session;
  switch (crypto_session.GetSecurityLevel()) {
    case kSecurityLevelL1:
      (*key_info)[QUERY_KEY_SECURITY_LEVEL] = QUERY_VALUE_SECURITY_LEVEL_L1;
      break;
    case kSecurityLevelL2:
      (*key_info)[QUERY_KEY_SECURITY_LEVEL] = QUERY_VALUE_SECURITY_LEVEL_L2;
      break;
    case kSecurityLevelL3:
      (*key_info)[QUERY_KEY_SECURITY_LEVEL] = QUERY_VALUE_SECURITY_LEVEL_L3;
      break;
    case kSecurityLevelUninitialized:
    case kSecurityLevelUnknown:
      (*key_info)[QUERY_KEY_SECURITY_LEVEL] =
          QUERY_VALUE_SECURITY_LEVEL_UNKNOWN;
      break;
    default:
      return KEY_ERROR;
  }

  std::string deviceId;
  bool success = crypto_session.GetDeviceUniqueId(&deviceId);
  if (success) {
    (*key_info)[QUERY_KEY_DEVICE_ID] = deviceId;
  }

  uint32_t system_id;
  success = crypto_session.GetSystemId(&system_id);
  if (success) {
    std::ostringstream system_id_stream;
    system_id_stream << system_id;
    (*key_info)[QUERY_KEY_SYSTEM_ID] = system_id_stream.str();
  }

  std::string provisioning_id;
  success = crypto_session.GetProvisioningId(&provisioning_id);
  if (success) {
    (*key_info)[QUERY_KEY_PROVISIONING_ID] = provisioning_id;
  }

  return NO_ERROR;
}

CdmResponseType CdmEngine::QuerySessionStatus(const CdmSessionId& session_id,
                                              CdmQueryMap* key_info) {
  LOGI("CdmEngine::QuerySessionStatus");
  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::QuerySessionStatus: session_id not found = %s",
         session_id.c_str());
    return KEY_ERROR;
  }
  return iter->second->QueryStatus(key_info);
}

CdmResponseType CdmEngine::QueryKeyStatus(
    const CdmSessionId& session_id,
    CdmQueryMap* key_info) {
  LOGI("CdmEngine::QueryKeyStatus");
  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::QueryKeyStatus: session_id not found = %s",
         session_id.c_str());
    return KEY_ERROR;
  }
  return iter->second->QueryKeyStatus(key_info);
}

CdmResponseType CdmEngine::QueryKeyControlInfo(
    const CdmSessionId& session_id,
    CdmQueryMap* key_info) {
  LOGI("CdmEngine::QueryKeyControlInfo");
  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::QueryKeyControlInfo: session_id not found = %s",
         session_id.c_str());
    return KEY_ERROR;
  }
  return iter->second->QueryKeyControlInfo(key_info);
}

/*
 * Composes a device provisioning request and output the request in JSON format
 * in *request. It also returns the default url for the provisioning server
 * in *default_url.
 *
 * Returns NO_ERROR for success and UNKNOWN_ERROR if fails.
 */
CdmResponseType CdmEngine::GetProvisioningRequest(
    CdmCertificateType cert_type,
    const std::string& cert_authority,
    CdmProvisioningRequest* request,
    std::string* default_url) {
  if (!request || !default_url) {
    LOGE("CdmEngine::GetProvisioningRequest: invalid input parameters");
    return UNKNOWN_ERROR;
  }
  if (NULL == cert_provisioning_.get()) {
    cert_provisioning_.reset(new CertificateProvisioning());
  }
  CdmResponseType ret = cert_provisioning_->GetProvisioningRequest(
      cert_provisioning_requested_security_level_,
      cert_type,
      cert_authority,
      request,
      default_url);
  if (ret != NO_ERROR) {
    cert_provisioning_.reset(NULL);  // Release resources.
  }
  return ret;
}

/*
 * The response message consists of a device certificate and the device RSA key.
 * The device RSA key is stored in the T.E.E. The device certificate is stored
 * in the device.
 *
 * Returns NO_ERROR for success and UNKNOWN_ERROR if fails.
 */
CdmResponseType CdmEngine::HandleProvisioningResponse(
    CdmProvisioningResponse& response,
    std::string* cert,
    std::string* wrapped_key) {
  if (response.empty()) {
    LOGE("CdmEngine::HandleProvisioningResponse: Empty provisioning response.");
    cert_provisioning_.reset(NULL);
    return UNKNOWN_ERROR;
  }
  if (NULL == cert) {
    LOGE("CdmEngine::HandleProvisioningResponse: invalid certificate "
         "destination");
    cert_provisioning_.reset(NULL);
    return UNKNOWN_ERROR;
  }
  if (NULL == wrapped_key) {
    LOGE("CdmEngine::HandleProvisioningResponse: invalid wrapped key "
         "destination");
    cert_provisioning_.reset(NULL);
    return UNKNOWN_ERROR;
  }
  if (NULL == cert_provisioning_.get()) {
    LOGE("CdmEngine::HandleProvisioningResponse: provisioning object missing.");
    return UNKNOWN_ERROR;
  }
  CdmResponseType ret =
      cert_provisioning_->HandleProvisioningResponse(response, cert,
                                                     wrapped_key);
  cert_provisioning_.reset(NULL);  // Release resources.
  return ret;
}

CdmResponseType CdmEngine::Unprovision(CdmSecurityLevel security_level) {
  DeviceFiles handle;
  if (!handle.Init(security_level)) {
    LOGE("CdmEngine::Unprovision: unable to initialize device files");
    return UNKNOWN_ERROR;
  }

  if (!handle.DeleteAllFiles()) {
    LOGE("CdmEngine::Unprovision: unable to delete files");
    return UNKNOWN_ERROR;
  }
  return NO_ERROR;
}

CdmResponseType CdmEngine::GetUsageInfo(CdmUsageInfo* usage_info) {
  usage_session_.reset(new CdmSession(NULL));

  CdmResponseType status = usage_session_->Init();
  if (NO_ERROR != status) {
    LOGE("CdmEngine::GetUsageInfo: session init error");
    return status;
  }

  DeviceFiles handle;
  if (!handle.Init(usage_session_->GetSecurityLevel())) {
    LOGE("CdmEngine::GetUsageInfo: unable to initialize device files");
    return status;
  }

  std::vector<std::pair<CdmKeyMessage, CdmKeyResponse> > license_info;
  if (!handle.RetrieveUsageInfo(&license_info)) {
    LOGE("CdmEngine::GetUsageInfo: unable to read usage information");
    return UNKNOWN_ERROR;
  }

  if (0 == license_info.size()) {
    usage_info->resize(0);
    return NO_ERROR;
  }

  std::string server_url;
  usage_info->resize(kUsageReportsPerRequest);

  uint32_t index = rand() % license_info.size();
  status = usage_session_->RestoreUsageSession(license_info[index].first,
                                                license_info[index].second);
  if (KEY_ADDED != status) {
    LOGE("CdmEngine::GetUsageInfo: restore usage session (%d) error %ld",
         index, status);
    usage_info->clear();
    return status;
  }

  status = usage_session_->GenerateReleaseRequest(&(*usage_info)[0], &server_url);

  if (KEY_MESSAGE != status) {
    LOGE("CdmEngine::GetUsageInfo: generate release request error: %ld",
         status);
    usage_info->clear();
    return status;
  }
  return KEY_MESSAGE;
}

CdmResponseType CdmEngine::ReleaseUsageInfo(
    const CdmUsageInfoReleaseMessage& message) {
  if (NULL == usage_session_.get()) {
    LOGE("CdmEngine::ReleaseUsageInfo: cdm session not initialized");
    return UNKNOWN_ERROR;
  }

  CdmResponseType status = usage_session_->ReleaseKey(message);
  if (NO_ERROR != status) {
    LOGE("CdmEngine::ReleaseUsageInfo: release key error: %ld", status);
  }
  return status;
}

CdmResponseType CdmEngine::Decrypt(
    const CdmSessionId& session_id,
    const CdmDecryptionParameters& parameters) {
  if (parameters.key_id == NULL) {
    LOGE("CdmEngine::Decrypt: no key_id");
    return KEY_ERROR;
  }

  if (parameters.encrypt_buffer == NULL) {
    LOGE("CdmEngine::Decrypt: no src encrypt buffer");
    return KEY_ERROR;
  }

  if (parameters.iv == NULL) {
    LOGE("CdmEngine::Decrypt: no iv");
    return KEY_ERROR;
  }

  if (parameters.decrypt_buffer == NULL) {
    if (!parameters.is_secure &&
        !Properties::Properties::oem_crypto_use_fifo()) {
      LOGE("CdmEngine::Decrypt: no dest decrypt buffer");
      return KEY_ERROR;
    } // else we must be level 1 direct and we don't need to return a buffer.
  }

  CdmSessionMap::iterator iter;
  if (session_id.empty()) {
    // Loop through the sessions to find the session containing the key_id.
    for (iter = sessions_.begin(); iter != sessions_.end(); ++iter) {
      if (iter->second->IsKeyLoaded(*parameters.key_id)) {
        break;
      }
    }
  } else {
    iter = sessions_.find(session_id);
  }
  if (iter == sessions_.end()) {
    LOGE("CdmEngine::Decrypt: session not found: id=%s, id size=%d",
         session_id.c_str(),
         session_id.size());
    return KEY_ERROR;
  }

  return iter->second->Decrypt(parameters);
}

bool CdmEngine::IsKeyLoaded(const KeyId& key_id) {
  for (CdmSessionMap::iterator iter = sessions_.begin();
       iter != sessions_.end(); ++iter) {
    if (iter->second->IsKeyLoaded(key_id)) {
      return true;
    }
  }
  return false;
}

bool CdmEngine::FindSessionForKey(
    const KeyId& key_id,
    CdmSessionId* session_id) {
  if (NULL == session_id) {
    LOGE("CdmEngine::FindSessionForKey: session id not provided");
    return false;
  }

  CdmSessionMap::iterator iter = sessions_.find(*session_id);
  if (iter != sessions_.end()) {
    if (iter->second->IsKeyLoaded(key_id)) {
      return true;
    }
  }

  uint32_t session_sharing_id = Properties::GetSessionSharingId(*session_id);

  for (iter = sessions_.begin(); iter != sessions_.end(); ++iter) {
    CdmSessionId local_session_id = iter->second->session_id();
    if (Properties::GetSessionSharingId(local_session_id) ==
        session_sharing_id) {
      if (iter->second->IsKeyLoaded(key_id)) {
        *session_id = local_session_id;
        return true;
      }
    }
  }
  return false;
}

bool CdmEngine::AttachEventListener(
    const CdmSessionId& session_id,
    WvCdmEventListener* listener) {

  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    return false;
  }

  return iter->second->AttachEventListener(listener);
}

bool CdmEngine::DetachEventListener(
    const CdmSessionId& session_id,
    WvCdmEventListener* listener) {

  CdmSessionMap::iterator iter = sessions_.find(session_id);
  if (iter == sessions_.end()) {
    return false;
  }

  return iter->second->DetachEventListener(listener);
}

bool CdmEngine::ValidateKeySystem(const CdmKeySystem& key_system) {
  return (key_system.find("widevine") != std::string::npos);
}

void CdmEngine::OnTimerEvent() {
  Clock clock;
  uint64_t current_time = clock.GetCurrentTime();
  bool update_usage_information = false;

  if (current_time - last_usage_information_update_time >
      kUpdateUsageInformationPeriod) {
    update_usage_information = true;
    last_usage_information_update_time = current_time;
  }

  for (CdmSessionMap::iterator iter = sessions_.begin();
      iter != sessions_.end(); ++iter) {
    iter->second->OnTimerEvent(update_usage_information);

    if (update_usage_information && iter->second->is_usage_update_needed()) {
      // usage is updated for all sessions so this needs to be
      // called only once per update usage information period
      CdmResponseType status = iter->second->UpdateUsageInformation();
      if (NO_ERROR != status) {
        LOGW("Update usage information failed: %u", status);
      } else {
        update_usage_information = false;
      }
    }
    iter->second->reset_is_usage_update_needed();
  }
}

void CdmEngine::OnKeyReleaseEvent(const CdmKeySetId& key_set_id) {

  for (CdmSessionMap::iterator iter = sessions_.begin();
       iter != sessions_.end(); ++iter) {
    iter->second->OnKeyReleaseEvent(key_set_id);
  }
}

}  // namespace wvcdm