/* Licensed to Stichting The Commons Conservancy (TCC) under one or more
 * contributor license agreements.  See the AUTHORS file distributed with
 * this work for additional information regarding copyright ownership.
 * TCC licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */



#include "crypto.h"
#include "message.h"

#include <stdio.h>
#include <assert.h>
#include <wincrypt.h>
#include <ncrypt.h>

void
crypto_start(SignTextData* signtext)
{
#if 1
	fprintf(stderr, "crypto_start\n");
#endif

#if 0
	volatile int done = 0;
	while (!done) sleep(1);
	/* inside gdb: set var done = 1 */
#endif

    HANDLE hTimer = NULL;
    LARGE_INTEGER liDueTime;

    liDueTime.QuadPart = -20000000LL;

    // Create an unnamed waitable timer.
    hTimer = CreateWaitableTimer(NULL, TRUE, NULL);
    if (NULL == hTimer)
    {
        assert(0);
    }

    // printf("Waiting for 2 seconds...\n");

    while (1) {

		SignTextCertificate* certificate = NULL;
		SignTextCertificate* certificates = NULL;

        HCERTSTORE sys = CertOpenStore(CERT_STORE_PROV_SYSTEM,
            X509_ASN_ENCODING, 0, CERT_SYSTEM_STORE_CURRENT_USER | CERT_STORE_READONLY_FLAG, L"MY");
        if (!sys) {
            assert(0);
        }

		PCCERT_CONTEXT cert = NULL;
		while ((cert = CertEnumCertificatesInStore(sys, cert)) != NULL) {

			DWORD flags = CRYPT_ACQUIRE_CACHE_FLAG | CRYPT_ACQUIRE_COMPARE_KEY_FLAG | CRYPT_ACQUIRE_SILENT_FLAG | CRYPT_ACQUIRE_PREFER_NCRYPT_KEY_FLAG;
			HCRYPTPROV_OR_NCRYPT_KEY_HANDLE key = 0;
			DWORD spec = 0;
			BOOL freeKey = FALSE;

			CryptAcquireCertificatePrivateKey(cert, flags, NULL, &key, &spec, &freeKey);

			/* ignore anything that doesn't have a key */
			if (!key) {
				continue;
			}

			DWORD size;

			LPWSTR subject;

			size = CertGetNameString(cert, CERT_NAME_FRIENDLY_DISPLAY_TYPE, CERT_NAME_STR_ENABLE_PUNYCODE_FLAG, NULL, NULL, 0);

			subject = malloc(size * sizeof(wchar_t));
			if (!subject) {
				assert(0);
			}

			CertGetNameString(cert, CERT_NAME_FRIENDLY_DISPLAY_TYPE, CERT_NAME_STR_ENABLE_PUNYCODE_FLAG, NULL, subject, size);

			LPWSTR issuer;

			size = CertGetNameString(cert, CERT_NAME_FRIENDLY_DISPLAY_TYPE, CERT_NAME_ISSUER_FLAG | CERT_NAME_STR_ENABLE_PUNYCODE_FLAG, NULL, NULL, 0);

			issuer = malloc(size * sizeof(wchar_t));
			if (!issuer) {
				assert(0);
			}

			CertGetNameString(cert, CERT_NAME_FRIENDLY_DISPLAY_TYPE, CERT_NAME_ISSUER_FLAG | CERT_NAME_STR_ENABLE_PUNYCODE_FLAG, NULL, issuer, size);

			size = swprintf(NULL, 0, L"%s (%s)", subject, issuer);

			LPWSTR name = malloc((size + 1) * sizeof(wchar_t));
			if (!name) {
				assert(0);
			}

			swprintf(name, size + 1, L"%s (%s)", subject, issuer);

			free(subject);
			free(issuer);

			switch (spec)
			{
			case CERT_NCRYPT_KEY_SPEC:
			{
				NCRYPT_PROV_HANDLE prov = 0;
				DWORD type = 0, prov_size = sizeof(prov);
				NCryptGetProperty(key, NCRYPT_PROVIDER_HANDLE_PROPERTY, (PBYTE) & prov, prov_size, &prov_size, 0);
				if (prov)
				{
					prov_size = sizeof(type);
					NCryptGetProperty(prov, NCRYPT_IMPL_TYPE_PROPERTY, (PBYTE) & type, prov_size, &prov_size, 0);
					NCryptFreeObject(prov);
				}

				// NCRYPT_EXPORT_POLICY_PROPERTY
				// NCRYPT_KEY_USAGE_PROPERTY

				SECRET_TYPE secretType = -1;
				PIN_INFO pin_info;
				DWORD pin_info_size = sizeof(PIN_INFO);
				NCryptGetProperty(key, NCRYPT_SCARD_PIN_INFO, (PBYTE)&pin_info, pin_info_size, &pin_info_size, 0);
				if (pin_info_size)
				{
					secretType = pin_info.PinType;
				}

				if (freeKey) {
					NCryptFreeObject(key);
				}

				certificate = signtext_certificate_new(signtext, name, cert, spec, type, secretType, certificate);

				if (!certificates) {
					certificates = certificate;
				}

				break;
			}
#if 0
			case AT_KEYEXCHANGE:
			case AT_SIGNATURE:
			{
				DWORD type = 0;
				DWORD size = sizeof(type);
				CryptGetProvParam(key, PP_IMPTYPE, (PBYTE) & type, &size, 0);

				if (freeKey) {
					CryptReleaseContext(key, 0);
				}

				certificate = signtext_certificate_new(signtext, name, cert, spec, type, certificate);

				if (!certificates) {
					certificates = certificate;
				}

				break;
			}
#endif
			}

		}
		CertCloseStore(sys, 0);

		PostMessage(signtext->hwnd, WM_CRYPTO_UPDATED, sizeof(certificates), (LPARAM)certificates);

		/* Set a timer to wait for 2 seconds. */
		if (!SetWaitableTimer(hTimer, &liDueTime, 0, NULL, NULL, 0))
		{
			//        printf("SetWaitableTimer failed (%d)\n", GetLastError());
			assert(0);
		}

		/* Wait for the timer. */

		if (WaitForSingleObject(hTimer, INFINITE) != WAIT_OBJECT_0) {
			assert(0);
		}

	}

}

int
crypto_certificate_is_hardware(SignTextCertificate* certificate)
{
	switch (certificate->spec)
	{
	case CERT_NCRYPT_KEY_SPEC:
	{
		return ((certificate->type & (NCRYPT_IMPL_HARDWARE_FLAG)));
	}
	case AT_KEYEXCHANGE:
	case AT_SIGNATURE:
	{
		return ((certificate->type & (CRYPT_IMPL_HARDWARE)));
	}
	default:
		assert(0);
		return 0;
	}
}

int
crypto_certificate_is_removable(SignTextCertificate* certificate)
{
	switch (certificate->spec)
	{
	case CERT_NCRYPT_KEY_SPEC:
	{
		return ((certificate->type & (NCRYPT_IMPL_REMOVABLE_FLAG)));
	}
	case AT_KEYEXCHANGE:
	case AT_SIGNATURE:
	{
		return ((certificate->type & (CRYPT_IMPL_REMOVABLE)));
	}
	default:
		assert(0);
		return 0;
	}
}

int
crypto_certificate_find(SignTextCertificate* certs, SignTextCertificate* cert)
{
	int found = 0;

	while (certs) {

		if (!crypto_certificate_compare(certs, cert)) {
			found = 1;
			break;
		}

		certs = certs->next;
	}

	return found;
}

size_t crypto_oidcmp(LPSTR oid1, LPSTR oid2, size_t oid2len)
{
	size_t oid1len = strlen(oid1);

	if (oid1len != oid2len)	{
		return (oid1len - oid2len);
	}

	return strcmp(oid1, oid2);
}

LPSTR
crypto_signature2digest(LPSTR signatureOid)
{
	if (!signatureOid) {
		return signatureOid;
	}

	size_t oidlen = strlen(signatureOid);

	/* some crypto libraries like bounceycastle are strict
	 * in what they accept.
	 * 
	 * map signature algorithms to digest algorithms
	 */
	if (!crypto_oidcmp(szOID_RSA_SHA256RSA, signatureOid, oidlen)) {
		return szOID_NIST_sha256;
	}

	if (!crypto_oidcmp(szOID_RSA_SHA384RSA, signatureOid, oidlen)) {
		return szOID_NIST_sha384;
	}

	if (!crypto_oidcmp(szOID_RSA_SHA512RSA, signatureOid, oidlen)) {
		return szOID_NIST_sha512;
	}

	if (!crypto_oidcmp(szOID_ECDSA_SHA256, signatureOid, oidlen)) {
		return szOID_NIST_sha256;
	}

	if (!crypto_oidcmp(szOID_ECDSA_SHA384, signatureOid, oidlen)) {
		return szOID_NIST_sha384;
	}

	if (!crypto_oidcmp(szOID_ECDSA_SHA512, signatureOid, oidlen)) {
		return szOID_NIST_sha512;
	}

	/* default to the signature OID for anything we don't recognise */
	return signatureOid;
}

int
crypto_certificate_compare(SignTextCertificate* c1, SignTextCertificate* c2)
{
	assert(c1);
	assert(c1->pbCertEncoded);
	assert(c2);
	assert(c2->pbCertEncoded);

	if (c1->cbCertEncoded != c2->cbCertEncoded) {
		return c1->cbCertEncoded - c2->cbCertEncoded;
	}

	return memcmp(c1->pbCertEncoded, c2->pbCertEncoded, c1->cbCertEncoded);
}

void
crypto_sign(SignTextInstance* instance)
{

	HCERTSTORE sys = CertOpenStore(CERT_STORE_PROV_SYSTEM,
		X509_ASN_ENCODING, 0, CERT_SYSTEM_STORE_CURRENT_USER | CERT_STORE_READONLY_FLAG, L"MY");
	if (!sys) {
		assert(0);
	}

	PCCERT_CONTEXT search;

	if (!(search = CertCreateCertificateContext(
		X509_ASN_ENCODING,
		instance->certificate->pbCertEncoded,
		instance->certificate->cbCertEncoded)))
	{
		assert(0);
	}

	PCCERT_CONTEXT cert = NULL;
	if (!(cert = CertFindCertificateInStore(sys, X509_ASN_ENCODING, 0, CERT_FIND_EXISTING, search, NULL))) {
		assert(0);
	}

	CertFreeCertificateContext(search);

	NCRYPT_HANDLE key;
	DWORD dwKeySpec = 0;
	if (!CryptAcquireCertificatePrivateKey(cert, CRYPT_ACQUIRE_CACHE_FLAG | CRYPT_ACQUIRE_COMPARE_KEY_FLAG | CRYPT_ACQUIRE_SILENT_FLAG | CRYPT_ACQUIRE_ONLY_NCRYPT_KEY_FLAG, NULL, &key, &dwKeySpec, NULL))
	{
		signtext_lasterror();

		CertCloseStore(sys, 0);

		PostMessage(instance->signtext->hwnd, WM_CRYPTO_NOTSIGNED, 0, 0);

		/* no need for the cert any more, forget it */
		signtext_certificate_unref(instance->certificate);

		instance->certificate = NULL;

		return;
	}

	switch (instance->certificate->secretType) {
	case AlphaNumericPinType:

		if (NCryptSetProperty(key, NCRYPT_PIN_PROPERTY, (BYTE*)instance->pinBuffer, instance->pinLen, 0) != ERROR_SUCCESS)
		{
			instance->is_signing = FALSE;

			CertCloseStore(sys, 0);

			/* wrong pin fails here */
			MessageBox(instance->signtext->hwnd, TEXT("The PIN was not accepted. Please enter it again carefully.\r\nEntering the PIN incorrectly multiple times will lock you out."), TEXT("Signing Error"), MB_OK);

			PostMessage(instance->signtext->hwnd, WM_CRYPTO_NOTSIGNED, 0, 0);

			/* no need for the cert any more, forget it */
			signtext_certificate_unref(instance->certificate);

			instance->certificate = NULL;

			return;
		}

		break;
	case ExternalPinType:
		break;
	case ChallengeResponsePinType:
		break;
	case EmptyPinType:
		break;
	default:
		break;
	}

	/* no need for the cert any more, forget it */
	signtext_certificate_unref(instance->certificate);

	instance->certificate = NULL;

	CCRYPT_OID_INFO* info;
	info = CryptFindOIDInfo(CRYPT_OID_INFO_OID_KEY, cert->pCertInfo->SignatureAlgorithm.pszObjId, 0);

	CRYPT_ALGORITHM_IDENTIFIER hashAlgorithm;
	memset(&hashAlgorithm, 0, sizeof(hashAlgorithm));

	if (instance->digestAlgorithm) {
		hashAlgorithm.pszObjId = instance->digestAlgorithm;
	}
	else {
		hashAlgorithm.pszObjId = crypto_signature2digest(cert->pCertInfo->SignatureAlgorithm.pszObjId);
	}

	CRYPT_SIGN_MESSAGE_PARA signParam;
	memset(&signParam, 0, sizeof(signParam));
	signParam.cbSize = sizeof(signParam);
	signParam.dwMsgEncodingType = PKCS_7_ASN_ENCODING;
	signParam.pSigningCert = cert;
	signParam.cMsgCert = 1;
	signParam.rgpMsgCert = &cert;
	signParam.HashAlgorithm = hashAlgorithm;
	signParam.dwFlags = CRYPT_MESSAGE_SILENT_KEYSET_FLAG;


	DWORD cbEncodedBlob = 0;
	const BYTE* rgpbToBeSigned = (BYTE*)instance->in_buffer;
	DWORD rgcbToBeSigned = (DWORD)instance->in_length;

	if (!CryptSignMessage(&signParam, FALSE, 1, &rgpbToBeSigned, &rgcbToBeSigned, NULL, &cbEncodedBlob))
	{
		signtext_lasterror();

		CertCloseStore(sys, 0);

		PostMessage(instance->signtext->hwnd, WM_CRYPTO_NOTSIGNED, 0, 0);

		return;
	}

	BYTE* pbEncodedBlob = (BYTE*)malloc(cbEncodedBlob);
	if (!CryptSignMessage(&signParam, FALSE , 1, &rgpbToBeSigned, &rgcbToBeSigned, pbEncodedBlob, &cbEncodedBlob))
	{
		signtext_lasterror();

		PostMessage(instance->signtext->hwnd, WM_CRYPTO_NOTSIGNED, 0, 0);

		free(pbEncodedBlob);

		instance->is_signing = FALSE;

		CertCloseStore(sys, 0);

		return;
	}
	else
	{

		DWORD cchString;
		CryptBinaryToStringA(pbEncodedBlob, cbEncodedBlob, CRYPT_STRING_BASE64 | CRYPT_STRING_NOCRLF, NULL, &cchString);
		LPSTR pszString = malloc((cchString + 1) * sizeof(char));
		CryptBinaryToStringA(pbEncodedBlob, cbEncodedBlob, CRYPT_STRING_BASE64 | CRYPT_STRING_NOCRLF, pszString, &cchString);

		free(pbEncodedBlob);

		instance->out_buffer = pszString;
		instance->out_length = cchString;

		message_send_response(instance);

		CertCloseStore(sys, 0);
	}

}