Jace/rebootinto
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

First working Linux version

Browse Source
This commit is contained in:
sippeangelo
2016-11-11 08:04:51 +01:00
parent fe19a81484
commit a5d238e7d9
9 changed files with 670 additions and 304 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -1 +1,3 @@
build/
bin/
.idea/
CMake/Findefivar.cmake
+17
View File
@@ -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
+2 -1
View File
@@ -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
+19
View File
@@ -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)
@@ -24,3 +39,7 @@ if(MSVC)
endif()
add_executable(RebootInto ${SOURCE_FILES})
target_link_libraries(RebootInto
${Boost_LIBRARIES}
${efivar_LIBRARIES}
)
src/UEFI.h
+31
View File
@@ -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
+111
View File
@@ -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
+2
View File
@@ -0,0 +1,2 @@
#include "UEFI.h"
src/main.cpp
+99 -238
View File
@@ -1,211 +1,57 @@
#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"
#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)
void printBootOrder(UEFI& uefi)
{
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;
auto order = uefi.ReadBootOrder();
for (auto& option : order) {
std::cout << option.ID << ": \"" << option.Description << "\"" << std::endl;
}
}
std::uint16_t GetCurrentBootOption()
// 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[])
{
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]);
}
UEFI uefi;
// 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);
@@ -226,45 +72,42 @@ int wmain(int argc, wchar_t* argv[])
cmd->add(arg_noreboot);
cmd->add(arg_verbose);
cmd->add(arg_entry);
cmd->parse(args);
cmd->parse(argc, argv);
} 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; }
VERBOSE = arg_verbose.getValue();
// Print boot order
if (arg_list.getValue() || Verbose) { // --list --verbose
PrintBootOrder();
if (!Verbose) {
if (arg_list.getValue()) { // --list
printBootOrder(uefi);
return 0;
}
if (VERBOSE) { // --verbose
printBootOrder(uefi);
}
// Fetch boot order
auto order = uefi.ReadBootOrder();
// Get option ID
std::uint16_t entry = INVALID_OPTION_ID;
std::vector<UEFI::BootOption>::iterator option = order.end();
if (arg_current.getValue()) { // --current
entry = GetCurrentBootOption();
if (entry == INVALID_OPTION_ID) {
return 1;
}
if (std::find(BootOrder.begin(), BootOrder.end(), entry) == BootOrder.end()) {
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) {
std::cout << "Currently booted UEFI entry is \"" << entry << ": " << IDToDescription.at(entry) << "\"" << std::endl;
if (VERBOSE) { // --verbose
std::cout << "Currently booted UEFI entry is \"" << (*option).ID << ": " << (*option).Description << "\"" << std::endl;
}
} else { // <UEFI entry>
// Find by ID
@@ -272,49 +115,67 @@ int wmain(int argc, wchar_t* argv[])
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) {
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) { }
// Find by description
if (entry == INVALID_OPTION_ID) {
auto it = DescriptionToID.find(arg_entry.getValue());
if (it == DescriptionToID.end()) {
// 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;
}
}
if (option == order.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; }
uefi.WriteBootNext(*option);
std::cout << "BootNext changed successfully." << std::endl;
} else {
// Move option to top of boot order
if (SetBootDefault(entry, arg_offset.getValue()) != ERROR_SUCCESS) { return 1; }
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;
// Apply boot order
// for (auto& option : order) {
// std::cout << option.ID << ": \"" << option.Description << "\"" << std::endl;
// }
uefi.WriteBootOrder(order);
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;
}
}
// 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;
sync();
if (reboot(RB_AUTOBOOT) != -1) {
std::cerr << "Failed to initiate reboot." << std::endl;
return 1;
}
}
src/main_win.cpp
+322
View File
@@ -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;
}