peb_lookup.cpp

Go to the documentation of this file.

#include "ntddk.h"
#include "peconv/logger.h"
#include <peconv/peb_lookup.h>
 
class SectionLocker {
public:
    SectionLocker(RTL_CRITICAL_SECTION &_section)
        : section(_section)
    {
        RtlEnterCriticalSection(&section);
    }
 
    ~SectionLocker()
    {
        RtlLeaveCriticalSection(&section);
    }
 
protected:
    RTL_CRITICAL_SECTION &section;
};
 
//here we don't want to use any functions imported form extenal modules
 
typedef struct _LDR_MODULE { 
    LIST_ENTRY  InLoadOrderModuleList;//   +0x00 
    LIST_ENTRY  InMemoryOrderModuleList;// +0x08   
    LIST_ENTRY  InInitializationOrderModuleList;// +0x10 
    void*   BaseAddress; // +0x18 
    void*   EntryPoint;  // +0x1c 
    ULONG   SizeOfImage; 
    UNICODE_STRING FullDllName; 
    UNICODE_STRING BaseDllName; 
    ULONG   Flags; 
    SHORT   LoadCount; 
    SHORT   TlsIndex; 
    HANDLE  SectionHandle; 
    ULONG   CheckSum; 
    ULONG   TimeDateStamp; 
} LDR_MODULE, *PLDR_MODULE;
 
inline PPEB get_peb()
{
#if defined(_M_AMD64)
    return (PPEB)__readgsqword(0x60);
#elif defined(_M_ARM64)
    PPEB peb = (PPEB)(*(__getReg(18) + 0x60));
    LOG_DEBUG("ARM64 TEB: %p PEB: %p.", (void*)__getReg(18), peb);
    return peb;
#else
    return (PPEB)__readfsdword(0x30);
/*
//alternative way to fetch it:
    LPVOID PEB = NULL;
    __asm {
        mov eax, fs:[30h]
        mov PEB, eax
    };
    return (PPEB)PEB;
 
    or:
    LPVOID PEB = RtlGetCurrentPeb();
*/
#endif
}
 
inline WCHAR to_lowercase(WCHAR c1)
{
    if (c1 <= L'Z' && c1 >= L'A') {
        c1 = (c1 - L'A') + L'a';
    }
    return c1;
}
 
bool is_wanted_module(LPCWSTR curr_name, LPCWSTR wanted_name)
{
    if (wanted_name == NULL || curr_name == NULL) return false;
 
    LPCWSTR curr_end_ptr = curr_name;
    while (*curr_end_ptr != L'\0') {
        curr_end_ptr++;
    }
    if (curr_end_ptr == curr_name) return false;
 
    LPCWSTR wanted_end_ptr = wanted_name;
    while (*wanted_end_ptr != L'\0') {
        wanted_end_ptr++;
    }
    if (wanted_end_ptr == wanted_name) return false;
 
    while ((curr_end_ptr != curr_name) && (wanted_end_ptr != wanted_name)) {
 
        if (to_lowercase(*wanted_end_ptr) != to_lowercase(*curr_end_ptr)) {
            return false;
        }
        wanted_end_ptr--;
        curr_end_ptr--;
    }
    return true;
}
 
HMODULE peconv::get_module_via_peb(IN OPTIONAL LPCWSTR module_name)
{
    PPEB peb = get_peb();
    if (!peb) {
        return NULL;
    }
    SectionLocker locker(*peb->LoaderLock);
    LIST_ENTRY head = peb->Ldr->InLoadOrderModuleList;
 
    const PLDR_MODULE first_module = *((PLDR_MODULE *)(&head));
    PLDR_MODULE curr_module = first_module;
    if (!module_name) {
        return (HMODULE)(curr_module->BaseAddress);
    }
 
    // it is a cyclic list, so if the next record links to the initial one, it means we went throught the full loop
    do {
        // this should also work as a terminator, because the BaseAddress of the last module in the cycle is NULL
        if (curr_module == NULL || curr_module->BaseAddress == NULL) {
            break;
        }
        if (is_wanted_module(curr_module->BaseDllName.Buffer, module_name)) {
            return (HMODULE)(curr_module->BaseAddress);
        }
        curr_module = (PLDR_MODULE)curr_module->InLoadOrderModuleList.Flink;
 
    } while (curr_module != first_module);
 
    return NULL;
}
 
size_t peconv::get_module_size_via_peb(IN OPTIONAL HMODULE hModule)
{
    PPEB peb = get_peb();
    if (!peb) {
        return 0;
    }
    SectionLocker locker(*peb->LoaderLock);
    LIST_ENTRY head = peb->Ldr->InLoadOrderModuleList;
 
    const PLDR_MODULE first_module = *((PLDR_MODULE *)(&head));
    PLDR_MODULE curr_module = first_module;
    if (!hModule) {
        return (size_t)(curr_module->SizeOfImage);
    }
 
    // it is a cyclic list, so if the next record links to the initial one, it means we went throught the full loop
    do {
        // this should also work as a terminator, because the BaseAddress of the last module in the cycle is NULL
        if (curr_module == NULL || curr_module->BaseAddress == NULL) {
            break;
        }
        if (hModule == (HMODULE)(curr_module->BaseAddress)) {
            return (size_t)(curr_module->SizeOfImage);
        }
        curr_module = (PLDR_MODULE)curr_module->InLoadOrderModuleList.Flink;
 
    } while (curr_module != first_module);
 
    return 0;
}
 
bool peconv::set_main_module_in_peb(HMODULE module_ptr)
{
    PPEB peb = get_peb();
    if (peb == NULL) {
        return false;
    }
    SectionLocker locker(*peb->FastPebLock);
    peb->ImageBaseAddress = module_ptr;
    return true;
}
 
HMODULE peconv::get_main_module_via_peb()
{
    PPEB peb = get_peb();
    if (peb == NULL) {
        return NULL;
    }
    SectionLocker locker(*peb->FastPebLock);
    return (HMODULE) peb->ImageBaseAddress;
}