/*
* THIS ENTIRE CODE COMES FROM THE RAPHAEL'S MODUGE GITHUB REPOSITORY NAMED "unhook-bof":
*   https://github.com/rsmudge/unhook-bof
* 
* Also, parts of this code contain functions implemented by Jeff Tang <jtang@cylance.com> in his apisetmap module.
* 
* I have only adapted it for needs of this project!
* 
*/

#pragma once

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <intrin.h>
#include <iostream>
#include <vector>

#define DEREF( name )*(UINT_PTR *)(name)
#define DEREF_64( name )*(DWORD64 *)(name)
#define DEREF_32( name )*(DWORD *)(name)
#define DEREF_16( name )*(WORD *)(name)
#define DEREF_8( name )*(BYTE *)(name)

class ModulesRefresher
{
    // Win 10
    typedef struct _API_SET_VALUE_ENTRY_V6
    {
        ULONG Flags;
        ULONG NameOffset;
        ULONG NameLength;
        ULONG ValueOffset;
        ULONG ValueLength;
    } API_SET_VALUE_ENTRY_V6, * PAPI_SET_VALUE_ENTRY_V6;

    typedef struct _API_SET_NAMESPACE_HASH_ENTRY_V6
    {
        ULONG Hash;
        ULONG Index;
    } API_SET_NAMESPACE_HASH_ENTRY_V6, * PAPI_SET_NAMESPACE_HASH_ENTRY_V6;

    typedef struct _API_SET_NAMESPACE_ENTRY_V6
    {
        ULONG Flags;
        ULONG NameOffset;
        ULONG Size;
        ULONG NameLength;
        ULONG DataOffset;
        ULONG Count;
    } API_SET_NAMESPACE_ENTRY_V6, * PAPI_SET_NAMESPACE_ENTRY_V6;

    typedef struct _API_SET_NAMESPACE_ARRAY_V6
    {
        ULONG Version;
        ULONG Size;
        ULONG Flags;
        ULONG Count;
        ULONG DataOffset;
        ULONG HashOffset;
        ULONG Multiplier;
        API_SET_NAMESPACE_ENTRY_V6 Array[ANYSIZE_ARRAY];
    } API_SET_NAMESPACE_ARRAY_V6, * PAPI_SET_NAMESPACE_ARRAY_V6;

    // Windows 8.1
    typedef struct _API_SET_VALUE_ENTRY_V4
    {
        ULONG Flags;
        ULONG NameOffset;
        ULONG NameLength;
        ULONG ValueOffset;
        ULONG ValueLength;
    } API_SET_VALUE_ENTRY_V4, * PAPI_SET_VALUE_ENTRY_V4;

    typedef struct _API_SET_VALUE_ARRAY_V4
    {
        ULONG Flags;
        ULONG Count;
        API_SET_VALUE_ENTRY_V4 Array[ANYSIZE_ARRAY];
    } API_SET_VALUE_ARRAY_V4, * PAPI_SET_VALUE_ARRAY_V4;

    typedef struct _API_SET_NAMESPACE_ENTRY_V4
    {
        ULONG Flags;
        ULONG NameOffset;
        ULONG NameLength;
        ULONG AliasOffset;
        ULONG AliasLength;
        ULONG DataOffset;
    } API_SET_NAMESPACE_ENTRY_V4, * PAPI_SET_NAMESPACE_ENTRY_V4;

    typedef struct _API_SET_NAMESPACE_ARRAY_V4
    {
        ULONG Version;
        ULONG Size;
        ULONG Flags;
        ULONG Count;
        API_SET_NAMESPACE_ENTRY_V4 Array[ANYSIZE_ARRAY];
    } API_SET_NAMESPACE_ARRAY_V4, * PAPI_SET_NAMESPACE_ARRAY_V4;

    // Windows 7/8
    typedef struct _API_SET_VALUE_ENTRY_V2
    {
        ULONG NameOffset;
        ULONG NameLength;
        ULONG ValueOffset;
        ULONG ValueLength;
    } API_SET_VALUE_ENTRY_V2, * PAPI_SET_VALUE_ENTRY_V2;

    typedef struct _API_SET_VALUE_ARRAY_V2
    {
        ULONG Count;
        API_SET_VALUE_ENTRY_V2 Array[ANYSIZE_ARRAY];
    } API_SET_VALUE_ARRAY_V2, * PAPI_SET_VALUE_ARRAY_V2;

    typedef struct _API_SET_NAMESPACE_ENTRY_V2
    {
        ULONG NameOffset;
        ULONG NameLength;
        ULONG DataOffset;
    } API_SET_NAMESPACE_ENTRY_V2, * PAPI_SET_NAMESPACE_ENTRY_V2;

    typedef struct _API_SET_NAMESPACE_ARRAY_V2
    {
        ULONG Version;
        ULONG Count;
        API_SET_NAMESPACE_ENTRY_V2 Array[ANYSIZE_ARRAY];
    } API_SET_NAMESPACE_ARRAY_V2, * PAPI_SET_NAMESPACE_ARRAY_V2;

    typedef struct _UNICODE_STR
    {
        USHORT Length;
        USHORT MaximumLength;
        PWSTR pBuffer;
    } UNICODE_STR, * PUNICODE_STR;

    // WinDbg> dt -v ntdll!_LDR_DATA_TABLE_ENTRY
    //__declspec( align(8) ) 
    typedef struct _LDR_DATA_TABLE_ENTRY
    {
        //LIST_ENTRY InLoadOrderLinks; // As we search from PPEB_LDR_DATA->InMemoryOrderModuleList we dont use the first entry.
        LIST_ENTRY InMemoryOrderModuleList;
        LIST_ENTRY InInitializationOrderModuleList;
        PVOID DllBase;
        PVOID EntryPoint;
        ULONG SizeOfImage;
        UNICODE_STR FullDllName;
        UNICODE_STR BaseDllName;
        ULONG Flags;
        SHORT LoadCount;
        SHORT TlsIndex;
        LIST_ENTRY HashTableEntry;
        ULONG TimeDateStamp;
    } LDR_DATA_TABLE_ENTRY, * PLDR_DATA_TABLE_ENTRY;

    // WinDbg> dt -v ntdll!_PEB_LDR_DATA
    typedef struct _PEB_LDR_DATA //, 7 elements, 0x28 bytes
    {
        DWORD dwLength;
        DWORD dwInitialized;
        LPVOID lpSsHandle;
        LIST_ENTRY InLoadOrderModuleList;
        LIST_ENTRY InMemoryOrderModuleList;
        LIST_ENTRY InInitializationOrderModuleList;
        LPVOID lpEntryInProgress;
    } PEB_LDR_DATA, * PPEB_LDR_DATA;

    // WinDbg> dt -v ntdll!_PEB_FREE_BLOCK
    typedef struct _PEB_FREE_BLOCK // 2 elements, 0x8 bytes
    {
        struct _PEB_FREE_BLOCK* pNext;
        DWORD dwSize;
    } PEB_FREE_BLOCK, * PPEB_FREE_BLOCK;

    // struct _PEB is defined in Winternl.h but it is incomplete
    // WinDbg> dt -v ntdll!_PEB
    typedef struct __PEB // 65 elements, 0x210 bytes
    {
        BYTE bInheritedAddressSpace;
        BYTE bReadImageFileExecOptions;
        BYTE bBeingDebugged;
        BYTE bSpareBool;
        LPVOID lpMutant;
        LPVOID lpImageBaseAddress;
        PPEB_LDR_DATA pLdr;
        LPVOID lpProcessParameters;
        LPVOID lpSubSystemData;
        LPVOID lpProcessHeap;
        PRTL_CRITICAL_SECTION pFastPebLock;
        LPVOID lpFastPebLockRoutine;
        LPVOID lpFastPebUnlockRoutine;
        DWORD dwEnvironmentUpdateCount;
        LPVOID lpKernelCallbackTable;
        DWORD dwSystemReserved;
        DWORD dwAtlThunkSListPtr32;
        PPEB_FREE_BLOCK pFreeList;
        DWORD dwTlsExpansionCounter;
        LPVOID lpTlsBitmap;
        DWORD dwTlsBitmapBits[2];
        LPVOID lpReadOnlySharedMemoryBase;
        LPVOID lpReadOnlySharedMemoryHeap;
        LPVOID lpReadOnlyStaticServerData;
        LPVOID lpAnsiCodePageData;
        LPVOID lpOemCodePageData;
        LPVOID lpUnicodeCaseTableData;
        DWORD dwNumberOfProcessors;
        DWORD dwNtGlobalFlag;
        LARGE_INTEGER liCriticalSectionTimeout;
        DWORD dwHeapSegmentReserve;
        DWORD dwHeapSegmentCommit;
        DWORD dwHeapDeCommitTotalFreeThreshold;
        DWORD dwHeapDeCommitFreeBlockThreshold;
        DWORD dwNumberOfHeaps;
        DWORD dwMaximumNumberOfHeaps;
        LPVOID lpProcessHeaps;
        LPVOID lpGdiSharedHandleTable;
        LPVOID lpProcessStarterHelper;
        DWORD dwGdiDCAttributeList;
        LPVOID lpLoaderLock;
        DWORD dwOSMajorVersion;
        DWORD dwOSMinorVersion;
        WORD wOSBuildNumber;
        WORD wOSCSDVersion;
        DWORD dwOSPlatformId;
        DWORD dwImageSubsystem;
        DWORD dwImageSubsystemMajorVersion;
        DWORD dwImageSubsystemMinorVersion;
        DWORD dwImageProcessAffinityMask;
        DWORD dwGdiHandleBuffer[34];
        LPVOID lpPostProcessInitRoutine;
        LPVOID lpTlsExpansionBitmap;
        DWORD dwTlsExpansionBitmapBits[32];
        DWORD dwSessionId;
        ULARGE_INTEGER liAppCompatFlags;
        ULARGE_INTEGER liAppCompatFlagsUser;
        LPVOID lppShimData;
        LPVOID lpAppCompatInfo;
        UNICODE_STR usCSDVersion;
        LPVOID lpActivationContextData;
        LPVOID lpProcessAssemblyStorageMap;
        LPVOID lpSystemDefaultActivationContextData;
        LPVOID lpSystemAssemblyStorageMap;
        DWORD dwMinimumStackCommit;
    } _PEB, * _PPEB;

    typedef struct
    {
        WORD	offset : 12;
        WORD	type : 4;
    } IMAGE_RELOC, * PIMAGE_RELOC;

public:
    ModulesRefresher();

    bool refresh(const std::vector<std::string>& skipModules);

private:

    BOOL IsBeaconDLL(const char* stomp, size_t beaconDllLength, PWSTR wszBaseDllName, USHORT BaseDllLength);
    HMODULE CustomLoadLibrary(const PWCHAR wszFullDllName, const PWCHAR wszBaseDllName, ULONG_PTR pDllBase);
    HMODULE CustomGetModuleHandleW(const PWSTR wszModule);
    FARPROC WINAPI CustomGetProcAddressEx(HMODULE hModule, const PCHAR lpProcName, PWSTR wszOriginalModule);
    VOID ScanAndFixModule(PCHAR pKnown, PCHAR pSuspect, PWCHAR wszBaseDllName);
    VOID ScanAndFixSection(PWCHAR wszBaseDllName, PCHAR pKnown, PCHAR pSuspect, size_t stLength);

    _PPEB GetProcessEnvironmentBlock();
    PLDR_DATA_TABLE_ENTRY GetInMemoryOrderModuleList();

    PWCHAR GetRedirectedName(const PWSTR wszImportingModule, const PWSTR wszVirtualModule, size_t* stSize);
    PWCHAR GetRedirectedName_V6(const PWSTR wszImportingModule, const PWSTR wszVirtualModule, size_t* stSize);
    PWCHAR GetRedirectedName_V4(const PWSTR wszImportingModule, const PWSTR wszVirtualModule, size_t* stSize);
    PWCHAR GetRedirectedName_V2(const PWSTR wszImportingModule, const PWSTR wszVirtualModule, size_t* stSize);
};