spicetools/script/api/control.cpp

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