Virtual Desktop: A Simple Desktop Management Tool

#include "StdAfx.h"
#include "DesktopManager.h"
#include "SHLWAPI.h"

int CDesktopManager::m_iDesktopCount = 0;
TCHAR *CDesktopManager::m_szDesktopNames[ARRAY_SIZE] = {0};

//Function enumerates all the desktops of the windows station.
BOOL CALLBACK EnumDesktopProc(LPTSTR lpszDesktopName,LPARAM lParam)
{
	try
	{
		//Adding Desktop Name To DesktopName List.
		CDesktopManager::AddDesktop(lpszDesktopName);
	}
	catch(...)
	{
		OutputDebugString(_T("\n\nException caught in EnumDesktopsProc."));
	}

	return TRUE;
}

CDesktopManager::CDesktopManager(void)
//: m_iDesktopCount(0)
{
	for(int iCount =0;iCount < ARRAY_SIZE;iCount ++)
		m_szDesktopNames[iCount] = NULL;

	//Populates all the desktop names.
	PopulateDesktopList();
}

CDesktopManager::~CDesktopManager(void)
{
	ReleaseMemory();
}

bool CDesktopManager::PopulateDesktopList(void)
{
	bool bReturn = false;
	try
	{
		HWINSTA hWindowsStation = GetProcessWindowStation();
		if(NULL == hWindowsStation)
			throw false;

		if(0 < m_iDesktopCount)
			ReleaseMemory();

		m_iDesktopCount = 0;
		bReturn = (FALSE != EnumDesktops(hWindowsStation,&EnumDesktopProc,NULL)); //m_szDesktopNames));
		//for(int i = 0;i< GetDesktopCount(); i++)
		//	OutputDebugString(m_szDesktopNames[i]);
	}
	catch(bool bThrowVal)
	{
		bReturn = bThrowVal;
	}
	catch(...)
	{
		OutputDebugString(_T("\n\nException caught in CDesktopManager::PopulateDesktopList"));
	}

	return bReturn;
}

//Returns the number of desktops in the window station.
int CDesktopManager::GetDesktopCount(void)
{
	//Refresh the list.
	PopulateDesktopList();
	return m_iDesktopCount;
}

//Release the memory to avoid memory leak.
void CDesktopManager::ReleaseMemory(void)
{
	for(int iCount =0;iCount < ARRAY_SIZE;iCount ++)
	{
		if(m_szDesktopNames[iCount] != NULL)
		{
			delete m_szDesktopNames[iCount];
			m_szDesktopNames[iCount] = NULL;
		}
	}
	m_iDesktopCount = 0;
}

//Switch between desktops.
bool CDesktopManager::SwitchDesktop(TCHAR * szDesktopName)
{
	bool bReturn = false;
	try
	{
		if (NULL == szDesktopName)
		{
			OutputDebugString(_T("\nNULL DesktopName in CDesktopManager::SwitchToDesktop"));
			throw false;
		}

		/*if( !_tcsicmp(_T("Winlogon"),szDesktopName) || !_tcsicmp(_T("Disconnect"),szDesktopName))
		{
			TCHAR szErrorMsg[ARRAY_SIZE] = {_T("You can not switch to ")};
			_tcscat_s(szErrorMsg,ARRAY_SIZE -1,szDesktopName);
			_tcscat_s(szErrorMsg,ARRAY_SIZE -1,_T(" Desktop."));
			MessageBox(NULL,szErrorMsg,_T("Virtual Desktop"),MB_ICONINFORMATION);
			throw false;
		}*/

		//Open desktop handle to switch to.
		HDESK hDesktopToSwitch = OpenDesktop(szDesktopName,DF_ALLOWOTHERACCOUNTHOOK,TRUE,GENERIC_ALL);
		if(NULL == hDesktopToSwitch)
		{
			TCHAR *pszError = NULL;
			TCHAR szErrorMsg[ARRAY_SIZE] = {0};
			int iErrorCode = GetLastError();
			if(5 == iErrorCode)
			{
				
				FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
								NULL,iErrorCode,0,(LPWSTR) &pszError,0,NULL);
				wsprintf(szErrorMsg,_T("Failed to switch to %s desktop.\n\t %s"),szDesktopName,pszError);
				MessageBox(NULL,szErrorMsg,_T("Virtual Desktop"),MB_ICONINFORMATION | MB_TOPMOST | MB_TASKMODAL);
				OutputDebugString(pszError);
			}

			wsprintf(szErrorMsg,_T("\nOpenDesktop Failed in CDesktopManager::SwitchToDesktop. Last Error : %s"),pszError);
			OutputDebugString(szErrorMsg);

			throw false;
		}

		//Switch the desktop.
		if(FALSE == ::SwitchDesktop(hDesktopToSwitch))
		{
			OutputDebugString(_T("\nSwitchDesktop Failed in CDesktopManager::SwitchToDesktop"));
			throw false;
		}
 
		//Close the desktop handle.
		CloseDesktop(hDesktopToSwitch);
		bReturn = true;
	}
	catch(bool bThrownVal)
	{
		bReturn = bThrownVal;
	}
	catch(...)
	{
		bReturn = false;
		OutputDebugString(_T("\nException caught in CDesktopManager::SwitchToDesktop"));
	}
	return bReturn;
}


//Creats new desktop.
bool CDesktopManager::CreateDesktop(TCHAR * szDesktopName)
{
	bool bReturn = false;
	try
	{
		if(NULL == szDesktopName)
		{
			OutputDebugString(_T("\nNULL DesktopName in CDesktopManager::CreateDesktop"));
			throw false;
		}

		//Setting the desktop security attributes.
		SECURITY_ATTRIBUTES sAttribute = {sizeof(SECURITY_ATTRIBUTES),NULL,TRUE};
		HDESK hNewDesktop = NULL;

		//Creating a new desktop with all access.
		hNewDesktop = ::CreateDesktop(szDesktopName,NULL,NULL,DF_ALLOWOTHERACCOUNTHOOK,GENERIC_ALL,&sAttribute);

		TCHAR *pszError = NULL;
		FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
						NULL,GetLastError(),0,(LPWSTR) &pszError,0,NULL);
		OutputDebugString(_T("\n\t\t"));
		OutputDebugString(pszError);

		int iCount = 0;
		//Checking whether desktop name already exists.
		for(iCount = 0;iCount < m_iDesktopCount;iCount++)
			if(_tcsicmp(szDesktopName,m_szDesktopNames[iCount]) == 0)
				break;

		if(iCount >= m_iDesktopCount)
		{
			OutputDebugString(_T("\nCreating New Desktop"));

			TCHAR szExplorerFile[ARRAY_SIZE]= {0};

			GetWindowsDirectory(szExplorerFile,ARRAY_SIZE -1);
			_tcscat_s(szExplorerFile,ARRAY_SIZE -1, _T("\\Explorer.Exe"));

			LaunchApplication(szExplorerFile,szDesktopName);
		}

		//Closing the desktop handle.
		if(NULL != hNewDesktop)
			CloseDesktop(hNewDesktop);

		bReturn = PopulateDesktopList();
	}
	catch(...)
	{
		OutputDebugString(_T("Exception Caught In CDesktopManager::CreateDesktop."));
	}
	return bReturn;
}

//Returns the desktop name at specified index.
void CDesktopManager::GetDesktopName(int iIndex, TCHAR * szDesktopName)
{
	try
	{
		if (iIndex < m_iDesktopCount)
			_tcscpy_s(szDesktopName,_tcslen(m_szDesktopNames[iIndex]) + 1, m_szDesktopNames[iIndex]);
	}
	catch(...)
	{
		OutputDebugString(_T("\n\nException caught in CDesktopManager::GetDesktopName."));
	}
}

//Adds the specified desktop name into the list.
void CDesktopManager::AddDesktop(TCHAR * lpszDesktopName)
{
	try
	{
		m_szDesktopNames[m_iDesktopCount] = new TCHAR[_tcslen(lpszDesktopName) + 1];
		ZeroMemory(m_szDesktopNames[m_iDesktopCount],_tcslen(lpszDesktopName) + 1);
		_tcscpy_s(m_szDesktopNames[m_iDesktopCount],_tcslen(lpszDesktopName) + 1,lpszDesktopName);
		m_iDesktopCount++;
	}
	catch(...)
	{
		OutputDebugString(_T("\n\nException caught in CDesktopManager::AddDesktop."));
	}
}

//Checks whether the specified desktop is current active desktop(deskop of the current thread).
bool CDesktopManager::IsCurrentDesktop(TCHAR * szDesktopName)
{
	bool bReturn = false;

	try
	{
		TCHAR szCurrentDesktopName[ARRAY_SIZE] = {0};
		DWORD iOutCount = 0;
		HDESK hCurrentDesktop = GetThreadDesktop(GetCurrentThreadId());
		GetUserObjectInformation(hCurrentDesktop,UOI_NAME,szCurrentDesktopName,ARRAY_SIZE - 1,&iOutCount);

		bReturn = (_tcsicmp(szDesktopName,szCurrentDesktopName) == 0);
	}
	catch(...)
	{
		OutputDebugString(_T("\n\nException caught in CDesktopManager::IsCurrentDesktop."));
		bReturn = false;
	}

	return bReturn;
}

bool CDesktopManager::LaunchApplication(TCHAR * szApplicationFilePath, TCHAR * szDesktopName)
{
	bool bReturn = false;

	try
	{
		if(!szApplicationFilePath || !szDesktopName || !_tcslen(szApplicationFilePath) || !_tcslen(szDesktopName))
			throw _T("Invalid Argument.");

		TCHAR szDirectoryName[ARRAY_SIZE] = {0};
		TCHAR szExplorerFile[ARRAY_SIZE]= {0};

		_tcscpy_s(szDirectoryName,_tcslen(szApplicationFilePath) + 1, szApplicationFilePath);
		
		if(!PathIsExe(szApplicationFilePath))
			throw _T("Invalid File Extension");

		PathRemoveFileSpec(szDirectoryName);

		STARTUPINFO sInfo = {0};
		PROCESS_INFORMATION pInfo = {0};

		sInfo.cb = sizeof(sInfo);
		sInfo.lpDesktop = szDesktopName;

		//Lanuching a application into dekstop
		BOOL bCreateProcessReturn = CreateProcess(szApplicationFilePath,
			NULL,
			NULL,
			NULL,
			TRUE,
			NORMAL_PRIORITY_CLASS,
			NULL,
			szDirectoryName,
			&sInfo,
			&pInfo);

		TCHAR *pszError = NULL;
		FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
			NULL,GetLastError(),0,(LPWSTR) &pszError,0,NULL);
		OutputDebugString(_T("\n\t\t"));
		OutputDebugString(pszError);

		if(bCreateProcessReturn)
			bReturn = true;
	}
	catch(TCHAR *pszError)
	{
		OutputDebugString(_T("\n\nCustom Exception caught in CDesktopManager::IsCurrentDesktop."));
		OutputDebugString(pszError);
		bReturn = false;
	}
	catch(bool)
	{
		OutputDebugString(_T("\n\nCustom Exception caught in CDesktopManager::IsCurrentDesktop."));
		bReturn = false;
	}
	catch(...)
	{
		OutputDebugString(_T("\n\nException caught in CDesktopManager::IsCurrentDesktop."));
		bReturn = false;
	}

	return bReturn;
}