#include <windows.h>
#include <stdio.h>
#include <stdlib.h>

#pragma hdrstop



BOOL SetCurrentPrivilege( LPCTSTR Privilege, BOOL bEnablePrivilege );
void DisplayError( DWORD err );



int main( int argc, char * argv[] )
{
	int i, errors = 0, local = 0, timeout = -1, force = 0, reboot = 0, oops = 0;
	char *remote = NULL, *message = NULL;
	BOOL rc;

	if ( argc == 1 )
	{
usage:
		if ( errors )
			putchar( '\n' );
		puts( "reboot {-l | -s servername} [-m \"message\"] [-t timeout] [-f] [-r] [-oops]\n" );
		puts( "-l             reboot local machine" );
		puts( "-s servername  reboot named computer" );
		puts( "-t timeout     show a logoff dialog for <timeout> seconds; 0 suppresses it" );
		puts( "-m \"message\"   display \"message\" in that dialog" );
		puts( "-f             force apps closed" );
		puts( "-r             reboot after shutdown" );
		puts( "-oops          aborts a previous shutdown request (only -l, -s count)" );
		return 1;
	}

	for ( i = 1; i < argc; i ++ )
	{
		if ( *argv[i] != '-' && *argv[i] != '/' )
		{
			printf( "\"%s\": switches must start with a '-' or '/' character.\n", argv[i] );
			errors ++;
			continue;
		}

		++ argv[i];
		switch ( *argv[i] )
		{
			case 'l':
				if ( remote != NULL || local != 0 )
				{
					puts( "Only one of -l and -s may be used." );
					++ errors;
				}
				else
					++ local;
				break;
			case 's':
				if ( remote != NULL || local != 0 )
				{
					puts( "Only one of -l and -s may be used." );
					++ errors;
				}
				else if ( i == argc - 1 )
				{
					puts( "-s must be followed by the name of the server to shut down." );
					++ errors;
				}
				else
				{
					++ i;
					remote = argv[i];
				}
				break;
			case 't':
				if ( timeout != -1 )
				{
					puts( "Only one timeout may be used." );
					++ errors;
				}
				else if ( i == argc - 1 )
				{
					puts( "-t must be followed by a timeout." );
					++ errors;
				}
				else
				{
					++ i;
					timeout = atoi( argv[i] );
				}
				break;
			case 'm':
				if ( message != NULL )
				{
					puts( "Only one message may be used." );
					++ errors;
				}
				else if ( i == argc - 1 )
				{
					puts( "-m must be followed by the message text." );
					++ errors;
				}
				else
				{
					++ i;
					message = argv[i];
				}
				break;
			case 'f':
				if ( force )
				{
					puts( "-f may only be used once." );
					++ errors;
				}
				else
					++ force;
				break;
			case 'r':
				if ( reboot )
				{
					puts( "-r may only be used once." );
					++ errors;
				}
				else
					++ reboot;
				break;
			case 'o':
				++ oops;
				break;
			default:
				printf( "\"%s\" is not a valid switch.\n", argv[i] - 1 );
				++ errors;
				break;
		}
	}

	if ( local == 0 && remote == NULL )
	{
		puts( "One of -l and -s must be used." );
		++ errors;
	}

	if ( ! oops && timeout == -1 ) // no timeout used
	{
		puts( "Anti-goof feature: you must specify a timeout (0 seconds is OK)" );
		++ errors;
	}

	if ( errors )
		goto usage;

	if ( local )
	{
		if ( ! SetCurrentPrivilege( SE_SHUTDOWN_NAME, TRUE ) )
		{
			DisplayError( GetLastError() );
			return 2;
		}
	}

	if ( oops )
	{
		printf( "\nTrying to abort the shutdown of %s ...", remote? remote: "this computer" );
		rc = AbortSystemShutdown( remote );
	}
	else
	{
		printf( "\nTrying to shut down %s in %d seconds ...", remote? remote: "this computer", timeout );
		rc = InitiateSystemShutdown( remote, message, (DWORD) timeout, force, reboot );
	}

	putchar( '\n' );

	if ( ! rc )
	{
		DisplayError( GetLastError() );
		return 1;
	}

	return 0;
}



BOOL SetCurrentPrivilege( LPCTSTR Privilege, BOOL bEnablePrivilege )
{
	HANDLE hToken;
	LUID luid;
	TOKEN_PRIVILEGES tp, tpPrevious;
	DWORD cbPrevious = sizeof( TOKEN_PRIVILEGES );
	BOOL bSuccess = FALSE;

	if ( ! LookupPrivilegeValue( NULL, Privilege, &luid ) )
		return FALSE;

	if( ! OpenProcessToken( GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken ) )
		return FALSE;

	tp.PrivilegeCount = 1;
	tp.Privileges[0].Luid = luid;
	tp.Privileges[0].Attributes = 0;

	AdjustTokenPrivileges( hToken, FALSE, &tp, sizeof( TOKEN_PRIVILEGES ), &tpPrevious, &cbPrevious );

	if ( GetLastError() == ERROR_SUCCESS )
	{
		tpPrevious.PrivilegeCount = 1;
		tpPrevious.Privileges[0].Luid = luid;

		if ( bEnablePrivilege )
			tpPrevious.Privileges[0].Attributes |= ( SE_PRIVILEGE_ENABLED );
		else
			tpPrevious.Privileges[0].Attributes &= ~( SE_PRIVILEGE_ENABLED );

		AdjustTokenPrivileges( hToken, FALSE, &tpPrevious, cbPrevious, NULL, NULL );

		if ( GetLastError() == ERROR_SUCCESS )
			bSuccess=TRUE;
	}

	CloseHandle( hToken );

	return bSuccess;
}



void DisplayError( DWORD err )
{
	char msgbuf[4096];

	FormatMessage( FORMAT_MESSAGE_FROM_SYSTEM, NULL, err,
		MAKELANGID( LANG_NEUTRAL, SUBLANG_DEFAULT),
		msgbuf, sizeof( msgbuf ), NULL );
	printf( "Error %d: %s\n", err, msgbuf );
}