// DeviceTreeAppDlg.cpp : implementation file
//

#include "stdafx.h"
#include "DeviceTreeApp.h"
#include "DeviceTreeAppDlg.h"
#include ".\devicetreeappdlg.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#endif


// CDeviceTreeAppDlg dialog



CDeviceTreeAppDlg::CDeviceTreeAppDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CDeviceTreeAppDlg::IDD, pParent)
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CDeviceTreeAppDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_TREE_DEVICE, m_ctrlDeviceTree);
}

BEGIN_MESSAGE_MAP(CDeviceTreeAppDlg, CDialog)
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	//}}AFX_MSG_MAP
END_MESSAGE_MAP()


// CDeviceTreeAppDlg message handlers

BOOL CDeviceTreeAppDlg::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

	m_hDeviceFile = CreateFile(_T("\\\\.\\DeviceTreeSys"),GENERIC_READ,0,0,OPEN_EXISTING,
							   FILE_ATTRIBUTE_NORMAL,0);

	if(m_hDeviceFile == INVALID_HANDLE_VALUE)
	{
		OnCancel();
		return TRUE;
	}

	DWORD dummy;

	DWORD len = static_cast<DWORD>(sizeof(GET_OBJ_BY_NAME) + sizeof(L"\\Driver\\PnpManager"));

	PGET_OBJ_BY_NAME pInput = reinterpret_cast<PGET_OBJ_BY_NAME>(new BYTE[len]);

	ZeroMemory(pInput,len);

	pInput->m_hRootHandle = 0;
	
	CopyMemory(pInput->m_strName,L"\\Driver\\PnpManager",sizeof(L"\\Driver\\PnpManager"));

	m_pRootDriver = 0;
	m_pRootDevice = 0;

	if(DeviceIoControl(m_hDeviceFile,DEVICE_TREE_GET_POINTER_FROM_NAME,pInput,len,
					   &m_pRootDriver,sizeof(PVOID),&dummy,0))
	{
		BuildDeviceTree();
	}
	
	delete pInput;

	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 CDeviceTreeAppDlg::OnPaint() 
{
	if (IsIconic())
	{
		CPaintDC dc(this); // device context for painting

		SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(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 CDeviceTreeAppDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}

VOID CDeviceTreeAppDlg::BuildDeviceTree()
{
	DWORD dummy;
	DRIVER_OBJECT driver;
	m_pRootDevice = 0;

	if(DeviceIoControl(m_hDeviceFile,DEVICE_TREE_READ_KERNEL_MEMORY,&m_pRootDriver,
					   sizeof(PVOID),&driver,sizeof(driver),&dummy,0))
	{
		PDEVICE_OBJECT pDevice = driver.DeviceObject;
		while(pDevice)
		{
			DEVICE_OBJECT device;
			if(!DeviceIoControl(m_hDeviceFile,DEVICE_TREE_READ_KERNEL_MEMORY,
								&pDevice,sizeof(PVOID),&device,sizeof(device),
								&dummy,0))
				return;

			m_pRootDevice = pDevice;

			pDevice = device.NextDevice;
		}
	}

	if(m_pRootDevice)
		BuildDeviceTree(m_pRootDevice,TVI_ROOT);
}

PWCHAR CDeviceTreeAppDlg::QueryObjectName(PVOID pObject)
{
	static WCHAR wName[255];

	PWCHAR pNameString = _T("(unnamed)");

	DWORD dummy;
	ZeroMemory(wName,sizeof(wName));

	if(DeviceIoControl(m_hDeviceFile,DEVICE_TREE_QUERY_OBJECT_NAME,&pObject,sizeof(PVOID),
					   wName,sizeof(wName),&dummy,0))
	{
		if(wName[0])
			pNameString = wName;
	}

	return pNameString;
}

VOID CDeviceTreeAppDlg::BuildDeviceTree(PDEVICE_OBJECT pDeviceObject,HTREEITEM hParent)
{
	DWORD dummy;

	DEVICE_OBJECT device;

	if(!DeviceIoControl(m_hDeviceFile,DEVICE_TREE_READ_KERNEL_MEMORY,&pDeviceObject,
					   sizeof(PVOID),&device,sizeof(device),&dummy,0))
	{
		return;
	}

	CString strDeviceName = QueryObjectName(pDeviceObject);
	CString strDriverName = QueryObjectName(device.DriverObject);

	CString strName;
	strName.Format(_T("PDO:%#x - "));
	strName += strDeviceName;
	strName += _T(" - ");
	strName += strDriverName;

	HTREEITEM hCurrent = m_ctrlDeviceTree.InsertItem(strName,hParent);

	PDEVICE_OBJECT pCurrentFdo = device.AttachedDevice;
	while(pCurrentFdo)
	{
		if(!DeviceIoControl(m_hDeviceFile,DEVICE_TREE_READ_KERNEL_MEMORY,&pCurrentFdo,
							sizeof(PVOID),&device,sizeof(device),&dummy,0))
		{
			return;
		}

		strDeviceName = QueryObjectName(pCurrentFdo);
		strDriverName = QueryObjectName(device.DriverObject);

		strName.Format(_T("FDO:%#x - "));
		strName += strDeviceName;
		strName += _T(" - ");
		strName += strDriverName;

		hCurrent = m_ctrlDeviceTree.InsertItem(strName,hCurrent);

		pCurrentFdo = device.AttachedDevice;
	}

	struct _dev_rel
	{
		DWORD m_dwCount;
		PDEVICE_OBJECT m_pDevices[0x400];
	};

	_dev_rel outRel;
	outRel.m_dwCount = 0;

	if(DeviceIoControl(m_hDeviceFile,DEVICE_TREE_QUERY_BUS_REL,&pDeviceObject,sizeof(PVOID),
					   &outRel,sizeof(outRel),&dummy,0))
	{
		for(DWORD i = 0; i < outRel.m_dwCount; i ++)
		{
			BuildDeviceTree(outRel.m_pDevices[i],hCurrent);
		}
	}
}

void CDeviceTreeAppDlg::OnOK()
{
}

void CDeviceTreeAppDlg::OnCancel()
{
	if(m_hDeviceFile != INVALID_HANDLE_VALUE)
		CloseHandle(m_hDeviceFile);

	CDialog::OnCancel();
}