#include "stdafx.h"
#include "WMIDeviceQuery.h"

#include <WinIoCtl.h>
#include <comutil.h>
#include <Wbemidl.h>
#include <tchar.h>
#include <strsafe.h>
#include <algorithm>
#include <atlconv.h>
#include <ntddndis.h>

#pragma comment (lib, "comsuppw.lib")
#pragma comment (lib, "wbemuuid.lib")


static const char *GetWQLQueryString(WMI_DEVICE_TYPE eDevType)
{
	switch (eDevType)
	{
	case Processor:
		return "SELECT * FROM Win32_Processor WHERE (ProcessorId IS NOT NULL)";
		break;
	case BaseBoard:
		return "SELECT * FROM Win32_BaseBoard WHERE (SerialNumber IS NOT NULL)";
		break;
	case DiskDrive:
		return "SELECT * FROM Win32_DiskDrive WHERE (SerialNumber IS NOT NULL) AND (MediaType LIKE 'Fixed hard disk%')";
		break;
	case NetworkAdapter:
		return "SELECT * FROM Win32_NetworkAdapter WHERE (PHYSICALADAPTER=TRUE) AND (MACAddress IS NOT NULL)";
		break;
	case BIOS:
		return "SELECT * FROM Win32_BIOS WHERE  (SerialNumber IS NOT NULL)";
		break;
	default:
		return NULL;
		break;
	}
}

static bool CheckHardDiskSerialNum(TCHAR *SerialNumber, int nBufLen)
{
	int nLen = _tcslen(SerialNumber);
	if (nLen == 40)	// InterfaceType = "IDE"
	{
		// 需要将16进制编码串转换为字符串
		TCHAR szBuf[32] = {};
		bool bSuc = true;
		for (int i = 0; i < 20; i++)
		{	
			// 将16进制字符转换为高4位
			BYTE b;
			TCHAR ch = SerialNumber[i * 2];
			if ((ch >= '0') && (ch <= '9'))
			{
				b = ch - '0';
			}
			else if ((ch >= 'A') && (ch <= 'F'))
			{
				b = ch - 'A' + 10;
			}
			else if ((ch >= 'a') && (ch <= 'f'))
			{
				b = ch - 'a' + 10;
			}
			else
			{	
				// 非法字符
				bSuc = false;
				break;
			}

			b <<= 4;

			// 将16进制字符转换为低4位
			ch = SerialNumber[i * 2 + 1];
			if ((ch >= '0') && (ch <= '9'))
			{
				b += ch - '0';
			}
			else if ((ch >= 'A') && (ch <= 'F'))
			{
				b += ch - 'A' + 10;
			}
			else if ((ch >= 'a') && (ch <= 'f'))
			{
				b += ch - 'a' + 10;
			}
			else
			{	
				// 非法字符
				bSuc = false;
				break;
			}

			szBuf[i] = b;
		}

		if (bSuc)
		{	
			// 转换成功
			szBuf[20] = 0;
			StringCchCopy(SerialNumber, nBufLen, szBuf);
			nBufLen = _tcslen(SerialNumber);
		}
	}

	// 每2个字符互换位置
	for (int i = 0; i < nBufLen; i += 2)
	{
		std::swap(SerialNumber[i], SerialNumber[i + 1]);
	}

	// 去掉空格
	std::remove(SerialNumber, SerialNumber + _tcslen(SerialNumber) + 1, ' ');

	return true;
}

static bool CheckWMIPropVal(WMI_DEVICE_TYPE eDevType, const char *szPropName, char *szPropVal, int nValSize)
{
	BOOL bOK = true;

	if (eDevType == DiskDrive && stricmp(szPropName, "SerialNumber") == 0)
		bOK = CheckHardDiskSerialNum(szPropVal, nValSize);

	// 去掉空格
	std::remove(szPropVal, szPropVal + _tcslen(szPropVal) + 1, ' ');
	return bOK;
}

// 基于Windows Management Instrumentation（Windows管理规范）
bool QueryWMIDevice(WMI_DEVICE_TYPE eDevType, const char *szPropName, char *pValBuf, int *pBufLen)
{	
	int nTotal = 0;
	
	// 初始化
	//HRESULT hResult = CoInitializeEx(NULL, COINIT_MULTITHREADED);
	HRESULT hResult = CoInitialize(NULL);
	if (FAILED(hResult))
		return false;

	// 设置COM的安全认证级别
	hResult = CoInitializeSecurity(
		NULL,
		-1,
		NULL,
		NULL,
		RPC_C_AUTHN_LEVEL_DEFAULT,
		RPC_C_IMP_LEVEL_IMPERSONATE,
		NULL,
		EOAC_NONE,
		NULL
		);
	
	if (FAILED(hResult) && RPC_E_TOO_LATE != hResult)
	{
		CoUninitialize();
		return false;
	}

	// 获得WMI连接COM接口
	IWbemLocator *pLoc = NULL;
	hResult = CoCreateInstance(
		CLSID_WbemLocator,
		NULL,
		CLSCTX_INPROC_SERVER,
		IID_IWbemLocator,
		reinterpret_cast<LPVOID*>(&pLoc)
		);
	if (FAILED(hResult))
	{
		CoUninitialize();
		return false;
	}

	// 通过连接接口连接WMI的内核对象名"ROOT//CIMV2"
	IWbemServices *pSvc = NULL;
	hResult = pLoc->ConnectServer(
		_bstr_t(L"ROOT\\CIMV2"),
		NULL,
		NULL,
		NULL,
		0,
		NULL,
		NULL,
		&pSvc
		);
	if (FAILED(hResult))
	{
		pLoc->Release();
		CoUninitialize();
		return false;
	}

	// 设置请求代理的安全级别
	hResult = CoSetProxyBlanket(
		pSvc,
		RPC_C_AUTHN_WINNT,
		RPC_C_AUTHZ_NONE,
		NULL,
		RPC_C_AUTHN_LEVEL_CALL,
		RPC_C_IMP_LEVEL_IMPERSONATE,
		NULL,
		EOAC_NONE
		);
	if (FAILED(hResult))
	{
		pSvc->Release();
		pLoc->Release();
		CoUninitialize();
		return false;
	}

	// 通过请求代理来向WMI发送请求
	IEnumWbemClassObject *pEnumerator = NULL;
	hResult = pSvc->ExecQuery(
		bstr_t("WQL"),
		bstr_t(GetWQLQueryString(eDevType)),
		WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
		NULL,
		&pEnumerator
		);
	if (FAILED(hResult))
	{
		pSvc->Release();
		pLoc->Release();
		CoUninitialize();
		return false;
	}

	// 循环枚举所有的结果对象  
	int nRetLen = 0;
	while (pEnumerator)
	{
		IWbemClassObject *pclsObj = NULL;
		ULONG uReturn = 0;

		pEnumerator->Next(
			WBEM_INFINITE,
			1,
			&pclsObj,
			&uReturn
			);

		if (uReturn == 0)
			break;

		// 获取属性值
		VARIANT vtProperty;
		VariantInit(&vtProperty);
		USES_CONVERSION;
		pclsObj->Get(A2W(szPropName), 0, &vtProperty, NULL, NULL);
		pclsObj->Release();

		if (vtProperty.vt == VT_EMPTY || vtProperty.vt == VT_NULL)
			continue;

		char szTmp[256] = {};
		StringCchCopyA(szTmp, sizeof(szTmp)-1, W2A(vtProperty.bstrVal));
		VariantClear(&vtProperty);

		// 对属性值做进一步的处理
		if (CheckWMIPropVal(eDevType, szPropName, szTmp, 256))
		{
			if (strlen(szTmp) + nRetLen +1 >= *pBufLen)
				break;
			
			nTotal++;
			if (nRetLen > 0)
			{
				nRetLen++;
				strcat(pValBuf, ";");
			}

			strcat(pValBuf, szTmp);
			nRetLen += strlen(szTmp);
		}
	} 

	// 释放资源
	pEnumerator->Release();
	pSvc->Release();
	pLoc->Release();
	CoUninitialize();

	return nTotal >0;
}