#include "NativeCode.h"

_NATIVE_CODE_LINK::_NATIVE_CODE_LINK()
{
	XedDecodedInstZero(&XedInst);
	XedDecodedInstSetMode(&XedInst, XED_MACHINE_MODE_LONG_64, XED_ADDRESS_WIDTH_64b);
	Flags = 0;
	Next = Prev = NULL;
	Label = 0;
	RawData = NULL;
	RawDataSize = 0UL;
}

_NATIVE_CODE_LINK::_NATIVE_CODE_LINK(ULONG LabelId)
	: _NATIVE_CODE_LINK()
{
	Label = LabelId;
	Flags = CODE_FLAG_IS_LABEL;
}

VOID NcInsertLinkAfter(PNATIVE_CODE_LINK Link1, PNATIVE_CODE_LINK Link2)
{
	if (Link1)
	{
		Link2->Prev = Link1;
		Link2->Next = Link1->Next;
		Link1->Next = Link2;
		if (Link2->Next)
			Link2->Next->Prev = Link2;
	}
}

VOID NcInsertLinkBefore(PNATIVE_CODE_LINK Link1, PNATIVE_CODE_LINK Link2)
{
	if (Link1)
	{
		Link2->Next = Link1;
		Link2->Prev = Link1->Prev;
		Link1->Prev = Link2;
		if (Link2->Prev)
			Link2->Prev->Next = Link2;
	}
}

VOID NcUnlink(PNATIVE_CODE_LINK Link)
{
	if (Link)
	{
		if (Link->Next)
			Link->Next->Prev = Link->Prev;
		if (Link->Prev)
			Link->Prev->Next = Link->Next;
	}
}

VOID NcConcat(PNATIVE_CODE_BLOCK Block1, PNATIVE_CODE_BLOCK Block2)
{
	//Cant simply concatinate blocks here, need to go through the second block and
	//update the label names so that there are no conflicts between the two blocks
}

VOID NcInsertBlockAfter(PNATIVE_CODE_LINK Link, PNATIVE_CODE_BLOCK Block)
{

}

VOID NcInsertBlockBfore(PNATIVE_CODE_LINK Link, PNATIVE_CODE_BLOCK Block)
{

}

BOOL NcCreateLabels(PNATIVE_CODE_BLOCK Block)
{
	ULONG CurrentLabelId = 0;
	for (PNATIVE_CODE_LINK T = Block->Start; T; T = T->Next)
	{
		XED_CATEGORY_ENUM Category = XedDecodedInstGetCategory(&T->XedInst);

		if (Category != XED_CATEGORY_COND_BR && Category != XED_CATEGORY_UNCOND_BR)
			continue;

		ULONG OperandCount = XedDecodedInstNumOperands(&T->XedInst);
		if (OperandCount < 1)
			continue;

		CONST XED_INST* Inst = XedDecodedInstInst(&T->XedInst);
		if (!Inst)
			continue;

		CONST XED_OPERAND* Operand = XedInstOperand(Inst, 0);
		if (!Operand)
			continue;

		XED_OPERAND_TYPE_ENUM OperandType = XedOperandType(Operand);
		if (OperandType != XED_OPERAND_TYPE_IMM && OperandType != XED_OPERAND_TYPE_IMM_CONST)
		{
			printf("Found jump to non immediate value. Category: %s\n", XedCategoryEnumToString(Category));
			continue;
		}

		INT32 BranchDisplacement = XedDecodedInstGetBranchDisplacement(&T->XedInst);
		PNATIVE_CODE_LINK JmpPos = NcValidateJmp(T, BranchDisplacement);
		if (!JmpPos)
		{
			printf("Failed to validate jump. Type: %s, Displacement: %d", XedCategoryEnumToString(Category), BranchDisplacement);
			return FALSE;
		}

		if (JmpPos->Prev && (JmpPos->Prev->Flags & CODE_FLAG_IS_LABEL))
		{
			T->Label = JmpPos->Prev->Label;
		}
		else
		{
			NcInsertLinkBefore(JmpPos, new NATIVE_CODE_LINK(CurrentLabelId));
			T->Label = CurrentLabelId;
			++CurrentLabelId;
		}
		T->Flags |= CODE_FLAG_IS_REL_JMP;
	}
}

PNATIVE_CODE_LINK NcValidateJmp(PNATIVE_CODE_LINK Jmp, INT32 Delta)
{
	PNATIVE_CODE_LINK T;
	if (Delta > 0)
	{
		T = Jmp->Next;
		while (Delta > 0 && T)
		{
			if (T->Flags & CODE_FLAG_IS_LABEL)
				continue;
			Delta -= XedDecodedInstGetLength(&T->XedInst);
			T = T->Next;
		}
		if (Delta != 0 || !T)
			return NULL;
		while (T && (T->Flags & CODE_FLAG_IS_LABEL))
			T = T->Next;
		return T;
	}
	else if (Delta < 0)
	{
		T = Jmp;
		while (T)
		{
			if (T->Flags & CODE_FLAG_IS_LABEL)
				continue;
			Delta += XedDecodedInstGetLength(&T->XedInst);
			if (Delta >= 0)
				break;
			T = T->Next;
		}
		if (Delta != 0 || !T)
			return NULL;
		while (T && (T->Flags & CODE_FLAG_IS_LABEL))
			T = T->Next;
		return T;
	}
	//return the jmp if that delta is zero
	return Jmp;
}

BOOL NcFromBuffer(PNATIVE_CODE_BLOCK Block, PVOID Buffer, ULONG BufferSize)
{
	Block->Start = new NATIVE_CODE_LINK;
	Block->End = Block->Start;
	
	PUCHAR Buf = (PUCHAR)Buffer;
	ULONG Offset = 0;

	while (Offset < BufferSize)
	{
		PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK;
		ULONG PossibleSize = min(15, BufferSize - Offset);
		
		XED_ERROR_ENUM DecodeError = XedDecode(&Link->XedInst, (Buf + Offset), PossibleSize);
		if (DecodeError != XED_ERROR_NONE)
		{
			printf("XedDecode failed with error %s\n", XedErrorEnumToString(DecodeError));
			NcDelete(Block);
			delete Link;
			return FALSE;
		}

		Link->Prev = Block->End;
		Block->End->Next = Link;
		Block->End = Link;
	}

	PNATIVE_CODE_LINK StartLink = Block->Start;
	Block->Start = Block->Start->Next;
	delete StartLink;

	NcCreateLabels(Block);

	return TRUE;
}

VOID NcDelete(PNATIVE_CODE_BLOCK Block)
{
	for (PNATIVE_CODE_LINK T = Block->Start; T;)
	{
		PNATIVE_CODE_LINK Next = T->Next;
		delete T;
		T = Next;
	}	
}