#include <cng_kex.h>
#include <dh.h>

typedef struct dh_st DH;

char	*strdelim(char **);
char *strsep(char **stringp, const char *delim);
// Prime in big-endian format.  (Group 2 prime from RFC 2409)
static
const
BYTE OakleyGroup2P[] =
{
	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xC9, 0x0F,
	0xDA, 0xA2, 0x21, 0x68, 0xC2, 0x34, 0xC4, 0xC6, 0x62, 0x8B,
	0x80, 0xDC, 0x1C, 0xD1, 0x29, 0x02, 0x4E, 0x08, 0x8A, 0x67,
	0xCC, 0x74, 0x02, 0x0B, 0xBE, 0xA6, 0x3B, 0x13, 0x9B, 0x22,
	0x51, 0x4A, 0x08, 0x79, 0x8E, 0x34, 0x04, 0xDD, 0xEF, 0x95,
	0x19, 0xB3, 0xCD, 0x3A, 0x43, 0x1B, 0x30, 0x2B, 0x0A, 0x6D,
	0xF2, 0x5F, 0x14, 0x37, 0x4F, 0xE1, 0x35, 0x6D, 0x6D, 0x51,
	0xC2, 0x45, 0xE4, 0x85, 0xB5, 0x76, 0x62, 0x5E, 0x7E, 0xC6,
	0xF4, 0x4C, 0x42, 0xE9, 0xA6, 0x37, 0xED, 0x6B, 0x0B, 0xFF,
	0x5C, 0xB6, 0xF4, 0x06, 0xB7, 0xED, 0xEE, 0x38, 0x6B, 0xFB,
	0x5A, 0x89, 0x9F, 0xA5, 0xAE, 0x9F, 0x24, 0x11, 0x7C, 0x4B,
	0x1F, 0xE6, 0x49, 0x28, 0x66, 0x51, 0xEC, 0xE6, 0x53, 0x81,
	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF
};

// Generator in big-endian format.  (Group 2 generator from RFC 2409)
static
const
BYTE OakleyGroup2G[] =
{
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02
};


// Prime in big-endian format.  (Group 14 prime from RFC 3526)
static
const
BYTE OakleyGroup14P[] =
{
	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xC9, 0x0F, 0xDA, 0xA2, 0x21, 0x68, 0xC2, 0x34, 0xC4, 0xC6, 0x62, 0x8B, 0x80, 0xDC, 0x1C, 0xD1,
	0x29, 0x02, 0x4E, 0x08, 0x8A, 0x67, 0xCC, 0x74, 0x02, 0x0B, 0xBE, 0xA6, 0x3B, 0x13, 0x9B, 0x22, 0x51, 0x4A, 0x08, 0x79, 0x8E, 0x34, 0x04, 0xDD,
	0xEF, 0x95, 0x19, 0xB3, 0xCD, 0x3A, 0x43, 0x1B, 0x30, 0x2B, 0x0A, 0x6D, 0xF2, 0x5F, 0x14, 0x37, 0x4F, 0xE1, 0x35, 0x6D, 0x6D, 0x51, 0xC2, 0x45,
	0xE4, 0x85, 0xB5, 0x76, 0x62, 0x5E, 0x7E, 0xC6, 0xF4, 0x4C, 0x42, 0xE9, 0xA6, 0x37, 0xED, 0x6B, 0x0B, 0xFF, 0x5C, 0xB6, 0xF4, 0x06, 0xB7, 0xED,
	0xEE, 0x38, 0x6B, 0xFB, 0x5A, 0x89, 0x9F, 0xA5, 0xAE, 0x9F, 0x24, 0x11, 0x7C, 0x4B, 0x1F, 0xE6, 0x49, 0x28, 0x66, 0x51, 0xEC, 0xE4, 0x5B, 0x3D,
	0xC2, 0x00, 0x7C, 0xB8, 0xA1, 0x63, 0xBF, 0x05, 0x98, 0xDA, 0x48, 0x36, 0x1C, 0x55, 0xD3, 0x9A, 0x69, 0x16, 0x3F, 0xA8, 0xFD, 0x24, 0xCF, 0x5F,
	0x83, 0x65, 0x5D, 0x23, 0xDC, 0xA3, 0xAD, 0x96, 0x1C, 0x62, 0xF3, 0x56, 0x20, 0x85, 0x52, 0xBB, 0x9E, 0xD5, 0x29, 0x07, 0x70, 0x96, 0x96, 0x6D,
	0x67, 0x0C, 0x35, 0x4E, 0x4A, 0xBC, 0x98, 0x04, 0xF1, 0x74, 0x6C, 0x08, 0xCA, 0x18, 0x21, 0x7C, 0x32, 0x90, 0x5E, 0x46, 0x2E, 0x36, 0xCE, 0x3B,
	0xE3, 0x9E, 0x77, 0x2C, 0x18, 0x0E, 0x86, 0x03, 0x9B, 0x27, 0x83, 0xA2, 0xEC, 0x07, 0xA2, 0x8F, 0xB5, 0xC5, 0x5D, 0xF0, 0x6F, 0x4C, 0x52, 0xC9,
	0xDE, 0x2B, 0xCB, 0xF6, 0x95, 0x58, 0x17, 0x18, 0x39, 0x95, 0x49, 0x7C, 0xEA, 0x95, 0x6A, 0xE5, 0x15, 0xD2, 0x26, 0x18, 0x98, 0xFA, 0x05, 0x10,
	0x15, 0x72, 0x8E, 0x5A, 0x8A, 0xAC, 0xAA, 0x68, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF
};

// Generator in big-endian format.  (Group 14 generator from RFC 3526)
static
const
BYTE OakleyGroup14G[] =
{
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
};



DH *
cng_dh_new_group_asc(const char *gen, const char *modulus)
{

	DH *dh;
	BIGNUM * p = NULL;
	BIGNUM * g = NULL;

	if (BN_hex2bn(&p, modulus) == 0 ||
		BN_hex2bn(&g, gen) == 0) {
		return NULL;
	}
	dh = cng_dh_new_group(g, p);
	BN_free(p);
	BN_free(g);
	return dh;

}



int cng_kex_supported(void)
{
	return IsWindows8Point1OrGreater() == FALSE ? 0 : 1;
}

DH * cng_setup_group1(void)
{
	PBYTE pBlob = NULL;
	DWORD keysize = sizeof(OakleyGroup2P);
	DWORD cbBlob = 0;

	pBlob = (PBYTE)malloc(sizeof(CNG_DH_CTX) + (3 * keysize));
	memset(pBlob, 0, sizeof(CNG_DH_CTX) + (3 * keysize));

	CNG_DH_CTX * pCtx = (CNG_DH_CTX *)pBlob;

	pCtx->size = keysize;
	pCtx->p = pBlob + sizeof(CNG_DH_CTX);
	pCtx->g = pCtx->p + keysize;
	pCtx->pub = pCtx->g + keysize;

	memcpy(pCtx->p, OakleyGroup2P, sizeof(OakleyGroup2P));
	memcpy(pCtx->g, OakleyGroup2G, sizeof(OakleyGroup2G));

	return (DH *)pCtx;
}

DH * cng_setup_group14(void)
{
	PBYTE pBlob = NULL;
	DWORD keysize = sizeof(OakleyGroup14P);
	DWORD cbBlob = 0;

	pBlob = (PBYTE)malloc(sizeof(CNG_DH_CTX) + (3 * keysize));
	memset(pBlob, 0, sizeof(CNG_DH_CTX) + (3 * keysize));

	CNG_DH_CTX * pCtx = (CNG_DH_CTX *)pBlob;

	pCtx->size = keysize;
	pCtx->p = pBlob + sizeof(CNG_DH_CTX);
	pCtx->g = pCtx->p + keysize;
	pCtx->pub = pCtx->g + keysize;

	memcpy(pCtx->p, OakleyGroup14P, sizeof(OakleyGroup14P));
	memcpy(pCtx->g, OakleyGroup14G, sizeof(OakleyGroup14G));

	return (DH *)pCtx;
}



//function to reverse the order of a byte array
static void 
cng_dh_swap_bytes(void *pv, size_t n)
{
	char *p = (char*)pv;
	size_t lo, hi;
	for (lo = 0, hi = n - 1; hi>lo; lo++, hi--)
	{
		char tmp = p[lo];
		p[lo] = p[hi];
		p[hi] = tmp;
	}
}


// import a bignum public key 
BCRYPT_KEY_HANDLE 
cng_dh_set_remote_pub_key(DH * dh, BIGNUM * b)
{
	BCRYPT_KEY_HANDLE hPub = NULL;
	CNG_DH_CTX				*pCtx = (CNG_DH_CTX *)dh;

	if (BN_num_bytes(b) > pCtx->size)
		return NULL;

	DWORD cbBlob = sizeof(BCRYPT_DH_KEY_BLOB) + (pCtx->size * 3);

	PBYTE pBlob = (PBYTE)malloc(cbBlob);

	BCRYPT_DH_KEY_BLOB *pKeyBlob = (BCRYPT_DH_KEY_BLOB *)pBlob;
	pKeyBlob->cbKey = pCtx->size;
	pKeyBlob->dwMagic = BCRYPT_DH_PUBLIC_MAGIC;
	PBYTE pModulus = pBlob + sizeof(BCRYPT_DH_KEY_BLOB);
	PBYTE pGenerator = pModulus + pKeyBlob->cbKey;
	PBYTE pPublic = pGenerator + pKeyBlob->cbKey;

	memcpy(pModulus, pCtx->p, pCtx->size);
	memcpy(pGenerator, pCtx->g, pCtx->size);
	memset(pPublic, 0, pCtx->size);

	BN_bn2bin(b, pPublic + pCtx->size - BN_num_bytes(b));


	HRESULT Status = 0;

	if (S_OK != (Status = BCryptImportKeyPair(pCtx->hAlg, NULL, BCRYPT_DH_PUBLIC_BLOB, &hPub, pBlob, cbBlob, 0)))
		goto cleanup;

cleanup:
	if (pBlob) free(pBlob);

	return hPub;
}


// generate KEX private/public keys for this side
int
cng_dh_gen_key(DH *dh, int need)
{
	HRESULT					Status;
	BCRYPT_DH_PARAMETER_HEADER  *DhParamHdrPointer = NULL;
	DWORD					DhParamBlobLength = 0;
	PBYTE					DhParamBlob = NULL;
	BCRYPT_ALG_HANDLE       hAlg = NULL;
	PBYTE                   pBlob = NULL;
	BCRYPT_KEY_HANDLE       hPriv = NULL;
	CNG_DH_CTX				*pCtx = (CNG_DH_CTX *)dh;
	DWORD				    cbBlob = 0;


	DhParamBlobLength = sizeof(BCRYPT_DH_PARAMETER_HEADER) + (pCtx->size * 2);

	if (NULL == (DhParamBlob = (PBYTE)malloc(DhParamBlobLength)))
		return -1;

	DhParamHdrPointer = (BCRYPT_DH_PARAMETER_HEADER *)DhParamBlob;
	DhParamHdrPointer->cbLength = DhParamBlobLength;
	DhParamHdrPointer->cbKeyLength = pCtx->size;
	DhParamHdrPointer->dwMagic = BCRYPT_DH_PARAMETERS_MAGIC;

	memcpy(DhParamBlob + sizeof(BCRYPT_DH_PARAMETER_HEADER), pCtx->p, pCtx->size);
	memcpy(DhParamBlob + sizeof(BCRYPT_DH_PARAMETER_HEADER) + pCtx->size, pCtx->g, pCtx->size);

	if (S_OK != (Status = BCryptOpenAlgorithmProvider(&hAlg, BCRYPT_DH_ALGORITHM, NULL, 0)))
		goto error;


	if (S_OK != (Status = BCryptGenerateKeyPair(hAlg, &hPriv, pCtx->size * 8, 0)))
		goto error;


	if (S_OK != (Status = BCryptSetProperty(hPriv, BCRYPT_DH_PARAMETERS, DhParamBlob, DhParamBlobLength, 0)))
		goto error;

	if (S_OK != (Status = BCryptFinalizeKeyPair(hPriv, 0)))
		goto error;

	if (S_OK != (Status = BCryptExportKey(hPriv, NULL, BCRYPT_DH_PUBLIC_BLOB, NULL, 0, &cbBlob, 0)))
		goto error;

	if (NULL == (pBlob = (PBYTE)malloc(cbBlob)))
	{
		Status = STATUS_NO_MEMORY;
		goto error;
	}

	if (S_OK != (Status = BCryptExportKey(hPriv, NULL, BCRYPT_DH_PUBLIC_BLOB, pBlob, cbBlob, &cbBlob, 0)))
		goto error;



	BCRYPT_DH_KEY_BLOB *pKeyBlob = (BCRYPT_DH_KEY_BLOB *)pBlob;
	PBYTE pModulus = pBlob + sizeof(BCRYPT_DH_KEY_BLOB);
	PBYTE pGenerator = pModulus + pKeyBlob->cbKey;
	PBYTE pPublic = pGenerator + pKeyBlob->cbKey;

	pCtx->hAlg = hAlg;
	pCtx->hPrivate = hPriv;
	memcpy(pCtx->pub, pPublic, pCtx->size);

	return 0;
error:

	return -1;
}



BIGNUM * 
cng_dh_export_secret(DH * dh, BCRYPT_SECRET_HANDLE hSecret)
{
	HRESULT hr = S_OK;
	CNG_DH_CTX				*pCtx = (CNG_DH_CTX *)dh;
	PBYTE pSecret = (PBYTE)malloc(pCtx->size);
	DWORD pSecret_len = pCtx->size;
	DWORD size = pCtx->size;
	BIGNUM * bn = NULL;

	hr = BCryptDeriveKey(hSecret, L"TRUNCATE", NULL, pSecret, pSecret_len, &size, 0);
	if (S_OK == hr)
	{
		cng_dh_swap_bytes(pSecret, size);
		bn = BN_bin2bn(pSecret, size, NULL);
		memset(pSecret, 0, size);
		free(pSecret);
	}
	return bn;
}

// calculate and return secret as a bignum
BIGNUM *
cng_dh_get_secret(DH * dh, BIGNUM * remote_pub)
{
	HRESULT hr = S_OK;
	CNG_DH_CTX				*pCtx = (CNG_DH_CTX *)dh;
	PBYTE pSecret = (PBYTE)malloc(pCtx->size);
	DWORD pSecret_len = pCtx->size;
	DWORD size = pCtx->size;
	BIGNUM * bn = NULL;
	BCRYPT_SECRET_HANDLE hSecret;

	BCRYPT_KEY_HANDLE  hRemotePub = cng_dh_set_remote_pub_key(dh, remote_pub);
	if (hRemotePub != NULL)
	{

		if (S_OK == (hr = BCryptSecretAgreement(pCtx->hPrivate, hRemotePub, &hSecret, 0)))
		{

			hr = BCryptDeriveKey(hSecret, L"TRUNCATE", NULL, pSecret, pSecret_len, &size, 0);
			if (S_OK == hr)
			{
				cng_dh_swap_bytes(pSecret, size);
				bn = BN_bin2bn(pSecret, size, NULL);
				memset(pSecret, 0, size);
				free(pSecret);
			}
			BCryptDestroySecret(hSecret);
		}
	}
	return bn;
}


// return our public key as a bignum (for sending to the other side)
BIGNUM *
cng_dh_get_local_pub_key(DH * dh)
{
	CNG_DH_CTX				*pCtx = (CNG_DH_CTX *)dh;
	BIGNUM * bn = NULL;

	bn = BN_bin2bn(pCtx->pub, pCtx->size, NULL);

	return bn;
}

BIGNUM *
cng_dh_get_p(DH * dh)
{
	CNG_DH_CTX				*pCtx = (CNG_DH_CTX *)dh;
	BIGNUM * bn = NULL;

	bn = BN_bin2bn(pCtx->p, pCtx->size, NULL);
	return bn;
}

BIGNUM *
cng_dh_get_g(DH * dh)
{
	CNG_DH_CTX				*pCtx = (CNG_DH_CTX *)dh;
	BIGNUM * bn = NULL;

	bn = BN_bin2bn(pCtx->g, pCtx->size, NULL);
	return bn;
}



int
cng_dh_pub_is_valid(DH * dh, BIGNUM *dh_pub)
{
	int i;
	int n = BN_num_bits(dh_pub);
	int bits_set = 0;
	BIGNUM *tmp;
	BIGNUM *p = cng_dh_get_p(dh);

	if (NULL == p) {
		logit("invalid DH prime value");
		return 0;
	}


	if (dh_pub->neg) {
		logit("invalid public DH value: negative");
		return 0;
	}
	if (BN_cmp(dh_pub, BN_value_one()) != 1) {	/* pub_exp <= 1 */
		logit("invalid public DH value: <= 1");
		return 0;
	}

	if ((tmp = BN_new()) == NULL) {
		error("%s: BN_new failed", __FUNCTION__);
		return 0;
	}
	if (!BN_sub(tmp, p, BN_value_one()) ||
		BN_cmp(dh_pub, tmp) != -1) {		/* pub_exp > p-2 */
		BN_clear_free(tmp);
		logit("invalid public DH value: >= p-1");
		return 0;
	}
	BN_clear_free(tmp);

	for (i = 0; i <= n; i++)
		if (BN_is_bit_set(dh_pub, i))
			bits_set++;
	debug2("bits set: %d/%d", bits_set, BN_num_bits(p));

	/* if g==2 and bits_set==1 then computing log_g(dh_pub) is trivial */
	if (bits_set > 1)
	{
		BN_clear_free(p);
		return 1;
	}

	logit("invalid public DH value (%d/%d)", bits_set, BN_num_bits(p));
	BN_clear_free(p);
	return 0;
}


DH *
cng_dh_new_group(BIGNUM *gen, BIGNUM *modulus)
{

	PBYTE pBlob = NULL;
	DWORD keysize = BN_num_bytes(modulus);;
	DWORD cbBlob = 0;


	pBlob = (PBYTE)malloc(sizeof(CNG_DH_CTX) + (3 * keysize));
	memset(pBlob, 0, sizeof(CNG_DH_CTX) + (3 * keysize));

	CNG_DH_CTX * pCtx = (CNG_DH_CTX *)pBlob;

	pCtx->size = keysize;
	pCtx->p = pBlob + sizeof(CNG_DH_CTX);
	pCtx->g = pCtx->p + keysize;
	pCtx->pub = pCtx->g + keysize;

	BN_bn2bin(gen, pCtx->g + keysize - BN_num_bytes(gen));
	BN_bn2bin(modulus, pCtx->p + keysize - BN_num_bytes(modulus));

	return (DH *)pCtx;
}



void cng_dh_free_context(DH * dh)
{
	CNG_DH_CTX * pCtx = (CNG_DH_CTX*)dh;

	if (pCtx->hAlg)
		BCryptCloseAlgorithmProvider(pCtx->hAlg, 0);

	if (pCtx->hPrivate)
		BCryptDestroyKey(pCtx->hPrivate);

	free(pCtx);

}