// Copyright 2021 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
//
//  Reference implementation of OEMCrypto APIs
//
#include "oemcrypto_oem_cert.h"

#include <string.h>

#include <openssl/pkcs7.h>
#include <openssl/rsa.h>
#include <openssl/x509.h>

#include "log.h"
#include "oemcrypto_rsa_key.h"
#include "scoped_object.h"

namespace wvoec_ref {
namespace {
using ScopedCertificate = ScopedObject<X509, X509_free>;
using ScopedEvpKey = ScopedObject<EVP_PKEY, EVP_PKEY_free>;
using ScopedPkcs7 = ScopedObject<PKCS7, PKCS7_free>;

constexpr size_t kExpectedCertCount = 2;  // Leaf and intermediate.
constexpr int kDeviceCertIndex = 0;

// Checks that the |public_key| from an X.509 certificate is the
// correct public key of the serialized |private_key_data|.
OEMCryptoResult VerifyRsaKey(const RSA* public_key,
                             const std::vector<uint8_t>& private_key_data) {
  if (public_key == nullptr) {
    LOGE("RSA key is null");
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  std::unique_ptr<RsaPrivateKey> private_key =
      RsaPrivateKey::Load(private_key_data);
  if (!private_key) {
    LOGE("Failed to parse provided RSA private key");
    return OEMCrypto_ERROR_INVALID_RSA_KEY;
  }
  if (!RsaKeysAreMatchingPair(public_key, private_key->GetRsaKey())) {
    LOGE("OEM certificate keys do not match");
    return OEMCrypto_ERROR_INVALID_RSA_KEY;
  }
  return OEMCrypto_SUCCESS;
}
}  // namespace

// This utility class encapsulates the minimum functionality of an
// OEM Public Certificate required to verify a device's OEM Public
// Certificate.
class OemPublicCertificate {
 public:
  // Loads a PKCS #7 signedData message with certificate chain.
  // Minimum validation is performed.  Only checks that the
  // device's public key is of a known type (RSA).
  static std::unique_ptr<OemPublicCertificate> Load(const uint8_t* public_cert,
                                                    size_t public_cert_size) {
    std::unique_ptr<OemPublicCertificate> oem_public_cert;
    if (public_cert == nullptr) {
      LOGE("Public cert buffer is null");
      return oem_public_cert;
    }
    if (public_cert_size == 0) {
      LOGE("Public cert buffer is empty");
      return oem_public_cert;
    }
    oem_public_cert.reset(new OemPublicCertificate());
    if (!oem_public_cert->InitFromBuffer(public_cert, public_cert_size)) {
      oem_public_cert.reset();
    }
    return oem_public_cert;
  }

  OemCertificate::KeyType key_type() const { return key_type_; }
  const std::vector<uint8_t>& cert_data() const { return cert_data_; }

  const RSA* GetPublicRsaKey() const {
    return EVP_PKEY_get0_RSA(device_public_key_.get());
  }

  ~OemPublicCertificate() = default;

  OemPublicCertificate(const OemPublicCertificate&) = delete;
  OemPublicCertificate(OemPublicCertificate&&) = delete;
  const OemPublicCertificate& operator=(const OemPublicCertificate&) = delete;
  OemPublicCertificate& operator=(OemPublicCertificate&&) = delete;

 private:
  OemPublicCertificate() {}

  bool InitFromBuffer(const uint8_t* public_cert, size_t public_cert_size) {
    // Step 1: Parse the PKCS7 certificate chain as signedData.
    const uint8_t* public_cert_ptr = public_cert;
    pkcs7_.reset(d2i_PKCS7(nullptr, &public_cert_ptr, public_cert_size));
    if (!pkcs7_) {
      LOGE("Failed to parse PKCS#7 certificate chain");
      return false;
    }
    if (!PKCS7_type_is_signed(pkcs7_.get())) {
      LOGE("OEM Public Certificate is not PKCS#7 signed data");
      return false;
    }
    PKCS7_SIGNED* signed_data = pkcs7_->d.sign;
    // Step 2: Get the leaf certificate.
    const size_t cert_count =
        static_cast<size_t>(sk_X509_num(signed_data->cert));
    if (cert_count != kExpectedCertCount) {
      LOGE("Unexpected number of certificates: expected = %zu, actual = %zu",
           kExpectedCertCount, cert_count);
      return false;
    }
    X509* leaf_cert = sk_X509_value(signed_data->cert, kDeviceCertIndex);
    // Step 3a: Get the device's public key.
    device_public_key_.reset(X509_get_pubkey(leaf_cert));
    if (!device_public_key_) {
      LOGE("Device X.509 certificate is missing a public key");
      return false;
    }
    // Step 3b: Check key type.
    if (EVP_PKEY_get0_RSA(device_public_key_.get()) == nullptr) {
      LOGE("Device public key is not RSA");
      return false;
    }
    key_type_ = OemCertificate::kRsa;
    cert_data_.assign(public_cert, public_cert + public_cert_size);
    return true;
  }

  OemCertificate::KeyType key_type_ = OemCertificate::kNone;
  // OpenSSL/BoringSSL's implementation of PKCS7 objects.
  ScopedPkcs7 pkcs7_;
  ScopedEvpKey device_public_key_;
  std::vector<uint8_t> cert_data_;
};

// ===== ===== ===== OEM Certificate ===== ===== =====

// static
std::unique_ptr<OemCertificate> OemCertificate::Create(
    const uint8_t* private_key_data, size_t private_key_size,
    const uint8_t* public_cert_data, size_t public_cert_size) {
  std::unique_ptr<OemCertificate> oem_cert;
  // Step 1: Verify public cert is well-formed.
  std::unique_ptr<OemPublicCertificate> oem_public_cert =
      OemPublicCertificate::Load(public_cert_data, public_cert_size);
  if (!oem_public_cert) {
    LOGE("Invalid OEM Public Certificate");
    return oem_cert;
  }
  // Step 2: Verify private key is well-formed.
  switch (oem_public_cert->key_type()) {
    case kRsa: {
      std::unique_ptr<RsaPrivateKey> oem_private_key =
          RsaPrivateKey::Load(private_key_data, private_key_size);
      if (!oem_private_key) {
        LOGE("Invalid OEM Private Key");
        return oem_cert;
      }
    } break;
    case kNone:  // Suppress compiler warnings.
      return oem_cert;
  }
  // Step 3: Copy over data.
  oem_cert.reset(new OemCertificate());
  oem_cert->private_key_.assign(private_key_data,
                                private_key_data + private_key_size);
  oem_cert->public_cert_ = std::move(oem_public_cert);
  return oem_cert;
}

// static
std::unique_ptr<OemCertificate> OemCertificate::Create(
    const std::vector<uint8_t>& private_key,
    const std::vector<uint8_t>& public_cert) {
  if (private_key.empty()) {
    LOGE("Private key buffer is empty");
    return std::unique_ptr<OemCertificate>();
  }
  if (public_cert.empty()) {
    LOGE("Public cert buffer is empty");
    return std::unique_ptr<OemCertificate>();
  }
  return Create(private_key.data(), private_key.size(), public_cert.data(),
                public_cert.size());
}

OemCertificate::KeyType OemCertificate::key_type() const {
  return public_cert_->key_type();
}

OEMCryptoResult OemCertificate::GetPublicCertificate(
    uint8_t* public_cert, size_t* public_cert_length) const {
  if (public_cert_length == nullptr) {
    LOGE("Output |public_cert_length| is null");
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }
  if (public_cert == nullptr && *public_cert_length > 0) {
    LOGE("Output |public_cert| is null");
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }
  const std::vector<uint8_t>& cert_data = public_cert_->cert_data();
  if (*public_cert_length < cert_data.size()) {
    *public_cert_length = cert_data.size();
    return OEMCrypto_ERROR_SHORT_BUFFER;
  }
  *public_cert_length = cert_data.size();
  memcpy(public_cert, cert_data.data(), cert_data.size());
  return OEMCrypto_SUCCESS;
}

const std::vector<uint8_t>& OemCertificate::GetPublicCertificate() const {
  return public_cert_->cert_data();
}

OEMCryptoResult OemCertificate::IsCertificateValid() const {
  switch (key_type()) {
    case kRsa:
      return VerifyRsaKey(public_cert_->GetPublicRsaKey(), private_key_);
    case kNone:  // Suppress compiler warnings.
      break;
  }
  LOGE("Unexpected error key type: type = %d", static_cast<int>(key_type()));
  return OEMCrypto_ERROR_UNKNOWN_FAILURE;
}

// Constructor and destructor do not perform anything special, but
// must be declared within a scope which defines OemPublicCertificate.
OemCertificate::OemCertificate() {}
OemCertificate::~OemCertificate() {}

}  // namespace wvoec_ref