#include "VirtualMachine.h"


XED_REG_ENUM VmOperandSizeToRegEnumBase(VM_OPERAND_SIZE_ENUM OperandSize)
{
	switch (OperandSize)
	{
	case VM_OPSIZE_8: return XED_REG_GPR8_FIRST;
	case VM_OPSIZE_16: return XED_REG_GPR16_FIRST;
	case VM_OPSIZE_32: return XED_REG_GPR32_FIRST;
	case VM_OPSIZE_64: return XED_REG_GPR64_FIRST;
	}
	return XED_REG_INVALID;
}
XED_REG_ENUM VmGetRegOfSize(XED_REG_ENUM Reg, VM_OPERAND_SIZE_ENUM OperandSize)
{
	if (Reg >= XED_REG_GPR8_FIRST && Reg <= XED_REG_GPR8_LAST)
		return (XED_REG_ENUM)(VmOperandSizeToRegEnumBase(OperandSize) + (Reg - XED_REG_GPR8_FIRST));
	else if (Reg >= XED_REG_GPR16_FIRST && Reg <= XED_REG_GPR16_LAST)
		return (XED_REG_ENUM)(VmOperandSizeToRegEnumBase(OperandSize) + (Reg - XED_REG_GPR16_FIRST));
	else if (Reg >= XED_REG_GPR32_FIRST && Reg <= XED_REG_GPR32_LAST)
		return (XED_REG_ENUM)(VmOperandSizeToRegEnumBase(OperandSize) + (Reg - XED_REG_GPR32_FIRST));
	else if (Reg >= XED_REG_GPR64_FIRST && Reg <= XED_REG_GPR64_LAST)
		return (XED_REG_ENUM)(VmOperandSizeToRegEnumBase(OperandSize) + (Reg - XED_REG_GPR64_FIRST));
	return XED_REG_INVALID;
}
VM_OPERAND_SIZE_ENUM VmXRegToOpSize(XED_REG_ENUM Reg)
{
	if (Reg >= XED_REG_GPR8_FIRST && Reg <= XED_REG_GPR8_LAST)
		return VM_OPSIZE_8;
	else if (Reg >= XED_REG_GPR16_FIRST && Reg <= XED_REG_GPR16_LAST)
		return VM_OPSIZE_16;
	else if (Reg >= XED_REG_GPR32_FIRST && Reg <= XED_REG_GPR32_LAST)
		return VM_OPSIZE_32;
	else if (Reg >= XED_REG_GPR64_FIRST && Reg <= XED_REG_GPR64_LAST)
		return VM_OPSIZE_64;
	return VM_OPSIZE_COUNT;
}
XED_REG_ENUM VmIRegToXReg(VM_IREG_ENUM IReg, VM_OPERAND_SIZE_ENUM OperandSize)
{
	switch (IReg)
	{
	case VM_IREG_1: return VmGetRegOfSize(XED_REG_RAX, OperandSize);
	case VM_IREG_2: return VmGetRegOfSize(XED_REG_RBX, OperandSize);
	case VM_IREG_3: return VmGetRegOfSize(XED_REG_RCX, OperandSize);
	}
	return XED_REG_INVALID;

	//Less portable version.
	/*if (OperandSize == VM_OPSIZE_8)
	{
		return (XED_REG_ENUM)(XED_REG_AL + IReg);
	}
	return (XED_REG_ENUM)(XED_REG_AX + (16 * (OperandSize - 1)) + IReg);*/
}
VM_REG_ENUM VmXRegToVRegId(XED_REG_ENUM Reg)
{
	return (VM_REG_ENUM)(VmGetRegOfSize(Reg, VM_OPSIZE_64) - XED_REG_RAX);
}

CONST CHAR* VmIClassToString(VM_ICLASS_ENUM IClass)
{
	if (IClass >= VM_ICLASS_COUNT)
		return "VM_ICLASS_INVALID";
	else if (IClass == VM_ICLASS_EXIT)
		return "VM_ICLASS_EXIT";
	else if (IClass >= VM_ICLASS_LD_IREG_MEM_START && IClass <= VM_ICLASS_LD_IREG_MEM_END)
		return "VM_ICLASS_LD_IREG_MEM";
	else if (IClass >= VM_ICLASS_ST_IREG_MEM_START && IClass <= VM_ICLASS_ST_IREG_MEM_END)
		return "VM_ICLASS_ST_IREG_MEM";
	else if (IClass >= VM_ICLASS_LD_IREG_MEM_START && IClass <= VM_ICLASS_LD_IREG_REG_END)
		return "VM_ICLASS_LD_IREG_REG";
	else if (IClass >= VM_ICLASS_ST_IREG_REG_START && IClass <= VM_ICLASS_ST_IREG_REG_END)
		return "VM_ICLASS_ST_IREG_REG";
	else if (IClass >= VM_ICLASS_LD_IREG_IMM_START && IClass <= VM_ICLASS_LD_IREG_IMM_END)
		return "VM_ICLASS_LD_IREG_IMM";
}

PUCHAR VmHandlerEpilogue(UINT InstructionSize, PUINT OutSize, XED_REG_ENUM Vip, XED_REG_ENUM HandlerTableReg)
{
	//	add rdx, InstructionSize + Prologue Size
	//	movzx r8,byte ptr[rdx]
	//	jmp qword ptr[rsi+r8*8h]

	XED_ENCODER_INSTRUCTION InstList[3];
	InstructionSize += VM_HANDLER_EPILOGUE_SIZE;
	XedInst2(&InstList[0], XedGlobalMachineState, XED_ICLASS_ADD, 64, XedReg(Vip), XedImm0(InstructionSize, 32));
	XedInst2(&InstList[1], XedGlobalMachineState, XED_ICLASS_MOVZX, 64, XedReg(XED_REG_R8), XedMemB(Vip, 16));
	XedInst1(&InstList[2], XedGlobalMachineState, XED_ICLASS_JMP, 64, XedMemBISD(HandlerTableReg, XED_REG_R8, 8, XedDisp(0, 0), 64));
	PUCHAR Ret = XedEncodeInstructions(InstList, 3, OutSize);
	if (*OutSize == VM_HANDLER_EPILOGUE_SIZE)
		return Ret;

	delete[] Ret;
	InstructionSize -= VM_HANDLER_EPILOGUE_SIZE;
	InstructionSize += *OutSize;
	XedInst2(&InstList[0], XedGlobalMachineState, XED_ICLASS_ADD, 64, XedReg(Vip), XedImm0(InstructionSize, 32));
	XedInst2(&InstList[1], XedGlobalMachineState, XED_ICLASS_MOVZX, 64, XedReg(XED_REG_R8), XedMemB(Vip, 16));
	XedInst1(&InstList[2], XedGlobalMachineState, XED_ICLASS_JMP, 64, XedMemBISD(HandlerTableReg, XED_REG_R8, 8, XedDisp(0, 0), 64));
	return XedEncodeInstructions(InstList, 3, OutSize);
}

PUCHAR VmHandlerEnter(PUINT OutSize)
{

}
PUCHAR VmHandlerExit(PUINT OutSize)
{

}

PUCHAR VmHandlerIRegMem_B(VM_IREG_ENUM IReg, VM_OPERAND_SIZE_ENUM OperandSize, BOOL Load, PUINT OutSize, XED_REG_ENUM XReg1)
{
	/*
	* movzx r8,byte ptr[rdx+2]
	* mov r8, qword ptr[rbp+r8*8]
	* mov (ireg), (size) ptr[r8]
	*/
	UINT OpSizeBits = VmOpSizeToBits(OperandSize);
	XED_ENCODER_INSTRUCTION InstList[3];
	XedInst2(&InstList[0], XedGlobalMachineState, XED_ICLASS_MOVZX, 64, XedReg(XReg1), XedMemBD(XED_REG_RDX, XedDisp(2, 8), 8));
	XedInst2(&InstList[1], XedGlobalMachineState, XED_ICLASS_MOV, 64, XedReg(XReg1), XedMemBISD(XED_REG_RBP, XReg1, 8, XedDisp(0,0), 64));
	if (Load)
		XedInst2(&InstList[2], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedReg(VmIRegToXReg(IReg, OperandSize)), XedMemB(XReg1, OpSizeBits));
	else
		XedInst2(&InstList[2], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedMemB(XReg1, OpSizeBits), XedReg(VmIRegToXReg(IReg, OperandSize)));

	return XedEncodeInstructions(InstList, 3, OutSize);
}
PUCHAR VmHandlerIRegMem_BD(VM_IREG_ENUM IReg, VM_OPERAND_SIZE_ENUM OperandSize, BOOL Load, PUINT OutSize, XED_REG_ENUM XReg1, XED_REG_ENUM XReg2)
{
	/*
	* movzx r8,byte ptr[rdx+2]
	* mov r8, qword ptr[rbp+r8*8]
	* mov r9, dword ptr[rdx+3]
	* mov (ireg), (size) ptr[r8+r9]
	*/
	UINT OpSizeBits = VmOpSizeToBits(OperandSize);
	XED_ENCODER_INSTRUCTION InstList[4];
	XedInst2(&InstList[0], XedGlobalMachineState, XED_ICLASS_MOVZX, 64, XedReg(XReg1), XedMemBD(XED_REG_RDX, XedDisp(2, 8), 8));
	XedInst2(&InstList[1], XedGlobalMachineState, XED_ICLASS_MOV, 64, XedReg(XReg1), XedMemBISD(XED_REG_RBP, XReg1, 8, XedDisp(0, 0), 64));
	XedInst2(&InstList[2], XedGlobalMachineState, XED_ICLASS_MOVSXD, 64, XedReg(XReg2), XedMemBD(XED_REG_RDX, XedDisp(3, 8), 32));
	if (Load)
		XedInst2(&InstList[3], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedReg(VmIRegToXReg(IReg, OperandSize)), XedMemBISD(XReg1, XReg2, 1, XedDisp(0, 0), OpSizeBits));
	else
		XedInst2(&InstList[3], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedMemBISD(XReg1, XReg2, 1, XedDisp(0, 0), OpSizeBits), XedReg(VmIRegToXReg(IReg, OperandSize)));
	return XedEncodeInstructions(InstList, 4, OutSize);
}
PUCHAR VmHandlerIRegMem_BIS(VM_IREG_ENUM IReg, VM_OPERAND_SIZE_ENUM OperandSize, BOOL Load, PUINT OutSize, XED_REG_ENUM XReg1, XED_REG_ENUM XReg2, XED_REG_ENUM XReg3)
{
	/*
	* movzx r8,byte ptr[rdx+2]
	* mov r8, qword ptr[rbp+r8*8] 
	* movzx r9, byte ptr[rdx+3]
	* movzx r10, byte ptr[rdx+4]	;load scale value(unsigned)
	* imul r10, qword ptr[rbp+r9*8]
	* mov (ireg), (size) ptr[r8+r10]
	*/

	UINT OpSizeBits = VmOpSizeToBits(OperandSize);
	XED_ENCODER_INSTRUCTION InstList[6];
	XedInst2(&InstList[0], XedGlobalMachineState, XED_ICLASS_MOVZX, 64, XedReg(XReg1), XedMemBD(XED_REG_RDX, XedDisp(2, 8), 8));
	XedInst2(&InstList[1], XedGlobalMachineState, XED_ICLASS_MOV, 64, XedReg(XReg1), XedMemBISD(XED_REG_RBP, XReg1, 8, XedDisp(0, 0), 64));
	XedInst2(&InstList[2], XedGlobalMachineState, XED_ICLASS_MOVSX, 64, XedReg(XReg2), XedMemBD(XED_REG_RDX, XedDisp(3, 8), 8));
	XedInst2(&InstList[3], XedGlobalMachineState, XED_ICLASS_MOVSX, 64, XedReg(XReg3), XedMemBD(XED_REG_RDX, XedDisp(4, 8), 8));
	XedInst2(&InstList[4], XedGlobalMachineState, XED_ICLASS_IMUL, 64, XedReg(XReg3), XedMemBISD(XED_REG_RBP, XReg2, 8, XedDisp(0, 0), 64));

	if (Load)
		XedInst2(&InstList[5], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedReg(VmIRegToXReg(IReg, OperandSize)), XedMemBISD(XReg1, XReg3, 1, XedDisp(0, 0), OpSizeBits));
	else
		XedInst2(&InstList[5], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedMemBISD(XReg1, XReg3, 1, XedDisp(0, 0), OpSizeBits), XedReg(VmIRegToXReg(IReg, OperandSize)));
	return XedEncodeInstructions(InstList, 6, OutSize);
}
PUCHAR VmHandlerIRegMem_BISD(VM_IREG_ENUM IReg, VM_OPERAND_SIZE_ENUM OperandSize, BOOL Load, PUINT OutSize, XED_REG_ENUM XReg1, XED_REG_ENUM XReg2, XED_REG_ENUM XReg3)
{
	/*
	* movzx r8, byte ptr[rdx+2]
	* mov r8, qword ptr[rbp+r8*8]
	* movzx r9, byte ptr[rdx+3]
	* movzx r10, byte ptr[rdx+4]	;load scale value(unsigned)
	* imul r10, qword ptr[rbp+r9*8]
	* movsxd r9, dword ptr[rdx+5]	;load immediate displacement
	* add r10, r9					;add immediate displacement
	* mov (ireg), (size) ptr[r8+r10]
	*/

	UINT OpSizeBits = VmOpSizeToBits(OperandSize);
	XED_ENCODER_INSTRUCTION InstList[8];
	XedInst2(&InstList[0], XedGlobalMachineState, XED_ICLASS_MOVZX, 64, XedReg(XReg1), XedMemBD(XED_REG_RDX, XedDisp(2, 8), 8));
	XedInst2(&InstList[1], XedGlobalMachineState, XED_ICLASS_MOV, 64, XedReg(XReg1), XedMemBISD(XED_REG_RBP, XReg1, 8, XedDisp(0, 0), 64));
	XedInst2(&InstList[2], XedGlobalMachineState, XED_ICLASS_MOVSX, 64, XedReg(XReg2), XedMemBD(XED_REG_RDX, XedDisp(3, 8), 8));
	XedInst2(&InstList[3], XedGlobalMachineState, XED_ICLASS_MOVSX, 64, XedReg(XReg3), XedMemBD(XED_REG_RDX, XedDisp(4, 8), 8));
	XedInst2(&InstList[4], XedGlobalMachineState, XED_ICLASS_IMUL, 64, XedReg(XReg3), XedMemBISD(XED_REG_RBP, XReg2, 8, XedDisp(0, 0), 64));
	XedInst2(&InstList[5], XedGlobalMachineState, XED_ICLASS_MOVSXD, 64, XedReg(XReg2), XedMemBD(XED_REG_RDX, XedDisp(5, 8), 32));
	XedInst2(&InstList[6], XedGlobalMachineState, XED_ICLASS_ADD, 64, XedReg(XReg3), XedReg(XReg2));
	if (Load)
		XedInst2(&InstList[7], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedReg(VmIRegToXReg(IReg, OperandSize)), XedMemBISD(XReg1, XReg3, 1, XedDisp(0, 0), OpSizeBits));
	else
		XedInst2(&InstList[7], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedMemBISD(XReg1, XReg3, 1, XedDisp(0, 0), OpSizeBits), XedReg(VmIRegToXReg(IReg, OperandSize)));

	return XedEncodeInstructions(InstList, 8, OutSize);
}
PUCHAR VmHandlerIRegMem(VM_MEMOP_TYPE_ENUM MemOpType, VM_IREG_ENUM IReg, VM_OPERAND_SIZE_ENUM OperandSize, BOOL Load, PUINT OutSize, XED_REG_ENUM XReg1, XED_REG_ENUM XReg2, XED_REG_ENUM XReg3)
{
	switch (MemOpType)
	{
	case VM_MEMOP_B: return VmHandlerIRegMem_B(IReg, OperandSize, Load, OutSize, XReg1);
	case VM_MEMOP_BD: return VmHandlerIRegMem_BD(IReg, OperandSize, Load, OutSize, XReg1, XReg2);
	case VM_MEMOP_BIS: return VmHandlerIRegMem_BIS(IReg, OperandSize, Load, OutSize, XReg1, XReg2, XReg3);
	case VM_MEMOP_BISD: return VmHandlerIRegMem_BISD(IReg, OperandSize, Load, OutSize, XReg1, XReg2, XReg3);
	}
	return NULL;
}
PUCHAR VmHandlerIRegReg(VM_IREG_ENUM IReg, VM_OPERAND_SIZE_ENUM OperandSize, BOOL Load, PUINT OutSize, XED_REG_ENUM XReg1)
{
	/*
	* movzx r8,byte ptr[rdx+2]
	* mov (ireg), (size) ptr[rbp+r8*8]
	*/

	UINT OpSizeBits = VmOpSizeToBits(OperandSize);
	XED_ENCODER_INSTRUCTION InstList[2];
	XedInst2(&InstList[0], XedGlobalMachineState, XED_ICLASS_MOVZX, 64, XedReg(XReg1), XedMemBD(XED_REG_RDX, XedDisp(2, 8), 8));
	if (Load)
		XedInst2(&InstList[1], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedReg(VmIRegToXReg(IReg, OperandSize)), XedMemBISD(XED_REG_RBP, XReg1, 8, XedDisp(0, 0), OpSizeBits));
	else
		XedInst2(&InstList[1], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedMemBISD(XED_REG_RBP, XReg1, 8, XedDisp(0, 0), OpSizeBits), XedReg(VmIRegToXReg(IReg, OperandSize)));
	return XedEncodeInstructions(InstList, 2, OutSize);
}
PUCHAR VmHandlerIRegImm(VM_IREG_ENUM IReg, VM_OPERAND_SIZE_ENUM OperandSize, PUINT OutSize)
{
	/*
	* mov (ireg), size ptr[rdx+2]
	*/

	UINT OpSizeBits = VmOpSizeToBits(OperandSize);
	XED_ENCODER_INSTRUCTION InstList[1];
	XedInst2(&InstList[0], XedGlobalMachineState, XED_ICLASS_MOV, OpSizeBits, XedReg(VmIRegToXReg(IReg, OperandSize)), XedMemBD(XED_REG_RDX, XedDisp(2, 8), OpSizeBits));
	return XedEncodeInstructions(InstList, 1, OutSize);
}

PUCHAR VmHandlerEncode0(XED_ICLASS_ENUM IClass, VM_OPERAND_SIZE_ENUM OperandSize, PUINT OutSize)
{
	XED_ENCODER_INSTRUCTION InstList[3];
	XedInst0(&InstList[0], XedGlobalMachineState, XED_ICLASS_POPFQ, 64);
	XedInst0(&InstList[1], XedGlobalMachineState, IClass, VmOpSizeToBits(OperandSize));
	XedInst0(&InstList[2], XedGlobalMachineState, XED_ICLASS_PUSHFQ, 64);
	return XedEncodeInstructions(InstList, 3, OutSize);
}
PUCHAR VmHandlerEncode1(XED_ICLASS_ENUM IClass, VM_OPERAND_SIZE_ENUM OperandSize, PUINT OutSize)
{
	XED_ENCODER_INSTRUCTION InstList[3];
	XedInst0(&InstList[0], XedGlobalMachineState, XED_ICLASS_POPFQ, 64);
	XedInst1(&InstList[1], XedGlobalMachineState, IClass, VmOpSizeToBits(OperandSize), XedReg(VmIRegToXReg(VM_IREG_1, OperandSize)));
	XedInst0(&InstList[2], XedGlobalMachineState, XED_ICLASS_PUSHFQ, 64);
	return XedEncodeInstructions(InstList, 3, OutSize);
}
PUCHAR VmHandlerEncode2(XED_ICLASS_ENUM IClass, VM_OPERAND_SIZE_ENUM OperandSize1, VM_OPERAND_SIZE_ENUM OperandSize2, PUINT OutSize)
{
	XED_ENCODER_INSTRUCTION InstList[3];
	XedInst0(&InstList[0], XedGlobalMachineState, XED_ICLASS_POPFQ, 64);
	XedInst2(&InstList[1], XedGlobalMachineState, IClass, VmOpSizeToBits(OperandSize1), XedReg(VmIRegToXReg(VM_IREG_1, OperandSize1)), XedReg(VmIRegToXReg(VM_IREG_2, OperandSize2)));
	XedInst0(&InstList[2], XedGlobalMachineState, XED_ICLASS_PUSHFQ, 64);
	return XedEncodeInstructions(InstList, 3, OutSize);
}
PUCHAR VmHandlerEncode3(XED_ICLASS_ENUM IClass, VM_OPERAND_SIZE_ENUM OperandSize1, VM_OPERAND_SIZE_ENUM OperandSize2, VM_OPERAND_SIZE_ENUM OperandSize3, PUINT OutSize)
{
	XED_ENCODER_INSTRUCTION InstList[3];
	XedInst0(&InstList[0], XedGlobalMachineState, XED_ICLASS_POPFQ, 64);
	XedInst3(&InstList[1], XedGlobalMachineState, IClass, VmOpSizeToBits(OperandSize1), XedReg(VmIRegToXReg(VM_IREG_1, OperandSize1)), XedReg(VmIRegToXReg(VM_IREG_2, OperandSize2)), XedReg(VmIRegToXReg(VM_IREG_3, OperandSize3)));
	XedInst0(&InstList[2], XedGlobalMachineState, XED_ICLASS_PUSHFQ, 64);
	return XedEncodeInstructions(InstList, 3, OutSize);
}