hwdetectdlg.cpp

一个VC编写的USB口的安全保护程序

					break;
				} else {
					// remove device
					delete pszData;
					m_ctrlTree.DeleteItem(hChildItem);
					if ( !m_ctrlTree.ItemHasChildren(hClass) ) {
						HeapFree(GetProcessHeap(), 0, (LPVOID)m_ctrlTree.GetItemData(hClass));
						m_ctrlTree.DeleteItem(hClass);
					}
					break;
				}
			} else {
				hChildItem = m_ctrlTree.GetNextItem(hChildItem, TVGN_NEXT);
			}
		}
		if ( DBT_DEVICEARRIVAL == wParam && !bFound ) {
			// Add device
			int n = GetClassImgIndex(&(spDevInfoData.ClassGuid));
			HTREEITEM hItem = m_ctrlTree.InsertItem(buf, n, n, hClass);
			m_ctrlTree.SetItemData(hItem,atoi(szDevId));
			// set pspDevInfoData to NULL, otherwise it will be HeapFree
			m_ctrlTree.Expand(hClass, TVE_EXPAND);
		}
	}

	szDevId.ReleaseBuffer();
	SetupDiDestroyDeviceInfoList(hDevInfo);
}

BOOL CHWDetectDlg::FindDevice(HDEVINFO& hDevInfo, 
							  CString& szDevId, 
							  SP_DEVINFO_DATA& spDevInfoData)
{
	spDevInfoData.cbSize = sizeof(SP_DEVINFO_DATA);
	for(int i=0; SetupDiEnumDeviceInfo(hDevInfo, i, &spDevInfoData); i++) {
		DWORD nSize=0 ;
		TCHAR buf[MAX_PATH];
		CString strDevIdTemp;

		if ( !SetupDiGetDeviceInstanceId(hDevInfo, &spDevInfoData, buf, sizeof(buf), &nSize) ) {
			TRACE(CString("SetupDiGetDeviceInstanceId(): ") 
				+ _com_error(GetLastError()).ErrorMessage());
			return FALSE;
		} 

		if ( szDevId == buf ) {
			// OK, device found
			return TRUE;
		}
	}
	return FALSE;
}

BOOL CHWDetectDlg::FindDeviceEx(HDEVINFO& hDevInfo,SP_DEVINFO_DATA& spDevInfoData)
{
	spDevInfoData.cbSize = sizeof(SP_DEVINFO_DATA);
	for(int i=0; SetupDiEnumDeviceInfo(hDevInfo, i, &spDevInfoData); i++) {
		DWORD nSize=0 ;
		TCHAR buf[MAX_PATH];
		CString strDevIdTemp;
		
		if ( !SetupDiGetDeviceInstanceId(hDevInfo, &spDevInfoData, buf, sizeof(buf), &nSize) ) {
			TRACE(CString("SetupDiGetDeviceInstanceId(): ") 
				+ _com_error(GetLastError()).ErrorMessage());
			return FALSE;
		} 
		
		strDevIdTemp = buf;
		strDevIdTemp = strDevIdTemp.Left(22);
		//STORAGE\REMOVABLEMEDIA
		if(0 == (strDevIdTemp.Compare(_T("STORAGE\\REMOVABLEMEDIA"))))
		{
			m_pszDevId[m_nStrID] = buf;
			m_nStrID++;
		}
	}
	if(!m_pszDevId[0].IsEmpty())
		return TRUE;
	return FALSE;
}

CString CHWDetectDlg::GetClassDesc(GUID* pGuid)
{
	TCHAR buf[MAX_PATH];
	DWORD size;
	if ( SetupDiGetClassDescription(pGuid, buf, sizeof(buf), &size) ) {
		return CString(buf);
	} else {
		TRACE("Can't get class description: %s", _com_error(GetLastError()).ErrorMessage());
		return _T("");
	}
}

int CHWDetectDlg::GetClassImgIndex(GUID* pGuid)
{
	int nIdx;
	if ( SetupDiGetClassImageIndex(&m_imgList, pGuid, &nIdx) ) {
		return nIdx;
	} else {
		TRACE("Can't get class image index: %s", _com_error(GetLastError()).ErrorMessage());
		return -1;
	}
}

HTREEITEM CHWDetectDlg::GetClassItem(GUID* pGuid, WPARAM wParam)
{
	if ( m_ctrlTree.ItemHasChildren(hRoot) ) {
		HTREEITEM hChildItem = m_ctrlTree.GetChildItem(hRoot);
		while (hChildItem != NULL) {
			GUID* tmp = (GUID*)m_ctrlTree.GetItemData(hChildItem);
			if ( IsEqualGUID(*tmp, *pGuid) ) {
				return hChildItem;
			} else {
				hChildItem = m_ctrlTree.GetNextItem(hChildItem, TVGN_NEXT);
			}
		}
	}
	
	if ( DBT_DEVICEARRIVAL == wParam ) {
		// no child or not found
		int n = GetClassImgIndex(pGuid);
 		GUID* tmp = (GUID*)HeapAlloc(GetProcessHeap(), 0, sizeof(GUID));
		CopyMemory(tmp, pGuid, sizeof(GUID));
		HTREEITEM hItem = m_ctrlTree.InsertItem(GetClassDesc(pGuid), n, n, hRoot);
		m_ctrlTree.SetItemData(hItem, (DWORD&)tmp);
		m_ctrlTree.Expand(hRoot, TVE_EXPAND);
//		hSelectedDev = hItem;
//		OnPopupDisable();
		return hItem;
	} else {
		return NULL;
	}
}

void CHWDetectDlg::OnContextMenu(CWnd* pWnd, CPoint point)
{
	// TODO: Add your message handler code here
	hSelectedDev = NULL;
	CPoint scPoint = point;
	m_ctrlTree.ScreenToClient(&scPoint);
	HTREEITEM hItem = m_ctrlTree.HitTest(scPoint);
	if ( NULL != hItem && hRoot != hItem ) {
		HTREEITEM hParent = m_ctrlTree.GetParentItem(hItem);
		if ( hParent != hRoot ) {
			m_ctrlTree.SelectItem(hItem);
			hSelectedDev = hItem;
			m_muPopup->TrackPopupMenu(TPM_LEFTALIGN | TPM_LEFTBUTTON, 
				point.x, point.y, this);
		}
	}
}

void CHWDetectDlg::OnPopupDisable()
{
	// TODO: Add your command handler code here
	//Modified by thit_chenwd ,the func is to disable usb device
	if ( NULL != hSelectedDev ) {
		HTREEITEM hClass = m_ctrlTree.GetParentItem(hSelectedDev);
		GUID* pGuid = (GUID*)m_ctrlTree.GetItemData(hClass);
		ASSERT(pGuid);

		HDEVINFO hDevInfo = SetupDiGetClassDevs(NULL,NULL,NULL,DIGCF_ALLCLASSES);
		if( INVALID_HANDLE_VALUE == hDevInfo ) {
			AfxMessageBox(CString("SetupDiGetClassDevs(): ") 
				+ _com_error(GetLastError()).ErrorMessage(), MB_ICONEXCLAMATION);
			return;
		}

		SP_PROPCHANGE_PARAMS spPropChangeParams;
		spPropChangeParams.ClassInstallHeader.cbSize = sizeof(SP_CLASSINSTALL_HEADER);
		spPropChangeParams.ClassInstallHeader.InstallFunction = DIF_PROPERTYCHANGE;
		spPropChangeParams.Scope = DICS_FLAG_GLOBAL;
		spPropChangeParams.StateChange = DICS_DISABLE;
		spPropChangeParams.HwProfile = 0;


		//thit_chenwd 
		//"USBSTOR\DISK&VEN_GENERIC&PROD_FLASH_DISK&REV_8.01\11C4EBE0&1"
		CString pszDevId = _T("STORAGE\\REMOVABLEMEDIA\\7&27081B85&0&RM");
		//String *pszDevId = &pszTemp;//(CString*)m_ctrlTree.GetItemData(hSelectedDev);
		
		SP_DEVINFO_DATA spDevInfoData;
		if ( FindDevice(hDevInfo,pszDevId, spDevInfoData) ) {
			if ( !SetupDiSetClassInstallParams(hDevInfo, &spDevInfoData, 
				(SP_CLASSINSTALL_HEADER*)&spPropChangeParams, sizeof(spPropChangeParams)) ) {
				AfxMessageBox(CString("SetupDiSetClassInstallParams(): ") 
					+ _com_error(GetLastError()).ErrorMessage(), MB_ICONEXCLAMATION);
			} else if ( !SetupDiCallClassInstaller(DIF_PROPERTYCHANGE, hDevInfo, &spDevInfoData) ) {
				AfxMessageBox(CString("SetupDiClassInstaller(): ") 
					+ _com_error(GetLastError()).ErrorMessage(), MB_ICONEXCLAMATION);
			} else {
				AfxMessageBox(_T("Dislabe OK"), MB_OK);
			}
		}

		SetupDiDestroyDeviceInfoList(hDevInfo);
	}
}


//single func to exec the use storage device be disabled.
void CHWDetectDlg::OnMyUsbDisable(BOOL bStateChange)
{
	// TODO: Add your command handler code here
	//Modified by thit_chenwd ,the func is to disable usb device
	//HTREEITEM hClass = m_ctrlTree.GetParentItem(hSelectedDev);
	//GUID* pGuid = (GUID*)m_ctrlTree.GetItemData(hClass);
	//ASSERT(pGuid);
	
	HDEVINFO hDevInfo = SetupDiGetClassDevs(NULL,NULL,NULL,DIGCF_ALLCLASSES);
	if( INVALID_HANDLE_VALUE == hDevInfo ) {
		AfxMessageBox(CString("SetupDiGetClassDevs(): ") 
			+ _com_error(GetLastError()).ErrorMessage(), MB_ICONEXCLAMATION);
		return;
	}
	
//	SP_PROPCHANGE_PARAMS spPropChangeParams;
//	spPropChangeParams.ClassInstallHeader.cbSize = sizeof(SP_CLASSINSTALL_HEADER);
//	spPropChangeParams.ClassInstallHeader.InstallFunction = DIF_PROPERTYCHANGE;
//	spPropChangeParams.Scope = DICS_FLAG_GLOBAL;
//	spPropChangeParams.StateChange = DICS_DISABLE;
//	spPropChangeParams.HwProfile = 0;
	
	
	//thit_chenwd 
	//"USBSTOR\DISK&VEN_GENERIC&PROD_FLASH_DISK&REV_8.01\11C4EBE0&1"
//	CString pszDevId[6] = {_T("STORAGE\\REMOVABLEMEDIA\\7&18F0EDB7&0&RM"),
//						_T("STORAGE\\REMOVABLEMEDIA\\7&24BD02C6&0&RM"),
//						_T("STORAGE\\REMOVABLEMEDIA\\7&27081B85&0&RM"),
//						_T("STORAGE\\REMOVABLEMEDIA\\7&27E0B682&0&RM"),
//						_T("STORAGE\\REMOVABLEMEDIA\\7&2E67401C&0&RM"),
//						 _T("STORAGE\\REMOVABLEMEDIA\\7&305B44FE&0&RM")
//	};
		
//	CString pszDevId3 = _T("USB\\Vid_04e8&Pid_503b\\0002F9A9828E0F06");
	
	//String *pszDevId = &pszTemp;//(CString*)m_ctrlTree.GetItemData(hSelectedDev);
	
//	for(int j = 0;j < 6;j++)


//	SP_PROPCHANGE_PARAMS spPropChangeParams;
//	spPropChangeParams.ClassInstallHeader.cbSize = sizeof(SP_CLASSINSTALL_HEADER);
//	spPropChangeParams.ClassInstallHeader.InstallFunction = DIF_PROPERTYCHANGE;
//	spPropChangeParams.Scope = DICS_FLAG_GLOBAL;
//	spPropChangeParams.StateChange = DICS_DISABLE;

//		spPropChangeParams.StateChange = DICS_ENABLE;
//	spPropChangeParams.HwProfile = 0;
	
	SP_DEVINFO_DATA spDevInfoData;
	if ( FindDeviceEx(hDevInfo,spDevInfoData) )
	{
		for(int j = 0; !m_pszDevId[j].IsEmpty() ; j++)
			DisableUsedev(hDevInfo,m_pszDevId[j],bStateChange);
	}
	
//	SetupDiDestroyDeviceInfoList(hDevInfo);
	SetupDiDestroyDeviceInfoList(hDevInfo);
}


void CHWDetectDlg::DisableUsedev(HDEVINFO& hDevInfo,CString& strDevId,BOOL bStateChange)
{
	
	SP_PROPCHANGE_PARAMS spPropChangeParams;
	spPropChangeParams.ClassInstallHeader.cbSize = sizeof(SP_CLASSINSTALL_HEADER);
	spPropChangeParams.ClassInstallHeader.InstallFunction = DIF_PROPERTYCHANGE;
	spPropChangeParams.Scope = DICS_FLAG_GLOBAL;
	if(TRUE == bStateChange)
		spPropChangeParams.StateChange = DICS_DISABLE;
	else if(FALSE == bStateChange)
		spPropChangeParams.StateChange = DICS_ENABLE;
	spPropChangeParams.HwProfile = 0;
	
	
	//thit_chenwd 
	//"USBSTOR\DISK&VEN_GENERIC&PROD_FLASH_DISK&REV_8.01\11C4EBE0&1"
	//CString pszDevId = _T("STORAGE\\REMOVABLEMEDIA\\7&27081B85&0&RM");
	//String *pszDevId = &pszTemp;//(CString*)m_ctrlTree.GetItemData(hSelectedDev);
	
	SP_DEVINFO_DATA spDevInfoData;
	if ( FindDevice(hDevInfo,strDevId, spDevInfoData) ) {
		if ( !SetupDiSetClassInstallParams(hDevInfo, &spDevInfoData, 
			(SP_CLASSINSTALL_HEADER*)&spPropChangeParams, sizeof(spPropChangeParams)) ) {
			AfxMessageBox(CString("SetupDiSetClassInstallParams(): ") 
				+ _com_error(GetLastError()).ErrorMessage(), MB_ICONEXCLAMATION);
		} else if ( !SetupDiCallClassInstaller(DIF_PROPERTYCHANGE, hDevInfo, &spDevInfoData) ) {
			AfxMessageBox(CString("SetupDiClassInstaller(): ") 
				+ _com_error(GetLastError()).ErrorMessage(), MB_ICONEXCLAMATION);
		} else {
			if(TRUE == bStateChange)
				AfxMessageBox(_T("Dislabe OK"), MB_OK);
			else if(FALSE == bStateChange)
				AfxMessageBox(_T("Enable oK"),MB_OK);
		}
	}
	
}

void CHWDetectDlg::OnDlgEnable()
{
//	AfxMessageBox("Dlg test");
	OnMyUsbDisable(FALSE);
}

void CHWDetectDlg::OnDlgDisable()
{
	OnMyUsbDisable(TRUE);
}