122 lines
3.1 KiB
C++
122 lines
3.1 KiB
C++
|
#include "control.h"
|
||
|
|
#include <csignal>
|
||
|
|
#include <cstring>
|
||
|
|
#include "external/LuaBridge.h"
|
||
|
|
#include "launcher/shutdown.h"
|
||
|
|
|
||
|
|
using namespace luabridge;
|
||
|
|
|
||
|
|
namespace script::api::control {
|
||
|
|
|
||
|
|
struct SignalMapping {
|
||
|
|
int signum;
|
||
|
|
const char* name;
|
||
|
|
};
|
||
|
|
static SignalMapping SIGNAL_MAPPINGS[] = {
|
||
|
|
{ SIGABRT, "SIGABRT" },
|
||
|
|
{ SIGFPE, "SIGFPE" },
|
||
|
|
{ SIGILL, "SIGILL" },
|
||
|
|
{ SIGINT, "SIGINT" },
|
||
|
|
{ SIGSEGV, "SIGSEGV" },
|
||
|
|
{ SIGTERM, "SIGTERM" },
|
||
|
|
};
|
||
|
|
|
||
|
|
static inline bool acquire_shutdown_privs() {
|
||
|
|
|
||
|
|
// check if already acquired
|
||
|
|
static bool acquired = false;
|
||
|
|
if (acquired)
|
||
|
|
return true;
|
||
|
|
|
||
|
|
// get process token
|
||
|
|
HANDLE hToken;
|
||
|
|
if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
|
||
|
|
return false;
|
||
|
|
|
||
|
|
// get the LUID for the shutdown privilege
|
||
|
|
TOKEN_PRIVILEGES tkp;
|
||
|
|
LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &tkp.Privileges[0].Luid);
|
||
|
|
tkp.PrivilegeCount = 1;
|
||
|
|
tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
|
||
|
|
|
||
|
|
// get the shutdown privilege for this process
|
||
|
|
AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, (PTOKEN_PRIVILEGES) NULL, 0);
|
||
|
|
|
||
|
|
// check for error
|
||
|
|
bool success = GetLastError() == ERROR_SUCCESS;
|
||
|
|
if (success)
|
||
|
|
acquired = true;
|
||
|
|
return success;
|
||
|
|
}
|
||
|
|
|
||
|
|
void raise(const std::string &signal) {
|
||
|
|
|
||
|
|
// get signal
|
||
|
|
int signal_val = -1;
|
||
|
|
for (auto mapping : SIGNAL_MAPPINGS) {
|
||
|
|
if (_stricmp(mapping.name, signal.c_str()) == 0) {
|
||
|
|
signal_val = mapping.signum;
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// raise if found
|
||
|
|
if (signal_val >= 0) {
|
||
|
|
::raise(signal_val);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void exit(int code) {
|
||
|
|
launcher::shutdown(code);
|
||
|
|
}
|
||
|
|
|
||
|
|
void restart() {
|
||
|
|
launcher::restart();
|
||
|
|
}
|
||
|
|
|
||
|
|
bool shutdown() {
|
||
|
|
|
||
|
|
// acquire privileges
|
||
|
|
if (!acquire_shutdown_privs())
|
||
|
|
return false;
|
||
|
|
|
||
|
|
// exit windows
|
||
|
|
if (!ExitWindowsEx(EWX_POWEROFF | EWX_FORCE,
|
||
|
|
SHTDN_REASON_MAJOR_APPLICATION |
|
||
|
|
SHTDN_REASON_MINOR_MAINTENANCE))
|
||
|
|
return false;
|
||
|
|
|
||
|
|
// terminate this process
|
||
|
|
launcher::shutdown(0);
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
bool reboot() {
|
||
|
|
|
||
|
|
// acquire privileges
|
||
|
|
if (!acquire_shutdown_privs())
|
||
|
|
return false;
|
||
|
|
|
||
|
|
// exit windows
|
||
|
|
if (!ExitWindowsEx(EWX_REBOOT | EWX_FORCE,
|
||
|
|
SHTDN_REASON_MAJOR_APPLICATION |
|
||
|
|
SHTDN_REASON_MINOR_MAINTENANCE))
|
||
|
|
return false;
|
||
|
|
|
||
|
|
// terminate this process
|
||
|
|
launcher::shutdown(0);
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
void init(lua_State *L) {
|
||
|
|
getGlobalNamespace(L)
|
||
|
|
.beginNamespace("control")
|
||
|
|
.addFunction("raise", raise)
|
||
|
|
.addFunction("exit", exit)
|
||
|
|
.addFunction("restart", restart)
|
||
|
|
.addFunction("shutdown", shutdown)
|
||
|
|
.addFunction("reboot", reboot)
|
||
|
|
.endNamespace();
|
||
|
|
}
|
||
|
|
}
|