#include "Common.h"

/*
 * This tells the loader to discard the image header,
 *   and sections .rsrc and .reloc.  If you need
 *   resources or you want to rebase again then
 *   undefine it.
 */
#define DISCARD_JUNK

static void rebase(
	PBYTE baseAddressLocal,
	PBYTE baseAddressRemote,
	ptrdiff_t offset)
{
	/*
	 * Credit to http://www.cultdeadcow.com/tools/pewrap.html
	 */
	DWORD i, numRelocs;
	PBYTE blockBase, *pPtr, ptr;
	PWORD pRelocs;
	PIMAGE_NT_HEADERS ntHeaders = (PIMAGE_NT_HEADERS) (baseAddressLocal + ((PIMAGE_DOS_HEADER) baseAddressLocal)->e_lfanew);
	PIMAGE_BASE_RELOCATION pBaseRelocation = (PIMAGE_BASE_RELOCATION) (baseAddressLocal +
			ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
	DWORD imageSize = ntHeaders->OptionalHeader.SizeOfImage;
	while (pBaseRelocation->VirtualAddress) {
		blockBase = baseAddressLocal + pBaseRelocation->VirtualAddress;
		pRelocs = (PWORD) ((PBYTE) pBaseRelocation + sizeof(IMAGE_BASE_RELOCATION));
		numRelocs = (pBaseRelocation->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
		for (i = 0; i < numRelocs; i++) {
			if ((pRelocs[i] >> 12) != IMAGE_REL_BASED_HIGHLOW)
				continue;
			pPtr = (PBYTE*) (blockBase + (pRelocs[i] & 0xFFFU));
			ptr = *pPtr + offset;
			if (ptr <  baseAddressRemote ||
				ptr >= baseAddressRemote + imageSize)
				continue;
			*pPtr = ptr;
		}
		pBaseRelocation = (PIMAGE_BASE_RELOCATION)
			((PBYTE) pBaseRelocation + pBaseRelocation->SizeOfBlock);
	}
	ntHeaders->OptionalHeader.ImageBase = (DWORD) baseAddressRemote;
}

static int copyBlock(
	HANDLE processHandle,
	PBYTE localAddress,
	PBYTE remoteAddress,
	DWORD blockSize,
	DWORD desiredProtection)
{
	DWORD d, writable;

	writable = desiredProtection & (PAGE_READWRITE | PAGE_EXECUTE_READWRITE);
	if (!VirtualAllocEx(
			processHandle,
			remoteAddress,
			blockSize,
			MEM_COMMIT,
			writable ? desiredProtection : PAGE_EXECUTE_READWRITE)) {
		strcpy(errorText, "Loader: remote memory allocation failed");
		return 0;
	}
	if (!WriteProcessMemory(
			processHandle,
			remoteAddress,
			localAddress,
			blockSize,
			&d) ||
		d < blockSize) {
		strcpy(errorText, "Loader: failed to copy image to remote process");
		return 0;
	}
	if (!writable)
		VirtualProtectEx(
			processHandle,
			remoteAddress,
			blockSize,
			desiredProtection,
			&d);
	return 1;
}

int loader(HANDLE processHandle, ptrdiff_t& addressOffset)
{
	MEMORY_BASIC_INFORMATION memoryBasicInformation;
	DWORD i, d, numSections, imageSize;
	PBYTE imageBase, baseAddressLocal, baseAddressRemote;
	ptrdiff_t offset;
	PIMAGE_NT_HEADERS ntHeaders;
	PIMAGE_SECTION_HEADER sections;

	imageBase = (PBYTE) GetModuleHandle(NULL);
	if (!imageBase) {
		strcpy(errorText, "Loader: unexpected fatal error");
		return 0;
	}
	ntHeaders = (PIMAGE_NT_HEADERS) (imageBase + ((PIMAGE_DOS_HEADER) imageBase)->e_lfanew);
	numSections = ntHeaders->FileHeader.NumberOfSections;
	sections = IMAGE_FIRST_SECTION(ntHeaders);
	/*
	 * Find out if we have relocation data
	 *   sets d if we don't
	 */
	for (i = 0; i < numSections; i++)
		if (!strcmp((const char*) sections[i].Name, ".reloc"))
			break;
	d = i == numSections ||
		ntHeaders->OptionalHeader.NumberOfRvaAndSizes <= IMAGE_DIRECTORY_ENTRY_BASERELOC ||
		!ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress ||
		!ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size;
	/*
	 * Allocate memory
	 *   try allocating at the same base address
	 *   so we don't have to rebase.
	 *   Note: in order to ask for an explicit
	 *   address in another process you need
	 *   to 1st reserve, then commit.  If you
	 *   try committing directly it fails.
	 */
	imageSize = ntHeaders->OptionalHeader.SizeOfImage;
	baseAddressRemote = (PBYTE)	VirtualAllocEx(
		processHandle,
		imageBase,
		imageSize,
		MEM_RESERVE,
		PAGE_NOACCESS);
	if (!baseAddressRemote) {
		if (d) {
			/*
			 * We need to rebase, but don't have a base relocation table.
			 */
			strcpy(errorText, "Loader: image missing base relocation table, use /FIXED:NO");
			return 0;
		}
		baseAddressRemote = (PBYTE) VirtualAllocEx(
			processHandle,
			NULL,
			imageSize,
			MEM_RESERVE,
			PAGE_NOACCESS);
		if (!baseAddressRemote) {
			strcpy(errorText, "Loader: remote memory reservation failed");
			return 0;
		}
	}
	if (baseAddressRemote == imageBase) {
		/*
		 * No need to rebase
		 */
		offset = 0;
		baseAddressLocal = imageBase;
	} else {
		baseAddressLocal = (PBYTE) VirtualAlloc(
			NULL,
			imageSize,
			MEM_COMMIT,
			PAGE_READWRITE);
		if (!baseAddressLocal) {
			VirtualFreeEx(processHandle, baseAddressRemote, 0, MEM_RELEASE);
			strcpy(errorText, "Loader: local memory allocation failed");
			return 0;
		}
		memcpy(baseAddressLocal, imageBase, imageSize);
		/*
		* Rebase the image
		*/
		offset = baseAddressRemote - imageBase;
		rebase(baseAddressLocal, baseAddressRemote, offset);
	}
	/*
	 * Copy to remote process
	 */
#ifndef DISCARD_JUNK
	if (!copyBlock(
			processHandle,
			baseAddressLocal,
			baseAddressRemote,
			ntHeaders->OptionalHeader.SizeOfHeaders,
			PAGE_READONLY) {
		if (offset)
			VirtualFree(baseAddressLocal, 0, MEM_RELEASE);
		VirtualFreeEx(processHandle, baseAddressRemote, 0, MEM_RELEASE);
		return 0;
	}
#endif
	for (i = 0; i < numSections; i++) {
#ifdef DISCARD_JUNK
		if (!strcmp((const char*) sections[i].Name, ".rsrc") ||
			!strcmp((const char*) sections[i].Name, ".reloc"))
			continue;
#endif
		if (VirtualQuery(
				imageBase + sections[i].VirtualAddress,
				&memoryBasicInformation,
				sizeof(MEMORY_BASIC_INFORMATION)) == sizeof(MEMORY_BASIC_INFORMATION)) {
			d = memoryBasicInformation.Protect;
			if ((d & 0xFF) == PAGE_WRITECOPY)
				d = (d & 0xFFFFFF00) | PAGE_READWRITE;
			else if ((d & 0xFF) == PAGE_EXECUTE_WRITECOPY)
				d = (d & 0xFFFFFF00) | PAGE_EXECUTE_READWRITE;
		} else
			d = PAGE_EXECUTE_READWRITE;
		if (!copyBlock(
				processHandle,
				baseAddressLocal + sections[i].VirtualAddress,
				baseAddressRemote + sections[i].VirtualAddress,
				sections[i].Misc.VirtualSize,
				d)) {
			if (offset)
				VirtualFree(baseAddressLocal, 0, MEM_RELEASE);
			VirtualFreeEx(processHandle, baseAddressRemote, 0, MEM_RELEASE);
			return 0;
		}
	}
	if (offset)
		VirtualFree(baseAddressLocal, 0, MEM_RELEASE);
	addressOffset = offset;
	return 1;
}

void unload(HANDLE processHandle, ptrdiff_t addressOffset)
{
	PBYTE imageBase;

	imageBase = (PBYTE) GetModuleHandle(NULL);
	if (!imageBase)
		return;
	VirtualFreeEx(processHandle, imageBase + addressOffset, 0, MEM_RELEASE);
}