//
//	GetQuerySystemInfoBuffer
//
//	Wrapper around ZwQuerySystemInformation
//
//	www.catch22.net
//
#include <windows.h>

typedef ULONG (WINAPI *ZWQUERYINFO)(ULONG, PVOID, ULONG, PULONG));

PVOID GetNativeCall(CHAR *szFunction)
{
	return (PVOID)GetProcAddress(GetModuleHandle("ntdll.dll"), szFunction);
}

//
//	Wrapper routine around ZwQuerySystemInformation
//	Return system information into a fixed-sized buffer
//
BOOL QuerySystemInfo(ULONG InfoClass, PVOID SysInfo, ULONG MemSize)
{
	ZWQUERYINFO ZwQuerySystemInformation;

	if((ZwQuerySystemInformation = (ZWQUERYINFO)GetNativeCall("ZwQuerySystemInformation")) == 0)
		return FALSE;

	return !ZwQuerySystemInformation(InfoClass, SysInfo, MemSize, &MemSize);
}

//
//	Wrapper routine around ZwQuerySystemInformation
//	Allocate and return system info into a malloc'd buffer.
//	Caller must free using FreeQuerySystemInfoBuffer
//
BOOL GetQuerySystemInfoBuffer(ULONG InfoClass, PVOID *SysInfo, PULONG MemSize)
{
	BYTE   *buf = 0;
	ULONG	len = 0;
	ULONG   status;

	ZWQUERYINFO ZwQuerySystemInformation;

	if((ZwQuerySystemInformation = (ZWQUERYINFO)GetNativeCall("ZwQuerySystemInformation")) == 0)
		return FALSE;

	do 
	{
		if(buf != 0)
			free(buf);

		len += 0x10000;

		if((buf = (BYTE *)malloc(len)) == 0)
			return FALSE;

		status = ZwQuerySystemInformation(InfoClass, buf, len, &len);
		
	} while(status == 0xC0000004L);	// STATUS_INFO_LENGTH_MISMATCH

	if(status == 0)
	{
		if(MemSize) *MemSize = len;
		*SysInfo = buf;
		return TRUE;
	}
	else
	{
		free(buf);
		return FALSE;
	}
}

VOID FreeQuerySystemInfoBuffer(PVOID buf)
{
	if(buf) free(buf);
}