MMap.h

#pragma once

#include "../Config.h"
#include "../Include/Winheaders.h"
#include "../Include/Macro.h"
#include "../PE/PEImage.h"

#include "../Process/MemBlock.h"
#include "MExcept.h"

#include <map>
#include <stack>

namespace blackbone
{

// Loader flags
enum eLoadFlags
{
    NoFlags         = 0x00,     // No flags
    ManualImports   = 0x01,     // Manually map import libraries
    CreateLdrRef    = 0x02,     // Create module references for native loader
    WipeHeader      = 0x04,     // Wipe image PE headers
    HideVAD         = 0x10,     // Make image appear as PAGE_NOACESS region
    MapInHighMem    = 0x20,     // Try to map image in address space beyond 4GB limit
    RebaseProcess   = 0x40,     // If target image is an .exe file, process base address will be replaced with mapped module value

    NoExceptions    = 0x01000,   // Do not create custom exception handler
    PartialExcept   = 0x02000,   // Only create Inverted function table, without VEH
    NoDelayLoad     = 0x04000,   // Do not resolve delay import
    NoSxS           = 0x08000,   // Do not apply SxS activation context
    NoTLS           = 0x10000,   // Skip TLS initialization and don't execute TLS callbacks
};

ENUM_OPS( eLoadFlags )

// Native loader flags callback
typedef enum LdrRefFlags( *LdrCallback )(void* context, const ModuleData& modInfo);


struct ImageContext
{
    typedef std::vector<ptr_t> vecPtr;

    pe::PEImage    peImage;                 // PE image data
    MemBlock       imgMem;                  // Target image memory region
    std::wstring   FilePath;                // path to image being mapped
    std::wstring   FileName;                // File name string
    vecPtr         tlsCallbacks;            // TLS callback routines
    ptr_t          pExpTableAddr = 0;       // Exception table address (amd64 only)
    ptr_t          EntryPoint = 0;          // Target image entry point
    eLoadFlags     flags = NoFlags;         // Image loader flags
    bool           initialized = false;     // Image entry point was called
};

typedef std::vector<std::unique_ptr<ImageContext>> vecImageCtx;

class MMap : public MExcept
{
        
public:
    BLACKBONE_API MMap( class Process& proc );
    BLACKBONE_API ~MMap( void );

    BLACKBONE_API const ModuleData* MapImage(
        const std::wstring& path,
        eLoadFlags flags = NoFlags,
        LdrCallback ldrCallback = nullptr,
        void* ldrContext = nullptr
        );

    BLACKBONE_API const ModuleData* MapImage(
        void* buffer, size_t size,
        bool asImage = false,
        eLoadFlags flags = NoFlags,
        LdrCallback ldrCallback = nullptr,
        void* ldrContext = nullptr
        );

    BLACKBONE_API bool UnmapAllModules();

    BLACKBONE_API void Cleanup();

    BLACKBONE_API inline void reset() { _images.clear(); _pAContext.Reset(); _usedBlocks.clear(); }

private:
    const ModuleData* MapImage(
        const std::wstring& path,
        void* buffer, size_t size,
        bool asImage = false,
        eLoadFlags flags = NoFlags,
        LdrCallback ldrCallback = nullptr,
        void* ldrContext = nullptr
        );

    const ModuleData* FindOrMapModule(
        const std::wstring& path,
        void* buffer, size_t size, bool asImage,
        eLoadFlags flags = NoFlags
        );

    bool RunModuleInitializers( ImageContext* pImage, DWORD dwReason );

    bool CopyImage( ImageContext* pImage );

    bool ProtectImageMemory( ImageContext* pImage );

    bool RelocateImage( ImageContext* pImage );

    bool ResolveImport( ImageContext* pImage, bool useDelayed = false );

    bool InitStaticTLS( ImageContext* pImage );

    bool EnableExceptions( ImageContext* pImage );

    bool DisableExceptions( ImageContext* pImage );

    bool CreateActx( const std::wstring& path, int id = 2, bool asImage = true );

    bool InitializeCookie( ImageContext* pImage );

    NTSTATUS ConcealVad( const MemBlock& imageMem );

    NTSTATUS AllocateInHighMem( MemBlock& imageMem, size_t size );

    const ModuleData* FindOrMapDependency( ImageContext* pImage, std::wstring& path );

    module_t MapPureManaged();

    DWORD GetSectionProt( DWORD characteristics );

private:
    vecImageCtx     _images;                // Mapped images
    class Process&  _process;               // Target process manager
    MemBlock        _pAContext;             // SxS activation context memory address
    LdrCallback     _ldrCallback = nullptr; // Loader callback for adding image into loader lists
    void*           _ldrContext = nullptr;  // user context for _ldrCallback       

    std::vector<std::pair<ptr_t, size_t>> _usedBlocks;   // Used memory blocks 
};

}