//==============================================================================
//
//     Copyright : (C) 2008 Videlli Limited
//     This file belongs to Videlli Limited
//
//==============================================================================
//
//     Purpose :
//     This file contains the win32 "buildtool attend" service source code.
//     This is a buildtool helper service, designed to "attend" to rogue processes
//     that insist on presenting a dialog on the console of unattended build machine
//     through termination.
//     This has the effect of holding a build thread to ransom.
//     The rules of engagement are:
//     1 a configured rogue process is the descendant of a configured ancestor process (the build daemon win32 service).
//     2 a configured rogue process has been running longer than a configured maximum run time.
//     Only when 1 and 2 are true will the rogue process be terminated.
//     The service is designed to be as lightweight as possible, checking for such a
//     scenario once every 60 seconds, or responding to a stop or shutdown request immediately.
//     The rogue tools are limited to a suite of borland tools at the time of writing,
//     though this list is configurable should the need arise to extend the list.
//     They are configured in the following registry key:
//     HKEY_LOCAL_MACHINE\\SYSTEM\\CurrentControlSet\\Services\\buildtool attend\\Parameters
//     AppParameters String Value.
//     This takes the form:
//     <fully qualified ancestor path> <max run time secs> <rogue child 1 process name> <rogue child 2 process name> ...
//     For example:
//     D:\\Work\\buildtool\\srvany.exe 60 bcc32.exe bpr2mak.exe brc32.exe brcc32.exe coff2omf.exe dcc32.exe ilink32.exe impdef.exe implib.exe make.exe tlib.exe tlibimp.exe
//     
//==============================================================================

#include <windows.h>
#include <iostream.h>
#include <vector>
#include <sstream>
#include "Psapi.h"
#include "Tlhelp32.h"
#include "Process.h"

void WINAPI ServiceMain( DWORD argc, LPTSTR * argv );
void WINAPI ControlHandler( DWORD dwControl );
void Start( DWORD argc, LPSTR * argv );
void Stop();
HANDLE hWaitEvent = 0L;
SERVICE_STATUS_HANDLE stat_hand = {0};
SERVICE_STATUS servstat = {0};
char *serviceName = "buildtool attend\0";

//------------------------------------------------------------------------------
//  log
//  Logs a message to the Application Event Log
//------------------------------------------------------------------------------
//  Parameters:
//      message IN - the message to log
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void log( const char * message )
{
	HANDLE hEventLog = RegisterEventSource( 0L, "Application" );

	if ( hEventLog != 0L )
	{
		ReportEvent(  hEventLog, EVENTLOG_INFORMATION_TYPE, 0, 0, 0L, 1, 0, &message, 0L );
		DeregisterEventSource( hEventLog );
	}
}

//------------------------------------------------------------------------------
//  isDescendedFrom
//  Determines if a child process is descended from an ancestor process
//------------------------------------------------------------------------------
//  Parameters:
//      ancestor IN - the ancestor process
//      child IN - the child process
//      processCollection IN - the collection of running processes
//------------------------------------------------------------------------------
//  Returns:
//      true if the child is descended from the ancestor
//------------------------------------------------------------------------------
bool isDescendedFrom( const std::string &ancestor, const Process &child, const std::vector< Process * > processCollection )
{
	bool retVal = false;

	if ( child.mPath == ancestor )
	{
		// child is descended from ancestor
		retVal = true;
	}
	else
	{
		// find the childs parent
		Process *parent = 0L;
		std::vector< Process * >::const_iterator processCollectionIterator = processCollection.begin();

		while ( processCollectionIterator != processCollection.end() )
		{
			Process *p = *processCollectionIterator;
			if ( child.mPpid == p->mPid  )
			{
				parent = p;
				break;
			}
		
			processCollectionIterator ++;
		}

		if ( parent != 0L )
		{
			retVal = isDescendedFrom( ancestor, *parent, processCollection );
		}
	}

	return retVal;
}

//------------------------------------------------------------------------------
//  main
//  This is the entry point for the process.
//  It sets up a service table entry and then calls StartServiceCtrlDispatcher.
//  This call does not return, but uses the thread which called it as a
//  control dispatcher for all the services implemented by this process.
//------------------------------------------------------------------------------
//  Parameters:
//      argc IN - the number of arguments on the command line
//      argv IN - the arguments on the command line
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void __cdecl main( DWORD argc, LPSTR * argv )
{
	// start the control dispatcher
	// this call gives the SCManager this thread for the entire period that this service is running,
	// so that it can call us back with service controls
	// it will spawn a new thread to run the service itself, starting at entrypoint ServiceMain
	SERVICE_TABLE_ENTRY stab[] = {
		{ serviceName, ServiceMain },
		{ NULL, NULL }
	};

	StartServiceCtrlDispatcher( stab );
}

//------------------------------------------------------------------------------
//  ServiceMain
//  This is the main entry point for the service.
//  The service control dispatcher creates a thread to start at this routine.
//------------------------------------------------------------------------------
//  Parameters:
//      argc IN - the number of arguments on the command line
//      argv IN - the arguments on the command line
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void WINAPI ServiceMain( DWORD argc, LPTSTR * argv )
{
	// register a service control handler for the service
	// this will be called by the control dispatcher when it has instructions for the service
	stat_hand = RegisterServiceCtrlHandler( serviceName, ControlHandler );

	if ( stat_hand == 0 )
	{
		std::string message = "ServiceMain RegisterServiceCrtlHandler failed with last error " + GetLastError();
		log( message.c_str() );
	}

	// create the handle used by the ControlHandler to signal to this method to return
	hWaitEvent = CreateEvent( 0L, FALSE, FALSE, 0L );

	if ( hWaitEvent == 0L )
	{
		std::string message = "ServiceMain CreateEvent failed with last error " + GetLastError();
		log( message.c_str() );
	}

	// inform the service control manager that the service is running
	servstat.dwServiceType = SERVICE_WIN32;
	servstat.dwCurrentState = SERVICE_RUNNING;
	servstat.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SHUTDOWN;
	servstat.dwWin32ExitCode = NO_ERROR;
	servstat.dwServiceSpecificExitCode = 0;
	servstat.dwCheckPoint = 0;
	servstat.dwWaitHint = 0;

	if ( SetServiceStatus( stat_hand, &servstat ) == 0 )
	{
		std::string message = "ServiceMain SetServiceStatus 1 failed with last error " + GetLastError();
		log( message.c_str() );
	}

	HKEY hKey;
	char appParams[255];
	DWORD dwBufLen;

	RegOpenKeyEx( HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\buildtool attend\\Parameters", 0, KEY_QUERY_VALUE, &hKey );
	RegQueryValueEx( hKey, "AppParameters", 0L, 0L, (LPBYTE) appParams, &dwBufLen);
	RegCloseKey( hKey );

	// convert appParams to args
	// calculate ac
	int ac = 0;

	// increment ac if a non space is read and either no previous character existed or previous character was space
	for ( int i = 0; i < strlen( appParams ); i++ )
	{
		if ( !isspace( appParams[ i ] ) && ( i == 0 || isspace( appParams[ i - 1] ) ) )
		{
			ac++;
		}
	}

	char ** args = new char * [ ac ];
	int arg = 0;

	for ( i = 0; i < strlen( appParams ); i++ )
	{
		if ( !isspace( appParams[ i ] ) && ( i == 0 || isspace( appParams[ i - 1] ) ) )
		{
			args[ arg ] = &appParams[ i ];
			arg++;
		}
	}

	int length = strlen( appParams );

	for ( i = 0; i < length; i++ )
	{
		if ( isspace( appParams[ i ] ) )
		{
			appParams[ i ] = '\0';
		}
	}

	for ( arg = 0; arg < ac; arg++ )
	{
		std::string message = "ServiceMain args[ ";
		std::string param;
		std::stringstream ss;
		ss << arg;
		ss >> param;
		message += param;
		message += "] = ";
		message += args[ arg ];
		log( message.c_str() );
	}

	// start the service
	Start( ac, args );

	// Start has returned because of a command
	// stop the service
	Stop();

	// inform the service control manager that the service is stopped
	servstat.dwCurrentState = SERVICE_STOPPED;

	if ( SetServiceStatus( stat_hand, &servstat ) == 0 )
	{
		std::string message = "ServiceMain SetServiceStatus 2 failed with last error " + GetLastError();
		log( message.c_str() );
	}

	return;
}

//------------------------------------------------------------------------------
//  ControlHandler
//  This is called by the control dispatcher when it has instructions for the service
//------------------------------------------------------------------------------
//  Parameters:
//      dwControl IN - the instruction
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void WINAPI ControlHandler( DWORD dwControl )
{
	switch( dwControl )
	{
	case SERVICE_CONTROL_PAUSE:
	case SERVICE_CONTROL_CONTINUE:
	case SERVICE_CONTROL_INTERROGATE:
		servstat.dwWaitHint = 0;
		break;
	case SERVICE_CONTROL_SHUTDOWN:
	case SERVICE_CONTROL_STOP:
		servstat.dwCurrentState = SERVICE_STOP_PENDING;
		servstat.dwWaitHint = 0;
		break;
	}

	// inform the service control manager of the service state
	if ( SetServiceStatus( stat_hand, &servstat ) == 0 )
	{
		std::string message = "ControlHandler SetServiceStatus failed with last error " + GetLastError();
		log( message.c_str() );
	}

	if ( ( dwControl == SERVICE_CONTROL_STOP ) || ( dwControl == SERVICE_CONTROL_SHUTDOWN ) )
	{
		log( "ControlHandler detected stop request" );

		// signal ServiceMain to return 
		if ( SetEvent( hWaitEvent ) == 0 )
		{
			std::string message = "ControlHandler SetEvent failed with last error " + GetLastError();
			log( message.c_str() );
		}

	}
}

//------------------------------------------------------------------------------
//  Start
//  The service logic, run when the service starts.
//------------------------------------------------------------------------------
//  Parameters:
//      argc IN - the number of arguments configured in the registry
//      argv IN - the arguments configured in the registry
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void Start( DWORD argc, LPSTR * argv )
{
	boolean doSomethingUseful = true;

	if ( argc < 3 )
	{
		log( "Start 1 configuration issue" );
		doSomethingUseful = false;
	}

	int maxRunTimeSecs = atoi( argv[ 1 ] );

	if ( maxRunTimeSecs == 0 )
	{
		log( "Start 2 configuration issue" );
		doSomethingUseful = false;
	}

	if ( doSomethingUseful )
	{
		std::vector< Process * > processCollection;
		std::string ancestor = argv[ 0 ];
		PROCESSENTRY32 entry;
		std::vector< Process * >::iterator processCollectionIterator = processCollection.begin();
		DWORD run = WAIT_TIMEOUT;

		// set the SeDebugPrivilege to prevent OpenProcess failing with access is denied
		LUID luid;

		if ( LookupPrivilegeValue( 0L, SE_DEBUG_NAME, &luid ) )
		{
			HANDLE hToken;
			if ( OpenProcessToken( GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken ) )
			{
				TOKEN_PRIVILEGES tp;
				tp.PrivilegeCount = 1;
				tp.Privileges[ 0 ].Luid = luid;
				tp.Privileges[ 0 ].Attributes = SE_PRIVILEGE_ENABLED;
				AdjustTokenPrivileges( hToken, FALSE, &tp, sizeof( TOKEN_PRIVILEGES ), 0L, 0L );
				CloseHandle( hToken );
			}

		}

		do
		{
			// clear processCollection
			processCollectionIterator = processCollection.begin();

			while ( processCollectionIterator != processCollection.end() )
			{
				Process *p = *processCollectionIterator;
				delete p;
				p = 0L;
				processCollectionIterator ++;
			}

			processCollection.clear();

			HANDLE snapshot = CreateToolhelp32Snapshot( TH32CS_SNAPPROCESS, 0 );

			if ( snapshot )
			{
				memset( &entry, 0, sizeof( entry ) );
				entry.dwSize = sizeof( entry );

				if ( Process32First( snapshot, &entry ) == TRUE )
				{
					do
					{
						// get a handle to the process
						HANDLE hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, entry.th32ProcessID );

						if ( hProcess )
						{
							// get processes fully qualified path
							HMODULE hMod;
							DWORD cbNeeded = 0;

							char path[MAX_PATH] = "";
							if ( EnumProcessModules( hProcess, &hMod, sizeof( hMod ), &cbNeeded ) )
							{
								FILETIME creationTime;
								FILETIME exitTime;
								FILETIME kernelTime;
								FILETIME userTime;
								GetModuleFileNameEx( hProcess, hMod, path, sizeof( path) );
								GetProcessTimes( hProcess, &creationTime, &exitTime, &kernelTime, &userTime );
								ULARGE_INTEGER nano;
								nano.LowPart = creationTime.dwLowDateTime;
								nano.HighPart = creationTime.dwHighDateTime;
								Process *p = new Process( path, entry.szExeFile, entry.th32ProcessID, entry.th32ParentProcessID, nano.QuadPart );
								processCollection.push_back( p );
							}

							CloseHandle( hProcess );
						}
						
					} while ( Process32Next( snapshot, &entry ) );
				}
				CloseHandle( snapshot );
			}

			// iterate through the processCollection
			// if its name matches a child from argv[2] to argv[argc-1], and it is a child of ancestor, and its runtime has expired, then terminate the process
			processCollectionIterator = processCollection.begin();

			while ( processCollectionIterator != processCollection.end() )
			{
				Process *p = *processCollectionIterator;
				int child = 2;

				while( child < argc )
				{
					if ( p->mName == argv[ child ] )
					{
						std::string message = "Start found configured process running ";
						message += p->mName;
						log( message.c_str() );
						// its name matches a child
						if ( isDescendedFrom( ancestor, *p, processCollection ) )
						{
							std::string message = "Start found configured process running and descendant ";
							message += p->mName;
							log( message.c_str() );
							// it is a child, directly or indirectly, of ancestor
							SYSTEMTIME st;
							FILETIME ft;
							// get current time
							GetSystemTime(&st);
							// convert to file time format
							SystemTimeToFileTime(&st, &ft);
							ULARGE_INTEGER nano;
							nano.LowPart = ft.dwLowDateTime;
							nano.HighPart = ft.dwHighDateTime;
							ULARGE_INTEGER runTime;
							runTime.QuadPart = nano.QuadPart - p->mTime;

							int runTimeSecs = runTime.QuadPart/10000000;

							if ( runTimeSecs > maxRunTimeSecs )
							{
								std::string message = "Start found configured process running too long ";
								message += p->mName;
								log( message.c_str() );
								// it has been running for too long, terminate
								HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_TERMINATE, FALSE, p->mPid );

								if ( hProcess )
								{
									std::string message = "Start terminating ";
									message += p->mName;
									log( message.c_str() );
									TerminateProcess( hProcess, CONTROL_C_EXIT );
									CloseHandle( hProcess );
								}

							}

						}
						break;
					}
					child++;
				}
			
				processCollectionIterator ++;
			}

			// this is effectively the ServiceMain routine
			// do not exit this loop and return until signalled
			// wait for 60 secs and continue to run if not signalled
			run = WaitForSingleObject( hWaitEvent, 60000 );

			if ( run == WAIT_FAILED )
			{
				std::string message = "Start WaitForSingleObject 1 failed with last error " + GetLastError();
				log( message.c_str() );
			}

		} while ( run == WAIT_TIMEOUT );

		// clear processCollection
		processCollectionIterator = processCollection.begin();

		while ( processCollectionIterator != processCollection.end() )
		{
			Process *p = *processCollectionIterator;
			delete p;
			p = 0L;
			processCollectionIterator ++;
		}

		processCollection.clear();

	}
	else
	{
		log( "Start calling WaitForSingleObject" );
		// this is effectively the ServiceMain routine
		// do not return until signalled
		if ( WaitForSingleObject( hWaitEvent, INFINITE ) == WAIT_FAILED )
		{
			std::string message = "Start WaitForSingleObject 2 failed with last error " + GetLastError();
			log( message.c_str() );
		}
		log( "Start called WaitForSingleObject" );
	}

}

//------------------------------------------------------------------------------
//  Stop
//  The service clean up, run when the service is to be stopped.
//------------------------------------------------------------------------------
//  Parameters:
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void Stop()
{
	// nothing to do
}