First working Linux version

.gitignore

@@ -1 +1,3 @@
-build/
+build/
+bin/
+.idea/

CMake/Findefivar.cmake

@@ -0,0 +1,17 @@
+# efivar_FOUND
+# efivar_INCLUDE_DIRS
+
+find_path(efivar_INCLUDE_DIR efivar.h
+    /usr/include/efivar
+    /usr/local/include/efivar
+)
+set(efivar_INCLUDE_DIRS ${efivar_INCLUDE_DIR})
+
+find_library(efivar_LIBRARIES libefivar.so
+    /usr/lib
+    /usr/local/lib
+)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(efivar DEFAULT_MSG efivar_INCLUDE_DIR efivar_LIBRARIES)
+mark_as_advanced(efivar_FOUND efivar_INCLUDE_DIR)

CMakeLists.txt

@@ -17,7 +17,8 @@ else()
 endif()
 
 if(CMAKE_COMPILER_IS_GNUCXX)
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14")
+    #set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14")
+	add_definitions(-std=c++14)
 elseif(MSVC)
 	set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /MP")
 	add_definitions(-D_CRT_SECURE_NO_WARNINGS)

src/CMakeLists.txt

@@ -2,17 +2,32 @@ project(RebootInto)
 find_package(TCLAP REQUIRED)
 find_package(Boost REQUIRED COMPONENTS locale)
 find_package(Nowide REQUIRED)
+if(UNIX)
+    find_package(efivar REQUIRED)
+endif()
 
 include_directories(
     ${INCLUDE_PATH}
     ${TCLAP_INCLUDE_DIRS}
     ${Boost_INCLUDE_DIRS}
     ${Nowide_INCLUDE_DIRS}
+    ${efivar_INCLUDE_DIRS}
 )
 
 set(SOURCE_FILES
+    UEFI.h
     main.cpp
 )
+if(WIN32)
+    set(SOURCE_FILES ${SOURCE_FILES}
+        UEFI_Windows.cpp
+    )
+endif()
+if(UNIX)
+    set(SOURCE_FILES ${SOURCE_FILES}
+        UEFI_Linux.cpp
+    )
+endif()
 source_group("" FILES ${SOURCE_FILES})
 
 if(MSVC)
@@ -23,4 +38,8 @@ if(MSVC)
     set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} /SUBSYSTEM:CONSOLE /level='requireAdministrator'")
 endif()
 
-add_executable(RebootInto ${SOURCE_FILES})
+add_executable(RebootInto ${SOURCE_FILES})
+target_link_libraries(RebootInto
+    ${Boost_LIBRARIES}
+    ${efivar_LIBRARIES}
+)

src/UEFI.h

@@ -0,0 +1,31 @@
+#include <cstdint>
+#include <limits>
+#include <string>
+#include <vector>
+#include <stdexcept>
+#include <memory>
+
+const std::size_t EFI_LOAD_OPTION_DESCRIPTION_OFFSET = sizeof(std::uint32_t) + sizeof(std::uint16_t);
+
+class UEFI
+{
+public:
+    struct BootOption
+    {
+        std::uint16_t ID;
+        std::string Description;
+    };
+
+    UEFI();
+
+    BootOption ReadBootCurrent();
+    BootOption ReadBootNext();
+    void WriteBootNext(const BootOption& option);
+    std::vector<BootOption> ReadBootOrder();
+    void WriteBootOrder(const std::vector<BootOption>& bootOrder);
+
+private:
+    std::string ReadDescription(std::uint16_t id);
+    std::string VarNameFromID(std::uint16_t id);
+};
+

src/UEFI_Linux.cpp

@@ -0,0 +1,111 @@
+#include "UEFI.h"
+#include <cstdlib>
+#include <sstream>
+#include <boost/locale.hpp>
+extern "C" {
+    #include <efivar.h>
+}
+
+template <typename T>
+std::pair<T*, std::size_t> ReadVariable(const std::string& name)
+{
+    std::uint8_t* data = nullptr;
+    std::size_t size;
+    std::uint32_t attributes;
+    if (efi_get_variable(EFI_GLOBAL_GUID, name.c_str(), &data, &size, &attributes) != 0) {
+        std::stringstream message;
+        message << "Failed to read EFI variable: " << name;
+        throw std::runtime_error(message.str());
+    }
+
+    return std::make_pair(reinterpret_cast<T*>(data), size);
+}
+
+UEFI::UEFI()
+{
+    efi_error_clear();
+    if (!efi_variables_supported()) {
+        throw std::runtime_error("System is not UEFI!");
+    }
+}
+
+UEFI::BootOption UEFI::ReadBootCurrent()
+{
+    auto data = ReadVariable<std::uint16_t>("BootCurrent");
+
+    UEFI::BootOption option;
+    option.ID = *data.first;
+    option.Description = ReadDescription(option.ID);
+
+    std::free(data.first);
+
+    return option;
+}
+
+UEFI::BootOption UEFI::ReadBootNext()
+{
+    auto data = ReadVariable<std::uint16_t>("BootNext");
+
+    UEFI::BootOption option;
+    option.ID = *data.first;
+    option.Description = ReadDescription(option.ID);
+
+    std::free(data.first);
+
+    return option;
+}
+
+void UEFI::WriteBootNext(const UEFI::BootOption& option)
+{
+    if (efi_set_variable(EFI_GLOBAL_GUID, "BootNext", reinterpret_cast<std::uint8_t*>(const_cast<std::uint16_t*>(&option.ID)), sizeof(std::uint16_t), EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_NON_VOLATILE, 0600) != 0) {
+        throw std::runtime_error("Failed to set EFI variable: BootNext");
+    }
+}
+
+std::vector<UEFI::BootOption> UEFI::ReadBootOrder()
+{
+    auto data = ReadVariable<std::uint16_t>("BootOrder");
+
+    int numOptions = data.second / sizeof(std::uint16_t);
+    std::vector<UEFI::BootOption> bootOrder(numOptions);
+    for (int i = 0; i < numOptions; ++i) {
+        auto& option = bootOrder[i];
+        option.ID = *(data.first + i);
+        option.Description = ReadDescription(option.ID);
+    }
+
+    std::free(data.first);
+
+    return bootOrder;
+}
+
+void UEFI::WriteBootOrder(const std::vector<UEFI::BootOption>& order)
+{
+    std::vector<std::uint16_t> data(order.size());
+    for (int i = 0; i < order.size(); ++i) {
+        data[i] = order[i].ID;
+    }
+
+    if (efi_set_variable(EFI_GLOBAL_GUID, "BootOrder", reinterpret_cast<std::uint8_t*>(data.data()), sizeof(std::uint16_t) * data.size(), EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_NON_VOLATILE, 0600) != 0) {
+        throw std::runtime_error("Failed to set EFI variable: BootOrder");
+    }
+}
+
+std::string UEFI::ReadDescription(std::uint16_t id)
+{
+    auto data = ReadVariable<char>(VarNameFromID(id));
+
+    char16_t* descUTF16 = reinterpret_cast<char16_t*>(data.first + EFI_LOAD_OPTION_DESCRIPTION_OFFSET);
+    std::string descUTF8 = boost::locale::conv::utf_to_utf<char, char16_t>(descUTF16);
+
+    std::free(data.first);
+
+    return descUTF8;
+}
+
+std::string UEFI::VarNameFromID(std::uint16_t id)
+{
+    char entry[9] = "Boot####";
+    std::snprintf(entry, 9, "Boot%04X", id);
+    return std::string(entry);
+}

src/UEFI_Windows.cpp

@@ -0,0 +1,2 @@
+#include "UEFI.h"
+

src/main.cpp

@@ -1,322 +1,183 @@
-#include <cstdint>
-#include <cstddef>
-#include <cstdio>
 #include <iostream>
-#include <map>
-#include <array>
-#include <vector>
-#include <memory>
-#define NOMINMAX
-#include <Windows.h>
-#include <boost/nowide/convert.hpp>
-//#define TCLAP_NAMESTARTSTRING "/"
-//#define TCLAP_FLAGSTARTSTRING "/"
 #include <tclap/CmdLine.h>
+#include <boost/optional.hpp>
+#include <unistd.h>
+#include <sys/reboot.h>
+#include "UEFI.h"
 
-using namespace boost::nowide;
+#define VERSION "2.0"
+bool VERBOSE = false;
 
-#define VERSION "1.0.1"
+void printBootOrder(UEFI& uefi)
-#define EFI_GLOBAL_VARIABLE L"{8BE4DF61-93CA-11D2-AA0D-00E098032B8C}"
-const std::size_t EFI_LOAD_OPTION_DESCRIPTION_OFFSET = sizeof(std::uint32_t) + sizeof(std::uint16_t);
-const std::uint16_t INVALID_OPTION_ID = std::numeric_limits<std::uint16_t>::max();
-
-bool Verbose = false;
-std::vector<std::uint16_t> BootOrder;
-std::map<std::string, std::uint16_t> DescriptionToID;
-std::map<std::uint16_t, std::string> IDToDescription;
-
-DWORD SetPrivilege(HANDLE hToken, LPCWSTR lpszPrivilege, bool bEnablePrivilege)
 {
-	LUID luid;
+    auto order = uefi.ReadBootOrder();
-	if (!LookupPrivilegeValueW(NULL, lpszPrivilege, &luid)) {
+    for (auto& option : order) {
-		std::cerr << "LookupPrivilegeValue failed: " << GetLastError() << std::endl;
+        std::cout << option.ID << ": \"" << option.Description << "\"" << std::endl;
-		return GetLastError();
+    }
-	}
-
-	TOKEN_PRIVILEGES tp;
-	tp.PrivilegeCount = 1;
-	tp.Privileges[0].Luid = luid;
-	if (bEnablePrivilege) {
-		tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
-	} else {
-		tp.Privileges[0].Attributes = 0;
-	}
-
-	if (!AdjustTokenPrivileges(
-		hToken,
-		false,
-		&tp,
-		sizeof(TOKEN_PRIVILEGES),
-		nullptr,
-		nullptr
-	)) {
-		std::cerr << "AdjustTokenPrivileges failed: " << GetLastError() << std::endl;
-		return GetLastError();
-	}
-
-	if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
-	{
-		std::cerr << "AdjustTokenPrivileges failed: The token does not have the specified privilege" << std::endl;
-		return GetLastError();
-	}
-
-	return ERROR_SUCCESS;
 }
 
-DWORD SetPrivileges()
+// Match a BootOption by numerical ID or Description
+//const UEFI::BootOption* fuzzyFindOption(const std::vector<UEFI::BootOption>& options, const std::string& identifier)
+//{
+//    // Find by ID
+//    try {
+//        std::size_t numNums;
+//        std::uint16_t id = std::stoi(identifier, &numNums);
+//        if (numNums == options.size()) {
+//            auto it = std::find_if(options.begin(), options.end(), [&id](auto& o) { return o.ID == id; });
+//            if (it != options.end()) {
+//                if (VERBOSE) { // --verbose
+//                    std::cout << "UEFI entry matched by ID." << std::endl;
+//                }
+//                return &(*it);
+//            }
+//        }
+//    } catch (std::invalid_argument& e) {
+//    } catch (std::out_of_range& e) { }
+//
+//    // Find by description
+//    {
+//        auto it = std::find_if(options.begin(), options.end(), [&identifier](auto& o) { return o.Description == identifier; });
+//        if (it != options.end()) {
+//            if (VERBOSE) {
+//                std::cout << "UEFI entry matched by description." << std::endl;
+//            }
+//            return &(*it);
+//        }
+//    }
+//
+//    return nullptr;
+//}
+
+int main(int argc, char* argv[])
 {
-	// Get access token
+    UEFI uefi;
-	HANDLE hInstance = GetCurrentProcess();
-	HANDLE hToken;
-	if (OpenProcessToken(hInstance, TOKEN_ADJUST_PRIVILEGES, &hToken) == 0) {
-		std::cerr << "OpenProcessToken failed: " << GetLastError() << std::endl;
-		return GetLastError();
-	}
 
-	DWORD result;
+    // Arguments
-	// Enable environment edit privileges
+    auto cmd = std::make_shared<TCLAP::CmdLine>("Tool to simplify dual-boot scenarios by changing the UEFI boot order to the desired entry and then reboot.", ' ', VERSION);
-	result = SetPrivilege(hToken, SE_SYSTEM_ENVIRONMENT_NAME, true);
+    TCLAP::SwitchArg arg_list("l", "list", "List current UEFI boot order and exit.");
-	if (result != ERROR_SUCCESS) {
+    TCLAP::SwitchArg arg_current("c", "current", "Move the currently booted UEFI entry to the top of the boot order and exit. Useful if you want your manual UEFI choices to stick, by automatically running this argument at startup.");
-		return result;
+    TCLAP::ValueArg<int> arg_offset("o", "offset", "Numerical offset from the top of the boot order that acts as a barrier to how high entries will be moved. Default is 0, which means boot entries will be moved to the very top. Useful if you want UEFI to try another media before proceeding with your chosen boot entry.", false, 0, "offset");
-	}
+    TCLAP::SwitchArg arg_once("n", "once", "Change boot order for the next boot only");
-	// Enable shutdown privileges
+    TCLAP::SwitchArg arg_noreboot("b", "noreboot", "Do not reboot.");
-	result = SetPrivilege(hToken, SE_SHUTDOWN_NAME, true);
+    TCLAP::SwitchArg arg_verbose("v", "verbose", "Verbose output.");
-	if (result != ERROR_SUCCESS) {
+    TCLAP::UnlabeledValueArg<std::string> arg_entry("entry", "The UEFI entry to reboot into.", false, "", "UEFI entry");
-		return result;
-	}
 
-	return ERROR_SUCCESS;
+    // Parse command line
-}
+    try {
+        cmd->add(arg_list);
+        cmd->add(arg_current);
+        cmd->add(arg_offset);
+        cmd->add(arg_once);
+        cmd->add(arg_noreboot);
+        cmd->add(arg_verbose);
+        cmd->add(arg_entry);
+        cmd->parse(argc, argv);
+    } catch (TCLAP::ArgException& e) {
+        std::cerr << "Error: " << e.error() << " for arg " << e.argId() << std::endl;
+        return 1;
+    }
+    // Positional argument requirement
+    if (arg_entry.getValue().empty() && !arg_list.getValue() && !arg_current.getValue()) {
+        auto output = cmd->getOutput();
+        output->usage(*cmd);
+        return 1;
+    }
+    VERBOSE = arg_verbose.getValue();
 
-DWORD LoadBootOrder()
+    // Print boot order
-{
+    if (arg_list.getValue()) { // --list
-	std::array<std::uint16_t, 32> buffer;
+        printBootOrder(uefi);
-	std::size_t size = GetFirmwareEnvironmentVariableW(L"BootOrder", EFI_GLOBAL_VARIABLE, buffer.data(), sizeof(std::uint16_t) * buffer.size());
+        return 0;
-	if (size == 0) {
+    }
-		std::cout << "LoadBootOrder failed: " << GetLastError() << std::endl;
+    if (VERBOSE) { // --verbose
-		return GetLastError();
+        printBootOrder(uefi);
-	}
+    }
 
-	int numEntries = size / sizeof(std::uint16_t);
+    // Fetch boot order
-	if (Verbose) {
+    auto order = uefi.ReadBootOrder();
-		std::cout << "Received " << numEntries << " UEFI entries." << std::endl;
-	}
 
-	BootOrder.resize(numEntries);
+    // Get option ID
-	std::copy_n(buffer.begin(), numEntries, BootOrder.begin());
+    std::vector<UEFI::BootOption>::iterator option = order.end();
+    if (arg_current.getValue()) { // --current
+        UEFI::BootOption current = uefi.ReadBootCurrent();
+        option = std::find_if(order.begin(), order.end(), [&current](auto& o) { return o.ID == current.ID; });
+        if (option == order.end()) {
+            std::cerr << "Error: Currently booted UEFI entry has been deleted since last boot!" << std::endl;
+            return 1;
+        }
+        if (VERBOSE) { // --verbose
+            std::cout << "Currently booted UEFI entry is \"" << (*option).ID << ": " << (*option).Description << "\"" << std::endl;
+        }
+    } else { // <UEFI entry>
+        // Find by ID
+        try {
+            std::size_t numNums;
+            std::uint16_t id = std::stoi(arg_entry.getValue(), &numNums);
+            if (numNums == arg_entry.getValue().size()) {
+                option = std::find_if(order.begin(), order.end(), [&id](auto& o) { return o.ID == id; });
+                if (VERBOSE && option != order.end()) { // --verbose
+                    std::cout << "UEFI entry matched by ID." << std::endl;
+                }
+            }
+        } catch (std::invalid_argument& e) {
+        } catch (std::out_of_range& e) { }
 
-	return ERROR_SUCCESS;
+        // Find by description if ID didn't match already
-}
+        if (option == order.end()) {
+            auto& identifier = arg_entry.getValue();
+            option = std::find_if(order.begin(), order.end(), [&identifier](auto& o) { return o.Description == identifier; });
+            if (VERBOSE && option != order.end()) { // --verbose
+                std::cout << "UEFI entry matched by description." << std::endl;
+            }
+        }
 
-DWORD LoadBootDescriptions()
+        if (option == order.end()) {
-{
+            std::cerr << "Error: No UEFI entry found with ID or description \"" << arg_entry.getValue() << "\"." << std::endl;
-	char* buffer = new char[1024];
+            std::cerr << "Use --list to list all available entries." << std::endl;
-	for (int i = 0; i < BootOrder.size(); i++) {
+            return 1;
-		std::uint16_t id = BootOrder.at(i);
+        }
-		char entry[9] = "Boot####";
+    }
-		std::snprintf(entry, 9, "Boot%04X", id);
 
-		std::size_t size = GetFirmwareEnvironmentVariableW(widen(entry).c_str(), EFI_GLOBAL_VARIABLE, buffer, sizeof(char) * 1024);
+    // Write boot commands to UEFI
-		if (size == 0) {
+    if (arg_once.getValue()) { // --once
-			std::cerr << "LoadBootDescriptions failed on " << entry << ": " << GetLastError() << std::endl;
+        // Set BootNext
-			return GetLastError();
+        uefi.WriteBootNext(*option);
-		}
+        std::cout << "BootNext changed successfully." << std::endl;
+    } else {
+        // Move option to top of boot order
+        auto currentDefault = order.begin();
+        std::advance(currentDefault, arg_offset.getValue()); // --offset=0
+        auto newDefault = option;
+        if (newDefault > currentDefault) {
+            // Shift entry to top
+            UEFI::BootOption tmp = *newDefault;
+            std::copy_backward(currentDefault, newDefault, std::next(newDefault));
+            *currentDefault = tmp;
 
-		std::string description(narrow(reinterpret_cast<wchar_t*>(buffer + EFI_LOAD_OPTION_DESCRIPTION_OFFSET)));
+            // Apply boot order
-		DescriptionToID[description] = id;
+//            for (auto& option : order) {
-		IDToDescription[id] = description;
+//                std::cout << option.ID << ": \"" << option.Description << "\"" << std::endl;
-	}
+//            }
-	delete[] buffer;
+            uefi.WriteBootOrder(order);
 
-	return ERROR_SUCCESS;
+            std::cout << "Boot order changed successfully." << std::endl;
-}
+        } else if (newDefault < currentDefault) {
+            std::cout << "Boot entry is already above provided offset! Nothing has been changed." << std::endl;
+        } else {
+            std::cout << "Boot entry is already the default! Nothing has been changed." << std::endl;
+        }
+    }
 
-void PrintBootOrder() 
+    // Reboot
-{
+    if (!arg_noreboot.getValue() && !arg_current.getValue()) { // --noreboot --current
-	for (auto& id : BootOrder) {
+        std::cout << "Rebooting..." << std::endl;
-		std::cout << id << ": \"" << IDToDescription.at(id) << "\"" << std::endl;
+        sync();
-	}
+        if (reboot(RB_AUTOBOOT) != -1) {
-}
+            std::cerr << "Failed to initiate reboot." << std::endl;
+            return 1;
+        }
+    }
 
-std::uint16_t GetCurrentBootOption()
+    return 0;
-{
-	std::uint16_t current;
-	std::size_t size = GetFirmwareEnvironmentVariableW(L"BootCurrent", EFI_GLOBAL_VARIABLE, &current, sizeof(std::uint16_t));
-	if (size == 0) {
-		std::cout << "GetCurrentBootOption failed: " << GetLastError() << std::endl;
-		return INVALID_OPTION_ID;
-	}
-	return current;
-}
-
-DWORD SetBootNext(std::uint16_t entry)
-{
-	bool result = SetFirmwareEnvironmentVariableW(L"BootNext", EFI_GLOBAL_VARIABLE, &entry, sizeof(std::uint16_t));
-	if (!result) {
-		std::cerr << "SetBootNext failed: " << GetLastError() << std::endl;
-		return GetLastError();
-	}
-
-	return ERROR_SUCCESS;
-}
-
-DWORD SetBootDefault(std::uint16_t entry, std::size_t offset = 0)
-{
-	auto currentDefault = BootOrder.begin();
-	std::advance(currentDefault, offset);
-	auto newDefault = std::find(BootOrder.begin(), BootOrder.end(), entry);
-	if (newDefault > currentDefault) {
-		// Shift entry to top
-		std::uint16_t tmp = *newDefault;
-		std::copy_backward(currentDefault, newDefault, std::next(newDefault));
-		*currentDefault = tmp;
-
-		// Apply boot order
-		bool result = SetFirmwareEnvironmentVariableW(L"BootOrder", EFI_GLOBAL_VARIABLE, BootOrder.data(), sizeof(std::uint16_t) * BootOrder.size());
-		if (!result) {
-			std::cerr << "Failed to apply boot order: " << GetLastError() << std::endl;
-			return GetLastError();
-		}
-
-		std::cout << "Boot order changed successfully." << std::endl;
-	} else if (newDefault < currentDefault) {
-		std::cout << "Boot entry is already above provided offset! Nothing has been changed." << std::endl;
-	} else {
-		std::cout << "Boot entry is already the default! Nothing has been changed." << std::endl;
-	}
-
-	return ERROR_SUCCESS;
-}
-
-int wmain(int argc, wchar_t* argv[])
-{
-	// Check for UEFI
-	FIRMWARE_TYPE firmwareType;
-	if (GetFirmwareType(&firmwareType) == 0) {
-		std::cerr << "GetFirmwareType failed: " << GetLastError() << std::endl;
-		return 1;
-	}
-	if (firmwareType != FirmwareTypeUefi) {
-		std::cerr << "Error: System is not UEFI!" << std::endl;
-		return 1;
-	}
-
-	// Create a UTF-8 argument vector for TCLAP
-	std::vector<std::string> args(argc);
-	for (int i = 0; i < argc; i++) {
-		args[i] = narrow(argv[i]);
-	}
-
-	// Arguments
-	auto cmd = std::make_shared<TCLAP::CmdLine>("Tool to simplify dual-boot scenarios by changing the UEFI boot order to the desired entry and then reboot.", ' ', VERSION);
-	TCLAP::SwitchArg arg_list("l", "list", "List current UEFI boot order and exit.");
-	TCLAP::SwitchArg arg_current("c", "current", "Move the currently booted UEFI entry to the top of the boot order and exit. Useful if you want your manual UEFI choices to stick, by automatically running this argument at startup.");
-	TCLAP::ValueArg<int> arg_offset("o", "offset", "Numerical offset from the top of the boot order that acts as a barrier to how high entries will be moved. Default is 0, which means boot entries will be moved to the very top. Useful if you want UEFI to try another media before proceeding with your chosen boot entry.", false, 0, "offset");
-	TCLAP::SwitchArg arg_once("n", "once", "Change boot order for the next boot only");
-	TCLAP::SwitchArg arg_noreboot("b", "noreboot", "Do not reboot.");
-	TCLAP::SwitchArg arg_verbose("v", "verbose", "Verbose output.");
-	TCLAP::UnlabeledValueArg<std::string> arg_entry("entry", "The UEFI entry to reboot into.", false, "", "UEFI entry");
-
-	// Parse command line
-	try {
-		cmd->add(arg_list);
-		cmd->add(arg_current);
-		cmd->add(arg_offset);
-		cmd->add(arg_once);
-		cmd->add(arg_noreboot);
-		cmd->add(arg_verbose);
-		cmd->add(arg_entry);
-		cmd->parse(args);
-	} catch (TCLAP::ArgException& e) {
-		std::cerr << "Error: " << e.error() << " for arg " << e.argId() << std::endl;
-		return 1;
-	}
-	Verbose = arg_verbose.getValue();
-	// Positional argument requirement
-	if (arg_entry.getValue().empty() && !arg_list.getValue() && !arg_current.getValue()) {
-		auto output = cmd->getOutput();
-		output->usage(*cmd);
-		return 1;
-	}
-
-	// Initialize
-	if (SetPrivileges() != ERROR_SUCCESS) { return 1; }
-	if (LoadBootOrder() != ERROR_SUCCESS) { return 1; }
-	if (LoadBootDescriptions() != ERROR_SUCCESS) { return 1; }
-
-	// Print boot order
-	if (arg_list.getValue() || Verbose) { // --list --verbose
-		PrintBootOrder();
-		if (!Verbose) {
-			return 0;
-		}
-	}
-
-	// Get option ID
-	std::uint16_t entry = INVALID_OPTION_ID;
-	if (arg_current.getValue()) { // --current
-		entry = GetCurrentBootOption();
-		if (entry == INVALID_OPTION_ID) {
-			return 1;
-		}
-		if (std::find(BootOrder.begin(), BootOrder.end(), entry) == BootOrder.end()) {
-			std::cerr << "Error: Currently booted UEFI entry has been deleted since last boot!" << std::endl;
-			return 1;
-		}
-		if (Verbose) {
-			std::cout << "Currently booted UEFI entry is \"" << entry << ": " << IDToDescription.at(entry) << "\"" << std::endl;
-		}
-	} else { // <UEFI entry>
-		// Find by ID
-		try {
-			std::size_t numNums;
-			std::uint16_t id = std::stoi(arg_entry.getValue(), &numNums);
-			if (numNums == arg_entry.getValue().size()) {
-				auto it = std::find(BootOrder.begin(), BootOrder.end(), id);
-				if (it != BootOrder.end()) {
-					entry = id;
-					if (Verbose) {
-						std::cout << "UEFI entry matched by ID." << std::endl;
-					}
-				}
-			}
-		} catch (std::invalid_argument& e) { 
-		} catch (std::out_of_range& e) { }
-		
-		// Find by description
-		if (entry == INVALID_OPTION_ID) {
-			auto it = DescriptionToID.find(arg_entry.getValue());
-			if (it == DescriptionToID.end()) {
-				std::cerr << "Error: No UEFI entry found with ID or description \"" << arg_entry.getValue() << "\"." << std::endl;
-				std::cerr << "Use --list to list all available entries." << std::endl;
-				return 1;
-			}
-
-			if (Verbose) {
-				std::cout << "UEFI entry matched by description." << std::endl;
-			}
-			entry = it->second;
-		}
-	}
-
-	// Write boot commands to UEFI
-	if (arg_once.getValue()) { // --once
-		// Set BootNext
-		if (SetBootNext(entry) != ERROR_SUCCESS) { return 1; }
-	} else {
-		// Move option to top of boot order
-		if (SetBootDefault(entry, arg_offset.getValue()) != ERROR_SUCCESS) { return 1; }
-	}
-
-	// Reboot
-	if (!arg_noreboot.getValue() && !arg_current.getValue()) { // --noreboot --current
-		if (InitiateSystemShutdown(nullptr, nullptr, 0, false, true) == 0) {
-			std::cerr << "Failed to initiate reboot: " << GetLastError() << std::endl;
-			return 1;
-		} else {
-			std::cout << "Rebooting..." << std::endl;
-		}
-	}
-
-	return 0;
 }

src/main_win.cpp

@@ -0,0 +1,322 @@
+#include <cstdint>
+#include <cstddef>
+#include <cstdio>
+#include <iostream>
+#include <map>
+#include <array>
+#include <vector>
+#include <memory>
+#define NOMINMAX
+#include <Windows.h>
+#include <boost/nowide/convert.hpp>
+//#define TCLAP_NAMESTARTSTRING "/"
+//#define TCLAP_FLAGSTARTSTRING "/"
+#include <tclap/CmdLine.h>
+
+using namespace boost::nowide;
+
+#define VERSION "1.0.1"
+#define EFI_GLOBAL_VARIABLE L"{8BE4DF61-93CA-11D2-AA0D-00E098032B8C}"
+const std::size_t EFI_LOAD_OPTION_DESCRIPTION_OFFSET = sizeof(std::uint32_t) + sizeof(std::uint16_t);
+const std::uint16_t INVALID_OPTION_ID = std::numeric_limits<std::uint16_t>::max();
+
+bool Verbose = false;
+std::vector<std::uint16_t> BootOrder;
+std::map<std::string, std::uint16_t> DescriptionToID;
+std::map<std::uint16_t, std::string> IDToDescription;
+
+DWORD SetPrivilege(HANDLE hToken, LPCWSTR lpszPrivilege, bool bEnablePrivilege)
+{
+	LUID luid;
+	if (!LookupPrivilegeValueW(NULL, lpszPrivilege, &luid)) {
+		std::cerr << "LookupPrivilegeValue failed: " << GetLastError() << std::endl;
+		return GetLastError();
+	}
+
+	TOKEN_PRIVILEGES tp;
+	tp.PrivilegeCount = 1;
+	tp.Privileges[0].Luid = luid;
+	if (bEnablePrivilege) {
+		tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
+	} else {
+		tp.Privileges[0].Attributes = 0;
+	}
+
+	if (!AdjustTokenPrivileges(
+		hToken,
+		false,
+		&tp,
+		sizeof(TOKEN_PRIVILEGES),
+		nullptr,
+		nullptr
+	)) {
+		std::cerr << "AdjustTokenPrivileges failed: " << GetLastError() << std::endl;
+		return GetLastError();
+	}
+
+	if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
+	{
+		std::cerr << "AdjustTokenPrivileges failed: The token does not have the specified privilege" << std::endl;
+		return GetLastError();
+	}
+
+	return ERROR_SUCCESS;
+}
+
+DWORD SetPrivileges()
+{
+	// Get access token
+	HANDLE hInstance = GetCurrentProcess();
+	HANDLE hToken;
+	if (OpenProcessToken(hInstance, TOKEN_ADJUST_PRIVILEGES, &hToken) == 0) {
+		std::cerr << "OpenProcessToken failed: " << GetLastError() << std::endl;
+		return GetLastError();
+	}
+
+	DWORD result;
+	// Enable environment edit privileges
+	result = SetPrivilege(hToken, SE_SYSTEM_ENVIRONMENT_NAME, true);
+	if (result != ERROR_SUCCESS) {
+		return result;
+	}
+	// Enable shutdown privileges
+	result = SetPrivilege(hToken, SE_SHUTDOWN_NAME, true);
+	if (result != ERROR_SUCCESS) {
+		return result;
+	}
+
+	return ERROR_SUCCESS;
+}
+
+DWORD LoadBootOrder()
+{
+	std::array<std::uint16_t, 32> buffer;
+	std::size_t size = GetFirmwareEnvironmentVariableW(L"BootOrder", EFI_GLOBAL_VARIABLE, buffer.data(), sizeof(std::uint16_t) * buffer.size());
+	if (size == 0) {
+		std::cout << "LoadBootOrder failed: " << GetLastError() << std::endl;
+		return GetLastError();
+	}
+
+	int numEntries = size / sizeof(std::uint16_t);
+	if (Verbose) {
+		std::cout << "Received " << numEntries << " UEFI entries." << std::endl;
+	}
+
+	BootOrder.resize(numEntries);
+	std::copy_n(buffer.begin(), numEntries, BootOrder.begin());
+
+	return ERROR_SUCCESS;
+}
+
+DWORD LoadBootDescriptions()
+{
+	char* buffer = new char[1024];
+	for (int i = 0; i < BootOrder.size(); i++) {
+		std::uint16_t id = BootOrder.at(i);
+		char entry[9] = "Boot####";
+		std::snprintf(entry, 9, "Boot%04X", id);
+
+		std::size_t size = GetFirmwareEnvironmentVariableW(widen(entry).c_str(), EFI_GLOBAL_VARIABLE, buffer, sizeof(char) * 1024);
+		if (size == 0) {
+			std::cerr << "LoadBootDescriptions failed on " << entry << ": " << GetLastError() << std::endl;
+			return GetLastError();
+		}
+
+		std::string description(narrow(reinterpret_cast<wchar_t*>(buffer + EFI_LOAD_OPTION_DESCRIPTION_OFFSET)));
+		DescriptionToID[description] = id;
+		IDToDescription[id] = description;
+	}
+	delete[] buffer;
+
+	return ERROR_SUCCESS;
+}
+
+void PrintBootOrder() 
+{
+	for (auto& id : BootOrder) {
+		std::cout << id << ": \"" << IDToDescription.at(id) << "\"" << std::endl;
+	}
+}
+
+std::uint16_t GetCurrentBootOption()
+{
+	std::uint16_t current;
+	std::size_t size = GetFirmwareEnvironmentVariableW(L"BootCurrent", EFI_GLOBAL_VARIABLE, &current, sizeof(std::uint16_t));
+	if (size == 0) {
+		std::cout << "GetCurrentBootOption failed: " << GetLastError() << std::endl;
+		return INVALID_OPTION_ID;
+	}
+	return current;
+}
+
+DWORD SetBootNext(std::uint16_t entry)
+{
+	bool result = SetFirmwareEnvironmentVariableW(L"BootNext", EFI_GLOBAL_VARIABLE, &entry, sizeof(std::uint16_t));
+	if (!result) {
+		std::cerr << "SetBootNext failed: " << GetLastError() << std::endl;
+		return GetLastError();
+	}
+
+	return ERROR_SUCCESS;
+}
+
+DWORD SetBootDefault(std::uint16_t entry, std::size_t offset = 0)
+{
+	auto currentDefault = BootOrder.begin();
+	std::advance(currentDefault, offset);
+	auto newDefault = std::find(BootOrder.begin(), BootOrder.end(), entry);
+	if (newDefault > currentDefault) {
+		// Shift entry to top
+		std::uint16_t tmp = *newDefault;
+		std::copy_backward(currentDefault, newDefault, std::next(newDefault));
+		*currentDefault = tmp;
+
+		// Apply boot order
+		bool result = SetFirmwareEnvironmentVariableW(L"BootOrder", EFI_GLOBAL_VARIABLE, BootOrder.data(), sizeof(std::uint16_t) * BootOrder.size());
+		if (!result) {
+			std::cerr << "Failed to apply boot order: " << GetLastError() << std::endl;
+			return GetLastError();
+		}
+
+		std::cout << "Boot order changed successfully." << std::endl;
+	} else if (newDefault < currentDefault) {
+		std::cout << "Boot entry is already above provided offset! Nothing has been changed." << std::endl;
+	} else {
+		std::cout << "Boot entry is already the default! Nothing has been changed." << std::endl;
+	}
+
+	return ERROR_SUCCESS;
+}
+
+int wmain(int argc, wchar_t* argv[])
+{
+	// Check for UEFI
+	FIRMWARE_TYPE firmwareType;
+	if (GetFirmwareType(&firmwareType) == 0) {
+		std::cerr << "GetFirmwareType failed: " << GetLastError() << std::endl;
+		return 1;
+	}
+	if (firmwareType != FirmwareTypeUefi) {
+		std::cerr << "Error: System is not UEFI!" << std::endl;
+		return 1;
+	}
+
+	// Create a UTF-8 argument vector for TCLAP
+	std::vector<std::string> args(argc);
+	for (int i = 0; i < argc; i++) {
+		args[i] = narrow(argv[i]);
+	}
+
+	// Arguments
+	auto cmd = std::make_shared<TCLAP::CmdLine>("Tool to simplify dual-boot scenarios by changing the UEFI boot order to the desired entry and then reboot.", ' ', VERSION);
+	TCLAP::SwitchArg arg_list("l", "list", "List current UEFI boot order and exit.");
+	TCLAP::SwitchArg arg_current("c", "current", "Move the currently booted UEFI entry to the top of the boot order and exit. Useful if you want your manual UEFI choices to stick, by automatically running this argument at startup.");
+	TCLAP::ValueArg<int> arg_offset("o", "offset", "Numerical offset from the top of the boot order that acts as a barrier to how high entries will be moved. Default is 0, which means boot entries will be moved to the very top. Useful if you want UEFI to try another media before proceeding with your chosen boot entry.", false, 0, "offset");
+	TCLAP::SwitchArg arg_once("n", "once", "Change boot order for the next boot only");
+	TCLAP::SwitchArg arg_noreboot("b", "noreboot", "Do not reboot.");
+	TCLAP::SwitchArg arg_verbose("v", "verbose", "Verbose output.");
+	TCLAP::UnlabeledValueArg<std::string> arg_entry("entry", "The UEFI entry to reboot into.", false, "", "UEFI entry");
+
+	// Parse command line
+	try {
+		cmd->add(arg_list);
+		cmd->add(arg_current);
+		cmd->add(arg_offset);
+		cmd->add(arg_once);
+		cmd->add(arg_noreboot);
+		cmd->add(arg_verbose);
+		cmd->add(arg_entry);
+		cmd->parse(args);
+	} catch (TCLAP::ArgException& e) {
+		std::cerr << "Error: " << e.error() << " for arg " << e.argId() << std::endl;
+		return 1;
+	}
+	Verbose = arg_verbose.getValue();
+	// Positional argument requirement
+	if (arg_entry.getValue().empty() && !arg_list.getValue() && !arg_current.getValue()) {
+		auto output = cmd->getOutput();
+		output->usage(*cmd);
+		return 1;
+	}
+
+	// Initialize
+	if (SetPrivileges() != ERROR_SUCCESS) { return 1; }
+	if (LoadBootOrder() != ERROR_SUCCESS) { return 1; }
+	if (LoadBootDescriptions() != ERROR_SUCCESS) { return 1; }
+
+	// Print boot order
+	if (arg_list.getValue() || Verbose) { // --list --verbose
+		PrintBootOrder();
+		if (!Verbose) {
+			return 0;
+		}
+	}
+
+	// Get option ID
+	std::uint16_t entry = INVALID_OPTION_ID;
+	if (arg_current.getValue()) { // --current
+		entry = GetCurrentBootOption();
+		if (entry == INVALID_OPTION_ID) {
+			return 1;
+		}
+		if (std::find(BootOrder.begin(), BootOrder.end(), entry) == BootOrder.end()) {
+			std::cerr << "Error: Currently booted UEFI entry has been deleted since last boot!" << std::endl;
+			return 1;
+		}
+		if (Verbose) {
+			std::cout << "Currently booted UEFI entry is \"" << entry << ": " << IDToDescription.at(entry) << "\"" << std::endl;
+		}
+	} else { // <UEFI entry>
+		// Find by ID
+		try {
+			std::size_t numNums;
+			std::uint16_t id = std::stoi(arg_entry.getValue(), &numNums);
+			if (numNums == arg_entry.getValue().size()) {
+				auto it = std::find(BootOrder.begin(), BootOrder.end(), id);
+				if (it != BootOrder.end()) {
+					entry = id;
+					if (Verbose) {
+						std::cout << "UEFI entry matched by ID." << std::endl;
+					}
+				}
+			}
+		} catch (std::invalid_argument& e) { 
+		} catch (std::out_of_range& e) { }
+		
+		// Find by description
+		if (entry == INVALID_OPTION_ID) {
+			auto it = DescriptionToID.find(arg_entry.getValue());
+			if (it == DescriptionToID.end()) {
+				std::cerr << "Error: No UEFI entry found with ID or description \"" << arg_entry.getValue() << "\"." << std::endl;
+				std::cerr << "Use --list to list all available entries." << std::endl;
+				return 1;
+			}
+
+			if (Verbose) {
+				std::cout << "UEFI entry matched by description." << std::endl;
+			}
+			entry = it->second;
+		}
+	}
+
+	// Write boot commands to UEFI
+	if (arg_once.getValue()) { // --once
+		// Set BootNext
+		if (SetBootNext(entry) != ERROR_SUCCESS) { return 1; }
+	} else {
+		// Move option to top of boot order
+		if (SetBootDefault(entry, arg_offset.getValue()) != ERROR_SUCCESS) { return 1; }
+	}
+
+	// Reboot
+	if (!arg_noreboot.getValue() && !arg_current.getValue()) { // --noreboot --current
+		if (InitiateSystemShutdown(nullptr, nullptr, 0, false, true) == 0) {
+			std::cerr << "Failed to initiate reboot: " << GetLastError() << std::endl;
+			return 1;
+		} else {
+			std::cout << "Rebooting..." << std::endl;
+		}
+	}
+
+	return 0;
+}