#include "stdafx.h"
#include "KeyStore.h"
#include "wcesecurity.h"
#include <jni.h>

static
_TCHAR* createAlias(PCCERT_CONTEXT pContext)
{
  // エイリアス名は
  // CERT_NAME_FRIENDLY_DISPLAY_TYPE + " [SerialNumber:" + シリアル番号 + "]\0" とする
  // シリアル番号を追加するのは、同一ストア内に同じ名前を持つ証明書が存在した場合にも
  // ユニークにする必要があるため
  DWORD name_count = ::CertGetNameString(pContext,
                                         CERT_NAME_FRIENDLY_DISPLAY_TYPE,
                                         0,
                                         NULL,
                                         NULL,
                                         0);
  int total_count = name_count // 名前の長さ と 末尾のNULL
                    + (pContext->pCertInfo->SerialNumber.cbData * 2) // シリアル番号 (1バイトにつき2文字）
                    + 2        // 左右括弧
                    + 13       // SerialNumber:
                    + 1;       // 名前とシリアル番号との間のスペース
  _TCHAR* alias = (_TCHAR*) ::malloc(sizeof(_TCHAR) * total_count);

  ::CertGetNameString(pContext,
                      CERT_NAME_FRIENDLY_DISPLAY_TYPE,
                      0,
                      NULL,
                      alias,
                      name_count);
  _tcscat(alias, _T(" [SerialNumber:"));

  // シリアル番号を追加していく
  // （little endianで格納されているので、高位バイトから低位バイトに向けて出力する）
  _TCHAR* work = alias + _tcslen(alias);
  for (int i = static_cast<int>(pContext->pCertInfo->SerialNumber.cbData) - 1; i >= 0 ; --i)
  {
    _stprintf(work, _T("%02x"), pContext->pCertInfo->SerialNumber.pbData[i]);
    work += 2;
  }
  _tcscat(work, _T("]"));
  assert(_tcslen(alias) == total_count - 1);

  return alias;
}

/**
 * コンストラクタ
 */
KeyStore::KeyStore(HCERTSTORE hCert) : hCertStore(hCert),
                             elementCount(0),
                             elements(NULL)
{
  loadAllCertificates();
}

/**
 * デストラクタ
 */
KeyStore::~KeyStore()
{
  freeAllCertificates();
  BOOL closed = ::CertCloseStore(this->hCertStore, CERT_CLOSE_STORE_CHECK_FLAG);
  if (! closed)
  {
#ifdef DEBUG
    DWORD lastError = GetLastError();
    OutputDebugString(_T("Failed to close store"));
#endif
  }
}


/**
 * 全証明書をロードする
 */
int
KeyStore::loadAllCertificates()
{
  if (! this->hCertStore)
  {
    return 0;
  }

  int size = 16;
  const int size_inc = 16;
  this->elementCount = 0;
  this->elements = reinterpret_cast<KeyStoreElement*>(::malloc(sizeof(KeyStoreElement) * size));
  if (! this->elements)
  {
    return -1;
  }

  int count = 0;
  PCCERT_CONTEXT pContext = NULL;
  while ((pContext = ::CertEnumCertificatesInStore(this->hCertStore, pContext)) != NULL)
  {
    if (count >= size)
    {
      size += size_inc;
      KeyStoreElement* tmp = reinterpret_cast<KeyStoreElement*>(::realloc(this->elements, sizeof(KeyStoreElement) * size));
      if (! tmp)
      {
        // ここで終了
        break;
      }
      this->elements = tmp;
    }

    // エイリアス名を決定する
    this->elements[count].alias = createAlias(pContext);

    // 複製を作成する
    // (CertEnumCertificatesInStore()は第2引数として渡したCERT_CONTEXTを解放するため）
    this->elements[count].certContext = ::CertDuplicateCertificateContext(pContext);
    assert(this->elements[count].certContext);

    count++;
  }

  // 取得した証明書の数を設定してリターン
  this->elementCount = count;
  return count;
}

/**
 * 全証明書を解放する
 */
void
KeyStore::freeAllCertificates()
{
  if (! this->hCertStore)
  {
    return;
  }
  for (int i = 0; i < this->elementCount; ++i)
  {
    assert(this->elements);
    ::free(this->elements[i].alias);
    ::CertFreeCertificateContext(this->elements[i].certContext);
  }
  ::free(this->elements);
  this->elementCount = 0;
  this->elements = NULL;
}

/**
 * 指定された証明書チェーンの内容に基づき、byte[][] を作成する
 */
jobjectArray
createEncodedCertificates(JNIEnv* env, PCCERT_CHAIN_CONTEXT chain)
{
  // 最初のチェーンのみを見る
  PCERT_SIMPLE_CHAIN simple_chain = chain->rgpChain[0];
  const DWORD cElement = simple_chain->cElement;
      
  jobjectArray encodedCerts = env->NewObjectArray(static_cast<jsize>(cElement), env->FindClass("[B"), NULL);

  // エンコードされた証明書情報を取得して、byte[][] に格納していく
  for (DWORD i = 0; i < cElement; ++i)
  {
    PCERT_CHAIN_ELEMENT elem = simple_chain->rgpElement[i];
    PCCERT_CONTEXT cert_context = elem->pCertContext;
    
    jbyteArray encoded = env->NewByteArray(static_cast<jsize>(cert_context->cbCertEncoded));
    if (encoded)
    {
      jbyte* p = static_cast<jbyte*>(env->GetPrimitiveArrayCritical(encoded, NULL));
      if (p)
      {
        memcpy(p, cert_context->pbCertEncoded, cert_context->cbCertEncoded);
      }
      env->ReleasePrimitiveArrayCritical(encoded, p, 0);

      env->SetObjectArrayElement(encodedCerts, static_cast<jsize>(i), encoded);
    }
  }
  return encodedCerts;
}
