////////////////////////////////////////////////////////////////////////////////
// Copyright 2019 Google LLC.
//
// 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 "common/ecies_crypto.h"

#include <memory>

#include "glog/logging.h"
#include "absl/strings/escaping.h"
#include "openssl/ec.h"
#include "common/aes_cbc_util.h"
#include "common/crypto_util.h"
#include "common/ec_key.h"
#include "common/ec_util.h"
#include "common/openssl_util.h"

namespace {

const char kZeroIv[] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

const size_t kMacSizeBytes = 32;
const size_t kMasterKeySizeBytes = 32;
const size_t kEncryptionKeySizeBytes = kMasterKeySizeBytes;
const size_t kSigningKeySizeBytes = kMasterKeySizeBytes;
const size_t kEncryptionKeySizeBits = kEncryptionKeySizeBytes * 8;
const size_t kSigningKeySizeBits = kSigningKeySizeBytes * 8;

const char kWrappingKeyLabel[] = "ECIESEncryptionLabel";
const char kSigningKeyLabel[] = "ECIESSigningLabel";

bool DeriveWrappingAndSigningKeys(const std::string& key,
                                  const std::string& context,
                                  std::string* wrapping_key,
                                  std::string* signing_key) {
  if (wrapping_key == nullptr || signing_key == nullptr) return false;
  if (key.size() != kMasterKeySizeBytes) {
    LOG(ERROR) << "Invalid master key size";
    return false;
  }

  *wrapping_key = widevine::crypto_util::DeriveKey(
      key, kWrappingKeyLabel, context, kEncryptionKeySizeBits);
  if (wrapping_key->empty()) {
    LOG(WARNING) << "Failed to derive encryption key.";
    return false;
  }

  *signing_key = widevine::crypto_util::DeriveKey(
      key, kSigningKeyLabel, context, kSigningKeySizeBits);
  if (signing_key->empty()) {
    LOG(WARNING) << "Failed to derive sigining key.";
    return false;
  }
  return true;
}

}  // anonymous namespace

namespace widevine {

EciesEncryptor::EciesEncryptor(std::unique_ptr<ECPublicKey> public_key,
                               ECKeySource* key_source)
    : public_key_(std::move(public_key)), key_source_(key_source) {}

std::unique_ptr<EciesEncryptor> EciesEncryptor::Create(
    const std::string& serialized_public_key, ECKeySource* key_source) {
  if (key_source == nullptr) {
    return nullptr;
  }
  std::unique_ptr<ECPublicKey> ec_key =
      ECPublicKey::Create(serialized_public_key);
  if (ec_key == nullptr) {
    return nullptr;
  }
  std::unique_ptr<EciesEncryptor> encryptor(
      new EciesEncryptor(std::move(ec_key), key_source));
  return encryptor;
}

bool EciesEncryptor::Encrypt(const std::string& plaintext,
                             const std::string& context,
                             std::string* ecies_message) const {
  if (ecies_message == nullptr) {
    LOG(ERROR) << "ecies_message cannot be null";
    return false;
  }

  std::string serialized_private_key;
  std::string serialized_public_key;
  if (!key_source_->GetECKey(public_key_->Curve(), &serialized_private_key,
                             &serialized_public_key)) {
    LOG(WARNING) << "Failed to get key pair from key source.";
    return false;
  }

  // Convert the ephemeral public key to an encoded EC point.
  std::unique_ptr<ECPublicKey> ephemeral_public_key =
      ECPublicKey::Create(serialized_public_key);
  std::string encoded_public_key;
  if (!ephemeral_public_key->GetPointEncodedKey(&encoded_public_key)) {
    LOG(ERROR) << "Could not encode the public key. ";
    return false;
  }

  // This condition is just an indication that the serialized_public_key doesn't
  // match our size expectations. It shouldn't block the encryption operation.
  size_t expected_size = ec_util::GetPublicKeyPointSize(public_key_->Curve());
  if (encoded_public_key.size() != expected_size) {
    LOG(WARNING) << "Unexpected key size for public key. "
                 << "Curve " << public_key_->Curve() << ". Expected "
                 << expected_size << ". Found " << encoded_public_key.size()
                 << ".";
  }

  std::unique_ptr<ECPrivateKey> ephemeral_private_key =
      ECPrivateKey::Create(serialized_private_key);

  std::string session_key;
  if (!ephemeral_private_key->DeriveSharedSessionKey(*public_key_,
                                                     &session_key)) {
    LOG(WARNING) << "Failed to derive shared session key.";
    return false;
  }

  std::string encryption_key;
  std::string signing_key;
  if (!DeriveWrappingAndSigningKeys(session_key, context, &encryption_key,
                                    &signing_key)) {
    return false;
  }

  std::string zero_iv(kZeroIv, sizeof(kZeroIv));
  std::string ciphertext =
      crypto_util::EncryptAesCbc(encryption_key, zero_iv, plaintext);
  if (ciphertext.empty()) {
    LOG(WARNING) << "Failed to encrypt plaintext.";
    return false;
  }

  std::string signature =
      crypto_util::CreateSignatureHmacSha256(signing_key, ciphertext);
  if (signature.empty()) {
    LOG(WARNING) << "Failed to sign plaintext.";
    return false;
  }

  *ecies_message = encoded_public_key + ciphertext + signature;

  return true;
}

std::unique_ptr<EciesDecryptor> EciesDecryptor::Create(
    const std::string& serialized_private_key) {
  std::unique_ptr<ECPrivateKey> ec_key =
      ECPrivateKey::Create(serialized_private_key);
  if (ec_key == nullptr) {
    return nullptr;
  }
  std::unique_ptr<EciesDecryptor> decryptor(
      new EciesDecryptor(std::move(ec_key)));
  if (decryptor == nullptr) {
    LOG(ERROR) << "Failed to create EciesDecryptor.";
  }
  return decryptor;
}

EciesDecryptor::EciesDecryptor(std::unique_ptr<ECPrivateKey> private_key)
    : private_key_(std::move(private_key)) {}

bool EciesDecryptor::Decrypt(const std::string& ecies_message,
                             const std::string& context,
                             std::string* plaintext) const {
  if (plaintext == nullptr) {
    LOG(ERROR) << "plaintext cannot be null";
    return false;
  }

  // Check the minimum std::string size.
  size_t key_size = ec_util::GetPublicKeyPointSize(private_key_->Curve());
  if (key_size + kMacSizeBytes > ecies_message.size()) {
    LOG(ERROR) << "The size of the message is too small.  Expected > "
               << key_size + kMacSizeBytes << ". found "
               << ecies_message.size();
    return false;
  }

  std::string ciphertext = ecies_message.substr(
      key_size, ecies_message.size() - kMacSizeBytes - key_size);
  std::string signature =
      ecies_message.substr(ecies_message.size() - kMacSizeBytes, kMacSizeBytes);

  std::unique_ptr<ECPublicKey> public_key = ECPublicKey::CreateFromKeyPoint(
      private_key_->Curve(), ecies_message.substr(0, key_size));
  if (public_key == nullptr) {
    LOG(WARNING) << "Failed to deserialize public key.";
    return false;
  }

  std::string session_key;
  if (!private_key_->DeriveSharedSessionKey(*public_key, &session_key)) {
    LOG(WARNING) << "Failed to derive shared session key.";
    return false;
  }

  std::string encryption_key;
  std::string signing_key;
  if (!DeriveWrappingAndSigningKeys(session_key, context, &encryption_key,
                                    &signing_key)) {
    LOG(WARNING) << "Failed to derive shared wrapping and signing keys.";
    return false;
  }

  if (!crypto_util::VerifySignatureHmacSha256(signing_key, signature,
                                              ciphertext)) {
    LOG(WARNING) << "Failed to verify signature on ciphertext.";
    return false;
  }

  std::string zero_iv(kZeroIv, sizeof(kZeroIv));
  // In theory, we should be able to decrypt a cipher block that was originally
  // the empty string. But DecryptAesCbc uses an empty std::string to indicate an
  // error. This means that we can't distinguish between an error and correctly
  // decrypted empty string.
  *plaintext = crypto_util::DecryptAesCbc(encryption_key, zero_iv, ciphertext);
  if (plaintext->empty()) {
    LOG(WARNING) << "Failed to decrypt plaintext.";
    return false;
  }

  return true;
}

}  // namespace widevine.