// Copyright (c) 2009 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/crypto/rsa_private_key.h"

#include <iostream>
#include <list>

#include "base/logging.h"
#include "base/scoped_ptr.h"
#include "base/string_util.h"

namespace {
  // Helper for error handling during key import.
#define READ_ASSERT(truth) \
  if (!(truth)) { \
  NOTREACHED(); \
  return false; \
  }
}  // namespace

namespace base {

// static
RSAPrivateKey* RSAPrivateKey::Create(uint16 num_bits) {
  scoped_ptr<RSAPrivateKey> result(new RSAPrivateKey);
  if (!result->InitProvider())
    return NULL;

  DWORD flags = CRYPT_EXPORTABLE;

  // The size is encoded as the upper 16 bits of the flags. :: sigh ::.
  flags |= (num_bits << 16);
  if (!CryptGenKey(result->provider_, CALG_RSA_SIGN, flags, &result->key_))
    return NULL;

  return result.release();
}

// static
RSAPrivateKey* RSAPrivateKey::CreateFromPrivateKeyInfo(
    const std::vector<uint8>& input) {
  scoped_ptr<RSAPrivateKey> result(new RSAPrivateKey);
  if (!result->InitProvider())
    return NULL;

  PrivateKeyInfoCodec pki(false);  // Little-Endian
  pki.Import(input);

  int blob_size = sizeof(PUBLICKEYSTRUC) +
                  sizeof(RSAPUBKEY) +
                  pki.modulus()->size() +
                  pki.prime1()->size() +
                  pki.prime2()->size() +
                  pki.exponent1()->size() +
                  pki.exponent2()->size() +
                  pki.coefficient()->size() +
                  pki.private_exponent()->size();
  scoped_array<BYTE> blob(new BYTE[blob_size]);

  uint8* dest = blob.get();
  PUBLICKEYSTRUC* public_key_struc = reinterpret_cast<PUBLICKEYSTRUC*>(dest);
  public_key_struc->bType = PRIVATEKEYBLOB;
  public_key_struc->bVersion = 0x02;
  public_key_struc->reserved = 0;
  public_key_struc->aiKeyAlg = CALG_RSA_SIGN;
  dest += sizeof(PUBLICKEYSTRUC);

  RSAPUBKEY* rsa_pub_key = reinterpret_cast<RSAPUBKEY*>(dest);
  rsa_pub_key->magic = 0x32415352;
  rsa_pub_key->bitlen = pki.modulus()->size() * 8;
  int public_exponent_int = 0;
  for (size_t i = pki.public_exponent()->size(); i > 0; --i) {
    public_exponent_int <<= 8;
    public_exponent_int |= (*pki.public_exponent())[i - 1];
  }
  rsa_pub_key->pubexp = public_exponent_int;
  dest += sizeof(RSAPUBKEY);

  memcpy(dest, &pki.modulus()->front(), pki.modulus()->size());
  dest += pki.modulus()->size();
  memcpy(dest, &pki.prime1()->front(), pki.prime1()->size());
  dest += pki.prime1()->size();
  memcpy(dest, &pki.prime2()->front(), pki.prime2()->size());
  dest += pki.prime2()->size();
  memcpy(dest, &pki.exponent1()->front(), pki.exponent1()->size());
  dest += pki.exponent1()->size();
  memcpy(dest, &pki.exponent2()->front(), pki.exponent2()->size());
  dest += pki.exponent2()->size();
  memcpy(dest, &pki.coefficient()->front(), pki.coefficient()->size());
  dest += pki.coefficient()->size();
  memcpy(dest, &pki.private_exponent()->front(), pki.private_exponent()->size());
  dest += pki.private_exponent()->size();

  READ_ASSERT(dest == blob.get() + blob_size);
  if (!CryptImportKey(
      result->provider_, reinterpret_cast<uint8*>(public_key_struc), blob_size,
      NULL, CRYPT_EXPORTABLE, &result->key_)) {
    return NULL;
  }

  return result.release();
}

RSAPrivateKey::RSAPrivateKey() : provider_(NULL), key_(NULL) {}

RSAPrivateKey::~RSAPrivateKey() {
  if (key_) {
    if (!CryptDestroyKey(key_))
      NOTREACHED();
  }

  if (provider_) {
    if (!CryptReleaseContext(provider_, 0))
      NOTREACHED();
  }
}

bool RSAPrivateKey::InitProvider() {
  return FALSE != CryptAcquireContext(&provider_, NULL, NULL,
                                      PROV_RSA_FULL, CRYPT_VERIFYCONTEXT);
}

bool RSAPrivateKey::ExportPrivateKey(std::vector<uint8>* output) {
  // Export the key
  DWORD blob_length = 0;
  if (!CryptExportKey(key_, NULL, PRIVATEKEYBLOB, 0, NULL, &blob_length)) {
    NOTREACHED();
    return false;
  }

  scoped_array<uint8> blob(new uint8[blob_length]);
  if (!CryptExportKey(key_, NULL, PRIVATEKEYBLOB, 0, blob.get(),
                      &blob_length)) {
    NOTREACHED();
    return false;
  }

  uint8* pos = blob.get();
  PUBLICKEYSTRUC *publickey_struct = reinterpret_cast<PUBLICKEYSTRUC*>(pos);
  pos += sizeof(PUBLICKEYSTRUC);

  RSAPUBKEY *rsa_pub_key = reinterpret_cast<RSAPUBKEY*>(pos);
  pos += sizeof(RSAPUBKEY);

  int mod_size = rsa_pub_key->bitlen / 8;
  int primes_size = rsa_pub_key->bitlen / 16;

  PrivateKeyInfoCodec pki(false);  // Little-Endian

  pki.modulus()->assign(pos, pos + mod_size);
  pos += mod_size;

  pki.prime1()->assign(pos, pos + primes_size);
  pos += primes_size;
  pki.prime2()->assign(pos, pos + primes_size);
  pos += primes_size;

  pki.exponent1()->assign(pos, pos + primes_size);
  pos += primes_size;
  pki.exponent2()->assign(pos, pos + primes_size);
  pos += primes_size;

  pki.coefficient()->assign(pos, pos + primes_size);
  pos += primes_size;

  pki.private_exponent()->assign(pos, pos + mod_size);
  pos += mod_size;

  pki.public_exponent()->assign(reinterpret_cast<uint8*>(&rsa_pub_key->pubexp),
      reinterpret_cast<uint8*>(&rsa_pub_key->pubexp) + 4);

  CHECK((pos - blob_length) == reinterpret_cast<BYTE*>(publickey_struct));

  return pki.Export(output);
}

bool RSAPrivateKey::ExportPublicKey(std::vector<uint8>* output) {
  DWORD key_info_len;
  if (!CryptExportPublicKeyInfo(
      provider_, AT_SIGNATURE, X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
      NULL, &key_info_len)) {
    NOTREACHED();
    return false;
  }

  scoped_array<uint8> key_info(new uint8[key_info_len]);
  if (!CryptExportPublicKeyInfo(
      provider_, AT_SIGNATURE, X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
      reinterpret_cast<CERT_PUBLIC_KEY_INFO*>(key_info.get()), &key_info_len)) {
    NOTREACHED();
    return false;
  }

  DWORD encoded_length;
  if (!CryptEncodeObject(
      X509_ASN_ENCODING | PKCS_7_ASN_ENCODING, X509_PUBLIC_KEY_INFO,
      reinterpret_cast<CERT_PUBLIC_KEY_INFO*>(key_info.get()), NULL,
      &encoded_length)) {
    NOTREACHED();
    return false;
  }

  scoped_array<BYTE> encoded(new BYTE[encoded_length]);
  if (!CryptEncodeObject(
      X509_ASN_ENCODING | PKCS_7_ASN_ENCODING, X509_PUBLIC_KEY_INFO,
      reinterpret_cast<CERT_PUBLIC_KEY_INFO*>(key_info.get()), encoded.get(),
      &encoded_length)) {
    NOTREACHED();
    return false;
  }

  for (size_t i = 0; i < encoded_length; ++i)
    output->push_back(encoded[i]);

  return true;
}

}  // namespace base