//***************************************************************
// From the book "Win32 System Services: The Heart of Windows 98
// and Windows 2000"
// by Marshall Brain
// Published by Prentice Hall
//
// Copyright 1995, by Prentice Hall.
//
// This code reboots the system. See also Chapter 16.
//***************************************************************

// sireboot.cpp

#include <windows.h>
#include <iostream.h>
#include <process.h>

void HandleError(char *s)
{
	cout << "Error in " << s << endl;
	cout << "Error code = " << GetLastError() 
		<< endl;
	exit(1);
}

BOOL EnablePrivilege(LPTSTR privilege)
{
	BOOL success;
	HANDLE token;
	LUID luid;
	TOKEN_PRIVILEGES tokenPrivileges;

	// Get token for this process
	success = OpenProcessToken(GetCurrentProcess(), 
		TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
		&token);
	if (!success) 
	{
		HandleError("OpenProcessToken");
		return FALSE;
	}

	// Gets the value for a privilege
	success = LookupPrivilegeValue(0, privilege, 
		&luid);
	if (!success) 
	{
		HandleError("LookupPrivilegeValue");
		return FALSE;
	}
	
	// Enable the privilege
	tokenPrivileges.PrivilegeCount = 1;
	tokenPrivileges.Privileges[0].Luid = luid;
	tokenPrivileges.Privileges[0].Attributes =
		SE_PRIVILEGE_ENABLED;
	success = AdjustTokenPrivileges(token, FALSE,
		&tokenPrivileges, 0, 0, 0);
	// Always returns true, so check GetLastError
	if (GetLastError() != ERROR_SUCCESS) 
	{
		HandleError("AdjustTokenPrivileges");
		return FALSE;
	}
	return TRUE;
}

BOOL DisablePrivilege(LPTSTR privilege)
{
	BOOL success;
	HANDLE token;
	LUID luid;
	TOKEN_PRIVILEGES tokenPrivileges;

	// Get tokens for this process
	success = OpenProcessToken(GetCurrentProcess(), 
		TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
		&token);
	if (!success) 
	{
		HandleError("OpenProcessToken");
		return FALSE;
	}

	// Gets the value for a privilege
	success = LookupPrivilegeValue(0, privilege, 
		&luid);
	if (!success) 
	{
		HandleError("LookupPrivilegeValue");
		return FALSE;
	}

	// Disable the privilege
	tokenPrivileges.PrivilegeCount = 1;
	tokenPrivileges.Privileges[0].Luid = luid;
	// disable the privilege
	tokenPrivileges.Privileges[0].Attributes = 0;
	success = AdjustTokenPrivileges(token, FALSE,
		&tokenPrivileges, 0, 0, 0);
	// Always returns true, so...
	if (GetLastError() != ERROR_SUCCESS) 
	{
		HandleError("AdjustTokenPrivileges");
		return FALSE;
	}
	return TRUE;
}

void main()
{
	BOOL success;

	if (EnablePrivilege(SE_SHUTDOWN_NAME))
	{
		success = ExitWindowsEx(EWX_REBOOT, 0);
		if (success)
			cout << "Success. Shutting down shortly." 
				<< endl;
		else
			HandleError("Shutting down");
	}
	DisablePrivilege(SE_SHUTDOWN_NAME);
}