imports_loader.cpp

Go to the documentation of this file.

#include "peconv/imports_loader.h"
 
#include <iostream>
 
using namespace peconv;
 
class FillImportThunks : public ImportThunksCallback
{
public:
    FillImportThunks(BYTE* _modulePtr, size_t _moduleSize, t_function_resolver* func_resolver)
        : ImportThunksCallback(_modulePtr, _moduleSize), funcResolver(func_resolver)
    {
    }
 
    virtual bool processThunks(LPSTR lib_name, ULONG_PTR origFirstThunkPtr, ULONG_PTR firstThunkPtr)
    {
        if (this->is64b) {
#ifdef _WIN64 // loader is 64 bit, allow to load imports for 64-bit payload:
            IMAGE_THUNK_DATA64* desc = reinterpret_cast<IMAGE_THUNK_DATA64*>(origFirstThunkPtr);
            ULONGLONG* call_via = reinterpret_cast<ULONGLONG*>(firstThunkPtr);
            return processThunks_tpl<ULONGLONG, IMAGE_THUNK_DATA64>(lib_name, desc, call_via, IMAGE_ORDINAL_FLAG64);
#else
            std::cerr << "[!] Cannot fill imports into 64 bit PE via 32 bit loader!\n";
            return false;
#endif
        }
        else {
#ifndef _WIN64 // loader is 32 bit, allow to load imports for 32-bit payload:
            IMAGE_THUNK_DATA32* desc = reinterpret_cast<IMAGE_THUNK_DATA32*>(origFirstThunkPtr);
            DWORD* call_via = reinterpret_cast<DWORD*>(firstThunkPtr);
            return processThunks_tpl<DWORD, IMAGE_THUNK_DATA32>(lib_name, desc, call_via, IMAGE_ORDINAL_FLAG32);
#else
            std::cerr << "[!] Cannot fill imports into 32 bit PE via 64 bit loader!\n";
            return false;
#endif 
        }
    }
 
protected:
    template <typename T_FIELD, typename T_IMAGE_THUNK_DATA>
    bool processThunks_tpl(LPSTR lib_name, T_IMAGE_THUNK_DATA* desc, T_FIELD* call_via, T_FIELD ordinal_flag)
    {
        if (!this->funcResolver) {
            return false;
        }
 
        bool is_by_ord = (desc->u1.Ordinal & ordinal_flag) != 0;
 
        FARPROC hProc = nullptr;
        if (is_by_ord) {
            T_FIELD raw_ordinal = desc->u1.Ordinal & (~ordinal_flag);
#ifdef _DEBUG
            std::cout << "raw ordinal: " << std::hex << raw_ordinal << std::endl;
#endif
            hProc = funcResolver->resolve_func(lib_name, MAKEINTRESOURCEA(raw_ordinal));
 
        }
        else {
            PIMAGE_IMPORT_BY_NAME by_name = (PIMAGE_IMPORT_BY_NAME)((ULONGLONG)modulePtr + desc->u1.AddressOfData);
            LPSTR func_name = reinterpret_cast<LPSTR>(by_name->Name);
#ifdef _DEBUG
            std::cout << "name: " << func_name << std::endl;
#endif
            hProc = this->funcResolver->resolve_func(lib_name, func_name);
        }
        if (!hProc) {
#ifdef _DEBUG
            std::cerr << "Could not resolve the function!" << std::endl;
#endif
            return false;
        }
        (*call_via) = reinterpret_cast<T_FIELD>(hProc);
        return true;
    }
 
    //fields:
    t_function_resolver* funcResolver;
};
 
//---
 
class CollectThunksCallback : public peconv::ImportThunksCallback
{
public:
    CollectThunksCallback(BYTE* _vBuf, size_t _vBufSize, std::set<DWORD>& _fields)
        : ImportThunksCallback(_vBuf, _vBufSize), vBuf(_vBuf), vBufSize(_vBufSize),
        fields(_fields)
    {
    }
 
    virtual bool processThunks(LPSTR libName, ULONG_PTR origFirstThunkPtr, ULONG_PTR va)
    {
        if (va == NULL) {
            return false; // invalid thunk
        }
        const ULONGLONG module_base = reinterpret_cast<ULONGLONG>(this->vBuf);
        if (va < module_base) {
            return false; // not this module
        }
        if (va > module_base + this->vBufSize) {
            return false; // not this module
        }
        DWORD thunk_rva = MASK_TO_DWORD(va - module_base);
        fields.insert(thunk_rva);
        return true;
    }
 
    std::set<DWORD>& fields;
    BYTE* vBuf;
    size_t vBufSize;
};
 
//---
 
class CollectImportsCallback : public ImportThunksCallback
{
public:
    CollectImportsCallback(BYTE* _modulePtr, size_t _moduleSize, std::map<DWORD, ExportedFunc*> &_thunkToFunc)
        : ImportThunksCallback(_modulePtr, _moduleSize), thunkToFunc(_thunkToFunc)
    {
    }
 
    virtual bool processThunks(LPSTR lib_name, ULONG_PTR origFirstThunkPtr, ULONG_PTR firstThunkPtr)
    {
        if (this->is64b) {
            IMAGE_THUNK_DATA64* desc = reinterpret_cast<IMAGE_THUNK_DATA64*>(origFirstThunkPtr);
            ULONGLONG* call_via = reinterpret_cast<ULONGLONG*>(firstThunkPtr);
            return processThunks_tpl<ULONGLONG, IMAGE_THUNK_DATA64>(lib_name, desc, call_via, IMAGE_ORDINAL_FLAG64);
 
        }
        else {
 
            IMAGE_THUNK_DATA32* desc = reinterpret_cast<IMAGE_THUNK_DATA32*>(origFirstThunkPtr);
            DWORD* call_via = reinterpret_cast<DWORD*>(firstThunkPtr);
            return processThunks_tpl<DWORD, IMAGE_THUNK_DATA32>(lib_name, desc, call_via, IMAGE_ORDINAL_FLAG32);
        }
    }
 
protected:
    template <typename T_FIELD, typename T_IMAGE_THUNK_DATA>
    bool processThunks_tpl(LPSTR lib_name, T_IMAGE_THUNK_DATA* desc, T_FIELD* call_via, T_FIELD ordinal_flag)
    {
        if (call_via == nullptr) {
            return false;
        }
 
        const std::string short_name = peconv::get_dll_shortname(lib_name);
        const bool is_by_ord = (desc->u1.Ordinal & ordinal_flag) != 0;
        ExportedFunc *func = nullptr;
 
        if (is_by_ord) {
            T_FIELD raw_ordinal = desc->u1.Ordinal & (~ordinal_flag);
            func = new ExportedFunc(short_name, IMAGE_ORDINAL64(raw_ordinal));
        }
        else {
            PIMAGE_IMPORT_BY_NAME by_name = (PIMAGE_IMPORT_BY_NAME)((ULONGLONG)modulePtr + desc->u1.AddressOfData);
            LPSTR func_name = reinterpret_cast<LPSTR>(by_name->Name);
            WORD ordinal = by_name->Hint;
            func = new ExportedFunc(short_name, func_name, ordinal);
        }
        if (!func) {
            return false;
        }
        const DWORD rva = MASK_TO_DWORD((ULONG_PTR)call_via - (ULONG_PTR)modulePtr);
        thunkToFunc[rva] = func;
        return true;
    }
 
    std::map<DWORD, ExportedFunc*> &thunkToFunc;
};
 
 
//---
 
template <typename T_FIELD, typename T_IMAGE_THUNK_DATA>
bool process_imp_functions_tpl(BYTE* modulePtr, size_t module_size, LPSTR lib_name, DWORD call_via, DWORD thunk_addr, IN ImportThunksCallback *callback)
{
    bool is_ok = true;
 
    T_FIELD *thunks = (T_FIELD*)((ULONGLONG)modulePtr + thunk_addr);
    T_FIELD *callers = (T_FIELD*)((ULONGLONG)modulePtr + call_via);
 
    for (size_t index = 0; true; index++) {
        if (!validate_ptr(modulePtr, module_size, &callers[index], sizeof(T_FIELD))) {
            break;
        }
        if (!validate_ptr(modulePtr, module_size, &thunks[index], sizeof(T_FIELD))) {
            break;
        }
        if (callers[index] == 0) {
            //nothing to fill, probably the last record
            return true;
        }
        LPVOID thunk_ptr = &thunks[index];
        T_IMAGE_THUNK_DATA* desc = reinterpret_cast<T_IMAGE_THUNK_DATA*>(thunk_ptr);
        if (!validate_ptr(modulePtr, module_size, desc, sizeof(T_IMAGE_THUNK_DATA))) {
            break;
        }
        if (desc->u1.Function == NULL) {
            break;
        }
        T_FIELD ordinal_flag = (sizeof(T_FIELD) == sizeof(ULONGLONG)) ? IMAGE_ORDINAL_FLAG64 : IMAGE_ORDINAL_FLAG32;
        bool is_by_ord = (desc->u1.Ordinal & ordinal_flag) != 0;
        if (!is_by_ord) {
            PIMAGE_IMPORT_BY_NAME by_name = (PIMAGE_IMPORT_BY_NAME)((ULONGLONG)modulePtr + desc->u1.AddressOfData);
            if (!validate_ptr(modulePtr, module_size, by_name, sizeof(IMAGE_IMPORT_BY_NAME))) {
                break;
            }
        }
        //when the callback is called, all the pointers should be already verified
        if (!callback->processThunks(lib_name, (ULONG_PTR)&thunks[index], (ULONG_PTR)&callers[index])) {
            is_ok = false;
        }
    }
    return is_ok;
}
 
//Walk through the table of imported DLLs (starting from the given descriptor) and execute the callback each time when the new record was found
bool process_dlls(BYTE* modulePtr, size_t module_size, IMAGE_IMPORT_DESCRIPTOR *first_desc, IN ImportThunksCallback *callback)
{
    bool isAllFilled = true;
#ifdef _DEBUG_EX
    std::cout << "---IMP---" << std::endl;
#endif
    const bool is64 = is64bit((BYTE*)modulePtr);
    IMAGE_IMPORT_DESCRIPTOR* lib_desc = nullptr;
 
    for (size_t i = 0; true; i++) {
        lib_desc = &first_desc[i];
        if (!validate_ptr(modulePtr, module_size, lib_desc, sizeof(IMAGE_IMPORT_DESCRIPTOR))) {
            break;
        }
        if (lib_desc->OriginalFirstThunk == NULL && lib_desc->FirstThunk == NULL) {
            break;
        }
        LPSTR lib_name = (LPSTR)((ULONGLONG)modulePtr + lib_desc->Name);
        if (!peconv::is_valid_import_name(modulePtr, module_size, lib_name)) {
            //invalid name
            return false;
        }
        DWORD call_via = lib_desc->FirstThunk;
        DWORD thunk_addr = lib_desc->OriginalFirstThunk;
        if (thunk_addr == NULL) {
            thunk_addr = lib_desc->FirstThunk;
        }
#ifdef _DEBUG_EX
        std::cout << "Imported Lib: " << std::hex << lib_desc->FirstThunk << " : " << std::hex << lib_desc->OriginalFirstThunk << " : " << lib_desc->Name << std::endl;
#endif
        size_t all_solved = false;
        if (is64) {
            all_solved = process_imp_functions_tpl<ULONGLONG, IMAGE_THUNK_DATA64>(modulePtr, module_size, lib_name, call_via, thunk_addr, callback);
        }
        else {
            all_solved = process_imp_functions_tpl<DWORD, IMAGE_THUNK_DATA32>(modulePtr, module_size, lib_name, call_via, thunk_addr, callback);
        }
        if (!all_solved) {
            isAllFilled = false;
        }
    }
#ifdef _DEBUG_EX
    printf("---------\n");
#endif
    return isAllFilled;
}
 
bool peconv::process_import_table(IN BYTE* modulePtr, IN SIZE_T moduleSize, IN ImportThunksCallback *callback)
{
    if (moduleSize == 0) { //if not given, try to fetch
        moduleSize = peconv::get_image_size((const BYTE*)modulePtr);
    }
    if (moduleSize == 0) return false;
 
    IMAGE_DATA_DIRECTORY *importsDir = get_directory_entry((BYTE*)modulePtr, IMAGE_DIRECTORY_ENTRY_IMPORT);
    if (!importsDir) {
        return true; //no import table
    }
    const DWORD impAddr = importsDir->VirtualAddress;
    IMAGE_IMPORT_DESCRIPTOR *first_desc = (IMAGE_IMPORT_DESCRIPTOR*)(impAddr + (ULONG_PTR)modulePtr);
    if (!peconv::validate_ptr(modulePtr, moduleSize, first_desc, sizeof(IMAGE_IMPORT_DESCRIPTOR))) {
        return false;
    }
    return process_dlls(modulePtr, moduleSize, first_desc, callback);
}
 
bool peconv::load_imports(BYTE* modulePtr, t_function_resolver* func_resolver)
{
    size_t moduleSize = peconv::get_image_size((const BYTE*)modulePtr);
    if (moduleSize == 0) return false;
 
    bool is64 = is64bit((BYTE*)modulePtr);
    bool is_loader64 = false;
#ifdef _WIN64
    is_loader64 = true;
#endif
    if (is64 != is_loader64) {
        std::cerr << "[ERROR] Loader/Payload bitness mismatch.\n";
        return false;
    }
 
    default_func_resolver default_res;
    if (!func_resolver) {
        func_resolver = (t_function_resolver*)&default_res;
    }
 
    FillImportThunks callback(modulePtr, moduleSize, func_resolver);
    return peconv::process_import_table(modulePtr, moduleSize, &callback);
}
 
// A valid name must contain printable characters. Empty name is also acceptable (may have been erased)
bool peconv::is_valid_import_name(const PBYTE modulePtr, const size_t moduleSize, LPSTR lib_name)
{
    while (true) {
        if (!peconv::validate_ptr(modulePtr, moduleSize, lib_name, sizeof(char))) {
            return false;
        }
        char next_char = *lib_name;
        if (next_char == '\0') break;
 
        if (next_char <= 0x20 || next_char >= 0x7E) {
            return false;
        }
        lib_name++;
    }
    return true;
}
 
bool peconv::has_valid_import_table(const PBYTE modulePtr, size_t moduleSize)
{
    IMAGE_DATA_DIRECTORY *importsDir = get_directory_entry((BYTE*)modulePtr, IMAGE_DIRECTORY_ENTRY_IMPORT);
    if (importsDir == NULL) return false;
 
    const DWORD impAddr = importsDir->VirtualAddress;
 
    IMAGE_IMPORT_DESCRIPTOR* lib_desc = NULL;
    DWORD parsedSize = 0;
    size_t valid_records = 0;
 
    while (true) { //size of the import table doesn't matter
        lib_desc = (IMAGE_IMPORT_DESCRIPTOR*)(impAddr + parsedSize + (ULONG_PTR)modulePtr);
        if (!peconv::validate_ptr(modulePtr, moduleSize, lib_desc, sizeof(IMAGE_IMPORT_DESCRIPTOR))) {
            return false;
        }
        parsedSize += sizeof(IMAGE_IMPORT_DESCRIPTOR);
 
        if (lib_desc->OriginalFirstThunk == NULL && lib_desc->FirstThunk == NULL) {
            break;
        }
        LPSTR lib_name = (LPSTR)((ULONGLONG)modulePtr + lib_desc->Name);
        if (!is_valid_import_name(modulePtr, moduleSize, lib_name)) return false;
 
        DWORD call_via = lib_desc->FirstThunk;
        DWORD thunk_addr = lib_desc->OriginalFirstThunk;
        if (thunk_addr == NULL) thunk_addr = lib_desc->FirstThunk;
 
        DWORD *thunks = (DWORD*)((ULONGLONG)modulePtr + thunk_addr);
        if (!peconv::validate_ptr(modulePtr, moduleSize, thunks, sizeof(DWORD))) return false;
 
        DWORD *callers = (DWORD*)((ULONGLONG)modulePtr + call_via);
        if (!peconv::validate_ptr(modulePtr, moduleSize, callers, sizeof(DWORD))) return false;
 
        valid_records++;
    }
 
    return (valid_records > 0);
}
 
bool peconv::collect_thunks(IN BYTE* modulePtr, IN SIZE_T moduleSize, OUT std::set<DWORD>& thunk_rvas)
{
    CollectThunksCallback collector(modulePtr, moduleSize, thunk_rvas);
    return peconv::process_import_table(modulePtr, moduleSize, &collector);
}
 
bool peconv::collect_imports(IN BYTE* modulePtr, IN SIZE_T moduleSize, OUT ImportsCollection &collection)
{
    CollectImportsCallback collector(modulePtr, moduleSize, collection.thunkToFunc);
    return peconv::process_import_table(modulePtr, moduleSize, &collector);
}