exports_mapper.h

Go to the documentation of this file.

 
#pragma once
 
#include <windows.h>
 
#include <string>
#include <map>
#include <set>
#include <sstream>
 
#include "pe_hdrs_helper.h"
#include "pe_raw_to_virtual.h"
#include "peconv/exported_func.h"
#include "peconv/file_util.h"
 
namespace peconv {
 
    struct DllInfo {
        DllInfo()
            : moduleBase(0), moduelSize(0), is64b(false)
        {
        }
 
        DllInfo(ULONGLONG _moduleBase, size_t _moduelSize, bool _is64b, std::string _moduleName)
        {
            moduleBase = _moduleBase;
            moduelSize = _moduelSize;
            moduleName = _moduleName;
            is64b = _is64b;
            shortName = get_dll_shortname(moduleName);
        }
 
        DllInfo(const DllInfo &other)
        {
            moduleBase = other.moduleBase;
            moduelSize = other.moduelSize;
            moduleName = other.moduleName;
            shortName = other.shortName;
            is64b = other.is64b;
        }
 
        bool operator<(const DllInfo &other) const
        {
            return this->moduleBase < other.moduleBase;
        }
 
    protected:
        ULONGLONG moduleBase;
        size_t moduelSize;
        std::string moduleName;
        std::string shortName;
        bool is64b;
 
        friend class ExportsMapper;
    };
 
    class ExportsMapper {
 
    public:
 
        size_t add_to_lookup(std::string moduleName, HMODULE modulePtr, size_t moduleSize, ULONGLONG moduleBase);
 
        size_t add_to_lookup(std::string moduleName, HMODULE modulePtr, ULONGLONG moduleBase)
        {
            return add_to_lookup(moduleName, modulePtr, 0, moduleBase);
        }
 
        size_t add_to_lookup(std::string moduleName, HMODULE modulePtr) 
        {
            return add_to_lookup(moduleName, modulePtr, reinterpret_cast<ULONGLONG>(modulePtr));
        }
 
        const std::set<ExportedFunc>* find_exports_by_va(ULONGLONG va) const
        {
            std::map<ULONGLONG, std::set<ExportedFunc>>::const_iterator itr = va_to_func.find(va);
            if (itr != va_to_func.end()) {
                const std::set<ExportedFunc> &fSet = itr->second;
                return &fSet;
            }
            return NULL;
        }
 
        ULONGLONG find_dll_base_by_func_va(ULONGLONG func_rva) const
        {
            // the first element that is greater than the start address
            std::map<ULONGLONG, DllInfo>::const_iterator firstGreater = dll_base_to_info.upper_bound(func_rva);
 
            std::map<ULONGLONG, DllInfo>::const_iterator itr;
            for (itr = dll_base_to_info.begin(); itr != firstGreater; ++itr) {
                const DllInfo& module = itr->second;
 
                if (func_rva >= module.moduleBase && func_rva <= (module.moduleBase + module.moduelSize)) {
                    // Address found in module:
                    return module.moduleBase;
                }
            }
            return 0;
        }
 
        std::string get_dll_path(ULONGLONG base) const
        {
            std::map<ULONGLONG, DllInfo>::const_iterator found =  this->dll_base_to_info.find(base);
            if (found == this->dll_base_to_info.end()) { // no DLL found at this base
                return "";
            }
            const DllInfo& info = found->second;
            return info.moduleName;
        }
 
        std::string get_dll_path(std::string short_name) const;
 
        size_t get_dll_paths(IN std::string short_name, OUT std::set<std::string>& paths) const;
 
        std::string get_dll_fullname(std::string short_name) const
        {
            std::string dll_path = get_dll_path(short_name);
            if (dll_path.length() == 0) return "";
 
            return get_file_name(dll_path);
        }
 
        const ExportedFunc* find_export_by_va(ULONGLONG va) const
        {
            const std::set<ExportedFunc>* exp_set = find_exports_by_va(va);
            if (exp_set == NULL) return NULL;
 
            std::set<ExportedFunc>::iterator fItr = exp_set->begin();
            const ExportedFunc* func = &(*fItr);
            return func;
        }
 
        void print_va_to_func(std::stringstream &stream) const;
        void print_func_to_va(std::stringstream &stream) const;
        
 
    private:
        enum ADD_FUNC_RES { RES_INVALID = 0, RES_MAPPED = 1, RES_FORWARDED = 2 };
        ADD_FUNC_RES add_function_to_lookup(HMODULE modulePtr, ULONGLONG moduleBase, size_t moduleSize, ExportedFunc &currFunc, DWORD callRVA);
 
        bool add_forwarded(ExportedFunc &currFunc, DWORD callRVA, PBYTE modulePtr, size_t moduleSize);
        bool add_to_maps(ULONGLONG va, ExportedFunc &currFunc);
 
        size_t resolve_forwarders(const ULONGLONG va, ExportedFunc &currFunc);
        size_t make_ord_lookup_tables(PVOID modulePtr, size_t moduleSize, std::map<PDWORD, DWORD> &va_to_ord);
 
    protected:
        void associateVaAndFunc(ULONGLONG va, const ExportedFunc& func)
        {
            va_to_func[va].insert(func);
            func_to_va[func] = va;
        }
 
        std::map<ULONGLONG, std::set<ExportedFunc>> va_to_func;
 
        std::map<ExportedFunc, std::set<ExportedFunc>> forwarders_lookup;
 
        std::map<ExportedFunc, ULONGLONG> func_to_va;
        
        std::map<std::string, std::set<ULONGLONG>> dll_shortname_to_base;
 
        std::map<ULONGLONG, DllInfo> dll_base_to_info;
    };
 
}; //namespace peconv