#include "UserAccountControl.h"

 #include <windows.h>

 #pragma comment (lib, "kernel32.lib")

 #pragma comment (lib, "advapi32.lib")

 #include <vcclr.h>

 #include <msclr\marshal.h>

 #include <msclr\marshal_windows.h>

 using namespace msclr::interop;

 using namespace System::ComponentModel;

 using namespace Microsoft::Win32;

 namespace UacHelpers {

 	Process^ UserAccountControl::CreateProcessAsAdmin(System::String^ exePath, System::String^ arguments)

     {

         ProcessStartInfo^ psi = gcnew ProcessStartInfo(exePath, arguments);

         psi->UseShellExecute = true;

         psi->Verb = "runas";

 		return Process::Start(psi);

     }

 	Process^ UserAccountControl::CreateProcessAsStandardUser(System::String^ exePath, System::String^ arguments)

 	{

 		marshal_context context;

 		//If the current process is not elevated, then there's no reason to go through the hassle --

 		//just use the standard System.Diagnostics.Process facilities.

 		if (!IsCurrentProcessElevated)

 		{

 			return Process::Start(exePath, arguments);

 		}

 		//The following implementation is roughly based on Aaron Margosis' post:

 		//http://blogs.msdn.com/aaron_margosis/archive/2009/06/06/faq-how-do-i-start-a-program-as-the-desktop-user-from-an-elevated-app.aspx

 		//Enable SeIncreaseQuotaPrivilege in this process.  (This requires administrative privileges.)

 		HANDLE hProcessToken = NULL;

 		if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hProcessToken))

 		{

 			throw gcnew Win32Exception(GetLastError());

 		}

 		else

 		{

 			TOKEN_PRIVILEGES tkp;

 			tkp.PrivilegeCount = 1;

 			LookupPrivilegeValueW(NULL, SE_INCREASE_QUOTA_NAME, &tkp.Privileges[0].Luid);

 			tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

 			AdjustTokenPrivileges(hProcessToken, FALSE, &tkp, 0, NULL, NULL);

 			DWORD dwLastErr = GetLastError();

 			CloseHandle(hProcessToken);

 			if (ERROR_SUCCESS != dwLastErr)

 			{

 				throw gcnew Win32Exception(dwLastErr);

 			}

 		}

 		//Get window handle representing the desktop shell.  This might not work if there is no shell window, or when

 		//using a custom shell.  Also note that we're assuming that the shell is not running elevated.

 		HWND hShellWnd = GetShellWindow();

 		if (hShellWnd == NULL)

 		{

 			throw gcnew System::InvalidOperationException("Unable to locate shell window; you might be using a custom shell");

 		}

 		//Get the ID of the desktop shell process.

 		DWORD dwShellPID;

 		GetWindowThreadProcessId(hShellWnd, &dwShellPID);

 		if (dwShellPID == 0)

 		{

 			throw gcnew Win32Exception(GetLastError());

 		}

 		//Open the desktop shell process in order to get the process token.

 		HANDLE hShellProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwShellPID);

 		if (hShellProcess == NULL)

 		{

 			throw gcnew Win32Exception(GetLastError());

 		}

 		HANDLE hShellProcessToken = NULL;

 		HANDLE hPrimaryToken = NULL;

 		try

 		{

 			//Get the process token of the desktop shell.

 			if (!OpenProcessToken(hShellProcess, TOKEN_DUPLICATE, &hShellProcessToken))

 			{

 				throw gcnew Win32Exception(GetLastError());

 			}

 			//Duplicate the shell's process token to get a primary token.

 			const DWORD dwTokenRights = TOKEN_QUERY | TOKEN_ASSIGN_PRIMARY | TOKEN_DUPLICATE | TOKEN_ADJUST_DEFAULT | TOKEN_ADJUST_SESSIONID;

 			if (!DuplicateTokenEx(hShellProcessToken, dwTokenRights, NULL, SecurityImpersonation, TokenPrimary, &hPrimaryToken))

 			{

 				throw gcnew Win32Exception(GetLastError());

 			}

 			//Start the target process with the new token.

 			STARTUPINFO si = {0}; si.cb = sizeof(si);

 			PROCESS_INFORMATION pi = {0};

 			if (!CreateProcessWithTokenW(hPrimaryToken, 0,

 				context.marshal_as<LPCWSTR>(exePath), context.marshal_as<LPWSTR>(exePath + " " + arguments),

 				0, NULL, NULL, &si, &pi))

 			{

 				throw gcnew Win32Exception(GetLastError());

 			}

 			CloseHandle(pi.hProcess);

 			CloseHandle(pi.hThread);

 			return Process::GetProcessById(pi.dwProcessId);

 		}

 		finally

 		{

 			if (hShellProcessToken != NULL)

 				CloseHandle(hShellProcessToken);

 			if (hPrimaryToken != NULL)

 				CloseHandle(hPrimaryToken);

 			if (hShellProcess != NULL)

 				CloseHandle(hShellProcess);

 		}

 	}

 	bool UserAccountControl::IsUserAdmin::get()

 	{

 		if (UserAccountControl::IsUacEnabled)

 			return GetProcessTokenElevationType() != TokenElevationTypeDefault;	//split token

 		//If UAC is off, we can't rely on the token; check for Admin group.

 		return WindowsPrincipal(WindowsIdentity::GetCurrent()).IsInRole("Administrators");

 	}

 	bool UserAccountControl::IsUacEnabled::get()

 	{

 		//Check the HKLM\Software\Microsoft\Windows\CurrentVersion\Policies\System\EnableLUA registry value.

 		RegistryKey^ key = Registry::LocalMachine->OpenSubKey(UacRegistryKey, false);

 		return key->GetValue(UacRegistryValue)->Equals(1);

     }

 	void UserAccountControl::DisableUac()

 	{

 		SetUacRegistryValue(false);

 	}

 	void UserAccountControl::DisableUacAndRestartWindows()

 	{

 		DisableUac();

 		RestartWindows();

 	}

 	void UserAccountControl::EnableUac()

 	{

 		SetUacRegistryValue(true);

 	}

 	void UserAccountControl::EnableUacAndRestartWindows()

 	{

 		EnableUac();

 		RestartWindows();

 	}

 	bool UserAccountControl::IsCurrentProcessElevated::get()

 	{

 		return GetProcessTokenElevationType() == TokenElevationTypeFull;	//elevated

 	}

 	bool UserAccountControl::IsCurrentProcessVirtualized::get()

 	{

 		HANDLE hToken;

         try

         {

             if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))

 				throw gcnew Win32Exception(GetLastError());

             DWORD virtualizationEnabled;

 			DWORD dwSize;

 			if (!GetTokenInformation(hToken, TokenVirtualizationEnabled, &virtualizationEnabled, sizeof(virtualizationEnabled), &dwSize))

                 throw gcnew Win32Exception(GetLastError());

 			return virtualizationEnabled != 0;

         }

         finally

         {

             CloseHandle(hToken);

         }

 	}

 	int UserAccountControl::GetProcessTokenElevationType()

 	{

 		HANDLE hToken;

         try

         {

             if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))

 				throw gcnew Win32Exception(GetLastError());

             TOKEN_ELEVATION_TYPE elevationType;

 			DWORD dwSize;

             if (!GetTokenInformation(hToken, TokenElevationType, &elevationType, sizeof(elevationType), &dwSize))

                 throw gcnew Win32Exception(GetLastError());

 			return elevationType;

         }

         finally

         {

             CloseHandle(hToken);

         }

 	}

 	void UserAccountControl::SetUacRegistryValue(bool enabled)

 	{

 		RegistryKey^ key = Registry::LocalMachine->OpenSubKey(UacRegistryKey, true);

 		key->SetValue(UacRegistryValue, enabled ? 1 : 0);

 	}

 	void UserAccountControl::RestartWindows()

 	{

 		InitiateSystemShutdownEx(NULL, NULL, 0/*Timeout*/,

 								 TRUE/*ForceAppsClosed*/, TRUE/*RebootAfterShutdown*/,

 								 SHTDN_REASON_MAJOR_OPERATINGSYSTEM | SHTDN_REASON_MINOR_RECONFIG | SHTDN_REASON_FLAG_PLANNED);

 		//This shutdown flag corresponds to: "Operating System: Reconfiguration (Planned)".

 	}

 }