#include <Windows.h>
#include <stdio.h>
#include "error.h"

#pragma comment(linker, "/DYNAMICBASE:NO")
#pragma comment(linker, "/BASE:0x02000000")

void * FileToMemory(LPTSTR file)
{
	HANDLE hFile = NULL;
	HANDLE hFileMap = NULL;
	LPVOID pMapView = NULL;

	hFile = CreateFile(file, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if (hFile == INVALID_HANDLE_VALUE)
	{
		CloseHandle(hFile);
		return NULL;
	}

	hFileMap = CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
	if (hFileMap == NULL)
	{
		CloseHandle(hFile);
		return NULL;
	}

	pMapView = MapViewOfFile(hFileMap, FILE_MAP_READ, 0, 0, 0);

	return pMapView;
}

void GetHeaders(PBYTE ibase, PIMAGE_FILE_HEADER *pfh, PIMAGE_OPTIONAL_HEADER *poh, PIMAGE_SECTION_HEADER *psh)
{
	PIMAGE_DOS_HEADER mzhead = (PIMAGE_DOS_HEADER)ibase;
	*pfh = (PIMAGE_FILE_HEADER)&ibase[mzhead->e_lfanew];
	*pfh = (PIMAGE_FILE_HEADER)((PBYTE)*pfh + sizeof(IMAGE_NT_SIGNATURE));
	*poh = (PIMAGE_OPTIONAL_HEADER)((PBYTE)*pfh + sizeof(IMAGE_FILE_HEADER));
	*psh = (PIMAGE_SECTION_HEADER)((PBYTE)*poh + sizeof(IMAGE_OPTIONAL_HEADER));
}

void * LoadPE(void *image)
{
	PIMAGE_DOS_HEADER dos_hdr = NULL;
	PIMAGE_FILE_HEADER file_hdr = NULL;
	PIMAGE_OPTIONAL_HEADER opt_hdr = NULL;
	PIMAGE_SECTION_HEADER section_hdr = NULL;
	LPVOID pImageBase = NULL;

	dos_hdr = (PIMAGE_DOS_HEADER) image;
	file_hdr = (PIMAGE_FILE_HEADER) ((DWORD)image + dos_hdr->e_lfanew + 4);
	opt_hdr = (PIMAGE_OPTIONAL_HEADER) (file_hdr + 1);
	section_hdr = (PIMAGE_SECTION_HEADER)((DWORD)opt_hdr + sizeof(IMAGE_OPTIONAL_HEADER));

	printf("Allocating memory for image\n");
	UnmapViewOfFile((LPVOID)opt_hdr->ImageBase);
	pImageBase = VirtualAlloc((LPVOID)opt_hdr->ImageBase, opt_hdr->SizeOfImage, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
	if (pImageBase == NULL)
	{
		FormatError(GetLastError());
		printf("Error while allocating memory. Possible conflict of base addresses\n");
		return NULL;
	}

	printf("Copying headers\n");
	CopyMemory(pImageBase, image, opt_hdr->SizeOfHeaders);

	printf("Copying sections\n");
	for (DWORD i = 0; i < file_hdr->NumberOfSections; i++)
	{
		LPVOID VirtualAddress = (LPVOID)((DWORD)pImageBase + section_hdr[i].VirtualAddress);
		CopyMemory(VirtualAddress, (void *)((DWORD)image + section_hdr[i].PointerToRawData), section_hdr[i].SizeOfRawData);
	}

	printf("Setting attributes\n");
	for (DWORD i = 0; i < file_hdr->NumberOfSections; i++)
	{
		LPVOID VirtualAddress = (LPVOID)((DWORD)pImageBase + section_hdr->VirtualAddress);
		DWORD VirtualSize = section_hdr[i].Misc.VirtualSize;

		DWORD attributes = 0;
		if(section_hdr[i].Characteristics & IMAGE_SCN_MEM_EXECUTE ||section_hdr[i].Characteristics & IMAGE_SCN_MEM_READ )
		{
			attributes |= PAGE_EXECUTE;
			if(section_hdr[i].Characteristics & IMAGE_SCN_MEM_WRITE )
				attributes |= PAGE_READWRITE;
			else
				attributes |= PAGE_READONLY;
		}
		else if(section_hdr[i].Characteristics & IMAGE_SCN_MEM_WRITE )
			attributes |= PAGE_READWRITE;

		VirtualProtect(VirtualAddress, VirtualSize, attributes, &attributes);
	}

	return pImageBase;
}

void CallPE(void *base)
{
	PIMAGE_FILE_HEADER pfh;
	PIMAGE_OPTIONAL_HEADER poh;
	PIMAGE_SECTION_HEADER psh;
	LPVOID entry = NULL;

    GetHeaders((PBYTE)base, &pfh, &poh, &psh);

	entry = (LPVOID)((DWORD)base + poh->AddressOfEntryPoint);
	__asm call dword ptr [entry];
}

int main()
{
	void *pImage = FileToMemory(TEXT("test.exe"));
	void *pModule = LoadPE(pImage);
	CallPE(pModule);
	return 0;
}