#include "stdafx.h"

ULONG_PTR RvaToOffset(IMAGE_NT_HEADERS *NT, ULONG_PTR Rva)
{
	ULONG_PTR Offset = Rva, Limit;
	IMAGE_SECTION_HEADER *Img;
	WORD i;

	Img = IMAGE_FIRST_SECTION(NT);

	if (Rva < Img->PointerToRawData)
		return Rva;

	for (i = 0; i < NT->FileHeader.NumberOfSections; i++)
	{
		if (Img[i].SizeOfRawData)
			Limit = Img[i].SizeOfRawData;
		else
			Limit = Img[i].Misc.VirtualSize;

		if (Rva >= Img[i].VirtualAddress &&
			Rva < (Img[i].VirtualAddress + Limit))
		{
			if (Img[i].PointerToRawData != 0)
			{
				Offset -= Img[i].VirtualAddress;
				Offset += Img[i].PointerToRawData;
			}

			return Offset;
		}
	}

	return 0;
}

BOOL GetExport(BYTE *PE, ULONG_PTR *EP, WORD *Ordinal, CHAR *FuncName) 
{
	IMAGE_DOS_HEADER *ET_DOS;
	IMAGE_NT_HEADERS *ET_NT;
	IMAGE_EXPORT_DIRECTORY *Export;
	ULONG_PTR ET, *Functions;
	PSTR *Names;
	WORD *Ordinals;
	CHAR *ApiName;

	__try
	{
		ET_DOS = (IMAGE_DOS_HEADER *)(ULONG_PTR) PE;

		ET_NT = (IMAGE_NT_HEADERS *)(ULONG_PTR) &PE[ET_DOS->e_lfanew];

		ET = ET_NT->OptionalHeader.DataDirectory
			[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;

		if (ET == NULL)
			return FALSE;

		ET = RvaToOffset(ET_NT, ET);

		if (ET == 0) return FALSE;

		Export = (IMAGE_EXPORT_DIRECTORY *)(ET + (ULONG_PTR) PE);

		Functions = (ULONG_PTR *)(RvaToOffset(ET_NT, 
			Export->AddressOfFunctions) + (ULONG_PTR) PE);
		Ordinals = (WORD *)(RvaToOffset(ET_NT, 
			Export->AddressOfNameOrdinals) + (ULONG_PTR) PE);
		Names = (PSTR *)(RvaToOffset(ET_NT, 
			Export->AddressOfNames) + (ULONG_PTR) PE);

		if (EP != NULL && *EP != 0)
		{
			for (WORD x = 0; x < Export->NumberOfFunctions; x++)
			{
				if (*EP == Functions[x])
				{
					if (Ordinal) *Ordinal = (WORD) (x + Export->Base);

					if (FuncName != NULL)
					{
						for (WORD i = 0; i < Export->NumberOfNames; i++)
						{
							if (Ordinals[i] == x)
							{
								ApiName = (char *) RvaToOffset(ET_NT, (ULONG_PTR) Names[i]);

								if (ApiName != NULL)
								{
									ApiName = (char *)((ULONG_PTR) ApiName + 
										(ULONG_PTR) PE);
									strcpy(FuncName, ApiName);
									break;
								}
							}
						}
					}

					return TRUE;
				}	
			}

			return FALSE;
		}
		else
		{
			if (FuncName == NULL || FuncName[0] == 0)
			{
				if (Ordinal == NULL || *Ordinal == 0)
					return FALSE;

				if (*Ordinal < Export->Base ||
					*Ordinal > (Export->Base + (Export->NumberOfFunctions - 1)))
					return FALSE;

				WORD FuncEntry = (WORD) (*Ordinal - Export->Base);

				if (EP) *EP = Functions[FuncEntry];

				if (FuncName != NULL)
				{
					for (DWORD i = 0; i < Export->NumberOfNames; i++)
					{
						if (Ordinals[i] == FuncEntry)
						{
							ApiName = (char *) RvaToOffset(ET_NT, (ULONG_PTR) Names[i]);

							if (ApiName != NULL)
							{
								ApiName = (char *)((ULONG_PTR) ApiName + (ULONG_PTR) PE);
								strcpy(FuncName, ApiName);
								break;
							}
						}
					}

				}

				return TRUE;
			}
			else
			{
				for (DWORD x = 0; x < Export->NumberOfFunctions; x++)
				{
					if (Functions[x] == 0)
						continue;

					for (DWORD i = 0; i < Export->NumberOfNames; i++)
					{
						if (Ordinals[i] == x)
						{
							ApiName = (char *) RvaToOffset(ET_NT, (ULONG_PTR) Names[i]);

							if (ApiName != NULL)
							{
								ApiName = (char *)((ULONG_PTR) ApiName + (ULONG_PTR) PE);

								if (strcmp(ApiName, FuncName) == 0)
								{
									if (Ordinal) *Ordinal = (WORD) (x + Export->Base);
									if (EP) *EP = Functions[x];
									return TRUE;
								}
							}
						}
					}

				}

				return FALSE;
			}
		}

	}

	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		return FALSE;
	}

	return FALSE;
}

ULONG_PTR GetExportRva(IN TCHAR *FileName, IN CHAR *FunctionName)
{
	HMODULE hModule = LoadLibrary(FileName);

	if (hModule == NULL)
		return 0;

	ULONG_PTR VA = (ULONG_PTR) GetProcAddress(hModule,
		FunctionName);

	FreeLibrary(hModule);

	return (ULONG_PTR) (VA - (ULONG_PTR) hModule);
}

BOOL InjectModule(IN ULONG_PTR ProcessID, IN TCHAR *ModuleName, 
				  OUT ULONG_PTR *BaseAddress OPTIONAL)
{
	HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE,
		ProcessID);

	if (hProcess == NULL)
		return FALSE;

	HANDLE hFile = CreateFile(ModuleName, GENERIC_READ, FILE_SHARE_READ,
		NULL, OPEN_EXISTING, 0, NULL);

	if (hFile == INVALID_HANDLE_VALUE)
	{
		CloseHandle(hProcess);
		return FALSE;
	}

	UINT PE_Size = GetFileSize(hFile, NULL);
	BYTE *PE_Buffer = (BYTE *) VirtualAlloc(NULL, PE_Size, 
		MEM_COMMIT, PAGE_READWRITE);

	if (PE_Buffer == NULL)
	{
		CloseHandle(hProcess);
		CloseHandle(hFile);
		return FALSE;
	}

	DWORD BR;

	if (!ReadFile(hFile, PE_Buffer, PE_Size, &BR, NULL))
	{
		VirtualFree(PE_Buffer, 0, MEM_RELEASE);
		CloseHandle(hProcess);
		CloseHandle(hFile);
		return FALSE;
	}

	CloseHandle(hFile);

	BYTE *PE_Image;

	IMAGE_DOS_HEADER *ImgDosHdr;
	IMAGE_NT_HEADERS *ImgNtHdrs;

	ULONG_PTR ImgBase, Delta, Reloc_Offset;
	IMAGE_BASE_RELOCATION *ImgBaseReloc;
	WORD *wData;
	UINT i, nItems;
	ULONG_PTR Offset;
	DWORD Type;
	ULONG_PTR *Block, BlockOffs;

	ULONG_PTR IT_Offset;
	IMAGE_IMPORT_DESCRIPTOR *ImgImpDescr;
	UINT x = 0, y = 0;
	CHAR *DllName;
	ULONG_PTR *Thunks, *FThunks;
	IMAGE_IMPORT_BY_NAME *ImgImpName;

	_try
	{
		ImgDosHdr = (IMAGE_DOS_HEADER *) PE_Buffer;

		ImgNtHdrs = (IMAGE_NT_HEADERS *) (ImgDosHdr->e_lfanew +
			(ULONG_PTR) PE_Buffer);

		if (ImgDosHdr->e_magic != IMAGE_DOS_SIGNATURE ||
			ImgNtHdrs->Signature != IMAGE_NT_SIGNATURE)
		{
			VirtualFree(PE_Buffer, 0, MEM_RELEASE);
			CloseHandle(hProcess);
			return FALSE;
		}

		ImgBase = (ULONG_PTR) VirtualAllocEx(hProcess, 
			(PVOID) 0, //ImgNtHdrs->OptionalHeader.ImageBase,
			ImgNtHdrs->OptionalHeader.SizeOfImage, MEM_COMMIT, 
			PAGE_EXECUTE_READWRITE);

		if (ImgBase == NULL)
		{
			VirtualFree(PE_Buffer, 0, MEM_RELEASE);
			CloseHandle(hProcess);
			return FALSE;
		}

		//
		// is file image base == memory image base?
		// if not, relocate
		//

		if (ImgNtHdrs->OptionalHeader.ImageBase != ImgBase)
		{
			Delta = ImgBase - ImgNtHdrs->OptionalHeader.ImageBase;

			if (!ImgNtHdrs->OptionalHeader.DataDirectory
				[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress ||
				!(Reloc_Offset = RvaToOffset(ImgNtHdrs, 
				ImgNtHdrs->OptionalHeader.DataDirectory
				[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress)))
			{
				VirtualFreeEx(hProcess, (LPVOID) ImgBase, 0, MEM_RELEASE);
				VirtualFree(PE_Buffer, 0, MEM_RELEASE);
				CloseHandle(hProcess);
				return FALSE;
			}

			ImgBaseReloc = (IMAGE_BASE_RELOCATION *) 
				(Reloc_Offset + (ULONG_PTR) PE_Buffer);

			do
			{
				if (!ImgBaseReloc->SizeOfBlock)
					break;

				nItems = (ImgBaseReloc->SizeOfBlock - 
					IMAGE_SIZEOF_BASE_RELOCATION) / sizeof (WORD);

				wData = (WORD *)(IMAGE_SIZEOF_BASE_RELOCATION + 
					(ULONG_PTR) ImgBaseReloc);

				for (i = 0; i < nItems; i++)
				{
					Offset = (*wData & 0xFFF) + ImgBaseReloc->VirtualAddress;

					Type = *wData >> 12;

					if (Type != IMAGE_REL_BASED_ABSOLUTE)
					{
						BlockOffs = RvaToOffset(ImgNtHdrs, Offset);

						if (BlockOffs == NULL)
						{
							VirtualFreeEx(hProcess, (LPVOID) ImgBase, 0, MEM_RELEASE);
							VirtualFree(PE_Buffer, 0, MEM_RELEASE);
							CloseHandle(hProcess);
							return FALSE;
						}

						Block = (DWORD *)(BlockOffs + (ULONG_PTR) PE_Buffer);

						*Block += Delta;
					}

					wData++;
				} 

				ImgBaseReloc = (PIMAGE_BASE_RELOCATION) wData;

			} while (*(DWORD *) wData);
		}

		//
		// fill the import addres table
		//

		if (ImgNtHdrs->OptionalHeader.DataDirectory
			[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress)
		{
			IT_Offset = RvaToOffset(ImgNtHdrs, 
				ImgNtHdrs->OptionalHeader.DataDirectory
				[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);

			ImgImpDescr = (IMAGE_IMPORT_DESCRIPTOR *) (IT_Offset +
				(ULONG_PTR) PE_Buffer);


			// for each descriptor
			while (ImgImpDescr[x].FirstThunk != 0)
			{
				DllName = (CHAR *) (RvaToOffset(ImgNtHdrs, 
					ImgImpDescr[x].Name) + (ULONG_PTR) PE_Buffer);

				Thunks = (ULONG_PTR *) (RvaToOffset(ImgNtHdrs,
					ImgImpDescr[x].OriginalFirstThunk != 0 ?
					ImgImpDescr[x].OriginalFirstThunk : 
				ImgImpDescr[x].FirstThunk) + (ULONG_PTR) PE_Buffer);

				FThunks = (ULONG_PTR *) (RvaToOffset(ImgNtHdrs,
					ImgImpDescr[x].FirstThunk) + (ULONG_PTR) PE_Buffer);

				y = 0;

				//
				// every imported function of the module
				//

				while (Thunks[y] != 0)
				{
					//
					// imported by ordinal?
					//

					if (Thunks[y] & IMAGE_ORDINAL_FLAG)
					{
						FThunks[y] = (ULONG_PTR) GetProcAddress(
							GetModuleHandle(DllName),
							(LPCSTR) (Thunks[y] - IMAGE_ORDINAL_FLAG));

						y++;

						continue;
					}

					ImgImpName = (IMAGE_IMPORT_BY_NAME *) (RvaToOffset(
						ImgNtHdrs, Thunks[y]) + (ULONG_PTR) PE_Buffer);

					FThunks[y] = (ULONG_PTR) GetProcAddress(GetModuleHandle(DllName), 
						(LPCSTR) &ImgImpName->Name);

					y++;
				}

				x++;
			}


		}
	}
	_except (EXCEPTION_EXECUTE_HANDLER)
	{
		VirtualFreeEx(hProcess, (LPVOID) ImgBase, 0, MEM_RELEASE);
		VirtualFree((LPVOID) PE_Buffer, 0, MEM_RELEASE);
		CloseHandle(hProcess);
		return FALSE;
	}

	//
	// create virtual image of the PE file
	//

	PE_Image = (BYTE *) VirtualAlloc(NULL, 
		ImgNtHdrs->OptionalHeader.SizeOfImage, 
		MEM_COMMIT, PAGE_READWRITE);

	if (PE_Image == NULL)
	{
		VirtualFreeEx(hProcess, (LPVOID) ImgBase, 0, MEM_RELEASE);
		VirtualFree(PE_Buffer, 0, MEM_RELEASE);
		CloseHandle(hProcess);
		return FALSE;
	}

	ZeroMemory(PE_Image, ImgNtHdrs->OptionalHeader.SizeOfImage);

	//
	// copy headers
	//

	IMAGE_SECTION_HEADER *Sect;

	Sect = IMAGE_FIRST_SECTION(ImgNtHdrs);

	RtlCopyMemory(PE_Image, PE_Buffer, Sect[0].PointerToRawData);

	//
	// map sections
	//

	for (UINT j = 0; j < ImgNtHdrs->FileHeader.NumberOfSections; j++)
	{
		BYTE *Source = (BYTE *)(Sect[j].PointerToRawData + 
			(ULONG_PTR) PE_Buffer);

		BYTE *Dest = (BYTE *)(Sect[j].VirtualAddress + 
			(ULONG_PTR) PE_Image);

		RtlCopyMemory(Dest, Source, Sect[j].SizeOfRawData);
	}

	BOOL bRet = WriteProcessMemory(hProcess, (LPVOID) ImgBase, 
		PE_Image, ImgNtHdrs->OptionalHeader.SizeOfImage, &BR);

	if (!bRet)
		VirtualFreeEx(hProcess, (LPVOID) ImgBase, 0, MEM_RELEASE);

	VirtualFree(PE_Image, 0, MEM_RELEASE);
	VirtualFree(PE_Buffer, 0, MEM_RELEASE);
	CloseHandle(hProcess);

	if (bRet == TRUE && BaseAddress != NULL)
		*BaseAddress = ImgBase; 

	return bRet;
}