Detecting a USB mounted hard drive

DetectorDlg.cpp

// DetectorDlg.cpp : implementation file
//
 
#include "stdafx.h"
#include "Detector.h"
#include "DetectorDlg.h"
#include "windows.h"
#include "WinDef.h"
#include "ntddstor.h"
#include "winioctl.h"
#include <initguid.h>
#include "dbt.h"
 
// Display installed devices
#include "setupapi.h"
#include "devguid.h"
#include "regstr.h"
 
// Drive type names
#define DRVUNKNOWN		0
#define DRVFIXED		1
#define DRVREMOTE		2
#define DRVRAM			3
#define DRVCD			4
#define DRVREMOVE		5
 
#ifdef _DEBUG
#define new DEBUG_NEW
#endif
 
//Function Prototypes
BOOL isItUsbDrive(TCHAR* a);
BOOL GetDisksProperty(HANDLE hDevice, PSTORAGE_DEVICE_DESCRIPTOR pDevDesc);
char GetDriveLetter(unsigned long a);
TCHAR* findUsbDrives(DWORD mask);
 
// CDetectorDlg dialog
 
CDetectorDlg::CDetectorDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CDetectorDlg::IDD, pParent)
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}
 
void CDetectorDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_EDIT1, m_text);
}
 
BEGIN_MESSAGE_MAP(CDetectorDlg, CDialog)
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	//}}AFX_MSG_MAP
	ON_BN_CLICKED(ID_DETECT, &CDetectorDlg::OnBnClickedDetect)
END_MESSAGE_MAP()
 
// CDetectorDlg message handlers
 
BOOL CDetectorDlg::OnInitDialog()
{
	CDialog::OnInitDialog();
 
	// Set the icon for this dialog.  The framework does this automatically
	//  when the application's main window is not a dialog
	SetIcon(m_hIcon, TRUE);			// Set big icon
	SetIcon(m_hIcon, FALSE);		// Set small icon
 
	// TODO: Add extra initialization here
 
	return TRUE;  // return TRUE  unless you set the focus to a control
}
 
//	If you add a minimize button to your dialog, you will need the code below
//  to draw the icon.  For MFC applications using the document/view model,
//  this is automatically done for you by the framework.
 
void CDetectorDlg::OnPaint()
{
	if (IsIconic())
	{
		CPaintDC dc(this); // device context for painting
 
		SendMessage(WM_ICONERASEBKGND, reinterpret_cast(dc.GetSafeHdc()), 0);
 
		// Center icon in client rectangle
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;
 
		// Draw the icon
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialog::OnPaint();
	}
}
 
// The system calls this function to obtain the cursor to display while the user drags
//  the minimized window.
HCURSOR CDetectorDlg::OnQueryDragIcon()
{
	return static_cast(m_hIcon);
}
 
/*********************************************************************
* Query a device handle and storage device descriptor for the
* IOCTL_STORAGE_QUERY_PROPERTY.
* Inputs:    HANDLE, PSTORAGE_DEVICE_DESCRIPTOR
* Outputs:   BOOL
*********************************************************************/
BOOL GetDisksProperty(HANDLE hDevice, PSTORAGE_DEVICE_DESCRIPTOR pDevDesc)
{
	STORAGE_PROPERTY_QUERY	Query;	// INPUT Parameter
	DWORD dwOutBytes;				// IOCTL Output Bytes
	BOOL bResult;					// IOCTL Return Value
 
	// specify the query type
	Query.PropertyId = StorageDeviceProperty;
	Query.QueryType = PropertyStandardQuery;
 
	// Query using IOCTL_STORAGE_QUERY_PROPERTY
	bResult = ::DeviceIoControl(hDevice,			// device handle
 
 
 
			IOCTL_STORAGE_QUERY_PROPERTY,			// info of device property
			&Query, sizeof(STORAGE_PROPERTY_QUERY),	// input data buffer
			pDevDesc, pDevDesc->Size,				// output data buffer
			&dwOutBytes,							// out's length
			(LPOVERLAPPED)NULL);					
 
	return bResult;
}
 
/*********************************************************************
* Checks the storage device descriptor for the corresponding BusType.
* Only BusTypeUsb is of interest in this case.
* Inputs:  TCHAR*
* Outputs: BOOL
*********************************************************************/
BOOL isItUsbDrive(TCHAR* a)
{
	PSTORAGE_DEVICE_DESCRIPTOR pDevDesc;
	HANDLE hDevice;
	int b;
	hDevice = CreateFile(a, GENERIC_READ,
					FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, NULL, NULL);
	 if(hDevice != INVALID_HANDLE_VALUE)
	 {
		 pDevDesc = (PSTORAGE_DEVICE_DESCRIPTOR)new BYTE[sizeof(STORAGE_DEVICE_DESCRIPTOR) + 512 - 1];
 
		 if(pDevDesc != NULL)
		 {
			 pDevDesc->Size = sizeof(STORAGE_DEVICE_DESCRIPTOR) + 512 -1;
 
			 if(GetDisksProperty(hDevice, pDevDesc))
			 {
				 if(pDevDesc->BusType == BusTypeUsb)		// BusTypeUSB is defined in ntddstor.h
				 {
					 //wcscat(finalString, CString(szDriveRoot));
					 //wcscat(finalString, CString(' '));
					 b = 1;
				 }
				 else
				 {
					 b = 0;
				 }
			 }
			 delete pDevDesc;
			 CloseHandle(hDevice);
			 return b;
		 }
		 else
		 {
			CloseHandle(hDevice);
			b = 0;
			return b;
		 }
	  }
}
 
void findUsbDrives(DWORD mask, TCHAR* firstString)
{
	TCHAR			szDriveRoot[5];
	TCHAR			szDriveCreate[7];
 
	TCHAR			cDrive;
	DWORD			drivetype;
	_stprintf(firstString, _T("USB disk(s) are: \r\n"));
	for ( cDrive = 'A'; cDrive <= 'Z'; cDrive++,mask>>=1 )
	{
		if ( !(mask&1) )
		{
			continue;
		}
		_stprintf(szDriveRoot, _T("%c:\\ "), cDrive);
		_stprintf(szDriveCreate, _T("\\\\.\\%c:"), cDrive);
		switch ( GetDriveType( szDriveRoot ) )
		{
			case 0:					// The drive type cannot be determined.
			case 1:					// The root directory does not exist.
				drivetype = DRVUNKNOWN;
			break;
 
			case DRIVE_REMOVABLE:	// The drive can be removed from the drive.
				drivetype = DRVREMOVE;
				if(isItUsbDrive(szDriveCreate))
				{
					wcscat(firstString, CString(szDriveRoot));
				}
			break;
 
			case DRIVE_CDROM:		// The drive is a CD-ROM drive.
				drivetype = DRVCD;
			break;
 
			case DRIVE_FIXED:		// The disk cannot be removed from the drive.
				drivetype = DRVFIXED;
				if(isItUsbDrive(szDriveCreate))
				{
					wcscat(firstString, CString(szDriveRoot));
				}
			break;
 
			case DRIVE_REMOTE:		// The drive is a remote (network) drive.
				drivetype = DRVREMOTE;
			break;
 
			case DRIVE_RAMDISK:		// The drive is a RAM disk.
				drivetype = DRVRAM;
			break;
		}
 
	}
	//return firstString;
}
 
/*********************************************************************
*	In order to determine USB HARD DISKS, we must first determine
*	the drive type using GetDriveType().  USB HD are determined as
*	FIX DISK.  Secondly, we determine the bus type of the Fix Drive.
*	This is done by calling the IOCTL_STORAGE_PROPERTY_QUERY.
*********************************************************************/
void CDetectorDlg::OnBnClickedDetect()
{
	DWORD			head;
	DWORD			drivetype;
	TCHAR			finalString[60];
 
	head = GetLogicalDrives();
	findUsbDrives(head, finalString);
	m_text.SetWindowTextW(CString(finalString));
}
 
char GetDriveLetter(unsigned long a)
{
	char c;
	for(c=0; c<26;c++)
	{
		if(a & 0x01)
		{
			break;
		}
		a = a >> 1;
	}
	return (c+'A');
}
 
LRESULT CDetectorDlg::WindowProc(UINT message, WPARAM wParam, LPARAM lParam)
{
 
	// TODO: Add your specialized code here and/or call the base class
	switch(message)
	{
		case WM_PAINT:
		break;
		case WM_DEVICECHANGE:
			PDEV_BROADCAST_HDR pHdr = (PDEV_BROADCAST_HDR) lParam;
			switch (wParam)
			{
				case DBT_DEVICEARRIVAL:
					AfxMessageBox(L"A device has been inserted.");
					//m_text.SetWindowText(CString("A device has been inserted"));
					if (pHdr->dbch_devicetype == DBT_DEVTYP_VOLUME)
					{
						PDEV_BROADCAST_VOLUME pVol = (PDEV_BROADCAST_VOLUME) pHdr;
						TCHAR szMessage[80];
						TCHAR cDriveLetter = GetDriveLetter(pVol->dbcv_unitmask);
						wsprintf(szMessage, L"Device '%c:' has been inserted.", cDriveLetter);
						m_text.SetWindowText(CString(szMessage));
					}
				break;
 
				case DBT_DEVICEREMOVECOMPLETE:
					//AfxMessageBox(L"A device has been removed");
					m_text.SetWindowText(CString("A device has been removed"));
				break;
			}
		break;
	}
 
	return CDialog::WindowProc(message, wParam, lParam);
}