DriverControl.h

#pragma once

#include "../Include/Winheaders.h"
#include "../Include/Types.h"
#include "../../BlackBoneDrv/BlackBoneDef.h"

#include <string>
#include <map>

namespace blackbone
{

// [Original ptr, size] <--> [Mapped ptr]
typedef std::map < std::pair<ptr_t, uint32_t>, ptr_t > mapMemoryMap;

struct MapMemoryResult
{
    ptr_t hostSharedPage;       // Shared page address in current process
    ptr_t targetSharedPage;     // Shared page address in target process
    HANDLE targetPipe;          // Hook pipe handle in the target process
            
    mapMemoryMap regions;       // Mapped regions info
};

struct MapMemoryRegionResult
{
    ptr_t originalPtr;          // Address of region in the target process
    ptr_t newPtr;               // Address of mapped region in the current process
    ptr_t removedPtr;           // Address of region unmapped because of address conflict 
    uint32_t size;              // Size of mapped region
    uint32_t removedSize;       // Size of unmapped region
};


class DriverControl
{
public:
    BLACKBONE_API DriverControl();
    BLACKBONE_API ~DriverControl();

    BLACKBONE_API static DriverControl& Instance();

    BLACKBONE_API NTSTATUS EnsureLoaded( const std::wstring& path = L"" );

    BLACKBONE_API NTSTATUS Reload( std::wstring path = L"" );

    BLACKBONE_API NTSTATUS Unload();

    BLACKBONE_API NTSTATUS DisableDEP( DWORD pid );

    BLACKBONE_API NTSTATUS ProtectProcess( DWORD pid, bool enable );

    BLACKBONE_API NTSTATUS PromoteHandle( DWORD pid, HANDLE handle, DWORD access );

    BLACKBONE_API NTSTATUS AllocateMem( DWORD pid, ptr_t& base, ptr_t& size, DWORD type, DWORD protection, bool physical = false );

    BLACKBONE_API NTSTATUS FreeMem( DWORD pid, ptr_t base, ptr_t size, DWORD type );

    BLACKBONE_API NTSTATUS ReadMem( DWORD pid, ptr_t base, ptr_t size, PVOID buffer );

    BLACKBONE_API NTSTATUS WriteMem( DWORD pid, ptr_t base, ptr_t size, PVOID buffer );

    BLACKBONE_API NTSTATUS ProtectMem( DWORD pid, ptr_t base, ptr_t size, DWORD protection );

    BLACKBONE_API NTSTATUS MapMemory( DWORD pid, const std::wstring& pipeName, bool mapSections, MapMemoryResult& result );

    BLACKBONE_API NTSTATUS MapMemoryRegion( DWORD pid, ptr_t base, uint32_t size, MapMemoryRegionResult& result );

    BLACKBONE_API NTSTATUS UnmapMemory( DWORD pid );

    BLACKBONE_API NTSTATUS UnmapMemoryRegion( DWORD pid, ptr_t base, uint32_t size );

    BLACKBONE_API NTSTATUS InjectDll(
        DWORD pid, 
        const std::wstring& path, 
        InjectType itype, 
        uint32_t initRVA = 0, 
        const std::wstring& initArg = L"", 
        bool wait = true
        );

    BLACKBONE_API NTSTATUS MMapDriver( const std::wstring& path );

    BLACKBONE_API NTSTATUS ConcealVAD( DWORD pid, ptr_t base, uint32_t size );

    BLACKBONE_API inline bool loaded() const { return _hDriver != INVALID_HANDLE_VALUE; }

private:
    DriverControl( const DriverControl& ) = delete;
    DriverControl& operator = (const DriverControl&) = delete;

    NTSTATUS LoadDriver( const std::wstring& svcName, const std::wstring& path );

    NTSTATUS UnloadDriver( const std::wstring& svcName );

    LSTATUS PrepareDriverRegEntry( const std::wstring& svcName, const std::wstring& path );

private:
    HANDLE _hDriver = INVALID_HANDLE_VALUE;
};

// Syntax sugar
inline DriverControl& Driver() { return DriverControl::Instance(); }

}