diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7cb0c99..4c1b4a9 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -15,16 +15,19 @@ include_directories( ) set(SOURCE_FILES + Platform.h UEFI.h main.cpp ) if(WIN32) set(SOURCE_FILES ${SOURCE_FILES} + Platform_Windows.cpp UEFI_Windows.cpp ) endif() if(UNIX) set(SOURCE_FILES ${SOURCE_FILES} + Platform_Linux.cpp UEFI_Linux.cpp ) endif() diff --git a/src/Platform.h b/src/Platform.h new file mode 100644 index 0000000..e160ca2 --- /dev/null +++ b/src/Platform.h @@ -0,0 +1,13 @@ +#ifndef Platform_h__ +#define Platform_h__ + +#include + +namespace Platform { + +void Initialize(); +void Reboot(); + +} + +#endif \ No newline at end of file diff --git a/src/Platform_Linux.cpp b/src/Platform_Linux.cpp new file mode 100644 index 0000000..76c2d96 --- /dev/null +++ b/src/Platform_Linux.cpp @@ -0,0 +1,16 @@ +#include "Platform.h" +#include +#include + +void Platform::Initialize() +{ + // TODO: root permission check here +} + +void Platform::Reboot() +{ + sync(); + if (reboot(RB_AUTOBOOT) != -1) { + throw std::runtime_error("Failed to initiate reboot."); + } +} \ No newline at end of file diff --git a/src/Platform_Windows.cpp b/src/Platform_Windows.cpp new file mode 100644 index 0000000..03f7d65 --- /dev/null +++ b/src/Platform_Windows.cpp @@ -0,0 +1,54 @@ +#include "Platform.h" +#include +#include + +void SetPrivilege(HANDLE hToken, LPCWSTR lpszPrivilege, bool bEnablePrivilege) +{ + LUID luid; + if (!LookupPrivilegeValueW(NULL, lpszPrivilege, &luid)) { + std::stringstream message; + message << "LookupPrivilegeValue failed: " << GetLastError(); + throw std::runtime_error(message.str()); + } + + 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::stringstream message; + message << "AdjustTockenPrivileges failed: " << GetLastError(); + throw std::runtime_error(message.str()); + } +} + +void Platform::Initialize() +{ + // Get access token + HANDLE hInstance = GetCurrentProcess(); + HANDLE hToken; + if (OpenProcessToken(hInstance, TOKEN_ADJUST_PRIVILEGES, &hToken) == 0) { + std::stringstream message; + message << "OpenProcessToken failed: " << GetLastError(); + throw std::runtime_error(message.str()); + } + + // Enable environment edit privileges + SetPrivilege(hToken, SE_SYSTEM_ENVIRONMENT_NAME, true); + // Enable shutdown privileges + SetPrivilege(hToken, SE_SHUTDOWN_NAME, true); +} + +void Platform::Reboot() +{ + if (InitiateSystemShutdown(nullptr, nullptr, 0, false, true) == 0) { + std::stringstream message; + message << "Failed to initiate reboot: " << GetLastError(); + throw std::runtime_error(message.str()); + } +} \ No newline at end of file diff --git a/src/main.cpp b/src/main.cpp index 80d0f32..9f303e9 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -1,8 +1,7 @@ #include #include #include -#include -#include +#include "Platform.h" #include "UEFI.h" #define VERSION "2.0" @@ -16,41 +15,9 @@ void printBootOrder(UEFI& uefi) } } -// Match a BootOption by numerical ID or Description -//const UEFI::BootOption* fuzzyFindOption(const std::vector& 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[]) { + Platform::Initialize(); UEFI uefi; // Arguments @@ -156,9 +123,6 @@ int main(int argc, char* argv[]) *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; @@ -172,11 +136,7 @@ int main(int argc, char* argv[]) // Reboot if (!arg_noreboot.getValue() && !arg_current.getValue()) { // --noreboot --current std::cout << "Rebooting..." << std::endl; - sync(); - if (reboot(RB_AUTOBOOT) != -1) { - std::cerr << "Failed to initiate reboot." << std::endl; - return 1; - } + Platform::Reboot(); } return 0;