#ifndef _VIRTUAL_MEMORY_TOOLS_HEADER
#define _VIRTUAL_MEMORY_TOOLS_HEADER

/*                                       ABOUT THIS LIBRARY                                            */
// If you get that the function you're using is deprecated, Add "Ex" at the end of the function name
// If you want to use the deprecated function anyway just define DEPRECATED_ALLOWED before the inclusion
// of this header file.
// If you want the debug mode with the run time and error messager just set the compiler to debug mode.

#pragma region DEFINES_AND_INCLUDES

#ifndef DEPRECATED_ALLOWED
   #define DEPRECATED_VM_METHOD __declspec(deprecated("Deprecated function, use: \"FuncName\" + \"Ex\" instead"))
#else
   #define DEPRECATED_VM_METHOD
#endif

#undef UNICODE

#include <Windows.h>
#include <TlHelp32.h>
#include <iostream>
#include <fstream>

#ifdef _DEBUG
   #undef NDEBUG
#else
   #define NDEBUG
   #define NODEBUG
#endif

#include <cassert>

//Flags relocation for manual mapping
#define RELOC_FLAG32(RelInfo)((RelInfo >> 0x0C) == IMAGE_REL_BASED_HIGHLOW)
#define RELOC_FLAG64(RelInfo)((RelInfo >> 0x0C) == IMAGE_REL_BASED_DIR64)

#ifdef _WIN64
   #define RELOC_FLAG RELOC_FLAG64
#else
   #define RELOC_FLAG RELOC_FLAG32
#endif

#pragma endregion

#pragma region OBJECTS_AND_FUCNTION_POINTERS

//Kernel32.dll function pointers
using pfLoadLibraryA   = HINSTANCE(WINAPI*)(const char* lpLibName);
using pfGetProcAddress = UINT_PTR(WINAPI*)(HINSTANCE hModule, const char* lpProcName);
using pfDllEntryPoint  = BOOL(WINAPI*)(LPVOID hDll, DWORD dwReason, LPVOID pReserved);

//Struct used for dll manual mapped dll
struct MANUAL_MAPPING_DATA
{
    pfLoadLibraryA    pLoadLibraryA;
    pfGetProcAddress  pGetProcAddress;
    HINSTANCE         hMod;
};


//Struct used to get Base address and size of a process in memory
#ifndef _SMODULE_DECLARED_
typedef struct SMODULE
{
    DWORD dwBase = 0;
    DWORD dwSize = 0;
} *PSMODULE, *LPSMODULE;
#endif


#pragma endregion

//Virtual memory managment class
typedef class VirtualMemory
{
public:

    //Base constructor
    VirtualMemory(_Null_ void)
    {
        this->StartEngine();
    }

    //Attach process by name constructor
    VirtualMemory(_In_ LPCWSTR pName, _In_opt_ DWORD rights = PROCESS_ALL_ACCESS, _In_opt_ DWORD SnapshotFlag = TH32CS_SNAPPROCESS)
    {
        this->StartEngine();
        this->AttachByName(pName, SnapshotFlag, rights);
    }

    //Attach process by pid constructor
    VirtualMemory(_In_ UINT pid, _In_opt_ DWORD rights = PROCESS_ALL_ACCESS)
    {
        this->StartEngine();
        this->AttachByPID(pid, rights);
    }

    //Destructor
    ~VirtualMemory(_Null_ void)
    {
        if (this->_handle != INVALID_HANDLE_VALUE) CloseHandle(this->_handle);
    }

    //Attach to a process by process name
    bool AttachByName(_In_ LPCWSTR processName, _In_opt_ DWORD rights = PROCESS_ALL_ACCESS, _In_opt_ DWORD SnapshotFlag = TH32CS_SNAPPROCESS)
    {
        HANDLE snapHandle = CreateToolhelp32Snapshot(SnapshotFlag, NULL);

        if (snapHandle == INVALID_HANDLE_VALUE)
        {
            this->_handle = NULL;
            return false;
        }

        if (Process32FirstW(snapHandle, &_entry))
        {
            do
            {
                if (!wcsicmp(_entry.szExeFile, processName))
                {
                    this->_pid = _entry.th32ProcessID;
                    CloseHandle(snapHandle);
                    this->_handle = OpenProcess(rights, false, this->_pid);
                    return true;
                }
            } while (Process32NextW(snapHandle, &_entry));
        }

        this->_handle = NULL;
        CloseHandle(snapHandle);
        return false;
    }

    //Attach to a process by existing pid
    bool AttachByPID(_In_ UINT pid, _In_opt_ DWORD rights = PROCESS_ALL_ACCESS)
    {
        this->_handle = OpenProcess(rights, false, pid);
        if (this->_handle == INVALID_HANDLE_VALUE) return false;
        return true;
    }

    //Get module informations
    LPSMODULE GetModule(_In_ LPCWSTR modulename, _In_opt_ DWORD ModuleFlag = TH32CS_SNAPMODULE)
    {
        HANDLE module = CreateToolhelp32Snapshot(ModuleFlag, _pid);

        if (module == INVALID_HANDLE_VALUE)return nullptr;

        if (Module32FirstW(module, &_mEntry))
        {
            do
            {
                if (!wcsicmp(_mEntry.szModule, modulename))
                {
                    CloseHandle(module);

                    LPSMODULE mod = new SMODULE;
                    mod->dwBase = (DWORD)_mEntry.hModule;
                    mod->dwSize = (DWORD)_mEntry.modBaseSize;

                    return mod;
                }
            } while (Module32NextW(module, &_mEntry));
        }

        CloseHandle(module);
        return nullptr;

    }

    //Find pattern in a range of bytes
    DWORD FindPattern(_In_ DWORD start, _In_ DWORD size, _In_ LPCSTR sig, _In_ LPCSTR mask)
    {
        BYTE* data = new BYTE[size];
        SIZE_T bytesread;

        if (!ReadProcessMemory(_handle, (LPVOID)start, data, size, &bytesread)) return NULL;

        for (DWORD offset = 0; offset < size; offset++)
        {
            if (DataCompare((CONST BYTE*)(data + offset), (CONST BYTE*) sig, mask))
            {
                delete[] data;
                return start + offset;
            }
        }

        delete[] data;
        return NULL;
    }

    //Read from memory
    template<typename TypeToRead = DWORD>
    TypeToRead Read(_In_ DWORD address, _Out_opt_ SIZE_T* BytesRead = NULL)
    {
        TypeToRead _Read;
        ReadProcessMemory(_handle, (LPVOID)address, &_Read, sizeof(TypeToRead), BytesRead);
        return _Read;
    }

    //Write in memory
    template<typename TypeToWrite = DWORD>
    void Write(_In_ DWORD address, _In_ TypeToWrite value, _Out_opt_ SIZE_T* BytesWrite = NULL)
    {
        WriteProcessMemory(_handle, (LPVOID)address, &value, sizeof(TypeToWrite), BytesWrite);
    }

    //Inject dll into a process
    bool InjectDll(_In_ LPCSTR dllPath)
    {
        LPVOID LoadLibAddr = (LPVOID)GetProcAddress(GetModuleHandleA("kernel32.dll"), "LoadLibraryA");
        LPVOID AllocatedMemory = VirtualAllocEx(_handle, NULL, strlen(dllPath), MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

        if (!WriteProcessMemory(_handle, AllocatedMemory, dllPath, strlen(dllPath), NULL)) return false;

        HANDLE thread = CreateRemoteThread(_handle, NULL, NULL, (LPTHREAD_START_ROUTINE)LoadLibAddr, AllocatedMemory, 0, NULL);

        WaitForSingleObject(thread, INFINITE);
        VirtualFreeEx(_handle, AllocatedMemory, strlen(dllPath), MEM_RELEASE);

        CloseHandle(thread);

        return true;

    }

    bool DEPRECATED_VM_METHOD InjectDllEx(_In_ LPCSTR szDllPath)
    {
        BYTE* pSrcData = nullptr;
        PIMAGE_NT_HEADERS pOldNtHeader = nullptr;
        PIMAGE_OPTIONAL_HEADER pOldOptHeader = nullptr;
        PIMAGE_FILE_HEADER pOldFileHeader = nullptr;
        BYTE* pTargetBase = nullptr;

        if (!GetFileAttributesA(szDllPath)) return false;

        std::ifstream File(szDllPath, std::ios::binary | std::ios::ate);
        if (File.fail())
        {
            File.close();
            File.clear();
            return false;
        }

        auto FileSize = File.tellg();
        if (FileSize < 0x1000)
        {
            File.close();
            File.clear();
            return false;
        }

        pSrcData = new BYTE[static_cast<UINT_PTR>(FileSize)];
        if (!pSrcData)
        {
            File.close();
            File.clear();
            return false;
        }

        File.seekg(0, std::ios::beg);
        File.read(reinterpret_cast<char*>(pSrcData), FileSize);
        File.close();
        File.clear();

        if (reinterpret_cast<PIMAGE_DOS_HEADER>(pSrcData)->e_magic != 0x54AD)
        {
            delete[] pSrcData;
            return false;
        }

        pOldNtHeader = reinterpret_cast<PIMAGE_NT_HEADERS>(pSrcData + reinterpret_cast<PIMAGE_DOS_HEADER>(pSrcData)->e_lfanew);
        pOldOptHeader = &pOldNtHeader->OptionalHeader;
        pOldFileHeader = &pOldNtHeader->FileHeader;

        #ifdef _WIN64
        if (pOldFileHeader->Machine != IMAGE_FILE_MACHINE_AMD64)
        {
            delete[] pSrcData;

            return false;
        }
        #else
        if (pOldFileHeader->Machine != IMAGE_FILE_MACHINE_I386)
        {
            delete[] pSrcData;

            return false;
        }
        #endif

        pTargetBase = reinterpret_cast<PBYTE>(VirtualAllocEx(_handle, reinterpret_cast<LPVOID>(pOldOptHeader->ImageBase), pOldOptHeader->SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE));
        if (!pTargetBase)
        {
            pTargetBase = reinterpret_cast<PBYTE>(VirtualAllocEx(_handle, nullptr, pOldOptHeader->SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE));
            if (!pTargetBase)
            {
                delete[] pSrcData;

                return false;
            }
        }

        MANUAL_MAPPING_DATA data{ 0 };
        data.pLoadLibraryA = LoadLibraryA;
        data.pGetProcAddress = reinterpret_cast<pfGetProcAddress>(GetProcAddress);

        auto* pSectionHeader = IMAGE_FIRST_SECTION(pOldNtHeader);
        for (UINT i = 0; i != pOldFileHeader->NumberOfSections; ++i, ++pSectionHeader)
        {
            if (pSectionHeader->SizeOfRawData)
            {
                if (!WriteProcessMemory(_handle, pTargetBase + pSectionHeader->VirtualAddress, pSrcData + pSectionHeader->PointerToRawData, pSectionHeader->SizeOfRawData, nullptr));
                {
                    delete[] pSrcData;
                    VirtualFreeEx(_handle, pTargetBase, 0, MEM_RELEASE);

                    return false;
                }
            }
        }

        memcpy(pSrcData, &data, sizeof(data));
        WriteProcessMemory(_handle, pTargetBase, pSrcData, 0x1000 , nullptr);

        delete[] pSrcData;

        LPVOID pShellCode = VirtualAllocEx(_handle, nullptr, 0x1000, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
        if (!pShellCode)
        {
            VirtualFreeEx(_handle, pTargetBase, 0, MEM_RELEASE);
            return false;
        }

        //WriteProcessMemory(_handle, pShellCode, this->ShellCode, 0x1000, nullptr);

        HANDLE hThread = CreateRemoteThread(_handle, nullptr, 0, reinterpret_cast<LPTHREAD_START_ROUTINE>(this->ShellCode), pTargetBase, 0, nullptr);
        if (!hThread)
        {
            VirtualFreeEx(_handle, pTargetBase, 0, MEM_RELEASE);
            VirtualFreeEx(_handle, pShellCode, 0, MEM_RELEASE);
        }

        //(hThread);

        HINSTANCE hCheck = NULL;
        while (!hCheck)
        {
            MANUAL_MAPPING_DATA data_checked{ 0 };
            ReadProcessMemory(_handle, pTargetBase, &data_checked, sizeof(data_checked), nullptr);
            hCheck = data_checked.hMod;
            Sleep(10);
        }

        VirtualFreeEx(_handle, pShellCode, 0, MEM_RELEASE);

        return true;

    }

    //Map file and get his bytes
    BYTE* MapFile(_In_ LPCSTR filename)
    {
        std::streampos size;
        std::fstream file(filename, std::ios::in | std::ios::binary | std::ios::ate);
        if (file.is_open())
        {
            size = file.tellg();

            BYTE* MemBlock = new BYTE[size];

            file.seekg(0, std::ios::beg);
            file.read((char*)MemBlock, size);
            file.close();
            file.clear();

            return MemBlock;
        }
        return 0;
    }

    //Run an exernal executable in a process (works only with x86 applications)
    #ifdef _X86_
    bool RunPortableExecutable(_In_ LPVOID Image)
    {
        IMAGE_DOS_HEADER* DOSHeader;
        IMAGE_NT_HEADERS* NtHeader;
        IMAGE_SECTION_HEADER* SectionHeader;

        PROCESS_INFORMATION PI;
        STARTUPINFOA SI;

        CONTEXT* CTX;

        DWORD* ImageBase;
        void* pImageBase;

        int count;
        char CurrentFilePath[1024];

        DOSHeader = PIMAGE_DOS_HEADER(Image);
        NtHeader = PIMAGE_NT_HEADERS(DWORD(Image) + DOSHeader->e_lfanew);

        GetModuleFileNameA(0, CurrentFilePath, 1024);

        if (NtHeader->Signature == IMAGE_NT_SIGNATURE)
        {
            ZeroMemory(&PI, sizeof(PI));
            ZeroMemory(&SI, sizeof(SI));

            if (CreateProcessA(CurrentFilePath, NULL, NULL, NULL, FALSE, CREATE_SUSPENDED, NULL, NULL, &SI, &PI))
            {
                CTX = LPCONTEXT(VirtualAlloc(NULL, sizeof(CTX), MEM_COMMIT, PAGE_READWRITE));
                CTX->ContextFlags = CONTEXT_FULL;

                if (GetThreadContext(PI.hThread, LPCONTEXT(CTX)))
                {
                    ReadProcessMemory(PI.hProcess, LPCVOID(CTX->Ebx + 8), LPVOID(&ImageBase), 4, 0);

                    pImageBase = VirtualAllocEx(PI.hProcess, LPVOID(NtHeader->OptionalHeader.ImageBase), NtHeader->OptionalHeader.SizeOfImage, 0x3000, PAGE_EXECUTE_READWRITE);

                    WriteProcessMemory(PI.hProcess, pImageBase, Image, NtHeader->OptionalHeader.SizeOfHeaders, NULL);

                    for (count = 0; count < NtHeader->FileHeader.NumberOfSections; count++)
                    {
                        SectionHeader = PIMAGE_SECTION_HEADER(DWORD(Image) + DOSHeader->e_lfanew + 248 + (count * 40));
                        WriteProcessMemory(PI.hProcess, LPVOID(DWORD(pImageBase) + SectionHeader->VirtualAddress), LPVOID(DWORD(Image) + SectionHeader->PointerToRawData), SectionHeader->SizeOfRawData, 0);
                    }
                    WriteProcessMemory(PI.hProcess, LPVOID(CTX->Ebx + 8), LPVOID(&NtHeader->OptionalHeader.ImageBase), 4, 0);

                    CTX->Eax = DWORD(pImageBase) + NtHeader->OptionalHeader.AddressOfEntryPoint;
                    SetThreadContext(PI.hThread, LPCONTEXT(CTX));
                    ResumeThread(PI.hThread);

                    return true;
                }
            }
        }

    }
    #endif

    //Get process id value
    DWORD GetPid(_Null_ void) const
    {
        return this->_pid;
    }

    //Get handle address
    HANDLE GetHandle(_Null_ void) const
    {
        return this->_handle;
    }

    //Get process information
    PROCESSENTRY32W GetEntry(_Null_ void) const
    {
        return this->_entry;
    }

    //Get dll module information
    MODULEENTRY32W GetModEntry(_Null_ void) const
    {
        return this->_mEntry;
    }

    /* Operators overloading */

    VirtualMemory& operator=(_In_ VirtualMemory& RightValue)
    {
        this->_handle = RightValue._handle;
        this->_pid = RightValue._pid;

        return *this;
    }

    inline bool operator==(_In_ const VirtualMemory& RightValue)
    {
        if (this->_handle != RightValue._handle) return false;
        if (this->_pid != RightValue._pid) return false;

        return true;
    }

    inline bool operator!(_Null_ void)
    {
        if (this->_handle == INVALID_HANDLE_VALUE) return true;
        else return false;
    }

private:

    //Constructor basic vars initializer
    void StartEngine(_Null_ void)
    {
        this->_entry.dwSize = sizeof(PROCESSENTRY32);
        this->_mEntry.dwSize = sizeof(MODULEENTRY32);
        this->_pid = 0;
        this->_handle = INVALID_HANDLE_VALUE;
    }

    //Compare data for pattern scan
    bool DataCompare(_In_ const BYTE* data, _In_ const BYTE* mask, _In_ LPCSTR szMask)
    {
        for (; *szMask; ++szMask, ++data, ++mask)
        {
            if (*szMask == 'x' && *data != *mask) return false;
        }
        return(*szMask == NULL);
    }

    //Shell code used in manual mapping
    void __stdcall ShellCode(_In_ MANUAL_MAPPING_DATA* pData)
    {
        if (!pData) return;

        PBYTE pBase = reinterpret_cast<PBYTE>(pData);
        auto* pOpt = &reinterpret_cast<PIMAGE_NT_HEADERS>(pBase + reinterpret_cast<PIMAGE_DOS_HEADER>(pData)->e_lfanew)->OptionalHeader;

        auto _LoadLibraryA = pData->pLoadLibraryA;
        auto _GetProcAddress = pData->pGetProcAddress;
        auto _DllMain = reinterpret_cast<pfDllEntryPoint>(pBase + pOpt->AddressOfEntryPoint);

        PBYTE LocationDelta = pBase - pOpt->ImageBase;
        if (LocationDelta)
        {
            if (!pOpt->DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size) return;

            auto* pRelocData = reinterpret_cast<PIMAGE_BASE_RELOCATION>(pBase + pOpt->DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
            while (pRelocData->VirtualAddress)
            {
                UINT AmountOfEntries = (pRelocData->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
                PWORD pRelativeInfo = reinterpret_cast<PWORD>(pRelocData + 1);

                for (UINT i = 0; i != AmountOfEntries; ++i, ++pRelativeInfo)
                {
                    if (RELOC_FLAG(*pRelativeInfo))
                    {
                        UINT_PTR* pPatch = reinterpret_cast<UINT_PTR*>(pBase + pRelocData->VirtualAddress + ((*pRelativeInfo) & 0xFFF));
                        *pPatch += reinterpret_cast<UINT_PTR>(LocationDelta);
                    }
                }

                pRelocData = reinterpret_cast<PIMAGE_BASE_RELOCATION>(reinterpret_cast<PBYTE>(pRelocData) + pRelocData->SizeOfBlock);
            }
        }

        if (pOpt->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size)
        {
            auto* pImportDescriptor = reinterpret_cast<PIMAGE_IMPORT_DESCRIPTOR>(pBase + pOpt->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
            while (pImportDescriptor->Name)
            {
                char* szMod = reinterpret_cast<char*>(pBase + pImportDescriptor->Name);
                HINSTANCE hDll = _LoadLibraryA(szMod);
                ULONG_PTR* pThunkRef = reinterpret_cast<ULONG_PTR*>(pBase + pImportDescriptor->OriginalFirstThunk);
                ULONG_PTR* pFuncRef = reinterpret_cast<ULONG_PTR*>(pBase + pImportDescriptor->FirstThunk);

                if (!pThunkRef) pThunkRef = pFuncRef;

                for (; *pThunkRef; ++pThunkRef, ++pFuncRef)
                {
                    if (IMAGE_SNAP_BY_ORDINAL(*pThunkRef))
                    {
                        *pFuncRef = _GetProcAddress(hDll, reinterpret_cast<PCHAR>(*pThunkRef & 0xFFFF));
                    }
                    else
                    {
                        auto* pImport = reinterpret_cast<PIMAGE_IMPORT_BY_NAME>(pBase + (*pThunkRef));
                        *pFuncRef = _GetProcAddress(hDll, pImport->Name);
                    }
                }
                pImportDescriptor++;
            }
        }

        if (pOpt->DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].Size)
        {
            auto* pTLS = reinterpret_cast<PIMAGE_TLS_DIRECTORY>(pBase + pOpt->DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].VirtualAddress);
            auto* pCallBack = reinterpret_cast<PIMAGE_TLS_CALLBACK*>(pTLS->AddressOfCallBacks);

            for (; pCallBack && *pCallBack; ++pCallBack)
                (*pCallBack)(pBase, DLL_PROCESS_ATTACH, nullptr);
        }

        _DllMain(pBase, DLL_PROCESS_ATTACH, nullptr);
        pData->hMod = reinterpret_cast<HINSTANCE>(pBase);

    }

    //Internal Variables
    PROCESSENTRY32W _entry;   //Store process informations to get pid and access
    MODULEENTRY32W _mEntry;   //Store last module informations found to get base address and size
    DWORD _pid;               //Process ID to get access
    HANDLE _handle;           //Process handle

} *pVirtualMemory, *lpVirtualMemory;

#endif