////////////////////////////////////////////////////////////////////////////////
// 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 <memory>
#include <string>

#include "gflags/gflags.h"
#include "glog/logging.h"
#include "common/aes_cbc_util.h"
#include "common/file_util.h"
#include "common/random_util.h"
#include "common/rsa_key.h"
#include "protos/public/certificate_provisioning.pb.h"
#include "protos/public/client_identification.pb.h"

DEFINE_string(service_public_key_path, "example_data/service.public",
              "Indicates the file path to service public key. If omitted, the "
              "token is not encrypted.");
DEFINE_string(
    certificate_path, "",
    "Indicates the file path to the certificate chain. Should not be empty.");
DEFINE_string(certificate_private_key_path, "",
              "Indicaets the file path to the certificate private key. Should "
              "not be empty.");
DEFINE_string(
    output_path, "",
    "Specifies where to write the output message. Should not be empty.");

namespace widevine {

bool GenerateProvisioningMessage(const std::string& service_public_key,
                                 const std::string& certificate,
                                 const std::string& private_key, std::string* message) {
  ClientIdentification client_id;
  client_id.set_token(certificate);
  client_id.set_type(ClientIdentification::OEM_DEVICE_CERTIFICATE);

  ProvisioningRequest request;
  if (service_public_key.empty()) {
    *request.mutable_client_id() = client_id;
  } else {
    auto encrypted_id = request.mutable_encrypted_client_id();
    CHECK(RandomBytes(16, encrypted_id->mutable_encrypted_client_id_iv()));
    std::string privacy_key;
    CHECK(RandomBytes(16, &privacy_key));
    encrypted_id->set_encrypted_client_id(crypto_util::EncryptAesCbc(
        privacy_key, encrypted_id->encrypted_client_id_iv(),
        client_id.SerializeAsString()));
    std::unique_ptr<RsaPublicKey> service_key(
        RsaPublicKey::Create(service_public_key));
    if (!service_key) return false;
    CHECK(service_key->Encrypt(privacy_key,
                               encrypted_id->mutable_encrypted_privacy_key()));
  }

  CHECK(RandomBytes(4, request.mutable_nonce()));

  SignedProvisioningMessage signed_message;
  signed_message.set_message(request.SerializeAsString());
  std::unique_ptr<RsaPrivateKey> signer(RsaPrivateKey::Create(private_key));
  CHECK(signer->GenerateSignature(signed_message.message(),
                                  signed_message.mutable_signature()));
  *message = signed_message.SerializeAsString();
  return true;
}

}  // namespace widevine

using widevine::GetContents;
using widevine::SetContents;
using widevine::GenerateProvisioningMessage;

int main(int argc, char** argv) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  std::string service_public_key;
  if (!FLAGS_service_public_key_path.empty()) {
    if (!GetContents(FLAGS_service_public_key_path, &service_public_key))
      return 1;
  }
  std::string certificate;
  if (!GetContents(FLAGS_certificate_path, &certificate)) return 1;
  std::string private_key;
  if (!GetContents(FLAGS_certificate_private_key_path, &private_key)) return 1;
  std::string message;
  if (!GenerateProvisioningMessage(service_public_key, certificate, private_key,
                                   &message))
    return 1;
  if (!SetContents(FLAGS_output_path, message)) return 1;
  return 0;
}