#include <com_mysaifu_jvm_java_security_provider_PFXKeyStoreSpi.h>
#include "wcesecurity.h"
#include "KeyStore.h"

/*
 * Class:     com_mysaifu_jvm_java_security_provider_PFXKeyStoreSpi
 * Method:    openNativeCertStore
 * Signature: ([B[C)I
 */
JNIEXPORT jint JNICALL Java_com_mysaifu_jvm_java_security_provider_PFXKeyStoreSpi_openNativeCertStore
  (JNIEnv *env, jobject, jbyteArray data, jcharArray password)
{
  // PFXImportCertStore()の有無を調べる
  typedef HCERTSTORE (*PFX_IMPORT_CERT_STORE)(CRYPT_DATA_BLOB* pPFX, LPCWSTR szPassword, DWORD dwFlags);
  PFX_IMPORT_CERT_STORE pPFXImportCertStore;

  HMODULE hDLL = GetModuleHandle(_T("crypt32"));
  if (! hDLL)
  {
    ::throwException(env, CERTIFICATE_EXCEPTION_NAME, GetLastError());
    return 0;
  }
  pPFXImportCertStore = (PFX_IMPORT_CERT_STORE) GetProcAddress(hDLL, _T("PFXImportCertStore"));
  if (! pPFXImportCertStore)
  {
    // PFXImportCertStore()が存在しない
    env->ThrowNew(env->FindClass(CERTIFICATE_EXCEPTION_NAME),
                  "Windows API 'PFXImportCertStore()' not found. this method works Windows Mobile 5.0 or later");
    return 0;
  }

  const size_t passwd_len = env->GetArrayLength(password);
  _TCHAR* passwd = new _TCHAR[passwd_len + 1];
  if (! passwd)
  {
    env->ThrowNew(env->FindClass(CERTIFICATE_EXCEPTION_NAME), "Failed to allocate a password array");
    return 0;
  }

  jchar* p = reinterpret_cast<jchar*>(env->GetPrimitiveArrayCritical(password, NULL));
  memcpy(passwd, p, passwd_len * sizeof(jchar));
  passwd[passwd_len] = _T('\0');

  CRYPT_DATA_BLOB pfxPacket;
  pfxPacket.cbData = env->GetArrayLength(data);
  pfxPacket.pbData = static_cast<BYTE*>(env->GetPrimitiveArrayCritical(data, NULL));

  // PFXImportCertStore()でインポート
  // フラグ CRYPT_EXPORTABLE を指定しておかないと、AcquireCredentialsHandle() が
  // エラーコード 0x8009030d (=SEC_E_UNKNOWN_CREDENTIALS) で失敗してしまう
  HCERTSTORE hCert = pPFXImportCertStore(&pfxPacket, passwd, CRYPT_EXPORTABLE);
  DWORD dwLastError = GetLastError();
  
  env->ReleasePrimitiveArrayCritical(password, p, 0);
  env->ReleasePrimitiveArrayCritical(data, pfxPacket.pbData, 0);

  if (! hCert)
  {
    ::throwException(env, CERTIFICATE_EXCEPTION_NAME, dwLastError);
    return 0;
  }

  KeyStore* pKeyStore = new KeyStore(hCert);
  if (! pKeyStore)
  {
    ::CertCloseStore(hCert, 0);
    env->ThrowNew(env->FindClass(CERTIFICATE_EXCEPTION_NAME), "Failed to create KeyStore object");
    return 0;
  }

  return reinterpret_cast<jint>(pKeyStore);
}

/*
 * Class:     com_mysaifu_jvm_java_security_provider_PFXKeyStoreSpi
 * Method:    closeNativeCertStore
 * Signature: (I)V
 */
JNIEXPORT void JNICALL Java_com_mysaifu_jvm_java_security_provider_PFXKeyStoreSpi_closeNativeCertStore
  (JNIEnv *env, jobject, jint nativePointer)
{
  KeyStore* pKeyStore = reinterpret_cast<KeyStore*>(nativePointer);
  if (! pKeyStore)
  {
    return;
  }

  for (int i = 0; i < pKeyStore->getElementCount(); ++i)
  {
    PCCERT_CONTEXT pCertContext = pKeyStore->getCertContext(i);

    // 証明書に関連付けられているキーペアの情報を取得する
    DWORD dwSize = 0;
    if (::CertGetCertificateContextProperty(pCertContext, CERT_KEY_PROV_INFO_PROP_ID, NULL, &dwSize))
    {
      PCRYPT_KEY_PROV_INFO pCryptKeyProvInfo;
      pCryptKeyProvInfo = reinterpret_cast<PCRYPT_KEY_PROV_INFO>(::malloc(dwSize));
      if (::CertGetCertificateContextProperty(pCertContext, CERT_KEY_PROV_INFO_PROP_ID, pCryptKeyProvInfo, &dwSize))
      {
        // 秘密鍵はストアをクローズしても削除されないため、アプリケーション側で消す必要がある
        HCRYPTPROV hCryptProv;
        BOOL deleted
          = ::CryptAcquireContext(&hCryptProv,
                                  pCryptKeyProvInfo->pwszContainerName,
                                  pCryptKeyProvInfo->pwszProvName,
                                  pCryptKeyProvInfo->dwProvType,
                                  CRYPT_DELETEKEYSET);
#ifdef DEBUG
        // デバッグ出力
        _TCHAR msg[256];
        _stprintf(msg, _T("%s key pair (CRYPT_KEY_PROV_INFO:%s %s 0x%x)\r\n"),
                      deleted ? _T("Deleted") : _T("Failed to delete"),
                      pCryptKeyProvInfo->pwszContainerName,
                      pCryptKeyProvInfo->pwszProvName,
                      pCryptKeyProvInfo->dwProvType);
        OutputDebugString(msg);
#endif
      }
      ::free(pCryptKeyProvInfo);
    }
    CertFreeCertificateContext(pCertContext);
  }

  // ストアを削除する
  delete pKeyStore;
}
