1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

#include <chrono>
#include <iostream>
#include <string>
#include <thread>
#include <Windows.h>

typedef enum _SHUTDOWN_ACTION
{
    ShutdownNoReboot,
    ShutdownReboot,
    ShutdownPowerOff
} SHUTDOWN_ACTION,
    *PSHUTDOWN_ACTION;

typedef LONG(__stdcall *PFN_NT_SHUTDOWN_SYSTEM)(DWORD SHUTDOWN_ACTION);

int DisplayError(const std::wstring &message);
std::wstring GetFormattedMessage(DWORD, LPCWSTR, DWORD, ...);

int main(int argc, char *argv[])
{
    std::wstring programName = std::wstring(argv[0], argv[0] + strlen(argv[0]));
    // std::wcout << L"Program Name: " << programName << std::endl;

    if (argc < 2 || strcmp(argv[1], "/?") == 0)
    {
        std::wcout << L"Usage: " << programName << L" [/r | /s] [/t seconds]\n"
                   << L"  /r : Reboot the system\n"
                   << L"  /s : Shutdown the system\n"
                   << L"  /t seconds : Set a countdown before executing (default 0)" << std::endl;
        return 0;
    }

    SHUTDOWN_ACTION action;
    int delay = 0;

    std::wcout << L"Processing command: " << argv[1] << std::endl;

    if (strcmp(argv[1], "/r") == 0)
    {
        action = ShutdownReboot;
        std::wcout << L"Action: Reboot" << std::endl;
    }
    else if (strcmp(argv[1], "/s") == 0)
    {
        action = ShutdownPowerOff;
        std::wcout << L"Action: Shutdown" << std::endl;
    }
    else
    {
        std::wcout << L"Invalid option. Use /? for help." << std::endl;
        return 1;
    }

    if (argc == 4 && strcmp(argv[2], "/t") == 0)
    {
        delay = std::stoi(argv[3]);
        std::wcout << L"Shutdown delay: " << delay << L" seconds" << std::endl;
    }

    if (delay > 0)
    {
        std::wcout << L"Shutdown will execute in " << delay << L" seconds..." << std::endl;
        std::this_thread::sleep_for(std::chrono::seconds(delay));
    }

    std::wcout << L"Loading ntdll.dll..." << std::endl;
    HINSTANCE hNtDll = GetModuleHandle(L"ntdll.dll");
    if (hNtDll == NULL)
        return DisplayError(L"GetModuleHandle() failed");

    std::wcout << L"Retrieving NtShutdownSystem function..." << std::endl;
    PFN_NT_SHUTDOWN_SYSTEM pNtShutdownSystem =
        (PFN_NT_SHUTDOWN_SYSTEM)GetProcAddress(hNtDll, "NtShutdownSystem");
    if (pNtShutdownSystem == NULL)
        return DisplayError(L"GetProcAddress() failed");

    std::wcout << L"Adjusting privileges..." << std::endl;
    HANDLE hToken;
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken))
        return DisplayError(L"OpenProcessToken() failed");

    TOKEN_PRIVILEGES tkp;
    if (!LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &tkp.Privileges[0].Luid))
        return DisplayError(L"LookupPrivilegeValue() failed");

    tkp.PrivilegeCount = 1;
    tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    if (!AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, NULL, 0))
        return DisplayError(L"AdjustTokenPrivileges() failed");

    CloseHandle(hToken);

    std::wcout << L"Calling NtShutdownSystem..." << std::endl;
    LONG status = pNtShutdownSystem(action);
    if (status)
    {
        std::wcout << L"NtShutdownSystem() returned 0x" << std::hex << status << std::endl;
        return 1;
    }

    std::wcout << L"System " << (action == ShutdownReboot ? L"reboot" : L"shutdown") << L" initiated successfully." << std::endl;
    return 0;
}

int DisplayError(const std::wstring &message)
{
    DWORD error = GetLastError();
    std::wstring errorMsg = GetFormattedMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, error);
    std::wcout << message << L": " << errorMsg << std::endl;
    return 1;
}

std::wstring GetFormattedMessage(DWORD dwFlags, LPCWSTR pMsg, DWORD dwMsgId, ...)
{
    LPWSTR pBuffer = NULL;
    va_list args;
    va_start(args, dwMsgId);

    FormatMessage(dwFlags | FORMAT_MESSAGE_ALLOCATE_BUFFER,
                  pMsg, dwMsgId,
                  MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL),
                  (LPWSTR)&pBuffer, 0, &args);
    va_end(args);

    std::wstring result = pBuffer ? pBuffer : L"Unknown error";
    LocalFree(pBuffer);
    return result;
}