// PRTKeyDerivation.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include "pch.h"
#include <iostream>
#include "ntstatus.h"
#include "windows.h"
#include "bcrypt.h"

int main(int argc, char* argv[], char* envp[])
{	
	if (argc < 3) {
		printf("Usage PRTKeyDerivation.exe <32-byte hex secretkey> <24-byte hex context>");
		return -1;
	}
	if (strlen(argv[1]) != 64) {
		printf("Size of secretkey != 32 hex bytes");
		return -1;
	}
	if (strlen(argv[2]) != 48) {
		printf("Size of context != 24 hex bytes");
		return -1;
	}
	UCHAR secret[32]{};

	UCHAR *j = 0;
	for (int i = 0; i < sizeof(secret); i++) {
		sscanf_s((argv[1]) + i * 2, "%02hhX", &j);
		secret[i] = (UCHAR)j;
	}
	UCHAR context[24]{};
	char *endptr;
	for (int  i = 0; i < sizeof(context); i++) {
		sscanf_s((argv[2])+ i * 2, "%02hhX", &j);
		context[i] = (UCHAR)j;
	}

	BCRYPT_ALG_HANDLE hProvider = NULL;
	BCRYPT_KEY_HANDLE hKey = NULL;

	NTSTATUS ret = BCryptOpenAlgorithmProvider(&hProvider, BCRYPT_SP800108_CTR_HMAC_ALGORITHM, MS_PRIMITIVE_PROVIDER, 0);
	
	if (!(BCRYPT_SUCCESS(ret)))
	{
		return -1;
	}
	NTSTATUS ret2 = BCryptGenerateSymmetricKey(hProvider, &hKey, NULL, 0, (PBYTE)secret, 32, 0);
	const char *text = "AzureAD-SecureConversation";
	static
		BCryptBuffer SP800108ParamBuffer[] = {
			{
				26,
				KDF_LABEL,
				(PBYTE)text,
			},
			{
				24,
				KDF_CONTEXT,
				(PBYTE)context,
			},
			{
				sizeof(BCRYPT_SHA256_ALGORITHM),
				KDF_HASH_ALGORITHM,
				(PBYTE) BCRYPT_SHA256_ALGORITHM,
			}
	};

	BCryptBufferDesc desc = { 0, 3, SP800108ParamBuffer };
	PBYTE   DerivedKey = NULL;
	DWORD   DerivedKeyLength = 0x20;
	DerivedKey = (PBYTE)HeapAlloc(GetProcessHeap(), 0, DerivedKeyLength);
	ULONG result = NULL;
	NTSTATUS ret3 = BCryptKeyDerivation(hKey, &desc, DerivedKey, DerivedKeyLength, &result, 0);
	int i = 0;

	for (i = 0; i < result; i++)
	{
		printf("%02X", DerivedKey[i]);
	}
}