#include "asm.h"

#ifndef WIN32

#define _GNU_SOURCE

#include <dlfcn.h>

#include <string.h>

#define REG_EAX			0

#define REG_ECX			1

#define REG_EDX			2

#define REG_EBX			3

#define IA32_MOV_REG_IMM		0xB8	// encoding is +r <imm32>

#endif

extern void Msg( const char *, ... );

/**

* Checks if a call to a fpic thunk has just been written into dest.

* If found replaces it with a direct mov that sets the required register to the value of pc.

*

* @param dest		Destination buffer where a call opcode + addr (5 bytes) has just been written.

* @param pc		The program counter value that needs to be set (usually the next address from the source).

* @noreturn

*/

void check_thunks(unsigned char *dest, unsigned char *pc)

{

#if defined WIN32

	return;

#else

	/* Step write address back 4 to the start of the function address */

	unsigned char *writeaddr = dest - 4;

	unsigned char *calloffset = *(unsigned char **)writeaddr;

	unsigned char *calladdr = (unsigned char *)(dest + (unsigned int)calloffset);

	/* Lookup name of function being called */

	if ((*calladdr == 0x8B) && (*(calladdr+2) == 0x24) && (*(calladdr+3) == 0xC3))

	{

		//a thunk maybe?

		char movByte = IA32_MOV_REG_IMM;

		/* Calculate the correct mov opcode */

		switch (*(calladdr+1))

		{

		case 0x04:

			{

				movByte += REG_EAX;

				break;

			}

		case 0x1C:

			{

				movByte += REG_EBX;

				break;

			}

		case 0x0C:

			{

				movByte += REG_ECX;

				break;

			}

		case 0x14:

			{

				movByte += REG_EDX;

				break;

			}

		default:

			{

				Msg("Unknown thunk: %c\n", *(calladdr+1));

				break;

			}

		}

		/* Move our write address back one to where the call opcode was */

		writeaddr--;

		/* Write our mov */

		*writeaddr = movByte;

		writeaddr++;

		/* Write the value - The provided program counter value */

		*(void **)writeaddr = (void *)pc;

		writeaddr += 4;

	}

	return;

#endif

}

//if dest is NULL, returns minimum number of bytes needed to be copied

//if dest is not NULL, it will copy the bytes to dest as well as fix CALLs and JMPs

//http://www.devmaster.net/forums/showthread.php?t=2311

int copy_bytes(unsigned char *func, unsigned char* dest, int required_len) {

	int bytecount = 0;

	while(bytecount < required_len && *func != 0xCC)

	{

		// prefixes F0h, F2h, F3h, 66h, 67h, D8h-DFh, 2Eh, 36h, 3Eh, 26h, 64h and 65h

		int operandSize = 4;

		int FPU = 0;

		int twoByte = 0;

		unsigned char opcode = 0x90;

		unsigned char modRM = 0xFF;

		while(*func == 0xF0 ||

			  *func == 0xF2 ||

			  *func == 0xF3 ||

			 (*func & 0xFC) == 0x64 ||

			 (*func & 0xF8) == 0xD8 ||

			 (*func & 0x7E) == 0x62)

		{

			if(*func == 0x66)

			{

				operandSize = 2;

			}

			else if((*func & 0xF8) == 0xD8)

			{

				FPU = *func;

				if (dest)

						*dest++ = *func++;

				else

						func++;

				bytecount++;

				break;

			}

			if (dest)

				*dest++ = *func++;

			else

				func++;

			bytecount++;

		}

		// two-byte opcode byte

		if(*func == 0x0F)

		{

			twoByte = 1;

			if (dest)

				*dest++ = *func++;

			else

				func++;

			bytecount++;

		}

		// opcode byte

		opcode = *func++;

		if (dest) *dest++ = opcode;

		bytecount++;

		// mod R/M byte

		modRM = 0xFF;

		if(FPU)

		{

			if((opcode & 0xC0) != 0xC0)

			{

				modRM = opcode;

			}

		}

		else if(!twoByte)

		{

			if((opcode & 0xC4) == 0x00 ||

			   ((opcode & 0xF4) == 0x60 && ((opcode & 0x0A) == 0x02 || (opcode & 0x09) == 0x09)) ||

			   (opcode & 0xF0) == 0x80 ||

			   ((opcode & 0xF8) == 0xC0 && (opcode & 0x0E) != 0x02) ||

			   (opcode & 0xFC) == 0xD0 ||

			   (opcode & 0xF6) == 0xF6)

			{

				modRM = *func++;

				if (dest) *dest++ = modRM;

				bytecount++;

			}

		}

		else

		{

			if(((opcode & 0xF0) == 0x00 && (opcode & 0x0F) >= 0x04 && (opcode & 0x0D) != 0x0D) ||

			   (opcode & 0xF0) == 0x30 ||

			   opcode == 0x77 ||

			   (opcode & 0xF0) == 0x80 ||

			   ((opcode & 0xF0) == 0xA0 && (opcode & 0x07) <= 0x02) ||

			   (opcode & 0xF8) == 0xC8)

			{

				// No mod R/M byte

			}

			else

			{

				modRM = *func++;

				if (dest) *dest++ = modRM;

				bytecount++;

			}

		}

		// SIB

		if((modRM & 0x07) == 0x04 &&

		   (modRM & 0xC0) != 0xC0)

		{

			if (dest)

				*dest++ = *func++;   //SIB

			else

				func++;

			bytecount++;

		}

		// mod R/M displacement

		// Dword displacement, no base

	if((modRM & 0xC5) == 0x05) {

		if (dest) {

			*(unsigned int*)dest = *(unsigned int*)func;

			dest += 4;

		}

		func += 4;

		bytecount += 4;

	}

		// Byte displacement

	if((modRM & 0xC0) == 0x40) {

		if (dest)

			*dest++ = *func++;

		else

			func++;

		bytecount++;

	}

		// Dword displacement

	if((modRM & 0xC0) == 0x80) {

		if (dest) {

			*(unsigned int*)dest = *(unsigned int*)func;

			dest += 4;

		}

		func += 4;

		bytecount += 4;

	}

		// immediate

		if(FPU)

		{

			// Can't have immediate operand

		}

		else if(!twoByte)

		{

			if((opcode & 0xC7) == 0x04 ||

			   (opcode & 0xFE) == 0x6A ||   // PUSH/POP/IMUL

			   (opcode & 0xF0) == 0x70 ||   // Jcc

			   opcode == 0x80 ||

			   opcode == 0x83 ||

			   (opcode & 0xFD) == 0xA0 ||   // MOV

			   opcode == 0xA8 ||			// TEST

			   (opcode & 0xF8) == 0xB0 ||   // MOV

			   (opcode & 0xFE) == 0xC0 ||   // RCL

			   opcode == 0xC6 ||			// MOV

			   opcode == 0xCD ||			// INT

			   (opcode & 0xFE) == 0xD4 ||   // AAD/AAM

			   (opcode & 0xF8) == 0xE0 ||   // LOOP/JCXZ

			   opcode == 0xEB ||

			   (opcode == 0xF6 && (modRM & 0x30) == 0x00))   // TEST

			{

				if (dest)

					*dest++ = *func++;

				else

					func++;

				bytecount++;

			}

		else if((opcode & 0xF7) == 0xC2) // RET

			{

				if (dest) {

					*(unsigned short*)dest = *(unsigned short*)func;

					dest += 2;

				}

				func += 2;

				bytecount += 2;

			}

			else if((opcode & 0xFC) == 0x80 ||

					(opcode & 0xC7) == 0x05 ||

					(opcode & 0xF8) == 0xB8 ||

					(opcode & 0xFE) == 0xE8 ||	  // CALL/Jcc

					(opcode & 0xFE) == 0x68 ||

					(opcode & 0xFC) == 0xA0 ||

					(opcode & 0xEE) == 0xA8 ||

					opcode == 0xC7 ||

					(opcode == 0xF7 && (modRM & 0x30) == 0x00))

			{

				if (dest) {

					//Fix CALL/JMP offset

					if ((opcode & 0xFE) == 0xE8) {

						if (operandSize == 4)

						{

							*(long*)dest = ((func + *(long*)func) - dest);

							//pRED* edit. func is the current address of the call address, +4 is the next instruction, so the value of $pc

							check_thunks(dest+4, func+4);

						}

						else

							*(short*)dest = ((func + *(short*)func) - dest);

					} else {

						if (operandSize == 4)

							*(unsigned long*)dest = *(unsigned long*)func;

						else

							*(unsigned short*)dest = *(unsigned short*)func;

					}

					dest += operandSize;

				}

				func += operandSize;

				bytecount += operandSize;

			}

		}

		else

		{

			if(opcode == 0xBA ||			// BT

			   opcode == 0x0F ||			// 3DNow!

			   (opcode & 0xFC) == 0x70 ||   // PSLLW

			   (opcode & 0xF7) == 0xA4 ||   // SHLD

			   opcode == 0xC2 ||

			   opcode == 0xC4 ||

			   opcode == 0xC5 ||

			   opcode == 0xC6)

			{

				if (dest)

					*dest++ = *func++;

				else

					func++;

			}

			else if((opcode & 0xF0) == 0x80) // Jcc -i

			{

				if (dest) {

					if (operandSize == 4)

						*(unsigned long*)dest = *(unsigned long*)func;

					else

						*(unsigned short*)dest = *(unsigned short*)func;

					dest += operandSize;

				}

				func += operandSize;

				bytecount += operandSize;

			}

		}

	}

	return bytecount;

}

//insert a specific JMP instruction at the given location

void inject_jmp(void* src, void* dest) {

	*(unsigned char*)src = OP_JMP;

	*(long*)((unsigned char*)src+1) = (long)((unsigned char*)dest - ((unsigned char*)src + OP_JMP_SIZE));

}

//fill a given block with NOPs

void fill_nop(void* src, unsigned int len) {

	unsigned char* src2 = (unsigned char*)src;

	while (len) {

		*src2++ = OP_NOP;

		--len;

	}

}

void* eval_jump(void* src) {

	unsigned char* addr = (unsigned char*)src;

	if (!addr) return 0;

	//import table jump

	if (addr[0] == OP_PREFIX && addr[1] == OP_JMP_SEG) {

		addr += 2;

		addr = *(unsigned char**)addr;

		//TODO: if addr points into the IAT

		return *(void**)addr;

	}

	//8bit offset

	else if (addr[0] == OP_JMP_BYTE) {

		addr = &addr[OP_JMP_BYTE_SIZE] + *(char*)&addr[1];

		//mangled 32bit jump?

		if (addr[0] == OP_JMP) {

			addr = addr + *(int*)&addr[1];

		}

		return addr;

	}

	/*

	//32bit offset

	else if (addr[0] == OP_JMP) {

		addr = &addr[OP_JMP_SIZE] + *(int*)&addr[1];

	}

	*/

	return addr;

}

/*

from ms detours package

static bool detour_is_imported(PBYTE pbCode, PBYTE pbAddress)

{

	MEMORY_BASIC_INFORMATION mbi;

	VirtualQuery((PVOID)pbCode, &mbi, sizeof(mbi));

	__try {

		PIMAGE_DOS_HEADER pDosHeader = (PIMAGE_DOS_HEADER)mbi.AllocationBase;

		if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE) {

			return false;

		}

		PIMAGE_NT_HEADERS pNtHeader = (PIMAGE_NT_HEADERS)((PBYTE)pDosHeader +

														  pDosHeader->e_lfanew);

		if (pNtHeader->Signature != IMAGE_NT_SIGNATURE) {

			return false;

		}

		if (pbAddress >= ((PBYTE)pDosHeader +

						  pNtHeader->OptionalHeader

						  .DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress) &&

			pbAddress < ((PBYTE)pDosHeader +

						 pNtHeader->OptionalHeader

						 .DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress +

						 pNtHeader->OptionalHeader

						 .DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].Size)) {

			return true;

		}

		return false;

	}

	__except(EXCEPTION_EXECUTE_HANDLER) {

		return false;

	}

}

*/