// Copyright 2012 Google Inc. All Rights Reserved.

#include "cdm_session.h"

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

#include "cdm_engine.h"
#include "clock.h"
#include "crypto_session.h"
#include "device_files.h"
#include "file_store.h"
#include "log.h"
#include "properties.h"
#include "string_conversions.h"
#include "wv_cdm_constants.h"
#include "wv_cdm_event_listener.h"

namespace {
const size_t kKeySetIdLength = 14;
}  // namespace

namespace wvcdm {

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

CdmSession::CdmSession(const CdmClientPropertySet* cdm_client_property_set)
    : session_id_(GenerateSessionId()),
      crypto_session_(NULL),
      license_received_(false),
      is_offline_(false),
      is_release_(false),
      is_usage_update_needed_(false),
      is_initial_decryption_(true) {
  if (cdm_client_property_set) {
    Properties::AddSessionPropertySet(session_id_, cdm_client_property_set);
  }
}

CdmSession::~CdmSession() { Properties::RemoveSessionPropertySet(session_id_); }

CdmResponseType CdmSession::Init() {
  scoped_ptr<CryptoSession> session(new CryptoSession());

  CdmResponseType sts = session->Open(GetRequestedSecurityLevel());
  if (NO_ERROR != sts) return sts;

  std::string token;
  if (Properties::use_certificates_as_identification()) {
    DeviceFiles handle;
    std::string wrapped_key;
    if (!handle.Init(session.get()->GetSecurityLevel()) ||
        !handle.RetrieveCertificate(&token, &wrapped_key) ||
        !session->LoadCertificatePrivateKey(wrapped_key)) {
      return NEED_PROVISIONING;
    }
  } else {
    if (!session->GetToken(&token)) return UNKNOWN_ERROR;
  }

  if (!license_parser_.Init(token, session.get(), &policy_engine_))
    return UNKNOWN_ERROR;

  crypto_session_.reset(session.release());
  license_received_ = false;
  is_initial_decryption_ = true;
  return NO_ERROR;
}

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 handle;
  if (!handle.Init(crypto_session_->GetSecurityLevel()))
    return UNKNOWN_ERROR;

  DeviceFiles::LicenseState license_state;

  if (!handle.RetrieveLicense(key_set_id, &license_state, &offline_init_data_,
                              &key_request_, &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 (!license_parser_.RestoreOfflineLicense(key_request_, key_response_,
                                             offline_key_renewal_response_)) {
    return UNKNOWN_ERROR;
  }

  license_received_ = true;
  is_offline_ = true;
  is_release_ = license_type == kLicenseTypeRelease;
  return KEY_ADDED;
}

CdmResponseType CdmSession::RestoreUsageSession(
    const CdmKeyMessage& key_request,
    const CdmKeyResponse& key_response) {

  key_request_ = key_request;
  key_response_ = key_response;
  if (!license_parser_.RestoreUsageLicense(key_request_, key_response_)) {
    return UNKNOWN_ERROR;
  }

  license_received_ = true;
  is_offline_ = false;
  is_release_ = true;
  return KEY_ADDED;
}

CdmResponseType CdmSession::GenerateKeyRequest(
    const InitializationData& init_data, const CdmLicenseType license_type,
    const CdmAppParameterMap& app_parameters, CdmKeyMessage* key_request,
    std::string* server_url) {
  if (crypto_session_.get() == NULL) {
    LOGW("CdmSession::GenerateKeyRequest: Invalid crypto session");
    return UNKNOWN_ERROR;
  }

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

  switch (license_type) {
    case kLicenseTypeStreaming: is_offline_ = false; break;
    case kLicenseTypeOffline: is_offline_ = true; break;
    case kLicenseTypeRelease: is_release_ = true; break;
    default:
      LOGE("CdmSession::GenerateKeyRequest: unrecognized license type: %ld",
           license_type);
      return UNKNOWN_ERROR;
  }

  if (is_release_) {
    return GenerateReleaseRequest(key_request, server_url);
  } else if (license_received_) {  // renewal
    return GenerateRenewalRequest(key_request, server_url);
  } else {
    if (!init_data.is_supported()) {
      LOGW("CdmSession::GenerateKeyRequest: unsupported init data type (%s)",
           init_data.type().c_str());
      return KEY_ERROR;
    }
    if (init_data.IsEmpty() && !license_parser_.HasInitData()) {
      LOGW("CdmSession::GenerateKeyRequest: init data absent");
      return KEY_ERROR;
    }

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

    key_request_ = *key_request;
    if (is_offline_) {
      offline_init_data_ = init_data.data();
      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_.get() == NULL) {
    LOGW("CdmSession::AddKey: Invalid crypto session");
    return UNKNOWN_ERROR;
  }

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

  if (is_release_) {
    CdmResponseType sts = ReleaseKey(key_response);
    return (NO_ERROR == sts) ? KEY_ADDED : sts;
  } else if (license_received_) {  // renewal
    return RenewKey(key_response);
  } else {
    CdmResponseType sts = license_parser_.HandleKeyResponse(key_response);

    if (sts != KEY_ADDED) return sts;

    license_received_ = true;
    key_response_ = key_response;

    if (is_offline_ || !license_parser_.provider_session_token().empty()) {
      sts = StoreLicense();
      if (sts != NO_ERROR) return sts;
    }

    *key_set_id = key_set_id_;
    return KEY_ADDED;
  }
}

CdmResponseType CdmSession::QueryStatus(CdmQueryMap* key_info) {
  if (crypto_session_.get() == NULL) {
    LOGE("CdmSession::QueryStatus: Invalid crypto session");
    return UNKNOWN_ERROR;
  }

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

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

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

CdmResponseType CdmSession::QueryKeyControlInfo(CdmQueryMap* key_info) {
  if (crypto_session_.get() == NULL) {
    LOGW("CdmSession::QueryKeyControlInfo: Invalid crypto session");
    return UNKNOWN_ERROR;
  }

  if (!crypto_session_->IsOpen()) {
    LOGW("CdmSession::QueryKeyControlInfo: Crypto session not open");
    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() {
  crypto_session_->Close();
  return NO_ERROR;
}

// Decrypt() - Accept encrypted buffer and return decrypted data.
CdmResponseType CdmSession::Decrypt(const CdmDecryptionParameters& params) {
  if (crypto_session_.get() == NULL || !crypto_session_->IsOpen())
    return UNKNOWN_ERROR;

  CdmResponseType status = crypto_session_->Decrypt(params);

  switch (status) {
    case NO_ERROR:
      if (is_initial_decryption_) {
        policy_engine_.BeginDecryption();
        is_initial_decryption_ = false;
      }
      if (!is_usage_update_needed_) {
        is_usage_update_needed_ =
          !license_parser_.provider_session_token().empty();
      }
      break;
    case UNKNOWN_ERROR:
      Clock clock;
      int64_t current_time = clock.GetCurrentTime();
      if (policy_engine_.IsLicenseDurationExpired(current_time) ||
          policy_engine_.IsPlaybackDurationExpired(current_time)) {
        return NEED_KEY;
      }
      break;
  }

  return status;
}

// 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)) {
    LOGE("CdmSession::GenerateRenewalRequest: ERROR on prepare");
    return KEY_ERROR;
  }

  if (is_offline_) {
    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 (is_offline_) {
    offline_key_renewal_response_ = key_response;
    if (!StoreLicense(DeviceFiles::kLicenseStateActive)) return UNKNOWN_ERROR;
  }
  return KEY_ADDED;
}

CdmResponseType CdmSession::GenerateReleaseRequest(CdmKeyMessage* key_request,
                                                   std::string* server_url) {
  is_release_ = true;
  if (!license_parser_.PrepareKeyUpdateRequest(false, key_request, server_url))
    return UNKNOWN_ERROR;

  if (is_offline_) {  // Mark license as being released
    if (!StoreLicense(DeviceFiles::kLicenseStateReleasing))
      return UNKNOWN_ERROR;
  }
  return KEY_MESSAGE;
}

// ReleaseKey() - Accept release response and  release license.
CdmResponseType CdmSession::ReleaseKey(const CdmKeyResponse& key_response) {
  CdmResponseType sts = license_parser_.HandleKeyUpdateResponse(false,
                                                                key_response);
  if (KEY_ADDED != sts)
    return sts;

  if (is_offline_ || !license_parser_.provider_session_token().empty()) {
    DeleteLicense();
  }
  return NO_ERROR;
}

bool CdmSession::IsKeyLoaded(const KeyId& key_id) {
  return license_parser_.IsKeyLoaded(key_id);
}

CdmSessionId CdmSession::GenerateSessionId() {
  static int session_num = 1;
  return SESSION_ID_PREFIX + IntToString(++session_num);
}

bool CdmSession::GenerateKeySetId(CdmKeySetId* key_set_id) {
  if (!key_set_id) {
    LOGW("CdmSession::GenerateKeySetId: key set id destination not provided");
    return false;
  }

  std::vector<uint8_t> random_data(
      (kKeySetIdLength - sizeof(KEY_SET_ID_PREFIX)) / 2, 0);

  DeviceFiles handle;
  if (!handle.Init(crypto_session_->GetSecurityLevel()))
    return false;

  while (key_set_id->empty()) {
    if (!crypto_session_->GetRandom(random_data.size(), &random_data[0]))
      return false;

    *key_set_id = KEY_SET_ID_PREFIX + b2a_hex(random_data);

    // key set collision
    if (handle.LicenseExists(*key_set_id)) {
      key_set_id->clear();
    }
  }
  return true;
}

CdmResponseType CdmSession::StoreLicense() {
  if (is_offline_) {
    if (!GenerateKeySetId(&key_set_id_)) {
      LOGE("CdmSession::StoreLicense: Unable to generate key set Id");
      return UNKNOWN_ERROR;
    }

    if (!StoreLicense(DeviceFiles::kLicenseStateActive)) {
      LOGE("CdmSession::StoreLicense: Unable to store license");
      CdmResponseType sts = Init();
      if (sts != NO_ERROR) {
        LOGW("CdmSession::StoreLicense: Reinitialization failed");
        return sts;
      }

      key_set_id_.clear();
      return UNKNOWN_ERROR;
    }
    return NO_ERROR;
  }

  std::string provider_session_token = license_parser_.provider_session_token();
  if (provider_session_token.empty()) {
    LOGE("CdmSession::StoreLicense: No provider session token and not offline");
    return UNKNOWN_ERROR;
  }

  DeviceFiles handle;
  if (!handle.Init(crypto_session_->GetSecurityLevel())) {
    LOGE("CdmSession::StoreLicense: Unable to initialize device files");
    return UNKNOWN_ERROR;
  }

  if (!handle.StoreUsageInfo(provider_session_token, key_request_,
                             key_response_)) {
    LOGE("CdmSession::StoreLicense: Unable to store usage info");
    return UNKNOWN_ERROR;
  }
  return NO_ERROR;
}

bool CdmSession::StoreLicense(DeviceFiles::LicenseState state) {
  DeviceFiles handle;
  if (!handle.Init(crypto_session_->GetSecurityLevel()))
    return false;

  return handle.StoreLicense(
      key_set_id_, state, offline_init_data_, key_request_,
      key_response_, offline_key_renewal_request_,
      offline_key_renewal_response_, offline_release_server_url_);
}

bool CdmSession::DeleteLicense() {
  if (!is_offline_ && license_parser_.provider_session_token().empty())
    return false;

  DeviceFiles handle;
  if (!handle.Init(crypto_session_->GetSecurityLevel())) {
    LOGE("CdmSession::DeleteLicense: Unable to initialize device files");
    return false;
  }

  if (is_offline_)
    return handle.DeleteLicense(key_set_id_);
  else
    return handle.DeleteUsageInfo(
               license_parser_.provider_session_token());
}

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) {
      (*iter)->OnEvent(session_id_, event);
    }
  }
}

void CdmSession::OnKeyReleaseEvent(const CdmKeySetId& key_set_id) {
  if (key_set_id_ == key_set_id) {
    for (CdmEventListenerIter iter = listeners_.begin();
         iter != listeners_.end(); ++iter) {
      (*iter)->OnEvent(session_id_, LICENSE_EXPIRED_EVENT);
    }
  }
}

SecurityLevel CdmSession::GetRequestedSecurityLevel() {
  std::string security_level;
  if (Properties::GetSecurityLevel(session_id_, &security_level) &&
      security_level == QUERY_VALUE_SECURITY_LEVEL_L3) {
    return kLevel3;
  }

  return kLevelDefault;
}

CdmSecurityLevel CdmSession::GetSecurityLevel() {
  if (NULL == crypto_session_.get())
    return kSecurityLevelUninitialized;

  return crypto_session_.get()->GetSecurityLevel();
}

CdmResponseType CdmSession::UpdateUsageInformation() {
  return crypto_session_->UpdateUsageInformation();
}

}  // namespace wvcdm