// Copyright 2012 Google Inc. All Rights Reserved.
// Author: jfore@google.com (Jeff Fore), rkuroiwa@google.com (Rintaro Kuroiwa)

#include "cdm_session.h"

#include <iostream>
#include <sstream>
#include <stdlib.h>

#include "clock.h"
#include "cdm_engine.h"
#include "crypto_engine.h"
#include "device_files.h"
#include "log.h"
#include "openssl/sha.h"
#include "properties.h"
#include "string_conversions.h"
#include "wv_cdm_constants.h"

namespace wvcdm {

typedef std::set<WvCdmEventListener*>::iterator CdmEventListenerIter;

CdmResponseType CdmSession::Init() {
  CryptoEngine* crypto_engine = CryptoEngine::GetInstance();
  if (!crypto_engine) {
    LOGE("CdmSession::Init failed to get CryptoEngine instance.");
    return UNKNOWN_ERROR;
  }

  crypto_session_ = crypto_engine->CreateSession(session_id_);
  if (!crypto_session_) {
    LOGE("CdmSession::Init crypto session creation failure");
    return UNKNOWN_ERROR;
  }

  std::string token;
  if (Properties::use_certificates_as_identification()) {
    if (!LoadDeviceCertificate(&token, &wrapped_key_)) {
      LOGE("CdmSession::Init provisioning needed");
      return NEED_PROVISIONING;
    }
  }
  else {
    if (!crypto_engine->GetToken(&token)) {
      LOGE("CdmSession::Init token retrieval failure");
      return UNKNOWN_ERROR;
    }
  }

  if (license_parser_.Init(token, crypto_session_, &policy_engine_))
    return NO_ERROR;
  else
    return UNKNOWN_ERROR;
}

CdmResponseType CdmSession::ReInit() {
  DestroySession();
  return Init();
}

bool CdmSession::DestroySession() {
  if (crypto_session_) {
    delete crypto_session_;
    crypto_session_ = NULL;
  }
  return true;
}

CdmResponseType CdmSession::RestoreOfflineSession(
    const CdmKeySetId& key_set_id,
    const CdmLicenseType license_type) {
  key_set_id_ = key_set_id;

  // Retrieve license information from persistent store
  DeviceFiles::LicenseState license_state;

  if (!DeviceFiles::RetrieveLicense(key_set_id, &license_state,
                                    &offline_pssh_data_,
                                    &offline_key_request_,
                                    &offline_key_response_,
                                    &offline_key_renewal_request_,
                                    &offline_key_renewal_response_,
                                    &offline_release_server_url_)) {
    LOGE("CdmSession::Init failed to retrieve license. key set id = %s",
        key_set_id.c_str());
    return UNKNOWN_ERROR;
  }

  if (license_state != DeviceFiles::kLicenseStateActive)  {
    LOGE("CdmSession::Init invalid offline license state = %s", license_state);
    return UNKNOWN_ERROR;
  }

  if (Properties::use_certificates_as_identification()) {
    if (!crypto_session_->LoadCertificatePrivateKey(wrapped_key_)) {
      return NEED_PROVISIONING;
    }
  }

  if (!license_parser_.RestoreOfflineLicense(offline_key_request_,
                                             offline_key_response_,
                                             offline_key_renewal_response_)) {
    return UNKNOWN_ERROR;
  }

  license_received_ = true;
  license_type_ = license_type;
  return KEY_ADDED;
}

bool CdmSession::VerifySession(const CdmKeySystem& key_system,
                               const CdmInitData& init_data) {
  // TODO(gmorgan): Compare key_system and init_data with value received
  // during session startup - they should be the same.
  return true;
}

CdmResponseType CdmSession::GenerateKeyRequest(
    const CdmInitData& init_data,
    const CdmLicenseType license_type,
    const CdmAppParameterMap& app_parameters,
    CdmKeyMessage* key_request,
    std::string* server_url) {

  if (reinitialize_session_) {
    CdmResponseType sts = ReInit();
    if (sts != NO_ERROR) {
      return sts;
    }
    reinitialize_session_ = false;
  }

  if (!crypto_session_) {
    LOGW("CdmSession::GenerateKeyRequest: Invalid crypto session");
    return UNKNOWN_ERROR;
  }

  if (!crypto_session_->IsOpen()) {
    LOGW("CdmSession::GenerateKeyRequest: Crypto session not open");
    return UNKNOWN_ERROR;
  }

  license_type_ = license_type;

  if (license_type_ == kLicenseTypeRelease) {
    return GenerateReleaseRequest(key_request, server_url);
  }
  else if (license_received_) {  // renewal
    return Properties::require_explicit_renew_request() ?
        UNKNOWN_ERROR : GenerateRenewalRequest(key_request, server_url);
  }
  else {
    CdmInitData pssh_data;
    if (!CdmEngine::ExtractWidevinePssh(init_data, &pssh_data)) {
      return KEY_ERROR;
    }

    if (Properties::use_certificates_as_identification()) {
      if (!crypto_session_->LoadCertificatePrivateKey(wrapped_key_)) {
        reinitialize_session_ = true;
        return NEED_PROVISIONING;
      }
    }

    if (!license_parser_.PrepareKeyRequest(pssh_data,
                                           license_type,
                                           app_parameters,
                                           key_request,
                                           server_url)) {
      return KEY_ERROR;
    }

    if (license_type_ == kLicenseTypeOffline) {
      offline_pssh_data_ = pssh_data;
      offline_key_request_ = *key_request;
      offline_release_server_url_ = *server_url;
    }

    return KEY_MESSAGE;
  }
}

// AddKey() - Accept license response and extract key info.
CdmResponseType CdmSession::AddKey(
    const CdmKeyResponse& key_response,
    CdmKeySetId* key_set_id) {
  if (!crypto_session_) {
    LOGW("CdmSession::AddKey: Invalid crypto session");
    return UNKNOWN_ERROR;
  }

  if (!crypto_session_->IsOpen()) {
    LOGW("CdmSession::AddKey: Crypto session not open");
    return UNKNOWN_ERROR;
  }

  if (license_type_ == kLicenseTypeRelease) {
    return ReleaseKey(key_response);
  }
  else if (license_received_) {  // renewal
    return Properties::require_explicit_renew_request() ?
        UNKNOWN_ERROR : RenewKey(key_response);
  }
  else {
    CdmResponseType sts = license_parser_.HandleKeyResponse(key_response);

    if (sts != KEY_ADDED)
      return sts;

    license_received_ = true;

    if (license_type_ == kLicenseTypeOffline) {
      offline_key_response_ = key_response;
      key_set_id_ = GenerateKeySetId(offline_pssh_data_);
      if (!StoreLicense(true)) {
        LOGE("CdmSession::AddKey: Unable to store license");
        ReInit();
        key_set_id_.clear();
        return UNKNOWN_ERROR;
      }
    }

    *key_set_id = key_set_id_;
    return KEY_ADDED;
  }
}

CdmResponseType CdmSession::QueryKeyStatus(CdmQueryMap* key_info) {
  return policy_engine_.Query(key_info);
}

CdmResponseType CdmSession::QueryKeyControlInfo(CdmQueryMap* key_info) {
  if ((!crypto_session_) || (!crypto_session_->IsOpen()))
    return UNKNOWN_ERROR;

  std::stringstream ss;
  ss << crypto_session_->oec_session_id();
  (*key_info)[QUERY_KEY_OEMCRYPTO_SESSION_ID] = ss.str();
  return NO_ERROR;
}

// CancelKeyRequest() - Cancel session.
CdmResponseType CdmSession::CancelKeyRequest() {
  // TODO(gmorgan): cancel and clean up session
  crypto_session_->Close();
  return NO_ERROR;
}

// Decrypt() - Accept encrypted buffer and return decrypted data.
CdmResponseType CdmSession::Decrypt(bool is_encrypted,
                                    bool is_secure,
                                    const KeyId& key_id,
                                    const uint8_t* encrypt_buffer,
                                    size_t encrypt_length,
                                    const std::vector<uint8_t>& iv,
                                    size_t block_offset,
                                    void* decrypt_buffer,
                                    size_t decrypt_buffer_offset,
                                    bool is_video) {
  if (!crypto_session_ || !crypto_session_->IsOpen())
    return UNKNOWN_ERROR;

  // Check if key needs to be selected
  if (is_encrypted) {
    if (key_id_.compare(key_id) != 0) {
      if (crypto_session_->SelectKey(key_id)) {
        key_id_ = key_id;
      }
      else {
        return NEED_KEY;
      }
    }
  }

  return crypto_session_->Decrypt(is_encrypted, is_secure, encrypt_buffer,
                                  encrypt_length, iv, block_offset,
                                  decrypt_buffer, decrypt_buffer_offset,
                                  is_video);
}

// License renewal
// GenerateRenewalRequest() - Construct valid renewal request for the current
// session keys.
CdmResponseType CdmSession::GenerateRenewalRequest(CdmKeyMessage* key_request,
                                                   std::string* server_url) {
  if (!license_parser_.PrepareKeyUpdateRequest(true, key_request, server_url))
    return KEY_ERROR;

  if (license_type_ == kLicenseTypeOffline) {
    offline_key_renewal_request_ = *key_request;
  }
  return KEY_MESSAGE;
}

// RenewKey() - Accept renewal response and update key info.
CdmResponseType CdmSession::RenewKey(const CdmKeyResponse& key_response) {
  CdmResponseType sts = license_parser_.HandleKeyUpdateResponse(true,
                                                                key_response);
  if (sts != KEY_ADDED)
    return sts;

  if (license_type_ == kLicenseTypeOffline) {
    offline_key_renewal_response_ = key_response;
    if (!StoreLicense(true))
      return UNKNOWN_ERROR;
  }
  return KEY_ADDED;
}

CdmResponseType CdmSession::GenerateReleaseRequest(CdmKeyMessage* key_request,
                                                   std::string* server_url) {
  if (license_parser_.PrepareKeyUpdateRequest(false, key_request,
                                              server_url)) {
    // Mark license as being released
    if (!StoreLicense(false))
      return UNKNOWN_ERROR;

    return KEY_MESSAGE;
  }
  return UNKNOWN_ERROR;
}

// ReleaseKey() - Accept release response and  release license.
CdmResponseType CdmSession::ReleaseKey(const CdmKeyResponse& key_response) {
  CdmResponseType sts = license_parser_.HandleKeyUpdateResponse(false,
                                                                key_response);
  DeviceFiles::DeleteLicense(key_set_id_);
  return sts;
}

bool CdmSession::IsKeyValid(const KeyId& key_id) {
  // TODO(gmorgan): lookup key and determine if valid.
  // return (session_keys_.find(key_id) != session_keys_.end());
  return true;
}

CdmSessionId CdmSession::GenerateSessionId() {
  static int session_num = 1;
  // TODO(rkuroiwa): Want this to be unique. Probably doing Hash(time+init_data)
  // to get something that is reasonably unique.
  return SESSION_ID_PREFIX + IntToString(++session_num);
}


CdmSessionId CdmSession::GenerateKeySetId(CdmInitData& pssh_data) {
  Clock clock;
  int64_t current_time = clock.GetCurrentTime();
  std::string key_set_id;

  while (key_set_id.empty()) {
    int random = rand();

    std::vector<uint8_t> hash(SHA256_DIGEST_LENGTH, 0);
    SHA256_CTX sha256;
    SHA256_Init(&sha256);
    SHA256_Update(&sha256, pssh_data.data(), pssh_data.size());
    SHA256_Update(&sha256, &current_time, sizeof(int64_t));
    SHA256_Update(&sha256, &random, sizeof(random));
    SHA256_Final(&hash[0], &sha256);
    for (int i = 0; i < SHA256_DIGEST_LENGTH; ++i) {
      hash[i%(SHA256_DIGEST_LENGTH/4)] ^= hash[i];
    }
    hash.resize(SHA256_DIGEST_LENGTH/4);
    key_set_id = KEY_SET_ID_PREFIX + b2a_hex(hash);

    if (DeviceFiles::LicenseExists(key_set_id)) {  // key set collision
      key_set_id.clear();
    }
  }
  return key_set_id;
}

bool CdmSession::LoadDeviceCertificate(std::string* certificate,
                                       std::string* wrapped_key) {
  return DeviceFiles::RetrieveCertificate(certificate,
                                          wrapped_key);
}

bool CdmSession::StoreLicense(bool active) {
  DeviceFiles::LicenseState state = DeviceFiles::kLicenseStateReleasing;
  if (active)
    state = DeviceFiles::kLicenseStateActive;

  return DeviceFiles::StoreLicense(key_set_id_, state, offline_pssh_data_,
                                   offline_key_request_, offline_key_response_,
                                   offline_key_renewal_request_,
                                   offline_key_renewal_response_,
                                   offline_release_server_url_);
}

bool CdmSession::AttachEventListener(WvCdmEventListener* listener) {
  std::pair<CdmEventListenerIter, bool> result = listeners_.insert(listener);
  return result.second;
}

bool CdmSession::DetachEventListener(WvCdmEventListener* listener) {
  return (listeners_.erase(listener) == 1);
}

void CdmSession::OnTimerEvent() {
  bool event_occurred = false;
  CdmEventType event;

  policy_engine_.OnTimerEvent(event_occurred, event);

  if (event_occurred) {
    for (CdmEventListenerIter iter = listeners_.begin();
         iter != listeners_.end(); ++iter) {
      CdmSessionId id = (*iter)->session_id();
      if (id.empty() || (id.compare(session_id_) == 0)) {
        (*iter)->onEvent(session_id_, event);
      }
    }
  }
}

void CdmSession::OnKeyReleaseEvent(CdmKeySetId key_set_id) {
  if (key_set_id_.compare(key_set_id) == 0) {
    for (CdmEventListenerIter iter = listeners_.begin();
      iter != listeners_.end(); ++iter) {
      CdmSessionId id = (*iter)->session_id();
      if (id.empty() || (id.compare(session_id_) == 0)) {
        (*iter)->onEvent(session_id_, LICENSE_EXPIRED_EVENT);
      }
    }
  }
}

}  // namespace wvcdm