#include "Virtualizer.h"
#include "VirtualMachine.h"


#define VmIRegLoadMemHandlerId(IReg, MemOpType) (VM_ICLASS_LD_IREG_MEM_START + VM_MEMOP_TYPE_COUNT * (IReg) + (MemOpType))
#define VmIRegStoreMemHandlerId(IReg, MemOpType) (VM_ICLASS_ST_IREG_MEM_START + VM_MEMOP_TYPE_COUNT * (IReg) + (MemOpType))
#define VmIRegLoadRegHandlerId(IReg) (VM_ICLASS_LD_IREG_REG_START + (IReg))
#define VmIRegStoreRegHandlerId(IReg) (VM_ICLASS_ST_IREG_REG_START + (IReg))
#define VmIRegLoadImmHandlerId(IReg) (VM_ICLASS_LD_IREG_IMM_START + (IReg))

PVM_HANDLER ViGetVmHandler(PVIRTUALIZER Virt, UINT Id)
{
	if (Virt->HandlerMap[Id] == 0)
	{
		Virt->Handlers.push_back(new VM_HANDLER());
		Virt->HandlerMap[Id] = (USHORT)(Virt->Handlers.size() - 1);
	}
	return Virt->Handlers[Virt->HandlerMap[Id]];
}

CONST CHAR* ViIClassToString(UINT IClass)
{
	if (IClass >= VM_ICLASS_ENTER)
		return VmIClassToString((VM_ICLASS_ENUM)IClass);
	return XedIClassEnumToString((XED_ICLASS_ENUM)IClass);
}
BOOL ViCanHandleInst(PNATIVE_CODE_LINK Link)
{
	switch (XedDecodedInstGetIClass(&Link->XedInstruction))
	{
	case XED_ICLASS_MOV: 
		return TRUE;
	}
	return FALSE;
}
BOOL ViValidateNativeCodeBlock(PVIRTUALIZER Virt, PNATIVE_CODE_BLOCK Block)
{
	for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
	{
		if (!ViCanHandleInst(T))
			return FALSE;

		//Cant handle RIP relative instructions.
		CONST XED_INST* Inst = XedDecodedInstInst(&T->XedInstruction);
		UINT OperandCount = XedDecodedInstNumOperands(&T->XedInstruction);
		for (UINT i = 0; i < OperandCount; i++)
		{
			XED_OPERAND_ENUM OperandName = XedOperandName(XedInstOperand(Inst, i));

			if ((OperandName == XED_OPERAND_MEM0 || OperandName == XED_OPERAND_MEM1)
				&& XedDecodedInstGetBaseReg(&T->XedInstruction, OperandName - XED_OPERAND_MEM0) == XED_REG_RIP)
				return FALSE;
		}

	}
	return TRUE;
}
BOOL ViVirtualizeInst(PVIRTUALIZER Virt, XED_DECODED_INST* DecodedInst, PVM_CODE_BLOCK Block)
{
	VM_CODE_BLOCK Prologue = { };
	VM_CODE_BLOCK Epilogue = { };

	CONST XED_INST* Inst = XedDecodedInstInst(DecodedInst);
	UINT OperandCount = XedDecodedInstNumOperands(DecodedInst);
	XED_ICLASS_ENUM IClass = XedDecodedInstGetIClass(DecodedInst);

	STDVECTOR<CONST XED_OPERAND*> Operands = {};
	VM_OPERAND_SIZE_ENUM OpSize[3] = { VM_OPSIZE_8 };

	for (UINT i = 0; i < OperandCount; i++)
	{
		CONST XED_OPERAND* Operand = XedInstOperand(Inst, i);
		XED_OPERAND_ENUM OperandName = XedOperandName(Operand);
		if (XedOperandIsRegister(OperandName) && XED_REG_RFLAGS == XedDecodedInstGetReg(DecodedInst, OperandName))
			continue;

		Operands.push_back(Operand);
		OpSize[Operands.size() - 1] = (VM_OPERAND_SIZE_ENUM)(log2(XedDecodedInstOperandLength(DecodedInst, i)));
	}

	for (UINT i = 0; i < Operands.size(); i++)
	{
		CONST XED_OPERAND* Operand = Operands[i];
		XED_OPERAND_ENUM OperandName = XedOperandName(Operand);
		VM_IREG_ENUM IReg = (VM_IREG_ENUM)(VM_IREG_1 + i);

		if (XedOperandIsRegister(OperandName))
		{
			if (XedOperandRead(Operand))
			{
				USHORT HandlerId = VmIRegLoadRegHandlerId(IReg);
				PVM_HANDLER_DESCRIPTOR Desc = ViGetVmHandler(Virt, HandlerId)->GetDesc(OpSize);
				if (!Desc->Data) Desc->Data = VmHandlerIRegReg(IReg, OpSize[i], TRUE, &Desc->DataSize);
				PVM_CODE_LINK Link = new VM_CODE_LINK();
				Link->HandlerId = HandlerId;
				Link->Flags |= CODE_FLAG_IREG_OPERATION;
				Link->Reg.Register = VmXRegToVRegId(XedDecodedInstGetReg(DecodedInst, OperandName));
				VcAppendToBlock(&Prologue, Link);
			}
			if (XedOperandWritten(Operand))
			{
				USHORT HandlerId = VmIRegStoreRegHandlerId(IReg);
				PVM_HANDLER_DESCRIPTOR Desc = ViGetVmHandler(Virt, HandlerId)->GetDesc(OpSize);
				if (!Desc->Data) Desc->Data = VmHandlerIRegReg(IReg, OpSize[i], FALSE, &Desc->DataSize);
				PVM_CODE_LINK Link = new VM_CODE_LINK();
				Link->HandlerId = HandlerId;
				Link->Flags |= CODE_FLAG_IREG_OPERATION;
				Link->Reg.Register = VmXRegToVRegId(XedDecodedInstGetReg(DecodedInst, OperandName));
				VcAppendToBlock(&Epilogue, Link);
			}
		}
		else if (OperandName == XED_OPERAND_MEM0/* || OperandName == XED_OPERAND_MEM1*/)
		{
			UINT MemOpIndex = OperandName - XED_OPERAND_MEM0;
			XED_REG_ENUM Base = XedDecodedInstGetBaseReg(DecodedInst, MemOpIndex);
			XED_REG_ENUM Index = XedDecodedInstGetIndexReg(DecodedInst, MemOpIndex);
			UINT Scale = XedDecodedInstGetScale(DecodedInst, MemOpIndex);
			LONGLONG Displacement = XedDecodedInstGetMemoryDisplacement(DecodedInst, MemOpIndex);
			UINT MemOpType = VM_MEMOP_TYPE_COUNT;
			if (Base != XED_REG_INVALID) MemOpType = VM_MEMOP_B;
			if (Index != XED_REG_INVALID) MemOpType = VM_MEMOP_BIS;
			if (Displacement != 0) MemOpType += 1;

			if (XedOperandRead(Operand))
			{
				USHORT HandlerId = VmIRegLoadMemHandlerId(IReg, MemOpType);
				PVM_HANDLER_DESCRIPTOR Desc = ViGetVmHandler(Virt, HandlerId)->GetDesc(OpSize);
				if (!Desc->Data) Desc->Data = VmHandlerIRegMem((VM_MEMOP_TYPE_ENUM)MemOpType, IReg, OpSize[i], TRUE, &Desc->DataSize);
				PVM_CODE_LINK Link = new VM_CODE_LINK();
				Link->HandlerId = HandlerId;
				Link->Flags |= CODE_FLAG_IREG_OPERATION;
				Link->Mem.Base = VmXRegToVRegId(Base);
				Link->Mem.Index = VmXRegToVRegId(Index);
				Link->Mem.Scale = Scale;
				Link->Mem.Displacement = Displacement;
				VcAppendToBlock(&Prologue, Link);
			}
			if (XedOperandWritten(Operand))
			{
				USHORT HandlerId = VmIRegStoreMemHandlerId(IReg, MemOpType);
				PVM_HANDLER_DESCRIPTOR Desc = ViGetVmHandler(Virt, HandlerId)->GetDesc(OpSize);
				if (!Desc->Data) Desc->Data = VmHandlerIRegMem((VM_MEMOP_TYPE_ENUM)MemOpType, IReg, OpSize[i], FALSE, &Desc->DataSize);
				PVM_CODE_LINK Link = new VM_CODE_LINK();
				Link->HandlerId = HandlerId;
				Link->Flags |= CODE_FLAG_IREG_OPERATION;
				Link->Mem.Base = VmXRegToVRegId(Base);
				Link->Mem.Index = VmXRegToVRegId(Index);
				Link->Mem.Scale = Scale;
				Link->Mem.Displacement = Displacement;
				VcAppendToBlock(&Epilogue, Link);
			}
		}
		else if (OperandName == XED_OPERAND_IMM0)
		{
			USHORT HandlerId = VmIRegLoadImmHandlerId(IReg);
			if (Operands.size() > 1)
				OpSize[i] = OpSize[i - 1];
			else
				OpSize[i] = VM_OPSIZE_64; //assume 16 for 1 operand imm instructions

			PVM_HANDLER_DESCRIPTOR Desc = ViGetVmHandler(Virt, HandlerId)->GetDesc(OpSize);
			if (!Desc->Data) Desc->Data = VmHandlerIRegImm(IReg, OpSize[i], &Desc->DataSize);

			PVM_CODE_LINK Link = new VM_CODE_LINK();
			Link->HandlerId = HandlerId;
			Link->Flags |= CODE_FLAG_IREG_OPERATION;
		
			if (XedDecodedInstGetImmediateIsSigned(DecodedInst))
				Link->Imm.Immediate.i64 = (INT64)XedDecodedInstGetSignedImmediate(DecodedInst);
			else
				Link->Imm.Immediate.u64 = XedDecodedInstGetUnsignedImmediate(DecodedInst);

			/*else
			{
				ULONG64 ImmRaw = xed_decoded_inst_get_unsigned_immediate(DecodedInst);
				memcpy(&Link->Imm.Immediate.Raw, &ImmRaw, VmOpSizeToBytes(OpSize[i]));
			}*/
			VcAppendToBlock(&Prologue, Link);
		}
	}

	switch (IClass)
	{
	default:
		PVM_CODE_LINK Link = new VM_CODE_LINK();
		Link->HandlerId = IClass;
		Link->Flags |= CODE_FLAG_NATIVE_OPERATION;
		PVM_HANDLER_DESCRIPTOR Desc = ViGetVmHandler(Virt, IClass)->GetDesc(OpSize);
		if (!Desc->Data)
		{
			switch (Operands.size())
			{
			case 1: Desc->Data = VmHandlerEncode1(IClass, OpSize[0], &Desc->DataSize); break;
			case 2: Desc->Data = VmHandlerEncode2(IClass, OpSize[0], OpSize[1], &Desc->DataSize); break;
			case 3: Desc->Data = VmHandlerEncode3(IClass, OpSize[0], OpSize[1], OpSize[2], &Desc->DataSize); break;
			}
		}
		VcAppendToBlock(&Prologue, Link);
	}

	if (Epilogue.Start)
	{
		Prologue.End->Next = Epilogue.Start;
		Epilogue.Start = Prologue.End;
		Prologue.End = Epilogue.End;
	}

	Block->Start = Prologue.Start;
	Block->End = Prologue.End;

	return TRUE;
}