policytoolgui/main.cpp

#include "PolicyToolGUI.h"
#include "incbin.h"

#include <string>
#include <string_view>

#include <windows.h>
#include <processenv.h>
#include <tlhelp32.h>
#include <psapi.h>

INCBIN(detoursDLL, "detours.dll");


static bool hasStarted = false,
            keepRunning = false;
static HMODULE hLibrary;
static STARTUPINFOA startInfo;
static PROCESS_INFORMATION processInfo;



static inline std::string LastErrorString() {
    auto errorMessageID = GetLastError();
    if (errorMessageID == 0) {
        return "No error";
    }
    LPSTR messageBuffer = nullptr;
    auto size = FormatMessageA(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                                 NULL, errorMessageID, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPSTR)&messageBuffer, 0, NULL);
    std::string message(messageBuffer, size);
    LocalFree(messageBuffer);
    return message;
}

bool SimpleIsCharValidForPath(char c) {
	if (c < 0x20)  return false;
	if (c > 0x7E)  return false;
	if (c == 0x60) return false;
	return true;
}
std::string StringSan(const std::string& str) {
	std::string fres;
	for (const auto c : str) {
		if (SimpleIsCharValidForPath(c)) fres.push_back(c);
	}
	return fres;
}
void Execute(std::string&& fileName, std::string&& arguments, std::string&& restartProcessName = "") {
	hasStarted = false;
	fileName = StringSan(fileName);
	restartProcessName = StringSan(restartProcessName);
	// Get temp folder
	wchar_t tempDir[MAX_PATH + 1];
	if (GetTempPathW(MAX_PATH, tempDir) == 0) {
		GetCurrentDirectoryW(MAX_PATH, tempDir);
	}
	// Make sure temp folder actually exists
	if (GetFileAttributesW(tempDir) == INVALID_FILE_ATTRIBUTES &&
		CreateDirectoryW(tempDir, NULL) == 0) {
		ErrorOK("Failed to access temp dir");
		return;
	}
	// Get temp file path
	auto tempFile = std::wstring(tempDir)+L"\\hooks_"+std::to_wstring(time(nullptr))+L".dll";
	// Write DLL
	{
		// Open file
		HANDLE hTempFile = CreateFileW(tempFile.c_str(), GENERIC_WRITE, 0,
	                                   NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
		if (!hTempFile) {
			ErrorOK("Failed to open temporary file for writing");
			return;
		}
		
		// Write data
		if (WriteFile(hTempFile, incbin_detoursDLL_start, incbin_detoursDLL_end - incbin_detoursDLL_start, NULL, NULL) == 0) {
			ErrorOK("Failed to write detours into temporary file");
			return;
		}
		// Close file
		CloseHandle(hTempFile);
	}
	// Kill process
	if (!restartProcessName.empty()) {
		// Create toolhelp snapshot.
		HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
		PROCESSENTRY32 process;
		ZeroMemory(&process, sizeof(process));
		process.dwSize = sizeof(process);
		// Walk through all processes.
		if (Process32First(snapshot, &process)) {
			do {
				// Compare process.szExeFile based on format of name, i.e., trim file path
				if (std::string_view(process.szExeFile) == restartProcessName) {
					// Attempt to terminate process
					HANDLE hProcess = OpenProcess(PROCESS_TERMINATE, 0,
                                                  (DWORD)process.th32ProcessID);
					if (hProcess != NULL) {
						TerminateProcess(hProcess, 1);
						// We've succeeded, clean up and break out of loop
						CloseHandle(hProcess);
						break;
					}
				}
			} while (Process32Next(snapshot, &process));
		}
		// We've succeeded, clean up
		CloseHandle(snapshot);
	}
	// Restart process with detours
	hasStarted = true;
	{
		ZeroMemory(&processInfo, sizeof(PROCESS_INFORMATION));
		ZeroMemory(&startInfo, sizeof(STARTUPINFOA));
		startInfo.cb = sizeof(STARTUPINFOA);
		// Load detours
		hLibrary = LoadLibraryW(tempFile.c_str());
		if (!hLibrary) {
			ErrorOK(("Failed to load detours: "+LastErrorString()).c_str());
			hasStarted = false;
		}
		// Mark temp file for deletion
		DeleteFileW(tempFile.c_str());
		// Create process
		bool createRes;
		if (fileName[2] == ':') {
			createRes = CreateProcessA(const_cast<char*>(fileName.c_str()), const_cast<char*>(arguments.c_str()), NULL, NULL, FALSE, 0, NULL, NULL,
					                   &startInfo, &processInfo);
		} else {
			createRes = CreateProcessA(NULL, const_cast<char*>((fileName+' '+arguments).c_str()), NULL, NULL, FALSE, 0, NULL, NULL,
					                   &startInfo, &processInfo);
		}
		if (createRes == 0) {
			ErrorOK(("Unable to invoke "+fileName).c_str());
			hasStarted = false;
			return;
		}
	}
}

MainWindow::MainWindow() {
	CtrlLayout(*this, "Windows Policy Tool");
	
	// initialize settings that will be inherited by subprocesses
	SetEnvironmentVariableW(L"__POLICYTOOL_ADMINIMPERSONATE", L"0");
	SetEnvironmentVariableW(L"__POLICYTOOL_REMOTELOCKBREAK", L"0");
	SetEnvironmentVariableW(L"__POLICYTOOL_POLICYDISABLE", L"1");
	OptDisablePolicies ^= [this] () {
		SetEnvironmentVariableW(L"__POLICYTOOL_POLICYDISABLE", OptDisablePolicies?L"1":L"0");
	};
	OptDisablePolicies.Set(1);
	OptHideProcess ^= [this] () {
		keepRunning = !OptHideProcess;
	};
	OptHideProcess.Set(1);
	OptAdminImpersonate ^= [this] () {
		bool enabled = OptAdminImpersonate;
		if (enabled) {
			SetEnvironmentVariableW(L"__COMPAT_LAYER", L"RunAsInvoker");
		} else {
			SetEnvironmentVariableW(L"__COMPAT_LAYER", nullptr);
		}
		SetEnvironmentVariableW(L"__POLICYTOOL_ADMINIMPERSONATE", enabled?L"1":L"0");
	};
	OptBreakRemoteLocks ^= [this] () {
		SetEnvironmentVariableW(L"__POLICYTOOL_REMOTELOCKBREAK", OptBreakRemoteLocks?L"1":L"0");
	};
	
	ApplyToExplorer ^= [this] () {
		// Execute
		::Execute("explorer.exe", "", "explorer.exe");
		// Continue in main()
		if (hasStarted) Close();
	};
	ApplyTo ^= [this] () {
		// File picker
		FileSel fileSel;
		fileSel.Multi(false)
		       .Type("Executable", "*.exe")
		       .AllFilesType()
		       .DefaultExt("exe")
		       .ExecuteOpen("Select executable");
		// Execute
		::Execute(DeQtf(~fileSel).ToStd(), "");
		// Continue in main()
		if (hasStarted) Close();
	};
	OpenControl ^= [this] () {
		// Put path into clipboard
		const static char path[] = R"(::{20D04FE0-3AEA-1069-A2D8-08002B30309D}\::{21EC2020-3AEA-1069-A2DD-08002B30309D})";
		HGLOBAL hMem = GlobalAlloc(GMEM_MOVEABLE, sizeof(path));
		memcpy(GlobalLock(hMem), path, sizeof(path));
		GlobalUnlock(hMem);
		OpenClipboard(0);
		EmptyClipboard();
		SetClipboardData(CF_TEXT, hMem);
		CloseClipboard();
		// Notify user
		MessageBoxA(nullptr, "Use the \"Apply\" button if you haven't done so already and paste in\r\nthe path that's been put into the clipboard for you.",
		                     "Windows Policy tool", MB_OK);
	};
}

GUI_APP_MAIN {
	// Run GUI
	MainWindow().Run();
	// Handle the rest
	if (hasStarted) {
		if (keepRunning) {
			// Wait for process to exit
			WaitForSingleObject(processInfo.hProcess, INFINITE);
		}
		// Clean up
		CloseHandle(processInfo.hProcess);
		CloseHandle(processInfo.hThread);
	}
}