Isolated platform-dependent functions from main

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()

src/Platform.h

@@ -0,0 +1,13 @@
+#ifndef Platform_h__
+#define Platform_h__
+
+#include <stdexcept>
+
+namespace Platform {
+
+void Initialize();
+void Reboot();
+
+}
+
+#endif

src/Platform_Linux.cpp

@@ -0,0 +1,16 @@
+#include "Platform.h"
+#include <unistd.h>
+#include <sys/reboot.h>
+
+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.");
+	}
+}

src/Platform_Windows.cpp

@@ -0,0 +1,54 @@
+#include "Platform.h"
+#include <sstream>
+#include <Windows.h>
+
+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());
+	}
+}

src/main.cpp

@@ -1,8 +1,7 @@
 #include <iostream>
 #include <tclap/CmdLine.h>
 #include <boost/optional.hpp>
-#include <unistd.h>
-#include <sys/reboot.h>
+#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<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[])
 {
+	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;