//-----------------------------------------------------------------------------
// ProcessEnvironmentBlock 
// (built upon examples in Bill Blunden's Rootkit Arsenal, 1st Ed)
//-----------------------------------------------------------------------------

#include "stdafx.h"
#include <stdio.h>
#include "ProcessEnvironmentBlock.h"

wchar_t debugBuffer[1024];

// The quick and hacky way to get a pointer to the PEB:
PEB* getPEBWithASM()
{
	PEB* peb;
	__asm
	{
		MOV EAX,FS:[30H]
		MOV peb,EAX
	}
	return(peb);
}

// The long and standard way to get a pointer to the PEB:
PEB* getPEB()
{
	HMODULE handleDLL;
	NtQueryInformationProcessPtr NtQueryInformationProcess;
	NTSTATUS ntStatus;
	PROCESS_BASIC_INFORMATION basicInfo;

	handleDLL = LoadLibraryA("ntdll.dll");
	if(handleDLL == NULL)
	{
		printf("[getPEB]: LoadLibraryA() failed.\n");
		return(NULL);
	}

	NtQueryInformationProcess = (NtQueryInformationProcessPtr)GetProcAddress(handleDLL, "NtQueryInformationProcess");
	if(NtQueryInformationProcess == NULL)
	{
		printf("[getPEB]: GetProcAddress() failed.\n");
		return(NULL);
	}

	ntStatus = NtQueryInformationProcess(GetCurrentProcess(),		// HANDLE ProcessHandle
										ProcessBasicInformation,	// PROCESSINFOCLASS ProcessInformationClass
										&basicInfo,					// PVOID ProcessInformation
										sizeof(PROCESS_BASIC_INFORMATION),	// ULONG ProcesInformationLength
										NULL);						// PULONG ReturnLength
	if(!NT_SUCCESS(ntStatus))
	{
		printf("[getPEB]: NtQueryInformationProcess() failed.\n");
		return(NULL);
	}

	return(basicInfo.PebBaseAddress);
}

MY_PLDR_DATA_TABLE_ENTRY getNextMemLdrDataTableEntry(MY_PLDR_DATA_TABLE_ENTRY ptr)
{
	BYTE* address;
	address = (BYTE*)((*ptr).InMemoryOrderLinks).Flink;
	address = address - LIST_ENTRY_OFFSET_MEM_ORDER;
	return((MY_PLDR_DATA_TABLE_ENTRY)address);
}

MY_PLDR_DATA_TABLE_ENTRY getNextLoadLdrDataTableEntry(MY_PLDR_DATA_TABLE_ENTRY ptr)
{
	BYTE* address;
	address = (BYTE*)((*ptr).InLoadOrderLinks).Flink;
	address = address - LIST_ENTRY_OFFSET_LOAD_ORDER;
	return((MY_PLDR_DATA_TABLE_ENTRY)address);
}

void printDLLInfo(MY_PLDR_DATA_TABLE_ENTRY ptr)
{
	printf("[printDLLInfo]: %S ", (*ptr).FullDllName.Buffer);
	printf("\t\tBase=%08X\n", (DWORD)(*ptr).DllBase);
	return;
}

void walkDLLList(MY_PEB* mpeb)
{
	PPEB_LDR_DATA loaderData;
	PRTL_USER_PROCESS_PARAMETERS procParams;

	BYTE* address;
	MY_PLDR_DATA_TABLE_ENTRY curr;
	MY_PLDR_DATA_TABLE_ENTRY first;
	DWORD nDLLs;

	procParams = (*mpeb).ProcessParameters;
	printf("[walkDLLList]: Image Path=%S\n", (*procParams).ImagePathName.Buffer);
	printf("[walkDLLList]: Command Line=%S\n", (*procParams).CommandLine.Buffer);

	loaderData = (*mpeb).LoaderData;
	address = (BYTE*)((*loaderData).InMemoryOrderModuleList).Flink;
	address = address - LIST_ENTRY_OFFSET_MEM_ORDER;
	first = (MY_PLDR_DATA_TABLE_ENTRY)address;
	curr = first;

	nDLLs = 0;
	do
	{
		nDLLs++;
		printDLLInfo(curr);
		curr = getNextMemLdrDataTableEntry(curr);

		// list is circular, but it does have a terminator to mark the end:
		if(((DWORD)(*curr).DllBase)==0)
			break;
	} while (curr != first);

	printf("[walkDLLList]: nDLLs=%u\n", nDLLs);
	return;
}

// The goal of this routine is to modify the PEB's "InLoadOrder" linked list 
// of loaded modules of this process, in order to move the list entry for our
// DLL from last (most recently loaded) to first. 
void promoteToFirstLoadedDLL(MY_PEB* mpeb, const DWORD ourBaseAddress)
{
	MY_PPEB_LDR_DATA loaderData;

	BYTE* address;
	MY_PLDR_DATA_TABLE_ENTRY curr;
	MY_PLDR_DATA_TABLE_ENTRY first;
	
	loaderData = (MY_PEB_LDR_DATA*)(*mpeb).LoaderData;
	address = (BYTE*)((*loaderData).InLoadOrderModuleList).Flink;
	address = address - LIST_ENTRY_OFFSET_LOAD_ORDER;
	first = (MY_PLDR_DATA_TABLE_ENTRY)address;
	curr = first;

	//printf("[promoteToFirstLoadedDLL]: Working with PEB at %08X\n", mpeb);
	//printf("[promoteToFirstLoadedDLL]: Working with PEB_LDR_DATA at %08X\n", loaderData);

	// First, I need to walk the linked list to the entry for our DLL.
	do
	{
		curr = getNextLoadLdrDataTableEntry(curr);

		// list is circular, but it does have a terminator to mark the end:
		if(((DWORD)(*curr).DllBase)==0)
			break;
		else if(((DWORD)(*curr).DllBase) == ourBaseAddress)
		{
			// If I want my DLL to appear first in the InLoadOrder list, I need to 
			// (1) remove its entry from the list; (2) add it back in at the beginning.

			// Remove:
			curr->InLoadOrderLinks.Blink->Flink = curr->InLoadOrderLinks.Flink;
			curr->InLoadOrderLinks.Flink->Blink = curr->InLoadOrderLinks.Blink;

			// Re-add at beginning:
			curr->InLoadOrderLinks.Flink = first->InLoadOrderLinks.Flink;
			curr->InLoadOrderLinks.Blink = first->InLoadOrderLinks.Flink->Blink;
			first->InLoadOrderLinks.Flink = (LIST_ENTRY*)(&(curr->InLoadOrderLinks.Flink));  //alternately, simply "curr"
			curr->InLoadOrderLinks.Flink->Blink = (LIST_ENTRY*)(&(curr->InLoadOrderLinks.Blink)); //alternately, simply "curr + 0x04"

			break;
		}
	} while (curr != first);

	return;
}

// The goal of this routine is to modify the PEB's "InMemoryOrder" linked list 
// of loaded modules of this process, in order to move the list entry for our
// DLL from last to first. 
void promoteToFirstInMemDLL(MY_PEB* mpeb, const DWORD ourBaseAddress)
{
	MY_PPEB_LDR_DATA loaderData;

	BYTE* address;
	MY_PLDR_DATA_TABLE_ENTRY curr;
	MY_PLDR_DATA_TABLE_ENTRY first;
	
	loaderData = (MY_PEB_LDR_DATA*)(*mpeb).LoaderData;
	address = (BYTE*)((*loaderData).InMemoryOrderModuleList).Flink;
	address = address - LIST_ENTRY_OFFSET_MEM_ORDER;
	first = (MY_PLDR_DATA_TABLE_ENTRY)address;
	curr = first;

	swprintf_s(debugBuffer, 1024, L"[promoteToFirstInMemDLL]: Working with PEB at %08X\n", mpeb);
	OutputDebugString(debugBuffer);
	swprintf_s(debugBuffer, 1024, L"[promoteToFirstInMemDLL]: Working with PEB_LDR_DATA at %08X\n", loaderData);
	OutputDebugString(debugBuffer);

	// First, I need to walk the linked list to the entry for our DLL.
	do
	{
		curr = getNextMemLdrDataTableEntry(curr);

		// list is circular, but it does have a terminator to mark the end:
		if(((DWORD)(*curr).DllBase)==0)
		{
			swprintf_s(debugBuffer, 1024, L"[promoteToFirstInMemDLL]: hit the end of the list. Exiting.\n");
			OutputDebugString(debugBuffer);

			break;
		}
		else if(((DWORD)(*curr).DllBase) == ourBaseAddress)
		{
			swprintf_s(debugBuffer, 1024, L"[promoteToFirstInMemDLL]: Found our DLL, at %08X, in list.\n", ourBaseAddress);
			OutputDebugString(debugBuffer);

			// If I want my DLL to appear first in the InMemOrder list, I need to 
			// (1) remove its entry from the list; (2) add it back in at the beginning.

			// Remove:
			curr->InMemoryOrderLinks.Blink->Flink = curr->InMemoryOrderLinks.Flink;
			curr->InMemoryOrderLinks.Flink->Blink = curr->InMemoryOrderLinks.Blink;

			// Re-add at beginning:
			curr->InMemoryOrderLinks.Flink = first->InMemoryOrderLinks.Flink;
			curr->InMemoryOrderLinks.Blink = first->InMemoryOrderLinks.Flink->Blink;
			first->InMemoryOrderLinks.Flink = (LIST_ENTRY*)(&(curr->InMemoryOrderLinks.Flink));  //a.k.a simply "curr"
			curr->InMemoryOrderLinks.Flink->Blink = (LIST_ENTRY*)(&(curr->InMemoryOrderLinks.Blink)); //a.k.a. simply "curr + 0x0C"

			swprintf_s(debugBuffer, 1024, L"[promoteToFirstInMemDLL]: List should be reordered now with our DLL as first in memory.\n");
			OutputDebugString(debugBuffer);

			break;
		}
	} while (curr != first);

	return;
}

/* 
//Example usage:
void main()
{
	PEB* peb;
	MY_PEB* mpeb;

	//peb = getPEB();
	peb = getPEBWithASM();

	mpeb = (MY_PEB*)peb;
	walkDLLList(mpeb);

	return;
}
*/