// Copyright 2005-2020 The Mumble Developers. All rights reserved.
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file at the root of the
// Mumble source tree or at <https://www.mumble.info/LICENSE>.

#include "win.h"

#include <string>

#include <shlwapi.h>
#include <stdio.h>
#include <sddl.h>

// Alert shows a fatal error dialog and waits for the user to click OK.
static void Alert(LPCWSTR title, LPCWSTR msg) {
	MessageBox(NULL, msg, title, MB_OK|MB_ICONERROR);
}

// GetExecutablePath returns the path to mumble.exe.
static const std::wstring GetExecutablePath() {
	wchar_t path[MAX_PATH];

	if (GetModuleFileNameW(NULL, path, MAX_PATH) == 0)
		return std::wstring();

	std::wstring exe_path(path);
	return exe_path;
}

// UIAccessDisabledViaConfig queries the configuration store, and returns
// whether UIAccess has been disabled by the user via the
// "shortcut/windows/uiaccess/enable" config option.
static bool UIAccessDisabledViaConfig() {
	// RegQueryValueEx only zero-terminates if the registry value's type is
	// one of the Windows registry's string types. Because of this, we have
	// to use a buffer of 7 elements in order to query for the string "false".
	//
	// If the value's type is string, we'll end up with "false\0\0", because we
	// always zero-pad.
	//
	// But that's better than ending up in an infinite loop because we failed
	// to zero-terminate.
	wchar_t buf[7];
	memset(&buf, 0, sizeof(buf));
	DWORD sz = sizeof(buf) - 1*sizeof(wchar_t);

	HKEY key = NULL;
	bool success = (RegOpenKeyExW(HKEY_CURRENT_USER, L"Software\\Mumble\\Mumble\\shortcut\\windows\\uiaccess", 0, KEY_READ, &key) == ERROR_SUCCESS) &&
	               (RegQueryValueExW(key, L"enable" , NULL, NULL, (LPBYTE)&buf, &sz) == ERROR_SUCCESS);
	if (success && _wcsicmp(buf, L"false") == 0) {
		return true;
	}

	return false;
}

// ProcessHasUIAccess returns true if the current process has UIAccess enabled.
static bool ProcessHasUIAccess() {
	HANDLE token = NULL;
	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token)) {
		return false;
	}

	DWORD ui_access = 0;
	DWORD ui_access_size = sizeof(ui_access);
	if (!GetTokenInformation(token, TokenUIAccess, &ui_access, sizeof(ui_access), &ui_access_size)) {
		CloseHandle(token);
		return false;
	}

	CloseHandle(token);
	return ui_access != 0;
}

// RelaunchWithoutUIAccessIfNecessary relaunches the process
// without UIAccess, if necessary.
//
// If a relaunch of Mumble is necessary, the new process is
// spawned and the existing process is terminated. In that
// situation, this function will not return to the caller.
static bool RelaunchWithoutUIAccessIfNecessary() {
	if (!ProcessHasUIAccess()) {
		return true;
	}
	if (!UIAccessDisabledViaConfig()) {
		return true;
	}

	STARTUPINFO startup_info;
	memset(&startup_info, 0, sizeof(startup_info));

	GetStartupInfo(&startup_info);

	PROCESS_INFORMATION process_info;
	memset(&process_info, 0, sizeof(process_info));

	// Store the original value of __COMPAT_LAYER so we can
	// restore it (which may be equivalent to not setting it all)
	// once we've been relaunched.
	const wchar_t *compat_layer_env = _wgetenv(L"__COMPAT_LAYER");
	if (compat_layer_env == NULL) {
		compat_layer_env = L"mumble_exe_none";
	}
	_wputenv_s(L"mumble_exe_prev_compat_layer", compat_layer_env);

	// The only way we have found thus far to ignore the executable's
	// manifest when calling CreateProcess, is to use the __COMPAT_LAYER
	// environment variable.
	//
	// Setting this to RunAsInvoker seemingly ignores the application's
	// manifest. RunAsInvoker is the default execution level for Windows
	// programs.
	_wputenv_s(L"__COMPAT_LAYER", L"RunAsInvoker");

	std::wstring exe_path = GetExecutablePath();
	if (exe_path.empty()) {
		return false;
	}

	if (!CreateProcessW(exe_path.c_str(),
	              GetCommandLineW(),
	              NULL,
	              NULL,
	              FALSE,
	              0,
	              NULL,
	              NULL,
	              &startup_info,
	              &process_info)) {
		return false;
	}

	exit(0);
}

// CleanupEnvironmentVariables cleans up any
// unnecessary environment variables.
//
// For now, this function cleans up any environment
// variables set by RelaunchWithoutUIAccessIfNecessary.
static void CleanupEnvironmentVariables() {
	wchar_t *prev_compat_layer_env = _wgetenv(L"mumble_exe_prev_compat_layer");
	if (prev_compat_layer_env != NULL) {
		if (wcscmp(prev_compat_layer_env, L"mumble_exe_none") == 0) {
			_wputenv_s(L"__COMPAT_LAYER", L"");
		} else {
			_wputenv_s(L"__COMPAT_LAYER", prev_compat_layer_env);
		}
	}
}

int os_early_init() {
	CleanupEnvironmentVariables();

	if (!RelaunchWithoutUIAccessIfNecessary()) {
		Alert(L"Mumble Early Init Error 50", L"Unable to complete RelaunchWithoutUIAccessIfNecessary");
		return 50;
	}

	return 0;
}