#include <windows.h>
#include <winternl.h>
#include <tlhelp32.h>
#include <stdio.h>

#ifdef _MSC_VER
#pragma comment (lib, "ntdll.lib")
#endif // _MSC_VER

#define ZwCurrentProcess() ((HANDLE)(LONG_PTR)-1)

NTSTATUS ZwOpenProcessToken(
    HANDLE ProcessHandle,
    ACCESS_MASK DesiredAccess,
    PHANDLE TokenHandle);

NTSTATUS ZwAdjustPrivilegesToken(
    HANDLE TokenHandle,
    BOOLEAN DisableAllPrivileges,
    PTOKEN_PRIVILEGES NewState,
    ULONG BufferLength,
    PTOKEN_PRIVILEGES PreviousState,
    PULONG ReturnLength);

NTSTATUS ZwOpenProcess(
    PHANDLE ProcessHandle,
    ACCESS_MASK DesiredAccess,
    POBJECT_ATTRIBUTES ObjectAttributes,
    CLIENT_ID* ClientId);

NTSTATUS ZwSuspendProcess(
    HANDLE ProcessHandle);

NTSTATUS ZwResumeProcess(
    HANDLE ProcessHandle);

NTSTATUS ZwTerminateProcess(
    HANDLE ProcessHandle,
    NTSTATUS ExitStatus);

NTSTATUS ZwClose(
    HANDLE Handle);

NTSTATUS EnablePrivilege(PWSTR lpName)
{
    HANDLE TokenHandle = NULL;
    LUID Luid;
    TOKEN_PRIVILEGES NewState;

    ZwOpenProcessToken(
        ZwCurrentProcess(),
        TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
        &TokenHandle);
    LookupPrivilegeValueW(NULL, lpName, &Luid);

    NewState.PrivilegeCount = 1;
    NewState.Privileges[0].Luid = Luid;
    NewState.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    NTSTATUS Status = ZwAdjustPrivilegesToken(
        TokenHandle, FALSE, &NewState,
        sizeof(TOKEN_PRIVILEGES), NULL, NULL);

    return Status;
}

void GetPid(PWSTR ProcessName, PULONG ProcessId)
{
    PROCESSENTRY32W entry;
    entry.dwSize = sizeof(PROCESSENTRY32W);

    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (Process32FirstW(snapshot, &entry))
    {
        while (Process32NextW(snapshot, &entry))
        {
            if (_wcsicmp(entry.szExeFile, ProcessName) == 0)
            {
                *ProcessId = entry.th32ProcessID;
            }
        }
    }
    ZwClose(snapshot);
}


int wmain(int wargc, wchar_t* wargv[])
{

    if (wargc < 2)
    {
        wprintf(L"No option provided\n");
        return 0;
    }

    ULONG ProcessId;
    HANDLE hProcess = NULL;
    OBJECT_ATTRIBUTES ObjectAttributes = { 0 };
    ObjectAttributes.Length = sizeof(OBJECT_ATTRIBUTES);
    CLIENT_ID ClientId = { 0 };

    NTSTATUS Status = EnablePrivilege(L"SeDebugPrivilege");

    if (!wcscmp(L"-s", wargv[1]))
    {
        GetPid(L"winlogon.exe", &ProcessId);
        ClientId.UniqueProcess = (HANDLE)(ULONG64)ProcessId;

        Status = ZwOpenProcess(&hProcess, PROCESS_ALL_ACCESS, &ObjectAttributes, &ClientId);
        Status = ZwSuspendProcess(hProcess);
        wprintf(L"winlogon: PID: %ld Handle: %p Status: %lu\n",
            ProcessId, hProcess, Status);
        ZwClose(hProcess);

        GetPid(L"dwm.exe", &ProcessId);
        ClientId.UniqueProcess = (HANDLE)(ULONG64)ProcessId;

        Status = ZwOpenProcess(&hProcess, PROCESS_ALL_ACCESS, &ObjectAttributes, &ClientId);
        Status = ZwTerminateProcess(hProcess, 0);
        wprintf(L"dwm: PID: %ld Handle: %p Status: %lu\n",
            ProcessId, hProcess, Status);
        ZwClose(hProcess);

        return 0;
    }
    else if (!wcscmp(L"-r", wargv[1]))
    {
        GetPid(L"winlogon.exe", &ProcessId);
        ClientId.UniqueProcess = (HANDLE)(ULONG64)ProcessId;

        Status = ZwOpenProcess(&hProcess, PROCESS_ALL_ACCESS, &ObjectAttributes, &ClientId);
        Status = ZwResumeProcess(hProcess);
        wprintf(L"winlogon: PID: %ld Handle: %p Status: %lu\n",
            ProcessId, hProcess, Status);
        ZwClose(hProcess);

        return 0;
    }
    else
    {
        wprintf(L"Only -r and -s option available\n");
        return 0;
    }
}