/*
	GrabCode code injection util-thingie

	By X-Calibre

	Yeah I know the source looks bad and not too much commented...
	but you should be lucky you got it at all :)
 */

#include <windows.h>
#include <stdio.h>
#include <winnt.h>
#include "resource.h"

// Our function prototypes
void BrowseSourceFile(HWND hDlg);
void BrowseTargetFile(HWND hDlg);
void OpenSourceFile(HWND hDlg);
void OpenTargetFile(HWND hDlg);
void GrabCode(HWND hDlg);
void DumpCode(HWND hDlg);
void CloseSourceFile();
void CloseTargetFile();
void OptimizePEHeader(HWND hDlg);

// Our file handles
HANDLE sourceFile, sourceFileMapping, targetFile, targetFileMapping, dumpFile;

// And filenames ofcourse
char sourceFileName[512] = {0};
char targetFileName[512] = {0};
char dumpFileName[512] = {0};

// Our header and data pointers
typedef struct
{
	IMAGE_DOS_HEADER *dos;
	IMAGE_NT_HEADERS *PE;
	IMAGE_SECTION_HEADER *section;
} HEADERS;

HEADERS source, target;

// Module handle ofcourse
HINSTANCE hInst;

// Buffer for wsprintf
char buffer[65535];

// Flags to check state of the app
#define	SOURCE_OPEN	1
#define	TARGET_OPEN	2

BYTE flags = 0;

WNDPROC OldSourceEditWndProc, OldTargetEditWndProc;

LRESULT CALLBACK SourceEditWndProc(HWND hEdit, UINT message, WPARAM wParam, LPARAM lParam)
{
	LRESULT retVal = CallWindowProc(OldSourceEditWndProc, hEdit, message, wParam, lParam);

	if(message == WM_KEYDOWN && (TCHAR)wParam == VK_RETURN)
	{
		// Get filename from EditBox
		GetWindowText(hEdit, sourceFileName, sizeof(sourceFileName));

		OpenSourceFile(GetParent(hEdit));
	}

	return retVal;
}

LRESULT CALLBACK TargetEditWndProc(HWND hEdit, UINT message, WPARAM wParam, LPARAM lParam)
{
	LRESULT retVal = CallWindowProc(OldTargetEditWndProc, hEdit, message, wParam, lParam);

	if(message == WM_KEYDOWN && (TCHAR)wParam == VK_RETURN)
	{
		// Get filename from EditBox
		GetWindowText(hEdit, targetFileName, sizeof(targetFileName));

		OpenTargetFile(GetParent(hEdit));
	}

	return retVal;
}

LRESULT CALLBACK DlgBoxProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
	switch (message)
	{
		case WM_COMMAND:
			switch (LOWORD(wParam))
			{
				case IDBROWSESOURCE:
					BrowseSourceFile(hDlg);
					OpenSourceFile(hDlg);
					return TRUE;
				case IDBROWSETARGET:
					BrowseTargetFile(hDlg);
					OpenTargetFile(hDlg);
					return TRUE;
				case IDINJECT:
					GrabCode(hDlg);
					return TRUE;
				case IDDUMPCODE:
					DumpCode(hDlg);
					return TRUE;
				case IDADDSECTION:
					OptimizePEHeader(hDlg);
					return TRUE;
				case IDABOUT:
					return MessageBox(hDlg, "GrabCode code injection utility\n\nVersion 1.0 By X-Calibre^PhrozenCrew", "About GrabCode...", MB_OK);
			}

			// New selection in ComboBox
			if (HIWORD(wParam) == CBN_SELCHANGE && LOWORD(wParam) == IDC_TARGETSECTION)
			{
				int i = SendDlgItemMessage(hDlg, IDC_TARGETSECTION, CB_GETCURSEL, 0, 0);
				
				// Update offset
				wsprintf(buffer, "%X", target.section[i].PointerToRawData);
				SetDlgItemText(hDlg, IDC_TARGETOFFSET, buffer);

				// Update size
				wsprintf(buffer, "%X", target.section[i].SizeOfRawData);
				SetDlgItemText(hDlg, IDC_SIZEOFBLOCK, buffer);
				return TRUE;
			}
			break;
			
		case WM_CLOSE:
			// The famous X in the top right corner is clicked!!
			ExitProcess(0);

		case WM_INITDIALOG:
			// Hook the EditBox return message.
			OldSourceEditWndProc = (WNDPROC)SetWindowLong(GetDlgItem(hDlg, IDC_SOURCE), GWL_WNDPROC, (LONG)SourceEditWndProc);
			OldTargetEditWndProc = (WNDPROC)SetWindowLong(GetDlgItem(hDlg, IDC_TARGET), GWL_WNDPROC, (LONG)TargetEditWndProc);
			return TRUE;
	}

	// Well, no more messages to be processed. Let's get outta here.
    return FALSE;
}

void GrabCode(HWND hDlg)
{
	DWORD bytesWritten, insertOffset, codeSize;

	// Parse user input
	GetDlgItemText(hDlg, IDC_TARGETOFFSET, buffer, sizeof(buffer));
	sscanf(buffer, "%X", &insertOffset);

	// Scan to correct offset in file
	SetFilePointer(targetFile, insertOffset, 0, FILE_BEGIN);

	// Inject code
	GetDlgItemText(hDlg, IDC_SIZEOFBLOCK, buffer, sizeof(buffer));
	sscanf(buffer, "%X", &codeSize);

	if (!WriteFile(targetFile, ((BYTE *)(source.dos) + (source.section[0].PointerToRawData + source.PE->OptionalHeader.AddressOfEntryPoint - source.PE->OptionalHeader.BaseOfCode)), codeSize, &bytesWritten, NULL))
	{
		MessageBox(hDlg, "Error injecting code into target!", "GrabCode failed", MB_OK);
		return;
	}

	// Notify user that injection is finished
	wsprintf(buffer, "Code portion of %0Xh bytes from %s injected at offset %Xh into %s", codeSize, sourceFileName, insertOffset, targetFileName);
	MessageBox(hDlg, buffer, "Finished", MB_OK);
}

OPENFILENAME openDumpFile = {sizeof(OPENFILENAME), NULL, NULL, "All files\0*.*\0", 0, 0, 0, dumpFileName, sizeof dumpFileName, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

void DumpCode(HWND hDlg)
{
	DWORD bytesWritten, codeSize;

	openDumpFile.hwndOwner = hDlg;
	openDumpFile.hInstance = hInst;

	if(!GetSaveFileName(&openDumpFile))
	{
		MessageBox(hDlg, "Error opening dump file!", "GrabCode failed", MB_OK);
		return;
	}

	dumpFile = CreateFile(dumpFileName, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, 0, NULL);

	// Inject code
	GetDlgItemText(hDlg, IDC_SIZEOFBLOCK, buffer, sizeof(buffer));
	sscanf(buffer, "%X", &codeSize);

	if (!WriteFile(dumpFile, ((BYTE *)(source.dos) + (source.section[0].PointerToRawData + source.PE->OptionalHeader.AddressOfEntryPoint - source.PE->OptionalHeader.BaseOfCode)), codeSize, &bytesWritten, NULL))
	{
		MessageBox(hDlg, "Error dumping code to file!", "GrabCode failed", MB_OK);
		return;
	}

	// Notify user that injection is finished
	wsprintf(buffer, "Code portion of %0Xh bytes from %s dumped to %s", codeSize, sourceFileName, dumpFileName);
	MessageBox(hDlg, buffer, "Finished", MB_OK);
}

OPENFILENAME openSourceFile = {sizeof(OPENFILENAME), NULL, NULL, "EXE files\0*.EXE\0", 0, 0, 0, sourceFileName, sizeof sourceFileName, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

void BrowseSourceFile(HWND hDlg)
{
	openSourceFile.hwndOwner = hDlg;
	openSourceFile.hInstance = hInst;

	// Check if file is open, and if so, close it
	if (flags & SOURCE_OPEN)
		CloseSourceFile();

	if(!GetOpenFileName(&openSourceFile))
		return;		

	// Set filename in editbox
	SetDlgItemText(hDlg, IDC_SOURCE, sourceFileName);
}

void OpenSourceFile(HWND hDlg)
{
	// Open file and memory-map it
	sourceFile = CreateFile(sourceFileName, GENERIC_READ, 0, NULL, OPEN_EXISTING, 0, NULL);

	sourceFileMapping = CreateFileMapping(sourceFile, NULL, PAGE_READONLY, 0, 0, NULL);
	
	if ((source.dos = MapViewOfFile(sourceFileMapping, FILE_MAP_READ, 0, 0, 0)) == 0)
	{
		MessageBox(hDlg, "Error opening source file", "Error", MB_OK);
		CloseSourceFile();
		return;
	}

	// Parse header
	if (source.dos->e_magic != 0x5A4D)
	{
		MessageBox(hDlg, "No MZ signature found. Incorrect file type", "Error", MB_OK);
		CloseSourceFile();
		return;
	}

	// Set pointer to start of PE header
	source.PE = (IMAGE_NT_HEADERS *)((source.dos->e_lfanew) + (DWORD)(source.dos));

	if (source.PE->Signature != 0x4550)
	{
		MessageBox(hDlg, "No PE signature found. Incorrect file type", "Error", MB_OK);
		CloseSourceFile();
		return;
	}

	// File is correct. Set flag to indicate it's opened
	flags |= SOURCE_OPEN;

	// Enable inject button, if possible
	EnableWindow(GetDlgItem(hDlg, IDINJECT), (flags == (SOURCE_OPEN | TARGET_OPEN)));

	// Enable dump code button
	EnableWindow(GetDlgItem(hDlg, IDDUMPCODE), TRUE);

	// Find section headers
	source.section = (IMAGE_SECTION_HEADER *)((DWORD)(source.PE) + (source.PE->FileHeader.SizeOfOptionalHeader + (sizeof(IMAGE_NT_HEADERS) - sizeof(IMAGE_OPTIONAL_HEADER))));

	// Determine code size
	wsprintf(buffer, "%X", source.section[0].SizeOfRawData - source.PE->OptionalHeader.AddressOfEntryPoint + source.section[0].VirtualAddress);
	SetDlgItemText(hDlg, IDC_SIZEOFBLOCK, buffer);
}

void CloseSourceFile()
{
	// Close handles
	UnmapViewOfFile(source.dos);
	CloseHandle(sourceFileMapping);
	CloseHandle(sourceFile);

	// Reset flag
	flags &= ~SOURCE_OPEN;
}

OPENFILENAME openTargetFile = {sizeof(OPENFILENAME), NULL, NULL, "EXE files\0*.EXE\0", 0, 0, 0, (char *)&targetFileName, sizeof targetFileName, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

void BrowseTargetFile(HWND hDlg)
{
	openTargetFile.hwndOwner = hDlg;
	openTargetFile.hInstance = hInst;

	// Check if file is open, and if so, close it
	if (flags & TARGET_OPEN)
	{
		// Remove sections from combobox
		SendDlgItemMessage(hDlg, IDC_TARGETSECTION, CB_RESETCONTENT, 0, 0);

		CloseTargetFile();
	}

	if(!GetOpenFileName(&openTargetFile))
		return;

	// Set filename in editbox
	SetDlgItemText(hDlg, IDC_TARGET, targetFileName);
}

void OpenTargetFile(HWND hDlg)
{
	unsigned int i;

	// Open file and memory-map it
	targetFile = CreateFile(targetFileName, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING, 0, NULL);

	targetFileMapping = CreateFileMapping(targetFile, NULL, PAGE_READWRITE, 0, 0, NULL);
	
	if ((target.dos = MapViewOfFile(targetFileMapping, FILE_MAP_WRITE, 0, 0, 0)) == 0)
	{
		MessageBox(hDlg, "Error opening target file", "Error", MB_OK);
		CloseTargetFile();
		return;
	}

	// Parse header
	if (target.dos->e_magic != 0x5A4D)
	{
		MessageBox(hDlg, "No MZ signature found. Incorrect file type", "Error", MB_OK);
		CloseTargetFile();
		return;
	}

	// Set pointer to start of PE header
	target.PE = (IMAGE_NT_HEADERS *)((target.dos->e_lfanew) + (DWORD)(target.dos));

	if (target.PE->Signature != 0x4550)
	{
		MessageBox(hDlg, "No PE signature found. Incorrect file type", "Error", MB_OK);
		CloseTargetFile();
		return;
	}

	// File is correct. Set flag to indicate it's opened
	flags |= TARGET_OPEN;

	// Enable inject button, if possible
	EnableWindow(GetDlgItem(hDlg, IDINJECT), (flags == (SOURCE_OPEN | TARGET_OPEN)));

	// Find section headers
	target.section = (IMAGE_SECTION_HEADER *)((DWORD)(target.PE) + (target.PE->FileHeader.SizeOfOptionalHeader + (sizeof(IMAGE_NT_HEADERS) - sizeof(IMAGE_OPTIONAL_HEADER))));

	// Parse sections
	for(i = 0; i < target.PE->FileHeader.NumberOfSections; i++)
		SendDlgItemMessage(hDlg, IDC_TARGETSECTION, CB_ADDSTRING, 0, (LPARAM)&target.section[i].Name);

	// Enable ComboBox
	EnableWindow(GetDlgItem(hDlg, IDC_TARGETSECTION), TRUE);
}

void CloseTargetFile()
{
	// Close handles
	UnmapViewOfFile(target.dos);
	CloseHandle(targetFileMapping);
	CloseHandle(targetFile);

	// Reset flag
	flags &= ~TARGET_OPEN;
}

void OptimizePEHeader(HWND hDlg)
{
	if (!(flags & TARGET_OPEN))
	{
		MessageBox(hDlg, "No target file opened!", "Error", MB_OK);
		return;
	}

	if (target.dos->e_lfanew == 0x40)
	{
		MessageBox(hDlg, "PE Header is already optimal", "Error", MB_OK);
		return;
	}

	MessageBox(hDlg, "Not supported yet", "Patience please", MB_OK);
}

void EntryPoint()
{
	hInst = GetModuleHandle(NULL);

	DialogBoxParam(hInst, MAKEINTRESOURCE(IDD_GRABCODE), NULL, (DLGPROC)DlgBoxProc, 0);
}

/*void EntryPoint()
{
	WinMain(GetModuleHandle(NULL), 0, 0, 0);
}

int APIENTRY WinMain(HINSTANCE hInstance,
                     HINSTANCE hPrevInstance,
                     LPSTR     lpCmdLine,
                     int       nCmdShow)
{
	hInst = hInstance;

	return DialogBoxParam(hInstance, MAKEINTRESOURCE(IDD_GRABCODE), NULL, (DLGPROC)DlgBoxProc, 0);
}*/