#include <Windows.h>

#include <iostream>
#include "ntdll_undoc.h"

PPEB get_default_peb()
{
#if defined(_WIN64)
    return (PPEB)__readgsqword(0x60);
#else
    return (PPEB)__readfsdword(0x30);
#endif
}

PPEB64 get_peb64(HANDLE hProcess, OUT PROCESS_BASIC_INFORMATION_WOW64 &pbi64)
{
    if (NtWow64QueryInformationProcess64 == nullptr) {
        return nullptr;
    }
    //reset structure:
    memset(&pbi64,0, sizeof(PROCESS_BASIC_INFORMATION_WOW64));
    
    ULONG outLength = 0;
    NTSTATUS status = NtWow64QueryInformationProcess64(
        hProcess,
        ProcessBasicInformation,
        &pbi64,
        sizeof(PROCESS_BASIC_INFORMATION_WOW64),
        &outLength
    );
    if (status != STATUS_SUCCESS) {
        return nullptr;
    }
    return (PPEB64) pbi64.PebBaseAddress;
}

int main()
{
    BOOL isWow64 = FALSE;
    IsWow64Process(GetCurrentProcess(), &isWow64);
    std::cout << "IsWow64" << " : " << isWow64 << std::endl;

    if (init_ntdll_func(isWow64) == false) {
        printf("Cannot load functions!\n");
        return -1;
    }

    PPEB myPeb = get_default_peb();
    std::cout << "PEB:  \t" ;
    std::cout << std::hex << myPeb << std::endl;

    PPEB64 pebWow64 = nullptr;
    if (isWow64) {
        PROCESS_BASIC_INFORMATION_WOW64 pbi64 = { 0 };
        pebWow64 = get_peb64(GetCurrentProcess(), pbi64);
        if (pebWow64 == nullptr) {
            std::cerr << "Fetching PEB64 failed!" << std::endl;
            return -1;
        }
        std::cout << "PEB64:\t" ;
        std::cout << std::hex << pebWow64 << std::endl;
    }
    std::cout << "ImageBaseAddress from PEB: \t" << std::hex << myPeb->ImageBaseAddress << std::endl;
    if (pebWow64 != nullptr) {
        std::cout << "ImageBaseAddress from PEB64: \t" << std::hex << pebWow64->ImageBaseAddress << std::endl;
    }
    system("pause");
    return 0;
}