////////////////////////////////////////////////////////////////////////////////
// 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.
////////////////////////////////////////////////////////////////////////////////
//
// Description:
//   Definition of elliptic curve public and private key classes.

#include "common/ec_key.h"

#include <vector>

#include "glog/logging.h"
#include "absl/memory/memory.h"
#include "openssl/base.h"
#include "openssl/bn.h"
#include "openssl/ec.h"
#include "openssl/ecdh.h"
#include "openssl/ecdsa.h"
#include "openssl/err.h"
#include "openssl/evp.h"
#include "openssl/sha.h"
#include "common/aes_cbc_util.h"
#include "common/ec_util.h"
#include "common/hash_algorithm.h"
#include "common/openssl_util.h"
#include "common/sha_util.h"

namespace widevine {

namespace {

void* SHA256KDF(const void* in, size_t inlen, void* out, size_t* outlen) {
  std::string shared_session_key(static_cast<const char*>(in), inlen);
  std::string derived_shared_session_key = Sha256_Hash(shared_session_key);
  if (*outlen != SHA256_DIGEST_LENGTH) {
    return nullptr;
  }
  memcpy(out, derived_shared_session_key.data(), *outlen);
  return out;
}

ECPrivateKey::EllipticCurve ECKeyToCurve(const EC_KEY* key) {
  if (key == nullptr) {
    return ECPrivateKey::UNDEFINED_CURVE;
  }
  return ec_util::NidToCurve(EC_GROUP_get_curve_name(EC_KEY_get0_group(key)));
}

std::string OpenSSLErrorString(uint32_t error) {
  char buf[ERR_ERROR_STRING_BUF_LEN];
  ERR_error_string_n(error, buf, sizeof(buf));
  return buf;
}

std::string GetMessageDigest(const std::string& message,
                             widevine::HashAlgorithm hash_algorithm) {
  switch (hash_algorithm) {
    case widevine::HashAlgorithm::kUnspecified:
    case widevine::HashAlgorithm::kSha256:
      return widevine::Sha256_Hash(message);
    case widevine::HashAlgorithm::kSha384:
      return widevine::Sha384_Hash(message);
    case widevine::HashAlgorithm::kSha1:
      LOG(ERROR) << "Unexpected hash algorithm: "
                 << static_cast<int>(hash_algorithm);
      return "";
  }
  LOG(FATAL) << "Unexpected hash algorithm: "
             << static_cast<int>(hash_algorithm);
  return "";
}

}  // namespace

ECPrivateKey::ECPrivateKey(EC_KEY* ec_key) : key_(ec_key) {
  CHECK(key() != nullptr);
  CHECK(EC_KEY_get0_private_key(key()) != nullptr);
  CHECK_NE(ECKeyToCurve(key()), ECPrivateKey::UNDEFINED_CURVE);
  CHECK_EQ(EC_KEY_check_key(key()), 1);
}

ECPrivateKey::ECPrivateKey(ScopedECKEY ec_key) : key_(std::move(ec_key)) {}

ECPrivateKey::ECPrivateKey(const ECPrivateKey& ec_key)
    : ECPrivateKey(EC_KEY_dup(ec_key.key())) {}

std::unique_ptr<ECPrivateKey> ECPrivateKey::Create(
    const std::string& serialized_key) {
  EC_KEY* key;
  if (!ec_util::DeserializeECPrivateKey(serialized_key, &key)) {
    return nullptr;
  }
  ScopedECKEY scoped_ec_key(key);
  if (EC_KEY_check_key(scoped_ec_key.get()) != 1) {
    LOG(ERROR) << "Invalid private EC key: "
               << OpenSSLErrorString(ERR_get_error());
    return nullptr;
  }
  if (ECKeyToCurve(scoped_ec_key.get()) == ECPrivateKey::UNDEFINED_CURVE) {
    LOG(ERROR) << "Key has unsupported curve";
    return nullptr;
  }
  return absl::make_unique<ECPrivateKey>(scoped_ec_key.release());
}

std::unique_ptr<ECPublicKey> ECPrivateKey::PublicKey() const {
  if (key_ == nullptr) {
    return nullptr;
  }
  return absl::make_unique<ECPublicKey>(ScopedECKEY(EC_KEY_dup(key_.get())));
}

bool ECPrivateKey::DeriveSharedSessionKey(
    const ECPublicKey& public_key,
    std::string* derived_shared_session_key) const {
  if (derived_shared_session_key == nullptr) {
    LOG(ERROR) << "|derived_shared_session_key| cannot be nullptr";
    return false;
  }
  const EC_GROUP* group1 = EC_KEY_get0_group(key());
  const EC_GROUP* group2 = EC_KEY_get0_group(public_key.key());
  if (EC_GROUP_cmp(group1, group2, nullptr) != 0) {
    LOG(ERROR) << "EC_GROUPs do not match";
    return false;
  }
  const EC_POINT* public_key_point = EC_KEY_get0_public_key(public_key.key());
  derived_shared_session_key->resize(SHA256_DIGEST_LENGTH, 0);
  int result = ECDH_compute_key(
      reinterpret_cast<void*>(
          const_cast<char*>(derived_shared_session_key->data())),
      derived_shared_session_key->size(), public_key_point, key(), SHA256KDF);
  if (result == -1) {
    LOG(ERROR) << "Could not write shared session key: "
               << OpenSSLErrorString(ERR_get_error());
    return false;
  }
  if (result != SHA256_DIGEST_LENGTH) {
    LOG(ERROR) << "Wrote " << result
               << " bytes to derived shared session key instead of "
               << SHA256_DIGEST_LENGTH << " bytes";
    return false;
  }
  return true;
}

bool ECPrivateKey::GenerateSignature(const std::string& message,
                                     HashAlgorithm hash_algorithm,
                                     std::string* signature) const {
  if (message.empty()) {
    LOG(ERROR) << "|message| cannot be empty";
    return false;
  }
  if (signature == nullptr) {
    LOG(ERROR) << "|signature| cannot be nullptr";
    return false;
  }

  // Hash the message using corresponding hash algorithm.
  std::string message_digest = GetMessageDigest(message, hash_algorithm);
  if (message_digest.empty()) {
    LOG(ERROR) << "Empty message digest";
    return false;
  }

  size_t max_signature_size = ECDSA_size(key());
  if (max_signature_size == 0) {
    LOG(ERROR) << "key_ does not have a group set";
    return false;
  }
  signature->resize(max_signature_size);
  unsigned int bytes_written = 0;
  int result = ECDSA_sign(
      0 /* unused type */,
      reinterpret_cast<const uint8_t*>(message_digest.data()),
      message_digest.size(),
      reinterpret_cast<uint8_t*>(const_cast<char*>(signature->data())),
      &bytes_written, key());
  if (result != 1) {
    LOG(ERROR) << "Could not calculate signature: "
               << OpenSSLErrorString(ERR_get_error());
    return false;
  }
  signature->resize(bytes_written);
  return true;
}

bool ECPrivateKey::MatchesPrivateKey(const ECPrivateKey& private_key) const {
  return BN_cmp(EC_KEY_get0_private_key(key()),
                EC_KEY_get0_private_key(private_key.key())) == 0;
}

bool ECPrivateKey::MatchesPublicKey(const ECPublicKey& public_key) const {
  return EC_POINT_cmp(EC_KEY_get0_group(key()), EC_KEY_get0_public_key(key()),
                      EC_KEY_get0_public_key(public_key.key()), nullptr) == 0;
}

ECPrivateKey::EllipticCurve ECPrivateKey::Curve() const {
  return ECKeyToCurve(key());
}

bool ECPrivateKey::SerializedKey(std::string* serialized_key) const {
  CHECK(serialized_key);
  return ec_util::SerializeECPrivateKey(key_.get(), serialized_key);
}

bool ECPrivateKey::GetRawPrivateKey(std::string* raw_private_key) const {
  if (raw_private_key == nullptr) {
    LOG(WARNING) << "|raw_private_key| cannot be nullptr.";
    return false;
  }
  const EC_GROUP* group = EC_KEY_get0_group(key());
  if (group == nullptr) {
    LOG(WARNING) << "Failed to load EC key group.";
    return false;
  }
  int key_size_bits = EC_GROUP_get_degree(group);
  if (key_size_bits == 0) {
    LOG(WARNING) << "Key size is 0.";
    return false;
  }
  const BIGNUM* key_bn = EC_KEY_get0_private_key(key());
  int key_size_bytes = (key_size_bits + 7) / 8;
  std::vector<uint8_t> tmp_raw_private_key_data(key_size_bytes);
  int num_bytes = BN_bn2bin(key_bn, tmp_raw_private_key_data.data());
  if (num_bytes != key_size_bytes) {
    LOG(WARNING) << "Failed to convert EC private key BIGNUM to binary data.";
    return false;
  }
  raw_private_key->assign(tmp_raw_private_key_data.begin(),
                          tmp_raw_private_key_data.end());
  return true;
}

ECPublicKey::ECPublicKey(EC_KEY* ec_key) : key_(ec_key) {
  CHECK(key() != nullptr);
  CHECK(EC_KEY_get0_private_key(key()) == nullptr);
  CHECK_NE(ECKeyToCurve(key()), ECPrivateKey::UNDEFINED_CURVE);
  CHECK_EQ(EC_KEY_check_key(key()), 1);
}

ECPublicKey::ECPublicKey(ScopedECKEY ec_key) : key_(std::move(ec_key)) {}

ECPublicKey::ECPublicKey(const ECPublicKey& ec_key)
    : ECPublicKey(EC_KEY_dup(ec_key.key())) {}

std::unique_ptr<ECPublicKey> ECPublicKey::Create(
    const std::string& serialized_key) {
  EC_KEY* key;
  if (!ec_util::DeserializeECPublicKey(serialized_key, &key)) {
    return nullptr;
  }
  ScopedECKEY scoped_ec_key(key);
  if (EC_KEY_check_key(scoped_ec_key.get()) != 1) {
    LOG(ERROR) << "Invalid public EC key: "
               << OpenSSLErrorString(ERR_get_error());
    return nullptr;
  }
  if (ECKeyToCurve(scoped_ec_key.get()) == ECPrivateKey::UNDEFINED_CURVE) {
    LOG(ERROR) << "Key has unsupported curve";
    return nullptr;
  }
  return absl::make_unique<ECPublicKey>(scoped_ec_key.release());
}

std::unique_ptr<ECPublicKey> ECPublicKey::CreateFromKeyPoint(
    ECPrivateKey::EllipticCurve curve, const std::string& key_point) {
  EC_KEY* key;
  if (!ec_util::GetPublicKeyFromKeyPoint(curve, key_point, &key)) {
    return nullptr;
  }

  ScopedECKEY scoped_ec_key(key);
  if (EC_KEY_check_key(scoped_ec_key.get()) != 1) {
    LOG(ERROR) << "Invalid public EC key: "
               << OpenSSLErrorString(ERR_get_error());
    return nullptr;
  }
  if (ECKeyToCurve(scoped_ec_key.get()) == ECPrivateKey::UNDEFINED_CURVE) {
    LOG(ERROR) << "Key has unsupported curve";
    return nullptr;
  }
  return absl::make_unique<ECPublicKey>(scoped_ec_key.release());
}

bool ECPublicKey::VerifySignature(const std::string& message,
                                  HashAlgorithm hash_algorithm,
                                  const std::string& signature) const {
  if (message.empty()) {
    LOG(ERROR) << "|message| cannot be empty";
    return false;
  }
  if (signature.empty()) {
    LOG(ERROR) << "|signature| cannot be empty";
    return false;
  }

  // Hash the message using corresponding hash algorithm.
  std::string message_digest = GetMessageDigest(message, hash_algorithm);
  if (message_digest.empty()) {
    LOG(ERROR) << "Empty message digest";
    return false;
  }

  int result = ECDSA_verify(
      0 /* unused type */,
      reinterpret_cast<const uint8_t*>(message_digest.data()),
      message_digest.size(),
      reinterpret_cast<uint8_t*>(const_cast<char*>(signature.data())),
      signature.size(), key());
  if (result != 1) {
    LOG(ERROR) << "Could not verify signature: "
               << OpenSSLErrorString(ERR_get_error());
    return false;
  }
  return true;
}

bool ECPublicKey::MatchesPrivateKey(const ECPrivateKey& private_key) const {
  return private_key.MatchesPublicKey(*this);
}

bool ECPublicKey::MatchesPublicKey(const ECPublicKey& public_key) const {
  return EC_POINT_cmp(EC_KEY_get0_group(key()), EC_KEY_get0_public_key(key()),
                      EC_KEY_get0_public_key(public_key.key()), nullptr) == 0;
}

ECPrivateKey::EllipticCurve ECPublicKey::Curve() const {
  return ECKeyToCurve(key());
}

bool ECPublicKey::SerializedKey(std::string* serialized_key) const {
  CHECK(serialized_key);
  return ec_util::SerializeECPublicKey(key_.get(), serialized_key);
}

bool ECPublicKey::GetPointEncodedKey(std::string* encoded_key) const {
  CHECK(encoded_key);
  return ec_util::GetPublicKeyPoint(key_.get(), encoded_key);
}

bool ECPublicKey::GetRawPublicKey(std::string* raw_public_key) const {
  if (raw_public_key == nullptr) {
    LOG(WARNING) << "|raw_public_key| cannot be nullptr.";
    return false;
  }
  std::string tmp_raw_public_key;
  if (!ec_util::GetPublicKeyPoint(key(), &tmp_raw_public_key) ||
      tmp_raw_public_key.empty()) {
    LOG(WARNING) << "Failed to encoded EC_KEY";
  }
  CHECK_EQ(tmp_raw_public_key[0], POINT_CONVERSION_UNCOMPRESSED);
  // Public key is X9.62 uncompressed format. Expected structure is <tag><X><Y>,
  // where tag = 0x04 indicating uncompressed, and X and Y lengths are equal but
  // dependent on ECC curve. Tag is removed because Intel Sigma 2.1.0 library
  // requires only X and Y bytes.
  raw_public_key->assign(tmp_raw_public_key.begin() + 1,
                         tmp_raw_public_key.end());
  return true;
}

}  // namespace widevine