#include "StdAfx.h"
#include ".\regquery.h"

#define VERSION_KEY		"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion"
#define OS_OPTIONS_KEY	"SYSTEM\\CurrentControlSet\\Control\\ProductOptions"

RegQuery::RegQuery(void)
{
}

RegQuery::~RegQuery(void)
{
}

bool RegQuery::GetOSVersion(char* szServer, char** lplpVersion)
{
	HKEY hRemote, hVersion, hOptions;
	DWORD dwSize; 
	char* szType = NULL;
	char* szVersion = NULL;
	char* szBuild = NULL;
	char* szServicePack = NULL;
	long nResult;

	char test[256];

	memset(test, 0, 256);
	_snprintf(test, 256, "\\\\%s", szServer);
	nResult = RegConnectRegistry(test, HKEY_LOCAL_MACHINE, &hRemote);
	if (ERROR_SUCCESS != nResult)
	{
		ErrorHandler("RegConnectRegistry", "GetOSVersion", nResult);
		return false;
	}

	nResult = RegOpenKey(hRemote, VERSION_KEY, &hVersion);
	if (ERROR_SUCCESS != nResult)
	{
		RegCloseKey(hRemote);
		ErrorHandler("RegOpenKey(Version)", "GetOSVersion", nResult);
		return false;
	}

	nResult = RegOpenKey(hRemote, OS_OPTIONS_KEY, &hOptions);
	if (ERROR_SUCCESS != nResult)
	{
		RegCloseKey(hVersion);
		RegCloseKey(hRemote);
		ErrorHandler("RegOpenKey(Options)", "GetOSVersion", nResult);
		return false;
	}

	dwSize = 0;
	RegQueryValueEx(hOptions, "ProductType", NULL, NULL, NULL, &dwSize);
	
	// Allocate memory for the buffer
	if (dwSize > 0)
	{
		szType = (char*)malloc(dwSize);
		nResult = RegQueryValueEx(hOptions, "ProductType", NULL, NULL, (BYTE*)szType, &dwSize);
		if (ERROR_SUCCESS != nResult)
		{
			RegCloseKey(hVersion);
			RegCloseKey(hOptions);
			RegCloseKey(hRemote);
			ErrorHandler("RegQueryValueEx(ProductType)", "GetOSVersion", nResult);
			return false;
		}
	}

	dwSize = 0;
	RegQueryValueEx(hVersion, "CurrentVersion", NULL, NULL, NULL, &dwSize);

	// Allocate memory for the buffer
	if (dwSize > 0)
	{
		szVersion = (char*)malloc(dwSize);
		nResult = RegQueryValueEx(hVersion, "CurrentVersion", NULL, NULL, (BYTE*)szVersion, &dwSize);
		if (ERROR_SUCCESS != nResult)
		{
			RegCloseKey(hVersion);
			RegCloseKey(hOptions);
			RegCloseKey(hRemote);
			ErrorHandler("RegQueryValueEx(CurrentVersion)", "GetOSVersion", nResult);
			return false;
		}
	}

	dwSize = 0;
	RegQueryValueEx(hVersion, "CurrentBuildNumber", NULL, NULL, NULL, &dwSize);

	// Allocate memory for the buffer
	if (dwSize > 0)
	{
		szBuild = (char*)malloc(dwSize);
		nResult = RegQueryValueEx(hVersion, "CurrentBuildNumber", NULL, NULL, (BYTE*)szBuild, &dwSize);
		if (ERROR_SUCCESS != nResult)
		{
			RegCloseKey(hVersion);
			RegCloseKey(hOptions);
			RegCloseKey(hRemote);
			ErrorHandler("RegQueryValueEx(BuildNumber)", "GetOSVersion", nResult);
			return false;
		}
	}

	dwSize = 0;
	RegQueryValueEx(hVersion, "CSDVersion", NULL, NULL, NULL, &dwSize);

	// Allocate memory for the buffer
	if (dwSize > 0)
	{
		szServicePack = (char*)malloc(dwSize);
		nResult = RegQueryValueEx(hVersion, "CSDVersion", NULL, NULL, (BYTE*)szServicePack, &dwSize);
		if (ERROR_SUCCESS != nResult)
		{
			RegCloseKey(hVersion);
			RegCloseKey(hOptions);
			RegCloseKey(hRemote);
			ErrorHandler("RegQueryValueEx(CSDVersion)", "GetOSVersion", nResult);
			return false;
		}
	}

	// Should have all the relevant information - build a string and return it to the caller
	const char* lpszFormat;
	int nArgSize;

	if (szServicePack == NULL)
	{
		lpszFormat = "Microsoft Windows %s %s (Build %s)";
		nArgSize = _scprintf(lpszFormat, MapRegValueToOS(szVersion), MapRegValueToOS(szType), szBuild);
		if (nArgSize > 0)
		{
			*lplpVersion = (char*)malloc(nArgSize + 1);
			memset(*lplpVersion, 0, nArgSize + 1);
			_snprintf(*lplpVersion, nArgSize, lpszFormat, MapRegValueToOS(szVersion), MapRegValueToOS(szType), szBuild);
		}
		else
		{
			*lplpVersion = NULL;
		}
	}
	else
	{
		lpszFormat = "Microsoft Windows %s %s %s (Build %s)";
		nArgSize = _scprintf(lpszFormat, MapRegValueToOS(szVersion), MapRegValueToOS(szType), szServicePack, szBuild);
		if (nArgSize > 0)
		{
			*lplpVersion = (char*)malloc(nArgSize + 1);
			memset(*lplpVersion, 0, nArgSize + 1);
			_snprintf(*lplpVersion, nArgSize, lpszFormat, MapRegValueToOS(szVersion), MapRegValueToOS(szType), szServicePack, szBuild);
		}
		else
		{
			*lplpVersion = NULL;
		}
	}

	if (szType != NULL)
		free(szType);
	if (szVersion != NULL)
		free(szVersion);
	if (szBuild != NULL)
		free(szBuild);
	if (szServicePack != NULL)
		free(szServicePack);

	RegCloseKey(hVersion);
	RegCloseKey(hOptions);
	RegCloseKey(hRemote);

	return true;
}

LPCTSTR RegQuery::MapRegValueToOS(LPCTSTR szRegValue)
{
	if (strncmp(szRegValue, "3.51", strlen(szRegValue)) == 0)
		return "3.51";
	else if (strncmp(szRegValue, "4.0", strlen(szRegValue)) == 0)
		return "4.0";
	else if (strncmp(szRegValue, "5.0", strlen(szRegValue)) == 0)
		return "2000";
	else if (strncmp(szRegValue, "5.1", strlen(szRegValue)) == 0)
		return "XP";
	else if (strncmp(szRegValue, "5.2", strlen(szRegValue)) == 0)
		return "2003";
	else if (strncmp(szRegValue, "6.0", strlen(szRegValue)) == 0)
		return "Vista";
	else if (strncmp(szRegValue, "ServerNT", strlen(szRegValue)) == 0)
		return "Server";
	else if (strncmp(szRegValue, "WinNT", strlen(szRegValue)) == 0)
		return "Professional";
	else
		return "Unknown";
}

