Task Manager Extension 2.0

// TaskManagerEx.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "LoadDll.h"
#include <locale.h>
#include "SplashWnd.h"
#include "Resource.h"

#include "../TaskManagerExDll/Localization.h"
#include "../TaskManagerExDll/TaskManagerEx.h"
#include "../TaskManagerExDll/Taskmgr.h"
#include "../TaskManagerExDll/Options.h"

DWORD	dwWindowsNTMajorVersion = 0;
HWND	hwndTaskManager = NULL;

HANDLE hAnotherTaskManagerExEvent = NULL;
LPCTSTR szGlobalKernelName = _T("Local\\TaskManagerExIsAlreadyRunning");

static char g_szStamp_Mark[] = "\r\n\r\nTaskManagerEx.cpp Timestamp: " __DATE__ ", " __TIME__ "\r\n\r\n";

//////////////////////////////////////////////////////////////////////////

BOOL IsAnotherTaskManagerEx_Capture( BOOL bInitialState )
{
	if( hAnotherTaskManagerExEvent != NULL )
	{
		return TRUE;
	}
	hAnotherTaskManagerExEvent = CreateEvent( NULL, TRUE, bInitialState, szGlobalKernelName );	

	DWORD dwer = GetLastError();	

	return (dwer==ERROR_ALREADY_EXISTS);
}

//////////////////////////////////////////////////////////////////////////

void ReleaseAnotherTaskManagerEx()
{
	if( hAnotherTaskManagerExEvent != NULL )
	{
		CloseHandle( hAnotherTaskManagerExEvent );
		hAnotherTaskManagerExEvent = NULL;
	}
}

//////////////////////////////////////////////////////////////////////////

DWORD GetTaskManagerProcessID()
{
	DWORD pID = 0;

	GetWindowThreadProcessId( hwndTaskManager, &pID );

	return pID;
}

//////////////////////////////////////////////////////////////////////////

//void MySleep( DWORD dwMilliseconds ) // for simplier debug
//{
//	Sleep(dwMilliseconds);
//}

DWORD MyWaitForSingleObject( HANDLE hHandle, DWORD dwMilliseconds ) // for simplier debug
{
	DWORD res = WaitForSingleObject( hHandle, dwMilliseconds );
	return res;
}

//////////////////////////////////////////////////////////////////////////

TCHAR szDllPath[_MAX_PATH] = _T(""); // TaskManagerEx.dll name

//////////////////////////////////////////////////////////////////////////

VOID CALLBACK ThreadTimerProc(
	HWND hwnd,
    UINT uMsg,
    UINT idEvent,
    DWORD dwTime
)
{
	hwnd = hwnd;
	uMsg = uMsg;
	idEvent = idEvent;
	dwTime = dwTime;

	hwndTaskManager = FindTaskManagerWindow();
	if( hwndTaskManager != NULL )
	{
		DWORD processId = GetTaskManagerProcessID();
		#ifdef _DEBUG
		TCHAR s[200] = _T("");
		wsprintf( s, _T("TaskManagerEx> LoadDllForRemoteThread for PID = %d\n"),
			processId );
		OutputDebugString( s );
		#endif

		DWORD dwRet = LoadDllForRemoteThread( 
				processId, 
				TRUE, 
				FALSE, 
				szDllPath,
				"Initialize",
				NULL,
				NULL,
				NULL,
				NULL,
				NULL,
				0,
				NULL
				);

		if( dwRet != 0 )
		{
			#ifdef _DEBUG
			TCHAR s[200] = _T("");
			wsprintf( s, _T("TaskManagerEx> LoadDllForRemoteThread for PID = %d returned %d (0x%X)\n"),
				processId, dwRet, dwRet );
			OutputDebugString( s );
			#endif
		}

		BOOL bRunOnceOnly = FALSE;

		{
			HKEY hKey = NULL;
			DWORD dwType = REG_DWORD;
			DWORD dwSize = sizeof(DWORD);
			LONG lRes;
			DWORD dwValue;

			::RegOpenKey( REG_ROOT, REG_KEY, &hKey );

			if( hKey != NULL )
			{
				dwType = REG_DWORD; dwSize = sizeof(DWORD);
				lRes = ::RegQueryValueEx( hKey, _T("LoaderStartOnceOnly"), NULL, &dwType, 
						(LPBYTE)&dwValue, &dwSize );

				if( lRes == ERROR_SUCCESS )
				{
					bRunOnceOnly = dwValue != 0;
				}

				::RegCloseKey( hKey );
			}
		}

		if( bRunOnceOnly )
		{
			ResetEvent( hAnotherTaskManagerExEvent ); // Stop process after one Dll load
		}
	}

	DWORD res = MyWaitForSingleObject( hAnotherTaskManagerExEvent, 0 );
	if( res != WAIT_OBJECT_0 )
	{
		PostQuitMessage( 0 );
	}
}

//////////////////////////////////////////////////////////////////////////

//LPCTSTR g_szLocale = NULL;

//////////////////////////////////////////////////////////////////////////

int MainFunction()
{
//	g_szLocale = _tsetlocale( LC_ALL, _T("") ); // it is commented because CRT was reduced from this project
												// so there is no need in localization of CRT library
	BOOL res;
	DWORD dwBufSize;

	TCHAR szCaption[200] = _T("");
	TCHAR szRequireNT[200] = _T("");
	TCHAR szGetModuleFilenameError[200] = _T("");
	TCHAR szCommandLineHelp[5000] = _T("");

	dwBufSize = SIZEOF_ARRAY(szCaption);
	LocLoadString(IDS_MESSAGEBOX_CAPTION, szCaption, &dwBufSize );
	dwBufSize = SIZEOF_ARRAY(szRequireNT);
	LocLoadString(IDS_REQUIRE_NT, szRequireNT, &dwBufSize );
	dwBufSize = SIZEOF_ARRAY(szGetModuleFilenameError);
	LocLoadString(IDS_GETMODULEFILNAME_ERROR, szGetModuleFilenameError, &dwBufSize );
	dwBufSize = SIZEOF_ARRAY(szCommandLineHelp);
	LocLoadString(IDS_COMMANDLINE_HELP, szCommandLineHelp, &dwBufSize );

	LPWSTR szCommandLineW = GetCommandLineW();
	if( szCommandLineW == NULL )
	{
		return -1;
	}
	int argc = 0;
	LPWSTR* argv = CommandLineToArgvW( szCommandLineW, &argc );
	if( argv == NULL )
	{
		MessageBox( NULL, szRequireNT, szCaption, MB_OK | MB_ICONERROR );
		return -1;
	}

	BOOL bAnother = FALSE;

	BOOL bUnload =	FALSE;
	BOOL bRestart =	FALSE;

	if( argc > 1 )
	{
		bUnload =	lstrcmpW( argv[1], L"/u" ) == 0 ||
					lstrcmpW( argv[1], L"/U" ) == 0;
		bRestart =	lstrcmpW( argv[1], L"/r" ) == 0 ||
					lstrcmpW( argv[1], L"/R" ) == 0;

		if( argc > 2 || ( !bUnload && !bRestart ) )
		{
			MessageBox( NULL, szCommandLineHelp, szCaption, MB_OK | MB_ICONINFORMATION );
			GlobalFree( argv );
			return 1;
		}

		if( bUnload )
		{
			bAnother = IsAnotherTaskManagerEx_Capture( TRUE );
			if( bAnother )
			{
				ResetEvent( hAnotherTaskManagerExEvent );
			}
			ReleaseAnotherTaskManagerEx();
			GlobalFree( argv );
			return 0;
		}
		else if( bRestart )
		{
			bAnother = IsAnotherTaskManagerEx_Capture( TRUE );
			if( bAnother )
			{
				ResetEvent( hAnotherTaskManagerExEvent );
				DWORD dwRet = MyWaitForSingleObject( hAnotherTaskManagerExEvent, 3000 );
				if( dwRet != WAIT_OBJECT_0 )
				{
					ASSERT( FALSE && "WaitForSingleObject error!" );
					GlobalFree( argv );
					return -1;
				}
			}
			bAnother = FALSE;
		}
		else
		{
			ASSERT(FALSE);
			// return -1; // warning C4702: unreachable code
		}
	}
	else
	{
		bAnother = IsAnotherTaskManagerEx_Capture( TRUE );
	}

	GlobalFree( argv );
	argv = NULL;

	if( bAnother != FALSE )
	{
		return 0;
	}

	if( hAnotherTaskManagerExEvent == NULL )
	{
		ASSERT(FALSE);
		return -1;
	}

	dwWindowsNTMajorVersion = IsWindowsNT();

	if ( dwWindowsNTMajorVersion < 4 )
	{
		MessageBox( NULL, szRequireNT, szCaption, MB_OK | MB_ICONERROR );
		return -1;
	}


	res = GetModuleFileName( NULL, szDllPath, _MAX_PATH );
	if ( !res )
	{
		MessageBox( NULL, szGetModuleFilenameError, szCaption, MB_OK | MB_ICONERROR );
		return -2;
	}

	LPTSTR p = NULL;

	for( int i = lstrlen(szDllPath)-1; i>=0; i-- )
	{
		if( szDllPath[i] == _T('\\') )
		{
			p = &szDllPath[i];
			break;
		}
	}

	if( p == NULL )
	{
	   p = szDllPath;
	}
	else
	{
	   p++;
	}

	lstrcpy( p, TMEX_INJECT_DLL );

	//////////////////////////////////////////////////////////////////////////
	//Let's work

	COptions options;
	BOOL bSplash = options.GetShowSplash();

	if( bSplash )
	{
		res = CreateSplashScreen();
	}

	SetTimer( NULL, 0, 400, ThreadTimerProc );

	MSG msg;

	while( (res = GetMessage( &msg, NULL, 0, 0 )) != 0)
	{ 
		if (res == -1)
		{
			// handle the error and possibly exit
			break;
		}
		else
		{
			TranslateMessage(&msg); 
			DispatchMessage(&msg); 
		}
	} 

	SetEvent( hAnotherTaskManagerExEvent );
	ReleaseAnotherTaskManagerEx();

	return 0;
}

///////////////////////////////////////////////////////////////////////////////////
///  CRT:

/*
#ifdef SUBSYSTEM_WINDOWS

int CALLBACK _tWinMain( HINSTANCE, HINSTANCE, LPTSTR, int )
{
	int res = MainFunction();
	return res;
}

#else

int _tmain( int, TCHAR ** ) //( int argc, TCHAR **argv )
{
	int res = MainFunction();
	return res;
}

#endif
//*/

//////////////////////////////////////////////////////////////////////////