#include "StdAfx.h"
#include <winioctl.h>
#include <strsafe.h>
#include <setupapi.h>
#include <ntddndis.h>
#include <ntdddisk.h>
#include <algorithm>
#include "WDK_DeviceQuery.h"

#ifdef __cplusplus
extern "C"
{
#endif

#include <hidsdi.h>

#ifdef __cplusplus
}
#endif

#pragma comment (lib, "Setupapi.lib")
#pragma comment (lib, "hid.lib")

const GUID GUID_QUERYSET[] = {
	// 人机交互设备VID-PID
	{0x4D1E55B2, 0xF16F, 0x11CF, 0x88, 0xCB, 0x00, 0x11, 0x11, 0x00, 0x00, 0x30},

	// 网卡原生MAC地址（包含USB网卡）
	{0xAD498944, 0x762F, 0x11D0, 0x8D, 0xCB, 0x00, 0xC0, 0x4F, 0xC3, 0x35, 0x8C},

	// 网卡原生MAC地址（剔除USB网卡）
	{0xAD498944, 0x762F, 0x11D0, 0x8D, 0xCB, 0x00, 0xC0, 0x4F, 0xC3, 0x35, 0x8C},

	// 硬盘序列号
	{0x53F56307, 0xB6BF, 0x11D0, 0x94, 0xF2, 0x00, 0xA0, 0xC9, 0x1E, 0xFB, 0x8B},
};

// 获取设备的VID-PID
static BOOL WDK_GetVIDPID( TCHAR* DevicePath, TCHAR* szProperty, UINT uSize )
{	
	HANDLE	hDeviceFile;
	BOOL	isOK = FALSE;

	// 获取设备句柄
	hDeviceFile = CreateFile( DevicePath,
		0,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		NULL,
		OPEN_EXISTING,
		0,
		NULL);

	if( hDeviceFile != INVALID_HANDLE_VALUE )
	{	// 获取设备属性
		HIDD_ATTRIBUTES	Attributes;

		Attributes.Size = sizeof(HIDD_ATTRIBUTES);
		isOK = HidD_GetAttributes( hDeviceFile, &Attributes );
		if( isOK )
		{
			if( szProperty != NULL )
			{
				StringCchPrintf( szProperty, uSize, TEXT("%04X-%04X"), Attributes.VendorID, Attributes.ProductID ); 
			}
		}

		CloseHandle( hDeviceFile );
	}

	return isOK;
}

// 获取网卡原生MAC地址
static BOOL WDK_GetMacAddress( TCHAR* DevicePath, TCHAR* szProperty, UINT uSize, BOOL isIncludeUSB )
{
	HANDLE	hDeviceFile;
	BOOL	isOK = FALSE;

	// 剔除虚拟网卡
	if( _tcsnicmp( DevicePath + 4, TEXT("root"), 4 ) == 0 )
	{
		return FALSE;
	}

	if( !isIncludeUSB )
	{	// 剔除USB网卡
		if( _tcsnicmp( DevicePath + 4, TEXT("usb"), 4 ) == 0 )
		{
			return FALSE;
		}
	}

	// 获取设备句柄
	hDeviceFile = CreateFile( DevicePath,
		0,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		NULL,
		OPEN_EXISTING,
		0,
		NULL);

	if( hDeviceFile != INVALID_HANDLE_VALUE )
	{	
		ULONG	dwID;
		BYTE	ucData[8];
		DWORD	dwByteRet;		

		// 获取原生MAC地址
		dwID = OID_802_3_PERMANENT_ADDRESS;
		isOK = DeviceIoControl( hDeviceFile, IOCTL_NDIS_QUERY_GLOBAL_STATS, &dwID, sizeof(dwID), ucData, sizeof(ucData), &dwByteRet, NULL );
		if( isOK )
		{	
			if( szProperty != NULL )
			{	// 将字节数组转换成16进制字符串
				for( DWORD i = 0; i < dwByteRet; i++ )
				{
					StringCchPrintf( szProperty + (i << 1), uSize - (i << 1), TEXT("%02X"), ucData[i] );
				}
			}
		}
	
		CloseHandle( hDeviceFile );
	}

	return isOK;
}

// 获取硬盘序列号（需要系统管理员权限）
static BOOL WDK_GetHarddiskSerialNumberWithAdminRights( TCHAR* DevicePath, TCHAR* szProperty, UINT uSize )
{	
	HANDLE	hDeviceFile;
	BOOL	isOK = FALSE;

	// 获取设备句柄
	hDeviceFile = CreateFile( DevicePath,
		GENERIC_READ | GENERIC_WRITE,	// 注意：需要系统管理员权限才能执行读写操作
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		NULL,
		OPEN_EXISTING,
		0,
		NULL);

	if( hDeviceFile != INVALID_HANDLE_VALUE )
	{	
		GETVERSIONINPARAMS VersionInParams;
		DWORD	dwByteRet;

		// 获取硬盘属性
		isOK = DeviceIoControl( hDeviceFile, SMART_GET_VERSION, NULL, 0, &VersionInParams, sizeof(GETVERSIONINPARAMS), &dwByteRet, NULL );
		if( isOK )
		{
			isOK = FALSE;
			if( VersionInParams.fCapabilities & CAP_ATA_ID_CMD )
			{	
				BYTE	abSCIP[sizeof(SENDCMDINPARAMS) - 1] = {0};
				BYTE	abSCOP[sizeof(SENDCMDOUTPARAMS) - 1 + 512] = {0};
				SENDCMDINPARAMS		*scip = reinterpret_cast<SENDCMDINPARAMS*>(abSCIP);
				SENDCMDOUTPARAMS	*scop = reinterpret_cast<SENDCMDOUTPARAMS*>(abSCOP);

				// 设置输入参数
				scip->irDriveRegs.bCommandReg = ID_CMD;
				
				// 设置输出参数
				scop->cBufferSize = 512;

				// 获取硬盘详细信息
				isOK = DeviceIoControl( hDeviceFile, SMART_RCV_DRIVE_DATA, scip, sizeof(abSCIP), scop, sizeof(abSCOP), &dwByteRet, NULL );
				if( isOK )
				{	
					if( szProperty != NULL )
					{	// 提取硬盘序列号
						for( UINT i = 0; i < 20; i += 2 )
						{	// 颠倒高低位字节
							szProperty[i] = scop->bBuffer[i+21];
							szProperty[i + 1] = scop->bBuffer[i+20];
						}

						szProperty[20] = TEXT('\0');

						// 去掉空格
						std::remove( szProperty, szProperty + _tcslen(szProperty) + 1, L' ' );
					}
				}
			}
		}
	
		CloseHandle( hDeviceFile );
	}

	return isOK;
}

static BOOL WDK_GetProperty( TCHAR* DevicePath, INT iQueryType, TCHAR* szProperty, UINT uSize )
{
	BOOL isOK = FALSE;

	switch( iQueryType )
	{
	case 0:
		// 获取设备的VID-PID
		isOK = WDK_GetVIDPID( DevicePath, szProperty, uSize );
		break;

	case 1:		// 网卡原生MAC地址（包含USB网卡）		
		isOK = WDK_GetMacAddress( DevicePath, szProperty, uSize, TRUE );
		break;

	case 2:		// 网卡原生MAC地址（剔除USB网卡）		
		isOK = WDK_GetMacAddress( DevicePath, szProperty, uSize, FALSE );
		break;

	case 3:		// 硬盘序列号（需要系统管理员权限）
		isOK = WDK_GetHarddiskSerialNumberWithAdminRights( DevicePath, szProperty, uSize );
		break;

	default:
		break;		
	}

	return isOK;
}

INT WDK_DeviceQuery_Property( INT iQueryType, T_DEVICE_PROPERTY *properties, INT iSize )
{	
	HDEVINFO	hDevInfo;
	DWORD		MemberIndex, RequiredSize;
	SP_DEVICE_INTERFACE_DATA			DeviceInterfaceData;
	PSP_DEVICE_INTERFACE_DETAIL_DATA	DeviceInterfaceDetailData;
	INT	iTotal = 0;

	// 判断查询类型是否支持
	if( (iQueryType < 0) || (iQueryType >= sizeof(GUID_QUERYSET)/sizeof(GUID)) )
	{
		return -2;	// 查询类型不支持
	}

	// 获取设备信息集
	hDevInfo = SetupDiGetClassDevs( GUID_QUERYSET + iQueryType, NULL, NULL, DIGCF_PRESENT | DIGCF_INTERFACEDEVICE );
	if( hDevInfo == INVALID_HANDLE_VALUE )
	{
		return -1;
	}

	// 枚举设备信息集中所有设备
	DeviceInterfaceData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);
	for( MemberIndex = 0; ((properties == NULL) || (iTotal < iSize)); MemberIndex++ )
	{	// 获取设备接口
		if( !SetupDiEnumDeviceInterfaces( hDevInfo, NULL, GUID_QUERYSET + iQueryType, MemberIndex, &DeviceInterfaceData ) )
		{	// 设备枚举完毕
			break;
		}

		// 获取接收缓冲区大小，函数返回值为FALSE，GetLastError()=ERROR_INSUFFICIENT_BUFFER
		SetupDiGetDeviceInterfaceDetail( hDevInfo, &DeviceInterfaceData, NULL, 0, &RequiredSize, NULL );
		
		// 申请接收缓冲区
		DeviceInterfaceDetailData = (PSP_DEVICE_INTERFACE_DETAIL_DATA)malloc( RequiredSize );
		DeviceInterfaceDetailData->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA);

		// 获取设备细节信息
		if( SetupDiGetDeviceInterfaceDetail( hDevInfo, &DeviceInterfaceData, DeviceInterfaceDetailData, RequiredSize, NULL, NULL ) )
		{
			if( properties != NULL )
			{
				if( WDK_GetProperty( DeviceInterfaceDetailData->DevicePath, iQueryType, properties[iTotal].szProperty, PROPERTY_MAX_LEN ) )
				{	
					if( iTotal > 0 )
					{	// 剔除重复设备
						if( _tcscmp( properties[iTotal].szProperty, properties[iTotal - 1].szProperty ) != 0 )
						{
							iTotal++;
						}
					}
					else
					{
						iTotal++;
					}
				}
			}
			else
			{
				iTotal++;
			}
		}

		free( DeviceInterfaceDetailData );
	}

	SetupDiDestroyDeviceInfoList( hDevInfo );

	return iTotal;
}

INT WDK_DeviceQuery_VIDPID( const GUID* pGuid, HIDD_VIDPID* pVIDPID, INT iSize )
{	
	HDEVINFO	hDevInfo;
	DWORD		MemberIndex, RequiredSize;
	SP_DEVICE_INTERFACE_DATA			DeviceInterfaceData;
	PSP_DEVICE_INTERFACE_DETAIL_DATA	DeviceInterfaceDetailData;
	INT	iTotal = 0;

	// 获取设备信息集
	hDevInfo = SetupDiGetClassDevs( pGuid, NULL, NULL, DIGCF_PRESENT | DIGCF_INTERFACEDEVICE );
	if( hDevInfo == INVALID_HANDLE_VALUE )
	{
		return -1;
	}

	// 枚举设备信息集中所有设备
	DeviceInterfaceData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);
	for( MemberIndex = 0; ((pVIDPID == NULL) || (iTotal < iSize)); MemberIndex++ )
	{	// 获取设备接口
		if( !SetupDiEnumDeviceInterfaces( hDevInfo, NULL, pGuid, MemberIndex, &DeviceInterfaceData ) )
		{	// 设备枚举完毕
			break;
		}

		// 获取接收缓冲区大小，函数返回值为FALSE，GetLastError()=ERROR_INSUFFICIENT_BUFFER
		SetupDiGetDeviceInterfaceDetail( hDevInfo, &DeviceInterfaceData, NULL, 0, &RequiredSize, NULL );
		
		// 申请接收缓冲区
		DeviceInterfaceDetailData = (PSP_DEVICE_INTERFACE_DETAIL_DATA)malloc( RequiredSize );
		DeviceInterfaceDetailData->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA);

		// 获取设备细节信息
		if( SetupDiGetDeviceInterfaceDetail( hDevInfo, &DeviceInterfaceData, DeviceInterfaceDetailData, RequiredSize, NULL, NULL ) )
		{	// 获取设备句柄
			HANDLE	hDeviceFile;				

			hDeviceFile = CreateFile( DeviceInterfaceDetailData->DevicePath,
				0,
				FILE_SHARE_READ | FILE_SHARE_WRITE,
				NULL,
				OPEN_EXISTING,
				0,
				NULL);

			if( hDeviceFile != INVALID_HANDLE_VALUE )
			{	// 获取设备属性
				HIDD_ATTRIBUTES	Attributes;

				Attributes.Size = sizeof(HIDD_ATTRIBUTES);
				if( HidD_GetAttributes( hDeviceFile, &Attributes ) )
				{
					if( pVIDPID != NULL )
					{
						pVIDPID[iTotal].VendorID = Attributes.VendorID;
						pVIDPID[iTotal].ProductID = Attributes.ProductID;
						if( iTotal > 0 )
						{	// 剔除重复设备
							if( (pVIDPID[iTotal].VendorID != pVIDPID[iTotal - 1].VendorID) || (pVIDPID[iTotal].ProductID != pVIDPID[iTotal - 1].ProductID) )
							{
								iTotal++;
							}
						}
						else
						{
							iTotal++;
						}
					}
					else
					{
						iTotal++;
					}					
				}				

				CloseHandle( hDeviceFile );
			}
		}

		free( DeviceInterfaceDetailData );
	}

	SetupDiDestroyDeviceInfoList( hDevInfo );

	return iTotal;
}