sniffusbdlg.cpp

USB sniffer for windows

                DWORD dwType = REG_SZ;
                DWORD dwSize = MAX_PATH;
                if(ERROR_SUCCESS == RegQueryValueEx(hInstKey, "DeviceDesc", NULL, &dwType, (LPBYTE) szData, &dwSize))
                {
                    sDescription = szData;
                }
                RegCloseKey(hInstKey);
                hInstKey = NULL;
            }
            dwIndex++;
        }
        RegCloseKey(hKey);
        hKey = NULL;
    }
}

BOOL CSniffUSBDlg::IsThereAFilter(LPCTSTR szVidPid)
{
    BOOL bThereIsAFilter = FALSE;
    TCHAR szEnumKey[MAX_PATH];
    _tcscpy(szEnumKey, "Enum\\USB\\");
    _tcscat(szEnumKey, szVidPid);
    HKEY hKey = NULL;
    if(ERROR_SUCCESS == RegOpenKeyEx(HKEY_LOCAL_MACHINE, szEnumKey, 0, KEY_ALL_ACCESS, &hKey))
    {
        DWORD dwIndex = 0;
        TCHAR sName[MAX_PATH];
        while(ERROR_SUCCESS == RegEnumKey(hKey, dwIndex, sName, MAX_PATH))
        {
            TRACE("Enumerated >%s<\n", sName);
            HKEY hInstKey = NULL;
            TCHAR szInstanceKey[MAX_PATH];
            _tcscpy(szInstanceKey, szEnumKey);
            _tcscat(szInstanceKey, _T("\\"));
            _tcscat(szInstanceKey, sName);
            if(ERROR_SUCCESS == RegOpenKeyEx(HKEY_LOCAL_MACHINE, szInstanceKey, 0, KEY_ALL_ACCESS, &hInstKey))
            {
                TCHAR szData[MAX_PATH];
                DWORD dwType = REG_SZ;
                DWORD dwSize = MAX_PATH;
                if(ERROR_SUCCESS == RegQueryValueEx(hInstKey, m_sLowerFilters, NULL, &dwType, (LPBYTE) szData, &dwSize))
                {
                    CString sData = szData;
                    sData.MakeLower();
                    if(NULL != _tcsstr(szData, m_sFilterName))
                    {
                        TRACE("Found filter!\n");
                        bThereIsAFilter = TRUE;
                    }
                }
                RegCloseKey(hInstKey);
                hInstKey = NULL;
            }
            dwIndex++;
        }
        RegCloseKey(hKey);
        hKey = NULL;
    }

    return bThereIsAFilter;
}

void CSniffUSBDlg::ModifyFilterOnVIDPID(LPCTSTR szVidPid, BOOL bAddFilter)
{
 	HDEVINFO hdev;
	
	hdev = SetupDiGetClassDevs(0,NULL,0,DIGCF_ALLCLASSES);
	if (hdev == INVALID_HANDLE_VALUE )
	{
		MessageBox("Unable to enumerate USB device");
		TRACE("SetupDiGetClassDevs = %d\n",GetLastError());
		return ;
	}

	for (DWORD idx=0;;idx++)
	{
		SP_DEVINFO_DATA  devinfo;
		devinfo.cbSize = sizeof(devinfo);

		CString sName, sFilter, sDescription;

		BYTE Buffer[200];
		DWORD BufferSize = 0;
		DWORD DataType;

		if (!SetupDiEnumDeviceInfo(hdev,idx,&devinfo))
		{
			if (GetLastError() != ERROR_NO_MORE_ITEMS)
			{
				MessageBox("Error while enumerating USB devices");
				TRACE("SetupDiEnumDeviceInfo = %d\n",GetLastError());
			}
			break;
		}

		if (SetupDiGetDeviceRegistryProperty(hdev,&devinfo,SPDRP_HARDWAREID  ,
			&DataType,Buffer,sizeof(Buffer),&BufferSize))
		{
			if (strcmp((const char *)Buffer,szVidPid)==0)
			{
				// gotcha !
				if (bAddFilter)
				{
					int len = strlen(SERVICE)+2;
					BYTE * LowerFilters = (BYTE *) malloc(len);
					memset(LowerFilters,0,len);
					strcpy((char *)LowerFilters,SERVICE);

					if (!SetupDiSetDeviceRegistryProperty(hdev,&devinfo,SPDRP_LOWERFILTERS ,
						LowerFilters,len))
					{
						MessageBox("Install failed!");
						TRACE("SetupDiSetDeviceRegistryProperty = %d\n",GetLastError());
					}
				}
				else
				{
					if (!SetupDiSetDeviceRegistryProperty(hdev,&devinfo,SPDRP_LOWERFILTERS ,
						NULL,0))
					{
						MessageBox("Uninstall failed!");
						TRACE("SetupDiSetDeviceRegistryProperty = %d\n",GetLastError());
					}
				}
			}
		}
    }

	SetupDiDestroyDeviceInfoList(hdev);
}

BOOL CSniffUSBDlg::GetSelectedVidPid(CString& sVidPid)
{
    UINT nSelected = m_cDevs.GetSelectedCount();
    if(0 == nSelected)
    {
        AfxMessageBox(IDS_SELECT_ITEM_FIRST);
        return FALSE;
    }

    int nIndex = m_cDevs.GetNextItem(-1, LVNI_SELECTED);
    sVidPid = m_cDevs.GetItemText(nIndex, 0);
    return TRUE;
}

void CSniffUSBDlg::OnInstall() 
{
    CString sVidPid;
    if(GetSelectedVidPid(sVidPid))
    {
        TRACE("Installing on %s\n", sVidPid);
        ModifyFilterOnVIDPID(sVidPid, TRUE);
        OnRefresh();
    }
}

void CSniffUSBDlg::OnUninstall() 
{
    CString sVidPid;
    if(GetSelectedVidPid(sVidPid))
    {
        TRACE("Installing on %s\n", sVidPid);
        ModifyFilterOnVIDPID(sVidPid, FALSE);
        OnRefresh();
    }
}

void CSniffUSBDlg::OnReplug() 
{
    CString HardwareId;
    if(GetSelectedVidPid(HardwareId))
    {
 		HDEVINFO hdev;
		hdev = SetupDiGetClassDevs(0,NULL,0,DIGCF_ALLCLASSES);
		if (hdev == INVALID_HANDLE_VALUE )
		{
			MessageBox("Unable to enumerate USB device");
			TRACE("SetupDiGetClassDevs = %d\n",GetLastError());
			return ;
		}

		for (DWORD idx=0;;idx++)
		{
			SP_DEVINFO_DATA  devinfo;
			devinfo.cbSize = sizeof(devinfo);

			CString sName, sFilter, sDescription;

			BYTE Buffer[200];
			DWORD BufferSize = 0;
			DWORD DataType;

			if (!SetupDiEnumDeviceInfo(hdev,idx,&devinfo))
			{
				if (GetLastError() != ERROR_NO_MORE_ITEMS)
				{
					MessageBox("Error while enumerating USB devices");
					TRACE("SetupDiEnumDeviceInfo = %d\n",GetLastError());
				}
				break;
			}

			if (SetupDiGetDeviceRegistryProperty(hdev,&devinfo,SPDRP_HARDWAREID  ,
				&DataType,Buffer,sizeof(Buffer),&BufferSize))
			{
				if (strcmp((const char *)Buffer,HardwareId)==0)
				{
					// gotcha !
	 				if (MessageBox("I will briefly remove/add this device "
								"and associated software "
								"from your system! Are you sure? ",HardwareId,MB_YESNO) == IDYES)
					{
						CWaitCursor waitcursor;

						if (SetupDiRemoveDevice(hdev,&devinfo))
						{
							if (SetupDiUnremoveDevice(hdev,&devinfo))
								MessageBox("Gotcha!");
							else
								TRACE("SetupDiUnremoveDevice = %d\n",GetLastError());
						}
						else
							TRACE("SetupDiRemoveDevice = %d\n",GetLastError());
					}
				}
			}
		}

		SetupDiDestroyDeviceInfoList(hdev);
		
		OnRefresh();
		TRACE("Back to life..\n");
	}
}

void CSniffUSBDlg::OnRclickUsbdevs(NMHDR* pNMHDR, LRESULT* pResult) 
{
    CMenu ctx;
    ctx.LoadMenu(IDR_SNOOPUSB);
    CMenu *popup = ctx.GetSubMenu(0);
    CPoint point;
    GetCursorPos(&point);
    CPoint pt = point;
    m_cDevs.ScreenToClient(&pt);
    UINT uFlags = 0;
    int nIndex = m_cDevs.HitTest(pt, &uFlags);
    if(LVHT_ONITEM & uFlags)
    {
        m_cDevs.SetItem(nIndex, 0, LVIF_STATE, NULL, 0, LVIS_SELECTED, LVIS_SELECTED, 0);
        CString sVidPid = m_cDevs.GetItemText(nIndex, 0);
        if(IsThereAFilter(sVidPid))
        {
            popup->EnableMenuItem(ID_SNOOPUSB_INSTALL, MF_BYCOMMAND | MF_GRAYED);
        }
        else
        {
            popup->EnableMenuItem(ID_SNOOPUSB_UNINSTALL, MF_BYCOMMAND | MF_GRAYED);
        }
        popup->TrackPopupMenu(TPM_LEFTALIGN | TPM_RIGHTBUTTON, point.x, point.y, this);
    }
	
	*pResult = 0;
}

void CSniffUSBDlg::OnSnoopusbInstall() 
{
    OnInstall();	
}

void CSniffUSBDlg::OnSnoopusbUninstall() 
{
    OnUninstall();
}

void CSniffUSBDlg::OnSnoopusbReplug() 
{
    OnReplug();	
}

void CSniffUSBDlg::OnFilterInstall() 
{
	// TODO: Add your control notification handler code here
	
}

void CSniffUSBDlg::CheckService()
{
	// check if "usbsnoop" service is installed

	SC_HANDLE hManager = OpenSCManager(NULL,NULL,
		SC_MANAGER_CREATE_SERVICE|SC_MANAGER_ENUMERATE_SERVICE);

	if (hManager == NULL)
	{
		MessageBox("Can't open service manager");
		return ;
	}

	SC_HANDLE hService = OpenService(hManager,"usbsnoop",DELETE);
	if (hService == NULL)
	{
		if (GetLastError() == ERROR_SERVICE_DOES_NOT_EXIST)
		{
			if (MessageBox("Service usbsnoop","Would you like to install?",MB_YESNO)==IDYES)
			{
				hService = CreateService(hManager,"usbsnoop","usbsnoop (display)",
					0,SERVICE_KERNEL_DRIVER,SERVICE_DEMAND_START,SERVICE_ERROR_NORMAL,
					"System32\\DRIVERS\\USBSNOOP.SYS",
					NULL,NULL,NULL,NULL,NULL);
				if (hService == NULL)
				{
					MessageBox("Can't create service");
				}
				else
					CloseServiceHandle(hService);
			}
		}
	}
	else
	{
/*
		if (!DeleteService(hService))
			MessageBox("Cannot remove existing usbsnoop service");
*/
		CloseServiceHandle(hService);
	}


	CloseServiceHandle(hManager);
}