Virtual Desktop: A Simple Desktop Management Tool

// Virtual DesktopDlg.cpp : implementation file
//

#include "stdafx.h"
#include "Virtual Desktop.h"
#include "Virtual DesktopDlg.h"
#include "DesktopManager.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#endif

const UINT WM_TRAYICON_NOTIFY_MESSAGE = RegisterWindowMessage(_T("WM_TRAYICON_NOTIFY_MESSAGE-{8DDBE93E-DFE8-4279-934E-05C39902F37D}"));
// CAboutDlg dialog used for App About

#define CONTEXT_MENU_IDS 600
#define MANAGE_DESKTOP_MENU_ID 500
#define VERIFY_SWITCH_MENU_ID 501
#define LAUNCH_APP_MENU_ID 502
#define EXIT_MENU_ID 503
#define SEPARATOR_MENU_ID 504

typedef bool(*InstallHook)(void);

class CAboutDlg : public CDialog
{
public:
	CAboutDlg();

// Dialog Data
	enum { IDD = IDD_ABOUTBOX };

	protected:
	virtual void DoDataExchange(CDataExchange* pDX);    // DDX/DDV support

// Implementation
protected:
	DECLARE_MESSAGE_MAP()
};

CAboutDlg::CAboutDlg() : CDialog(CAboutDlg::IDD)
{
}

void CAboutDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
}

BEGIN_MESSAGE_MAP(CAboutDlg, CDialog)
END_MESSAGE_MAP()


// CVirtualDesktopDlg dialog




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

void CVirtualDesktopDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_DESKTOP_LIST, m_DesktopListControl);
	DDX_Control(pDX, IDC_DESKTOP_NAME, m_DesktopNameControl);
	DDX_Control(pDX, IDC_ADD_NEW_DESKTOP, m_AddNewDesktop);
	DDX_Control(pDX, IDC_SWITCH_TO_DESKTOP, m_SwitchToDesktop);
	DDX_Control(pDX, IDC_VERIFY_CHECK, m_ChkVerifyDesktopSwitch);
}

BEGIN_MESSAGE_MAP(CVirtualDesktopDlg, CDialog)
	ON_WM_SYSCOMMAND()
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	ON_REGISTERED_MESSAGE(WM_TRAYICON_NOTIFY_MESSAGE,OnTrayMessage)
	//}}AFX_MSG_MAP
	ON_WM_DESTROY()
	ON_LBN_SELCHANGE(IDC_DESKTOP_LIST, &CVirtualDesktopDlg::OnLbnSelchangeDesktopList)
	ON_BN_CLICKED(IDC_ADD_NEW_DESKTOP, &CVirtualDesktopDlg::OnBnClickedAddNewDesktop)
	ON_BN_CLICKED(IDC_SWITCH_TO_DESKTOP, &CVirtualDesktopDlg::OnBnClickedSwitchToDesktop)
	ON_BN_CLICKED(IDC_LAUNCH_APPLICATION, &CVirtualDesktopDlg::OnBnClickedLaunchApplication)
END_MESSAGE_MAP()


// CVirtualDesktopDlg message handlers

BOOL CVirtualDesktopDlg::OnInitDialog()
{
	CDialog::OnInitDialog();

	// Add "About..." menu item to system menu.

	// IDM_ABOUTBOX must be in the system command range.
	ASSERT((IDM_ABOUTBOX & 0xFFF0) == IDM_ABOUTBOX);
	ASSERT(IDM_ABOUTBOX < 0xF000);

	CMenu* pSysMenu = GetSystemMenu(FALSE);
	if (pSysMenu != NULL)
	{
		CString strAboutMenu;
		strAboutMenu.LoadString(IDS_ABOUTBOX);
		if (!strAboutMenu.IsEmpty())
		{
			pSysMenu->AppendMenu(MF_SEPARATOR);
			pSysMenu->AppendMenu(MF_STRING, IDM_ABOUTBOX, strAboutMenu);
		}
	}

	// 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
	//Adding try notify icon for the application.
	NOTIFYICONDATA nData;
	nData.cbSize = sizeof(NOTIFYICONDATA);
	nData.hIcon = LoadIcon(AfxGetInstanceHandle(),MAKEINTRESOURCE(IDR_MAINFRAME));
	nData.hWnd = m_hWnd;
	wcscpy_s(nData.szTip,127,_T("Virtual Desktop"));
	nData.uCallbackMessage = WM_TRAYICON_NOTIFY_MESSAGE;
	nData.uFlags = NIF_ICON | NIF_TIP | NIF_MESSAGE ;
	nData.uID = 1;

	Shell_NotifyIcon(NIM_ADD,&nData);


	//Load the Event hooker dll
	HMODULE hModule = LoadLibrary(_T("Event Hooker Dll.dll"));
	InstallHook fpInstallHook = NULL;

	if(NULL != hModule)
	{
		//Get the windows procedure hook installer function address.
		fpInstallHook = (InstallHook) GetProcAddress(hModule,("InstallWinProcHook"));

		if(NULL != fpInstallHook)
			if(!fpInstallHook()) //Install the windows procedure hook.
				MessageBox(_T("Failed to install hooks."),_T("Virtual Desktop"),MB_ICONINFORMATION | MB_TOPMOST | MB_TASKMODAL);

		//Get the message hook installer function address.
		fpInstallHook = (InstallHook) GetProcAddress(hModule,("InstallMessageHook"));

		if(NULL != fpInstallHook)
			if(!fpInstallHook()) //Install the message hook.
				MessageBox(_T("Failed to install hooks."),_T("Virtual Desktop"),MB_ICONINFORMATION | MB_TOPMOST | MB_TASKMODAL);
	}
	else
		MessageBox(_T("Failed to library."),_T("Virtual Desktop"),MB_ICONINFORMATION | MB_TOPMOST | MB_TASKMODAL);


	return TRUE;  // return TRUE  unless you set the focus to a control
}

void CVirtualDesktopDlg::OnSysCommand(UINT nID, LPARAM lParam)
{
	if ((nID & 0xFFF0) == IDM_ABOUTBOX)
	{
		CAboutDlg dlgAbout;
		dlgAbout.DoModal();
	}
	else
	{
		CDialog::OnSysCommand(nID, lParam);
	}
}

// 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 CVirtualDesktopDlg::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 CVirtualDesktopDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}


void CVirtualDesktopDlg::OnDestroy()
{
	CDialog::OnDestroy();

	// TODO: Add your message handler code here
	NOTIFYICONDATA nData;
	nData.cbSize = sizeof(NOTIFYICONDATA);
	nData.hWnd = m_hWnd;
	nData.uID = 1;

	Shell_NotifyIcon(NIM_DELETE,&nData);

	//Load the event hooker dll.
	HMODULE hModule = LoadLibrary(_T("Event Hooker Dll.Dll"));

	if(NULL != hModule)
	{
		InstallHook fpUninstallHook = NULL;
		//Get the hook uninstaller function address.
		fpUninstallHook = (InstallHook) GetProcAddress(hModule,("UnInstallMsgHook"));
		if(NULL == fpUninstallHook || !fpUninstallHook()) //Uninstall the hooks.
			OutputDebugString(_T("Failed to uninstall Msg Hook in OnDestroy()"));

		fpUninstallHook = NULL;
		//Get the hook uninstaller function address.
		fpUninstallHook = (InstallHook) GetProcAddress(hModule,("UnInstallWinProcHook"));
		if(NULL == fpUninstallHook || !fpUninstallHook()) //Uninstall the hooks.
			OutputDebugString(_T("Failed to uninstall WinProc Hook in OnDestroy()"));
	}

	//Release the memory to avoid the memory leaks.
	CDesktopManager::ReleaseMemory();
}

LRESULT CVirtualDesktopDlg::OnTrayMessage(WPARAM wParam, LPARAM lParam)
{
	UINT uMsg = (UINT) lParam;

	if(uMsg == WM_RBUTTONDOWN || uMsg == WM_CONTEXTMENU)
	{
		POINT pt;
		GetCursorPos(&pt);

		//Display a tray menu with all the desktop names as menu items.
		int iDesktopCounts = CDesktopManager::GetDesktopCount();
		int iMenuItemCount = 0;

		HMENU hContextMenu = CreatePopupMenu();

		//Iterate to add all the desktop names as menu items into the tray menu.
		for(iMenuItemCount = 0;iMenuItemCount < iDesktopCounts ;iMenuItemCount++)
		{
			TCHAR szDesktopName[ARRAY_SIZE] = {0};
			CDesktopManager::GetDesktopName(iMenuItemCount,szDesktopName);
			
			if(CDesktopManager::IsCurrentDesktop(szDesktopName))
				AppendMenu(hContextMenu,MF_STRING | MF_ENABLED | MF_CHECKED,CONTEXT_MENU_IDS + iMenuItemCount ,szDesktopName);
			else
				AppendMenu(hContextMenu,MF_STRING | MF_ENABLED,CONTEXT_MENU_IDS + iMenuItemCount ,szDesktopName);
		}


		if(iMenuItemCount > 0)
			AppendMenu(hContextMenu, MF_ENABLED | MF_SEPARATOR, SEPARATOR_MENU_ID,_T("Virtual Desktop Separator"));

		if(m_ChkVerifyDesktopSwitch.GetCheck())
			AppendMenu(hContextMenu, MF_ENABLED | MF_STRING | MF_CHECKED, VERIFY_SWITCH_MENU_ID,_T("&Confirm \"Desktop Switch\""));
		else
			AppendMenu(hContextMenu, MF_ENABLED | MF_STRING, VERIFY_SWITCH_MENU_ID,_T("&Confirm \"Desktop Switch\""));

		AppendMenu(hContextMenu, MF_ENABLED | MF_STRING, MANAGE_DESKTOP_MENU_ID,_T("&Manage Desktops"));
		AppendMenu(hContextMenu, MF_ENABLED | MF_STRING, LAUNCH_APP_MENU_ID,_T("&Launch Application"));
		AppendMenu(hContextMenu, MF_ENABLED | MF_SEPARATOR, SEPARATOR_MENU_ID,_T("Virtual Desktop Separator"));
		AppendMenu(hContextMenu, MF_ENABLED | MF_STRING, EXIT_MENU_ID,_T("&Exit"));
		
		SetForegroundWindow();
		//Display the context menu.
		int iSelectedIndex = TrackPopupMenu(hContextMenu,TPM_TOPALIGN | TPM_VERPOSANIMATION | TPM_RETURNCMD,pt.x,pt.y,0,m_hWnd,NULL);

		switch(iSelectedIndex)
		{
			case EXIT_MENU_ID:
				//Exit the application.
				PostQuitMessage(0);
				break;
			case MANAGE_DESKTOP_MENU_ID:
				//Show the Manage Desktop dialog.
				ShowManageDesktopsDialog();
				break;
			case VERIFY_SWITCH_MENU_ID:
				//Toggle the check.
				m_ChkVerifyDesktopSwitch.SetCheck(!m_ChkVerifyDesktopSwitch.GetCheck());
				break;
			case LAUNCH_APP_MENU_ID:
				OnBnClickedLaunchApplication();
				break;
			default:
				if(iSelectedIndex >= CONTEXT_MENU_IDS)
				{
					TCHAR szSwitchToDesktopName[ARRAY_SIZE] = {0};
					//Get the desktop name to be switched to.
					CDesktopManager::GetDesktopName(iSelectedIndex - CONTEXT_MENU_IDS,szSwitchToDesktopName);
					//Switch to the selected desktop.
					if(_tcslen(szSwitchToDesktopName))
						SwitchDesktopTo(szSwitchToDesktopName);
				}
		}
	}

	return LRESULT();
}

//Display the Manage Desktops dialog.
void CVirtualDesktopDlg::ShowManageDesktopsDialog(void)
{
	try
	{
		//CDesktopManager objDeskManager;
		int iDeskCount = CDesktopManager::GetDesktopCount();
		//Clear the list box items.
		m_DesktopListControl.ResetContent();

		//Insert all the desktop names into the list box.
		for(int i = 0;i < iDeskCount;i++)
		{
			TCHAR szTempDeskName[ARRAY_SIZE]  = {0};
			CDesktopManager::GetDesktopName(i,szTempDeskName);
			m_DesktopListControl.AddString(szTempDeskName);
		}

		TCHAR szSelectedDesktopName[ARRAY_SIZE] = {0};

		m_DesktopNameControl.GetWindowText(szSelectedDesktopName,ARRAY_SIZE - 1);

		if(0 != _tcslen(szSelectedDesktopName) && -1 != m_DesktopListControl.SelectString(0,szSelectedDesktopName))
			m_DesktopNameControl.EnableWindow(FALSE);
		else
		{
			m_DesktopNameControl.SetWindowText(NULL);
			m_DesktopNameControl.EnableWindow(TRUE);
		}

		ShowWindow(SW_SHOW);
	}
	catch(...)
	{
		OutputDebugString(_T("\nException caught in CVirtualDesktopDlg::ShowManageDesktopsDialog."));
	}
}

void CVirtualDesktopDlg::OnLbnSelchangeDesktopList()
{
	// TODO: Add your control notification handler code here
	TCHAR szSelectedDesktopName[ARRAY_SIZE] = {0};
	m_DesktopListControl.GetText(m_DesktopListControl.GetCurSel(),szSelectedDesktopName);
	m_DesktopNameControl.SetWindowText(szSelectedDesktopName);

	m_AddNewDesktop.SetWindowText(_T("&New"));
	m_SwitchToDesktop.EnableWindow(TRUE);
	m_DesktopNameControl.EnableWindow(FALSE);
}

//Creates the new desktop (Adds new desktop)
void CVirtualDesktopDlg::OnBnClickedAddNewDesktop()
{
	// TODO: Add your control notification handler code here
	TCHAR szCaption[ARRAY_SIZE] = {0};
	m_AddNewDesktop.GetWindowText(szCaption,ARRAY_SIZE - 1);
	if(_tcsicmp(szCaption,_T("&New")))
	{
		m_DesktopNameControl.GetWindowText(szCaption,ARRAY_SIZE -1);

		//Left trimming the string.
		while(' ' == szCaption[0])
			_tcscpy_s(szCaption,ARRAY_SIZE - 1, szCaption + 1);
		
		int iLen = (int) _tcslen(szCaption);
		//Right trimming the string.
		while(' ' == szCaption[--iLen])
			szCaption[iLen] = '\0';

		if(_tcslen(szCaption))
		{
			//Check the desktop name in the list.
			if(-1 != m_DesktopListControl.SelectString(0,szCaption))
			{
				MessageBox(_T("Desktop already created !"),_T("Virtual Desktop"),MB_ICONEXCLAMATION | MB_TOPMOST | MB_TASKMODAL);
			}

			m_DesktopListControl.AddString(szCaption);
			m_DesktopListControl.SelectString(0,szCaption);
			OnLbnSelchangeDesktopList();
			
			//Create the desktop
			CDesktopManager::CreateDesktop(szCaption);
			if(IDYES == MessageBox(_T("New Desktop is been created.\nWould you like to switch to new desktop ?"),_T("Virtual Desktop"),MB_YESNO | MB_ICONINFORMATION | MB_TOPMOST | MB_TASKMODAL))
				SwitchDesktopTo(szCaption);

		}
		else
		{
			MessageBox(_T("Please enter Desktop Name"),_T("Virtual Desktop"),MB_ICONEXCLAMATION | MB_TOPMOST | MB_TASKMODAL);
			m_DesktopNameControl.SetWindowText(_T("")); 
			m_DesktopNameControl.SetFocus();
		}
	}
	else
	{
		m_AddNewDesktop.SetWindowText(_T("&Add"));
		m_DesktopNameControl.SetWindowText(_T(""));
		m_DesktopNameControl.EnableWindow(TRUE);
		m_SwitchToDesktop.EnableWindow(FALSE);
		m_DesktopNameControl.SetFocus();
	}
}

//Switch desktop
void CVirtualDesktopDlg::OnBnClickedSwitchToDesktop()
{
	// TODO: Add your control notification handler code here

	TCHAR szSwitchToDesktopName[ARRAY_SIZE] = {0};

	m_DesktopListControl.GetText(m_DesktopListControl.GetCurSel(),szSwitchToDesktopName);
	//Switching the desktop to specified one.
	SwitchDesktopTo(szSwitchToDesktopName);
}

void CVirtualDesktopDlg::SwitchDesktopTo(TCHAR * szDesktopName)
{
	try
	{
		//CDesktopManager objDeskManager; 
		if(NULL == szDesktopName)
			return;

		//Checking whether we're in the same desktop.
		if(CDesktopManager::IsCurrentDesktop(szDesktopName))
		{
			//If we're in the same specified desktop, just return.
			MessageBox(_T("You are currently on the same Desktop."),_T("Virtual Desktop"),MB_ICONINFORMATION  | MB_TOPMOST | MB_TASKMODAL);
			return;
		}

		if(m_ChkVerifyDesktopSwitch.GetCheck())
			if(IDNO == MessageBox(_T("Are you sure to switch to selected Desktop ?"),_T("Virtual Desktop"),MB_YESNO | MB_ICONINFORMATION | MB_TOPMOST | MB_TASKMODAL))
				return;

		//Exit the application from current desktop.
		if(CDesktopManager::SwitchDesktop(szDesktopName))
		{
			TCHAR szAppName[ARRAY_SIZE] = {0};
			//Get the application full path, so that it can be launch into the switching desktop.
			GetModuleFileName(GetModuleHandle(NULL),szAppName,ARRAY_SIZE - 1);

			//Launch the same application into the switching desktop.
			CDesktopManager::LaunchApplication(szAppName,szDesktopName);
			PostQuitMessage(0);
		}
	}
	catch(...)
	{
		OutputDebugString(_T("\nException caught in CVirtualDesktopDlg::SwitchDesktopTo."));
	}
}
void CVirtualDesktopDlg::OnBnClickedLaunchApplication()
{
	// TODO: Add your control notification handler code here

	TCHAR szDesktopName[ARRAY_SIZE] = {0};
	m_DesktopListControl.GetText(m_DesktopListControl.GetCurSel(), szDesktopName);

	//Checking the Desktop name.
	if(!_tcslen(szDesktopName))
	{
		MessageBox(_T("Please select the desktop name."), _T("Virtual Desktop"));
		//Making the Manage Desktop Dialog Visible. (If selected "Launch Application" From context menu.
		ShowManageDesktopsDialog();
		return;
	}

	if(!_tcsicmp(szDesktopName, _T("WinLogon")) || !_tcsicmp(szDesktopName, _T("Disconnect")) ) 
	{
		MessageBox(_T("Application cann't be launched in this Desktop."), _T("Virtual Desktop"));
		//Making the Manage Desktop Dialog Visible. (If selected "Launch Application" From context menu.
		ShowManageDesktopsDialog();
		return;
	}

	CFileDialog dlgOpen(TRUE, _T("*.Exe|"), NULL, 4|2, _T("Applications (*.Exe)|*.Exe|"), this);

	if(IDOK == dlgOpen.DoModal())
	{
		CString szFileName = dlgOpen.GetPathName();

		//Launching the selected application.
		if(CDesktopManager::LaunchApplication(szFileName.GetBuffer(), szDesktopName))
			MessageBox(_T("Application is launched into the specified Desktop."), _T("Virtual Desktop"), MB_ICONINFORMATION);
		else
			MessageBox(_T("Failed to launch application into the specified Desktop."), _T("Virtual Desktop"), MB_ICONERROR);

		szFileName.ReleaseBuffer();
	}
}