////////////////////////////////////////////////////////////////////////////////
// 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.
////////////////////////////////////////////////////////////////////////////////

//
// Description:
//   RSA utility functions for serializing and deserializing RSA keys,
//   encryption, and signing.

#include "common/rsa_util.h"

#include <cstring>

#include "glog/logging.h"
#include "openssl/bio.h"
#include "openssl/evp.h"
#include "openssl/pem.h"
#include "openssl/x509.h"

namespace widevine {
namespace rsa_util {

static bool SerializeRsaKey(const RSA* key, std::string* serialized_key,
                            bool serialize_private_key) {
  if (key == NULL) {
    LOG(ERROR) << (serialize_private_key ? "Private" : "Public")
               << " RSA key is NULL.";
    return false;
  }
  if (serialized_key == NULL) {
    LOG(ERROR) << "Pointer to hold serialized RSA" << (serialize_private_key ?
        "Private" : "Public") << "Key is NULL.";
    return false;
  }
  BIO* bio = BIO_new(BIO_s_mem());
  if (bio == NULL) {
    LOG(ERROR) << "BIO_new returned NULL";
    return false;
  }
  bool success = false;
  if ((serialize_private_key ?
       i2d_RSAPrivateKey_bio(bio, const_cast<RSA*>(key)) :
       i2d_RSAPublicKey_bio(bio, const_cast<RSA*>(key))) != 0) {
    int serialized_size = BIO_pending(bio);
    serialized_key->assign(serialized_size, 0);
    if (BIO_read(bio, &(*serialized_key)[0], serialized_size) ==
        serialized_size) {
      success = true;
    } else {
      LOG(ERROR) << "BIO_read failure";
    }
  } else {
    LOG(ERROR) << (serialize_private_key ? "Private" : "Public") <<
        " key serialization failure";
  }
  BIO_free(bio);
  return success;
}

static bool DeserializeRsaKey(const std::string& serialized_key, RSA** key,
                              bool deserialize_private_key) {
  if (serialized_key.empty()) {
    LOG(ERROR) << "Serialized RSA" << (deserialize_private_key ?
                  "Private" : "Public") << "Key is empty.";
    return false;
  }
  if (key == NULL) {
    LOG(ERROR) << "Pointer to hold new RSA " << (deserialize_private_key ?
                  "private" : "public") << " key is NULL.";
    return false;
  }
  BIO* bio = BIO_new_mem_buf(const_cast<char*>(serialized_key.data()),
                             serialized_key.size());
  if (bio == NULL) {
    LOG(ERROR) << "BIO_new_mem_buf returned NULL";
    return false;
  }
  *key = deserialize_private_key ? d2i_RSAPrivateKey_bio(bio, NULL) :
      d2i_RSAPublicKey_bio(bio, NULL);
  BIO_free(bio);
  if (*key == NULL) {
    LOG(ERROR) << (deserialize_private_key ? "Private" : "Public") <<
        " RSA key deserialization failure";
  }
  return *key != NULL;
}

bool SerializeRsaPrivateKey(const RSA* private_key,
                            std::string* serialized_private_key) {
  return SerializeRsaKey(private_key, serialized_private_key, true);
}

bool DeserializeRsaPrivateKey(const std::string& serialized_private_key,
                              RSA** private_key) {
  return DeserializeRsaKey(serialized_private_key, private_key, true);
}

bool SerializeRsaPublicKey(const RSA* public_key,
                           std::string* serialized_public_key) {
  return SerializeRsaKey(public_key, serialized_public_key, false);
}

bool DeserializeRsaPublicKey(const std::string& serialized_public_key,
                             RSA** public_key) {
  return DeserializeRsaKey(serialized_public_key, public_key, false);
}

bool SerializePrivateKeyInfo(const RSA* private_key,
                             std::string* serialized_private_key) {
  if (private_key == NULL) {
    LOG(ERROR) << "Private RSA key is NULL.";
    return false;
  }
  if (serialized_private_key == NULL) {
    LOG(ERROR) << "Pointer to hold serialized PrivateKeyInfo is NULL.";
    return false;
  }
  // The following method of serializing a PKCS#8 PrivateKeyInfo object
  // was obtained from analyzing the openssl utility code, as the official
  // mechanism via i2d_PKCS8PrivateKey_bio is broken in the current openssl
  // version (1.0.0c). Please refer to b/8560683.
  EVP_PKEY* evp = EVP_PKEY_new();
  if (evp == NULL) {
    LOG(ERROR) << "EVP_PKEY_new returned NULL.";
    return false;
  }
  bool success = false;
  PKCS8_PRIV_KEY_INFO *pkcs8_pki = NULL;
  BIO* bio = NULL;
  if (EVP_PKEY_set1_RSA(evp, const_cast<RSA*>(private_key)) == 0) {
    LOG(ERROR) << "EVP_PKEY_set1_RSA failed.";
    goto cleanup;
  }
  pkcs8_pki = EVP_PKEY2PKCS8(evp);
  if (pkcs8_pki == NULL) {
    LOG(ERROR) << "EVP_PKEY2PKCS8 returned NULL.";
    goto cleanup;
  }
  bio = BIO_new(BIO_s_mem());
  if (bio == NULL) {
    LOG(ERROR) << "BIO_new returned NULL.";
    goto cleanup;
  }
  if (i2d_PKCS8_PRIV_KEY_INFO_bio(bio, pkcs8_pki) == 0) {
    LOG(ERROR) << "i2d_PKCS8_PRIV_KEY_INFO_bio failed.";
    goto cleanup;
  }
  {
    int serialized_size = BIO_pending(bio);
    serialized_private_key->assign(serialized_size, 0);
    if (BIO_read(bio, &(*serialized_private_key)[0], serialized_size) !=
        serialized_size) {
      LOG(ERROR) << "BIO_read failed.";
      goto cleanup;
    }
  }
  success = true;

cleanup:
  if (bio != NULL) {
    BIO_free(bio);
  }
  if (pkcs8_pki != NULL) {
    PKCS8_PRIV_KEY_INFO_free(pkcs8_pki);
  }
  EVP_PKEY_free(evp);
  return success;
}

bool DeserializePrivateKeyInfo(const std::string& serialized_private_key,
                               RSA** private_key) {
  if (serialized_private_key.empty()) {
    LOG(ERROR) << "Serialized PrivateKeyInfo is empty.";
    return false;
  }
  if (private_key == NULL) {
    LOG(ERROR) << "Pointer to hold new RSA private key is NULL.";
    return false;
  }
  // The following method of deserializing a PKCS#8 PrivateKeyInfo object
  // was obtained from analyzing the openssl utility code, as the official
  // mechanism via d2i_PKCS8PrivateKey_bio is broken in the current openssl
  // version (1.0.0c). Please refer to b/8560683.
  BIO* bio = BIO_new_mem_buf(const_cast<char*>(serialized_private_key.data()),
                             serialized_private_key.size());
  if (bio == NULL) {
    LOG(ERROR) << "BIO_new_mem_buf returned NULL";
    return false;
  }
  bool success = false;
  EVP_PKEY* evp = NULL;
  PKCS8_PRIV_KEY_INFO *pkcs8_pki = d2i_PKCS8_PRIV_KEY_INFO_bio(bio, NULL);
  if (pkcs8_pki == NULL) {
    LOG(ERROR) << "d2i_PKCS8_PRIV_KEY_INFO_bio returned NULL.";
    goto cleanup;
  }
  evp = EVP_PKCS82PKEY(pkcs8_pki);
  if (evp == NULL) {
    LOG(ERROR) << "EVP_PKCS82PKEY returned NULL.";
    goto cleanup;
  }
  *private_key = EVP_PKEY_get1_RSA(evp);
  if (*private_key == NULL) {
    LOG(ERROR) << "PrivateKeyInfo did not contain an RSA key.";
    goto cleanup;
  }
  success = true;

cleanup:
  if (evp != NULL) {
    EVP_PKEY_free(evp);
  }
  if (pkcs8_pki != NULL) {
    PKCS8_PRIV_KEY_INFO_free(pkcs8_pki);
  }
  BIO_free(bio);
  return success;
}

bool RsaPrivateKeyToPrivateKeyInfo(const std::string& rsa_private_key,
                                   std::string* private_key_info) {
  RSA* key = NULL;
  if (DeserializeRsaPrivateKey(rsa_private_key, &key)) {
    bool success = SerializePrivateKeyInfo(key, private_key_info);
    RSA_free(key);
    return success;
  }
  return false;
}

bool PrivateKeyInfoToRsaPrivateKey(const std::string& private_key_info,
                                   std::string* rsa_private_key) {
  RSA* key = NULL;
  if (DeserializePrivateKeyInfo(private_key_info, &key)) {
    bool success = SerializeRsaPrivateKey(key, rsa_private_key);
    RSA_free(key);
    return success;
  }
  return false;
}

bool SerializeEncryptedPrivateKeyInfo(const RSA* private_key,
                                      const std::string& passphrase,
                                      std::string* serialized_private_key) {
  if (private_key == NULL) {
    LOG(ERROR) << "Private RSA key is NULL.";
    return false;
  }
  if (passphrase.empty()) {
    LOG(ERROR) << "Passphrase for RSA key encryption is empty.";
    return false;
  }
  if (serialized_private_key == NULL) {
    LOG(ERROR) << "Pointer to hold serialized EncryptedPrivateKeyInfo is NULL.";
    return false;
  }
  EVP_PKEY* evp = EVP_PKEY_new();
  if (evp == NULL) {
    LOG(ERROR) << "EVP_PKEY_new returned NULL.";
    return false;
  }
  bool success = false;
  BIO* bio = NULL;
  if (EVP_PKEY_set1_RSA(evp, const_cast<RSA*>(private_key)) == 0) {
    LOG(ERROR) << "EVP_PKEY_set1_RSA failed.";
    goto cleanup;
  }
  bio = BIO_new(BIO_s_mem());
  if (bio == NULL) {
    LOG(ERROR) << "BIO_new returned NULL.";
    goto cleanup;
  }
  if (i2d_PKCS8PrivateKey_bio(bio, evp, EVP_aes_256_cbc(),
                              const_cast<char*>(passphrase.data()),
                              passphrase.size(), NULL, NULL) == 0) {
    LOG(ERROR) << "i2d_PKCS8PrivateKey_bio failed.";
    goto cleanup;
  }
  {
    int serialized_size = BIO_pending(bio);
    serialized_private_key->assign(serialized_size, 0);
    if (BIO_read(bio, &(*serialized_private_key)[0], serialized_size) !=
        serialized_size) {
      LOG(ERROR) << "BIO_read failed.";
      goto cleanup;
    }
  }
  success = true;

cleanup:
  if (bio != NULL) {
    BIO_free(bio);
  }
  EVP_PKEY_free(evp);
  return success;
}

namespace {
// Password retrieval function used by DeserializeEncryptedPrivateKeyInfo below.
int get_password(char *buf, int size, int rwflag, void *u) {
  CHECK(buf);
  CHECK(u);
  const std::string* pass(static_cast<const std::string*>(u));
  if (!pass->empty() && size >= static_cast<int>(pass->size())) {
    memcpy(buf, pass->data(), pass->size());
    return pass->size();
  }
  return 0;
}
}  // namespace

bool DeserializeEncryptedPrivateKeyInfo(const std::string& serialized_private_key,
                                        const std::string& passphrase,
                                        RSA** private_key) {
  if (serialized_private_key.empty()) {
    LOG(ERROR) << "Serialized RSAEncryptedPrivateKeyInfo is empty.";
    return false;
  }
  if (passphrase.empty()) {
    LOG(ERROR) << "Passphrase for RSA key decryption is empty.";
    return false;
  }
  if (private_key == NULL) {
    LOG(ERROR) << "Pointer to hold new RSA private key is NULL.";
    return false;
  }
  BIO* bio = BIO_new_mem_buf(const_cast<char*>(serialized_private_key.data()),
                             serialized_private_key.size());
  if (bio == NULL) {
    LOG(ERROR) << "BIO_new_mem_buf returned NULL";
    return false;
  }
  bool success = false;
  EVP_PKEY* evp = d2i_PKCS8PrivateKey_bio(bio, NULL, get_password,
                                          const_cast<std::string*>(&passphrase));
  if (evp == NULL) {
    LOG(ERROR) << "d2i_PKCS8PrivateKey_bio returned NULL.";
    goto cleanup;
  }
  *private_key = EVP_PKEY_get1_RSA(evp);
  if (*private_key == NULL) {
    LOG(ERROR) << "EncryptedPrivateKeyInfo did not contain an RSA key.";
    goto cleanup;
  }
  success = true;

cleanup:
  if (evp != NULL) {
    EVP_PKEY_free(evp);
  }
  BIO_free(bio);
  return success;
}

bool RsaPrivateKeyToEncryptedPrivateKeyInfo(const std::string& rsa_private_key,
                                            const std::string& passphrase,
                                            std::string* private_key_info) {
  RSA* key = NULL;
  if (DeserializeRsaPrivateKey(rsa_private_key, &key)) {
    bool success = SerializeEncryptedPrivateKeyInfo(key,
                                                    passphrase,
                                                    private_key_info);
    RSA_free(key);
    return success;
  }
  return false;
}

bool EncryptedPrivateKeyInfoToRsaPrivateKey(const std::string& private_key_info,
                                            const std::string& passphrase,
                                            std::string* rsa_private_key) {
  RSA* key = NULL;
  if (DeserializeEncryptedPrivateKeyInfo(private_key_info, passphrase, &key)) {
    bool success = SerializeRsaPrivateKey(key, rsa_private_key);
    RSA_free(key);
    return success;
  }
  return false;
}

}  // namespace rsa_util
}  // namespace widevine