// Copyright 2012 Google Inc. All Rights Reserved.

#include "license.h"

#include <vector>

#include "crypto_engine.h"
#include "crypto_session.h"
#include "log.h"
#include "policy_engine.h"
#include "properties.h"
#include "string_conversions.h"
#include "wv_cdm_constants.h"

namespace {
  std::string kCompanyNameKey = "company_name";
  std::string kCompanyNameValueGoogle = "Google";
  std::string kModelNameKey = "model_name";
  std::string kArchitectureNameKey = "architecture_name";
  std::string kDeviceNameKey = "device_name";
  std::string kProductNameKey = "product_name";
  std::string kBuildInfoKey = "build_info";
  std::string kDeviceIdKey = "device_id";
}

namespace wvcdm {

// Protobuf generated classes.
using video_widevine_server::sdk::ClientIdentification;
using video_widevine_server::sdk::ClientIdentification_NameValue;
using video_widevine_server::sdk::LicenseRequest;
using video_widevine_server::sdk::LicenseRequest_ContentIdentification;
using video_widevine_server::sdk::LicenseRequest_ContentIdentification_CENC;
using video_widevine_server::sdk::LicenseRequest_ContentIdentification_ExistingLicense;
using video_widevine_server::sdk::License;
using video_widevine_server::sdk::License_KeyContainer;
using video_widevine_server::sdk::LicenseError;
using video_widevine_server::sdk::SignedMessage;


static std::vector<CryptoKey> ExtractContentKeys(const License& license) {
  std::vector<CryptoKey> key_array;

  // Extract content key(s)
  for (int i = 0; i < license.key_size(); ++i) {
    CryptoKey key;
    size_t length;
    switch (license.key(i).type()) {
      case License_KeyContainer::CONTENT:
      case License_KeyContainer::OPERATOR_SESSION:
        key.set_key_id(license.key(i).id());
        // Strip off PKCS#5 padding - since we know the key is 16 or 32 bytes,
        // the padding will always be 16 bytes.
        length = license.key(i).key().size() - 16;
        key.set_key_data(license.key(i).key().substr(0, length));
        key.set_key_data_iv(license.key(i).iv());
        if (license.key(i).has_key_control()) {
          key.set_key_control(license.key(i).key_control().key_control_block());
          key.set_key_control_iv(license.key(i).key_control().iv());
        }
        key_array.push_back(key);
        break;
      case License_KeyContainer::KEY_CONTROL:
        if (license.key(i).has_key_control()) {
          key.set_key_control(license.key(i).key_control().key_control_block());
          key.set_key_control_iv(license.key(i).key_control().iv());
          key_array.push_back(key);
        }
        break;
      default:
        // Ignore SIGNING key types as they are not content related
        break;
    }
  }

  return key_array;
}

CdmLicense::CdmLicense(): session_(NULL) {}

CdmLicense::~CdmLicense() {}

bool CdmLicense::Init(const std::string& token,
                      CryptoSession* session,
                      PolicyEngine* policy_engine) {
  if (token.size() == 0)
    return false;
  if (session == NULL || !session->IsValid() || !session->IsOpen())
    return false;
  token_ = token;
  session_ = session;
  policy_engine_ = policy_engine;
  return true;
}

bool CdmLicense::PrepareKeyRequest(const CdmInitData& init_data,
                                   const CdmLicenseType license_type,
                                   const CdmAppParameterMap& app_parameters,
                                   CdmKeyMessage* signed_request,
                                   std::string* server_url) {
  if (!session_ ||
      token_.empty()) {
    return false;
  }
  if (init_data.empty()) {
    LOGE("CdmLicense::PrepareKeyRequest : No init data provided;");
    return false;
  }
  if (!signed_request) {
    LOGE("CdmLicense::PrepareKeyRequest : No signed request provided.");
    return false;
  }
  if (!server_url) {
    LOGE("CdmLicense::PrepareKeyRequest : No server url provided.");
    return false;
  }

  // TODO(gmorgan): Request ID owned by session?
  std::string request_id;
  session_->GenerateRequestId(request_id);

  LicenseRequest license_request;
  ClientIdentification* client_id = license_request.mutable_client_id();

  if (Properties::use_certificates_as_identification())
    client_id->set_type(ClientIdentification::DEVICE_CERTIFICATE);
  else
    client_id->set_type(ClientIdentification::KEYBOX);
  client_id->set_token(token_);

  ClientIdentification_NameValue* client_info;
  CdmAppParameterMap::const_iterator iter;
  for (iter = app_parameters.begin(); iter != app_parameters.end(); iter++) {
    client_info = client_id->add_client_info();
    client_info->set_name(iter->first);
    client_info->set_value(iter->second);
  }
  std::string value;
  client_info = client_id->add_client_info();
  client_info->set_name(kCompanyNameKey);
  client_info->set_value(kCompanyNameValueGoogle);
  if (Properties::GetModelName(value)) {
    client_info = client_id->add_client_info();
    client_info->set_name(kModelNameKey);
    client_info->set_value(value);
  }
  if (Properties::GetArchitectureName(value)) {
    client_info = client_id->add_client_info();
    client_info->set_name(kArchitectureNameKey);
    client_info->set_value(value);
  }
  if (Properties::GetDeviceName(value)) {
    client_info = client_id->add_client_info();
    client_info->set_name(kDeviceNameKey);
    client_info->set_value(value);
  }
  if (Properties::GetProductName(value)) {
    client_info = client_id->add_client_info();
    client_info->set_name(kProductNameKey);
    client_info->set_value(value);
  }
  if (Properties::GetBuildInfo(value)) {
    client_info = client_id->add_client_info();
    client_info->set_name(kBuildInfoKey);
    client_info->set_value(value);
  }

  if (CryptoEngine::GetInstance()->GetDeviceUniqueId(&value)) {
    client_info = client_id->add_client_info();
    client_info->set_name(kDeviceIdKey);
    client_info->set_value(value);
  }

  // Content Identification may be a cenc_id, a webm_id or a license_id
  LicenseRequest_ContentIdentification* content_id =
      license_request.mutable_content_id();

  LicenseRequest_ContentIdentification_CENC* cenc_content_id =
      content_id->mutable_cenc_id();
  cenc_content_id->add_pssh(init_data);

  switch (license_type) {
    case kLicenseTypeOffline:
      cenc_content_id->set_license_type(video_widevine_server::sdk::OFFLINE);
      break;
    case kLicenseTypeStreaming:
      cenc_content_id->set_license_type(video_widevine_server::sdk::STREAMING);
      break;
    default:
      LOGD("CdmLicense::PrepareKeyRequest: Unknown license type = %u",
          (int)license_type);
      return false;
      break;
  }
  cenc_content_id->set_request_id(request_id);

  // TODO(jfore): The time field will be updated once the cdm wrapper
  // has been updated to pass us in the time.
  license_request.set_request_time(0);

  license_request.set_type(LicenseRequest::NEW);

  // Get/set the nonce.  This value will be reflected in the Key Control Block
  // of the license response.
  uint32_t nonce;
  if (!session_->GenerateNonce(&nonce)) {
    return false;
  }
  license_request.set_key_control_nonce(UintToString(nonce));
  LOGD("PrepareKeyRequest: nonce=%u", nonce);
  license_request.set_protocol_version(video_widevine_server::sdk::VERSION_2_1);

  // License request is complete. Serialize it.
  std::string serialized_license_req;
  license_request.SerializeToString(&serialized_license_req);

  if (Properties::use_certificates_as_identification())
    key_request_ = serialized_license_req;

  // Derive signing and encryption keys and construct signature.
  std::string license_request_signature;
  if (!session_->PrepareRequest(serialized_license_req,
                                &license_request_signature, false)) {
    signed_request->clear();
    return false;
  }

  if (license_request_signature.empty()) {
    LOGE("CdmLicense::PrepareKeyRequest: License request signature empty");
    signed_request->clear();
    return false;
  }

  // Put serialize license request and signature together
  SignedMessage signed_message;
  signed_message.set_type(SignedMessage::LICENSE_REQUEST);
  signed_message.set_signature(license_request_signature);
  signed_message.set_msg(serialized_license_req);

  signed_message.SerializeToString(signed_request);

  *server_url = server_url_;
  return true;
}

bool CdmLicense::PrepareKeyUpdateRequest(bool is_renewal,
                                         CdmKeyMessage* signed_request,
                                         std::string* server_url) {
  if (!session_) {
    LOGE("CdmLicense::PrepareKeyUpdateRequest: Invalid crypto session");
    return false;
  }
  if (!signed_request) {
    LOGE("CdmLicense::PrepareKeyUpdateRequest: No signed request provided");
    return false;
  }
  if (!server_url) {
    LOGE("CdmLicense::PrepareKeyUpdateRequest: No server url provided");
    return false;
  }

  LicenseRequest license_request;
  if (is_renewal)
    license_request.set_type(LicenseRequest::RENEWAL);
  else
    license_request.set_type(LicenseRequest::RELEASE);

  LicenseRequest_ContentIdentification_ExistingLicense* current_license =
      license_request.mutable_content_id()->mutable_license();
  current_license->mutable_license_id()->CopyFrom(license_id_);

  // Get/set the nonce.  This value will be reflected in the Key Control Block
  // of the license response.
  uint32_t nonce;
  if (!session_->GenerateNonce(&nonce)) {
    return false;
  }
  license_request.set_key_control_nonce(UintToString(nonce));
  LOGD("PrepareKeyUpdateRequest: nonce=%u", nonce);
  license_request.set_protocol_version(video_widevine_server::sdk::VERSION_2_1);

  // License request is complete. Serialize it.
  std::string serialized_license_req;
  license_request.SerializeToString(&serialized_license_req);

  // Construct signature.
  std::string license_request_signature;
  if (!session_->PrepareRenewalRequest(serialized_license_req,
                                       &license_request_signature))
    return false;

  if (license_request_signature.empty()) {
    LOGE("CdmLicense::PrepareKeyUpdateRequest: empty license request"
        " signature");
    return false;
  }

  // Put serialize license request and signature together
  SignedMessage signed_message;
  signed_message.set_type(SignedMessage::LICENSE_REQUEST);
  signed_message.set_signature(license_request_signature);
  signed_message.set_msg(serialized_license_req);

  signed_message.SerializeToString(signed_request);
  *server_url = server_url_;
  return true;
}

CdmResponseType CdmLicense::HandleKeyResponse(
    const CdmKeyResponse& license_response) {
  if (!session_) {
    return KEY_ERROR;
  }
  if (license_response.empty()) {
    LOGE("CdmLicense::HandleKeyResponse : Empty license response.");
    return KEY_ERROR;
  }

  SignedMessage signed_response;
  if (!signed_response.ParseFromString(license_response))
    return KEY_ERROR;

  if (signed_response.type() == SignedMessage::ERROR) {
    return HandleKeyErrorResponse(signed_response);
  }

  if (!signed_response.has_signature())
    return KEY_ERROR;

  License license;
  if (!license.ParseFromString(signed_response.msg()))
    return KEY_ERROR;

  if (Properties::use_certificates_as_identification()) {
    if (!signed_response.has_session_key())
      return KEY_ERROR;

    if (!session_->GenerateDerivedKeys(key_request_,
                                       signed_response.session_key()))
      return KEY_ERROR;
  }

  // Extract mac key
  std::string mac_key_iv;
  std::string mac_key;
  if (license.policy().can_renew())
  {
    for (int i = 0; i < license.key_size(); ++i) {
      if (license.key(i).type() == License_KeyContainer::SIGNING) {
        mac_key_iv.assign(license.key(i).iv());

        // Strip off PKCS#5 padding
        mac_key.assign(license.key(i).key().data(), MAC_KEY_SIZE);
      }
    }

    if (mac_key_iv.size() != KEY_IV_SIZE ||
        mac_key.size() != MAC_KEY_SIZE) {
      return KEY_ERROR;
    }

    // License Id should not be empty for renewable license
    if (!license.has_id()) return KEY_ERROR;

    license_id_.CopyFrom(license.id());
  }

  std::vector<CryptoKey> key_array = ExtractContentKeys(license);
  if (!key_array.size()) {
    LOGE("CdmLicense::HandleKeyResponse : No content keys.");
    return KEY_ERROR;
  }

  if (license.policy().has_renewal_server_url()) {
    server_url_ = license.policy().renewal_server_url();
  }

  policy_engine_->SetLicense(license);

  if (session_->LoadKeys(signed_response.msg(),
                         signed_response.signature(),
                         mac_key_iv,
                         mac_key,
                         key_array.size(),
                         &key_array[0])) {
    return KEY_ADDED;
  }
  else {
    return KEY_ERROR;
  }
}

CdmResponseType CdmLicense::HandleKeyUpdateResponse(
    bool is_renewal,
    const CdmKeyResponse& license_response) {
  if (!session_) {
    return KEY_ERROR;
  }
  if (license_response.empty()) {
    LOGE("CdmLicense::HandleKeyUpdateResponse : Empty license response.");
    return KEY_ERROR;
  }

  SignedMessage signed_response;
  if (!signed_response.ParseFromString(license_response)) {
    LOGE("CdmLicense::HandleKeyUpdateResponse: Unable to parse signed message");
    return KEY_ERROR;
  }

  if (signed_response.type() == SignedMessage::ERROR) {
    return HandleKeyErrorResponse(signed_response);
  }

  if (!signed_response.has_signature()) {
    LOGE("CdmLicense::HandleKeyUpdateResponse: signature missing");
    return KEY_ERROR;
  }

  License license;
  if (!license.ParseFromString(signed_response.msg())) {
    LOGE("CdmLicense::HandleKeyUpdateResponse: Unable to parse license"
        " from signed message");
    return KEY_ERROR;
  }

  if (!license.has_id()) {
    LOGE("CdmLicense::HandleKeyUpdateResponse: license id not present");
    return KEY_ERROR;
  }

  if (license.id().version() > license_id_.version()) {
    // This is the normal case.
    license_id_.CopyFrom(license.id());

    if (is_renewal) {
      if (license.policy().has_renewal_server_url() &&
          license.policy().renewal_server_url().size() > 0) {
        server_url_ = license.policy().renewal_server_url();
      }
    }

    policy_engine_->UpdateLicense(license);

    if (!is_renewal)
      return KEY_ADDED;

    std::vector<CryptoKey> key_array = ExtractContentKeys(license);

    if (session_->RefreshKeys(signed_response.msg(),
                              signed_response.signature(),
                              key_array.size(),
                              &key_array[0])) {
      return KEY_ADDED;
    }
    else {
      return KEY_ERROR;
    }
  }

  // This isn't supposed to happen.
  // TODO(jfore): Handle wrap? We can miss responses and that should be
  // considered normal until retries are exhausted.
  LOGE("CdmLicense::HandleKeyUpdateResponse: license version: expected > %u,"
      " actual = %u", license_id_.version(), license.id().version());
  return KEY_ERROR;
}

bool CdmLicense::RestoreOfflineLicense(
    CdmKeyMessage& license_request,
    CdmKeyResponse& license_response,
    CdmKeyResponse& license_renewal_response) {

  if (license_request.empty() || license_response.empty()) {
    LOGE("CdmLicense::RestoreOfflineLicense: key_request or response empty: "
         "%u %u", license_request.size(), license_response.size());
    return false;
  }

  SignedMessage signed_request;
  if (!signed_request.ParseFromString(license_request)) {
    LOGE("CdmLicense::RestoreOfflineLicense: license_request parse failed");
    return false;
  }

  if (signed_request.type() != SignedMessage::LICENSE_REQUEST) {
    LOGE("CdmLicense::RestoreOfflineLicense: license request type: expected = "
        "%d, actual = %d",
        SignedMessage::LICENSE_REQUEST,
        signed_request.type());
    return false;
  }

  if (Properties::use_certificates_as_identification()) {
    key_request_ = signed_request.msg();
  }
  else {
    if (!session_->GenerateDerivedKeys(signed_request.msg()))
      return false;
  }

  CdmResponseType sts = HandleKeyResponse(license_response);

  if (sts != KEY_ADDED)
    return false;

  if (!license_renewal_response.empty()) {
    sts = HandleKeyUpdateResponse(true, license_renewal_response);

    if (sts != KEY_ADDED)
      return false;
  }

  return true;
}

CdmResponseType CdmLicense::HandleKeyErrorResponse(
    const SignedMessage& signed_message) {

  LicenseError license_error;
  if (!license_error.ParseFromString(signed_message.msg())) {
    LOGE("CdmLicense::HandleKeyErrorResponse: Unable to parse license error");
    return KEY_ERROR;
  }

  switch (license_error.error_code()) {
    case LicenseError::INVALID_CREDENTIALS:
      return NEED_PROVISIONING;
    case LicenseError::REVOKED_CREDENTIALS:
      return DEVICE_REVOKED;
    case LicenseError::SERVICE_UNAVAILABLE:
    default:
      LOGW("CdmLicense::HandleKeyErrorResponse: Unknwon error type = %d",
          license_error.error_code());
      return KEY_ERROR;
  }
}

}  // namespace wvcdm