////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 Google Inc.
//
// This software is licensed under the terms defined in the Widevine Master
// License Agreement. For a copy of this agreement, please contact
// widevine-licensing@google.com.
////////////////////////////////////////////////////////////////////////////////

#include "provisioning_sdk/internal/provisioning_session_impl.h"

#include <stddef.h>

#include "gflags/gflags.h"
#include "glog/logging.h"
#include "common/aes_cbc_util.h"
#include "common/random_util.h"
#include "common/rsa_key.h"
#include "common/sha_util.h"
#include "provisioning_sdk/public/provisioning_status.h"

DEFINE_int32(provisioning_log_every_n, 1,
             "parameter for LOG_EVERY_N to help abate log spamming.");

#define LOG_EVERY_N_WITH_PROTO(message, proto)         \
  LOG_EVERY_N(WARNING, FLAGS_provisioning_log_every_n) \
      << (message) << " [proto: " << (proto).ShortDebugString() << "]"

namespace widevine {

ProvisioningSessionImpl::ProvisioningSessionImpl(
    const ProvisioningEngineImpl& engine, const OemDeviceCert& oem_device_cert,
    const RsaPrivateKey& service_private_key)
    : engine_(engine),
      oem_device_cert_(oem_device_cert),
      service_private_key_(service_private_key),
      rsa_key_factory_(new RsaKeyFactory) {}

ProvisioningSessionImpl::~ProvisioningSessionImpl() {}

ProvisioningStatus ProvisioningSessionImpl::Initialize(
    const std::string& device_public_key, const std::string& device_private_key) {
  auto rsa_public_key =
      rsa_key_factory_->CreateFromPkcs1PublicKey(device_public_key);
  if (!rsa_public_key) return INVALID_DEVICE_PUBLIC_KEY;
  auto rsa_private_key =
      rsa_key_factory_->CreateFromPkcs1PrivateKey(device_private_key);
  if (!rsa_private_key) return INVALID_DEVICE_PRIVATE_KEY;
  if (!rsa_public_key->MatchesPrivateKey(*rsa_private_key)) {
    LOG(WARNING) << "Device public key and private key do not match.";
    return INVALID_DEVICE_PRIVATE_KEY;
  }

  device_public_key_ = device_public_key;
  device_private_key_ = device_private_key;
  return OK;
}

ProvisioningStatus ProvisioningSessionImpl::ProcessMessage(
    const std::string& message, std::string* response) {
  SignedProvisioningMessage signed_request;
  ProvisioningRequest request;
  if (!ValidateAndDeserializeRequest(message, &signed_request, &request))
    return INVALID_REQUEST_MESSAGE;

  ClientIdentification client_id;
  if (request.has_encrypted_client_id()) {
    if (!DecryptClientIdentification(request.encrypted_client_id(), &client_id))
      return INVALID_REQUEST_MESSAGE;
  } else {
    DCHECK(request.has_client_id());
    client_id.Swap(request.mutable_client_id());
  }

  if (client_id.type() != ClientIdentification::OEM_DEVICE_CERTIFICATE) {
    LOG_EVERY_N_WITH_PROTO("Invalid client_id type", client_id);
    return INVALID_REQUEST_MESSAGE;
  }
  if (client_id.token().empty()) {
    LOG_EVERY_N_WITH_PROTO("Missing client_id.token", client_id);
    return INVALID_REQUEST_MESSAGE;
  }

  std::unique_ptr<RsaPublicKey> cert_public_key;
  uint32_t system_id;
  std::string oem_ca_serial_number;
  if (!oem_device_cert_.VerifyCertificateChain(client_id.token(),
                                               &cert_public_key, &system_id,
                                               &oem_ca_serial_number)) {
    LOG_EVERY_N_WITH_PROTO("Invalid token", client_id);
    return INVALID_REQUEST_MESSAGE;
  }
  if (!cert_public_key->VerifySignature(signed_request.message(),
                                        signed_request.signature())) {
    LOG_EVERY_N_WITH_PROTO("Signature verification failed", client_id);
    return INVALID_REQUEST_MESSAGE;
  }

  // Save device_info for query later.
  device_info_ = engine_.GetDeviceInfo(system_id);

  std::string certificate_serial_number;
  if (request.has_spoid()) {
    certificate_serial_number = request.spoid();
  } else {
    // Generate stable serial number.
    const std::string stable_data(client_id.token() + request.stable_id() +
                             request.provider_id() +
                             engine_.secret_spoid_sauce());
    const std::string hash = Sha256_Hash(stable_data);
    const size_t kCertificateSerialNumberSize = 16;
    certificate_serial_number = hash.substr(0, kCertificateSerialNumberSize);
  }

  ProvisioningResponse provisioning_response;
  ProvisioningStatus status = GenerateProvisioningResponse(
      system_id, oem_ca_serial_number, request.provider_id(),
      certificate_serial_number, *cert_public_key, &provisioning_response);
  if (status != OK) return status;
  provisioning_response.set_nonce(request.nonce());

  // Sign the response.
  SignedProvisioningMessage signed_message;
  if (!provisioning_response.SerializeToString(
          signed_message.mutable_message())) {
    LOG(WARNING) << "Error serializing ProvisioningResponse.";
    return INTERNAL_ERROR;
  }
  if (!service_private_key_.GenerateSignature(
          signed_message.message(), signed_message.mutable_signature())) {
    LOG(WARNING) << "Failed to sign ProvisioningResponse.";
    return INTERNAL_ERROR;
  }
  if (!signed_message.SerializeToString(response)) {
    LOG(WARNING) << "Error serializing SignedProvisioningMessage.";
    return INTERNAL_ERROR;
  }
  return OK;
}

bool ProvisioningSessionImpl::ValidateAndDeserializeRequest(
    const std::string& message, SignedProvisioningMessage* signed_request,
    ProvisioningRequest* request) const {
  if (!signed_request->ParseFromString(message)) {
    LOG_EVERY_N(WARNING, FLAGS_provisioning_log_every_n)
        << "Failed to parse SignedProvisioningMessage.";
    return false;
  }
  VLOG(1) << "signed_request: " << signed_request->ShortDebugString();

  if (signed_request->message().empty()) {
    LOG_EVERY_N_WITH_PROTO("Missing message", *signed_request);
    return false;
  }
  if (signed_request->signature().empty()) {
    LOG_EVERY_N_WITH_PROTO("Missing signature", *signed_request);
    return false;
  }
  if (!request->ParseFromString(signed_request->message())) {
    LOG_EVERY_N_WITH_PROTO("Failed to parse ProvisioningRequest",
                           *signed_request);
    return false;
  }

  if (request->has_encrypted_client_id()) {
    const EncryptedClientIdentification& encrypted_client_id =
        request->encrypted_client_id();
    if (encrypted_client_id.encrypted_client_id().empty()) {
      LOG_EVERY_N_WITH_PROTO("Missing encrypted_client_id",
                             encrypted_client_id);
      return false;
    }
    if (encrypted_client_id.encrypted_client_id_iv().empty()) {
      LOG_EVERY_N_WITH_PROTO("Missing encrypted_client_id_iv",
                             encrypted_client_id);
      return false;
    }
    if (encrypted_client_id.encrypted_privacy_key().empty()) {
      LOG_EVERY_N_WITH_PROTO("Missing encrypted_privacy_key",
                             encrypted_client_id);
      return false;
    }
  } else if (!request->has_client_id()) {
    LOG_EVERY_N_WITH_PROTO("Missing clear_or_encrypted_client_id", *request);
    return false;
  }

  const size_t kMinimumRequiredNonceLength = 4;
  if (request->nonce().size() < kMinimumRequiredNonceLength) {
    LOG_EVERY_N_WITH_PROTO("Missing or invalid nonce", *request);
    return false;
  }
  return true;
}

bool ProvisioningSessionImpl::DecryptClientIdentification(
    const EncryptedClientIdentification& encrypted_client_id,
    ClientIdentification* client_id) {
  std::string privacy_key;
  if (!service_private_key_.Decrypt(encrypted_client_id.encrypted_privacy_key(),
                                    &privacy_key)) {
    LOG_EVERY_N_WITH_PROTO("Failed to decrypt encrypted_privacy_key",
                           encrypted_client_id);
    return false;
  }
  std::string serialized_client_id(crypto_util::DecryptAesCbc(
      privacy_key, encrypted_client_id.encrypted_client_id_iv(),
      encrypted_client_id.encrypted_client_id()));
  if (serialized_client_id.empty()) {
    LOG_EVERY_N_WITH_PROTO("Failed to decrypt client_id", encrypted_client_id);
    return false;
  }
  if (!client_id->ParseFromString(serialized_client_id)) {
    LOG_EVERY_N_WITH_PROTO("Failed to parse client_id", encrypted_client_id);
    return false;
  }
  return true;
}

ProvisioningStatus ProvisioningSessionImpl::GenerateProvisioningResponse(
    uint32_t system_id, const std::string& oem_ca_serial_number,
    const std::string& provider_id, const std::string& certificate_serial_number,
    const RsaPublicKey& cert_public_key, ProvisioningResponse* response) {
  ProvisioningStatus status = engine_.GenerateProviderDeviceDrmCertificate(
      system_id, oem_ca_serial_number, provider_id, device_public_key_,
      certificate_serial_number, response->mutable_device_certificate());
  if (status != OK) return status;

  const size_t kAesKeySize = 16;
  const size_t kIvSize = 16;

  // Encrypt private key.
  std::string message_key;
  if (!RandomBytes(kAesKeySize, &message_key)) {
    LOG(WARNING) << "Failed to generate message_key.";
    return INTERNAL_ERROR;
  }
  std::string iv;
  if (!RandomBytes(kIvSize, &iv)) {
    LOG(WARNING) << "Failed to generate iv.";
    return INTERNAL_ERROR;
  }
  response->set_device_rsa_key_iv(iv);
  response->set_device_rsa_key(
      crypto_util::EncryptAesCbc(message_key, iv, device_private_key_));
  if (response->device_rsa_key().empty()) {
    LOG(WARNING) << "Failed to encrypt device_rsa_key";
    return INTERNAL_ERROR;
  }
  if (!cert_public_key.Encrypt(message_key, response->mutable_wrapping_key())) {
    LOG(WARNING) << "Failed to encrypt wrapping_key";
    return INTERNAL_ERROR;
  }
  return OK;
}

}  // namespace widevine