#include "NativeCode.h"

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

_NATIVE_CODE_LINK::_NATIVE_CODE_LINK(ULONG LabelId, _NATIVE_CODE_BLOCK* B)
	: _NATIVE_CODE_LINK()
{
	Block = B;
	Label = LabelId;
	Flags = CODE_FLAG_IS_LABEL;
}

_NATIVE_CODE_LINK::_NATIVE_CODE_LINK(ULONG F, PVOID Rd, ULONG Rds)
	: _NATIVE_CODE_LINK()
{
	Flags = F;
	RawDataSize = Rds;
	RawData = new UCHAR[Rds];
	if (Rd)
		memcpy(RawData, Rd, Rds);
}

_NATIVE_CODE_LINK::~_NATIVE_CODE_LINK()
{
	if (RawData)
		delete RawData;
}

_NATIVE_CODE_BLOCK::_NATIVE_CODE_BLOCK()
{
	Start = End = NULL;
	HasRelativeJumps = FALSE;
}

VOID NcAppendToBlock(PNATIVE_CODE_BLOCK Block, PNATIVE_CODE_LINK Link)
{
	if (!Link)
		return;

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

	if (!Block->End || !Block->Start)
	{
		Block->Start = Block->End = Link;
	}
	else
	{
		Block->End->Next = Link;
		Block->End = Link;
	}
}

VOID NcPrependToBlock(PNATIVE_CODE_BLOCK Block, PNATIVE_CODE_LINK Link)
{
	if (!Link)
		return;

	Link->Block = Block;
	Link->Next = Block->Start;
	Link->Prev = NULL;

	if (!Block->End || !Block->Start)
	{
		Block->Start = Block->End = Link;
	}
	else
	{
		Block->Start->Prev = Link;
		Block->Start = Link;
	}
}

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;
	}
}

ULONG NcGenUnusedLabelId(STDVECTOR<ULONG> CONST& LabelIds)
{
	ULONG ReturnLabelId = rand();
	while (StdFind(LabelIds.begin(), LabelIds.end(), ReturnLabelId) != LabelIds.end())
		ReturnLabelId = rand();
	return ReturnLabelId;
}

VOID NcChangeLabelId(PNATIVE_CODE_BLOCK Block1, ULONG Original, ULONG New)
{
	for (PNATIVE_CODE_LINK T = Block1->Start; T; T = T->Next)
	{
		if (((T->Flags & CODE_FLAG_IS_LABEL) || (T->Flags & CODE_FLAG_IS_REL_JMP)) && T->Label == Original)
			T->Label = New;
	}
}

VOID NcFixLabelsForBlocks(PNATIVE_CODE_BLOCK Block1, PNATIVE_CODE_BLOCK Block2)
{
	STDVECTOR<ULONG> BlockOneLabels;
	for (PNATIVE_CODE_LINK T = Block1->Start; T; T = T->Next)
	{
		if ((T->Flags & CODE_FLAG_IS_LABEL) && StdFind(BlockOneLabels.begin(), BlockOneLabels.end(), T->Label) != BlockOneLabels.end())
			BlockOneLabels.push_back(T->Label);
	}

	for (PNATIVE_CODE_LINK T = Block2->Start; T; T = T->Next)
	{
		if ((T->Flags & CODE_FLAG_IS_LABEL) && StdFind(BlockOneLabels.begin(), BlockOneLabels.end(), T->Label) != BlockOneLabels.end())
			NcChangeLabelId(Block2, T->Label, NcGenUnusedLabelId(BlockOneLabels));
	}
}

BOOL NcInsertBlockAfter(PNATIVE_CODE_LINK Link, PNATIVE_CODE_BLOCK Block)
{
	if (!Link || !Link->Block || !Block || !Block->Start || !Block->End || Link->Block == Block)
		return FALSE;

	if (Block->HasRelativeJumps && Link->Block->HasRelativeJumps)
		NcFixLabelsForBlocks(Link->Block, Block);

	if (Link->Next)
		Link->Next->Prev = Block->End;
	Block->End->Next = Link->Next;
	Block->Start->Prev = Link;
	Link->Next = Block->Start;

	for (PNATIVE_CODE_LINK T = Block->Start; T; T = T->Next)
		T->Block = Link->Block;

	return TRUE;
}

BOOL NcInsertBlockBefore(PNATIVE_CODE_LINK Link, PNATIVE_CODE_BLOCK Block)
{
	if (!Link || !Link->Block || !Block || !Block->Start || !Block->End)
		return FALSE;

	if (Block->HasRelativeJumps && Link->Block->HasRelativeJumps)
		NcFixLabelsForBlocks(Link->Block, Block);
	
	if (Link->Prev)
		Link->Prev->Next = Block->Start;
	Block->Start->Prev = Link->Prev;
	Block->End->Next = Link;
	Link->Prev = Block->End;

	for (PNATIVE_CODE_LINK T = Block->Start; T; T = T->Next)
		T->Block = Link->Block;

	return TRUE;
}

BOOL NcCreateLabels(PNATIVE_CODE_BLOCK Block)
{
	ULONG CurrentLabelId = 0;
	for (PNATIVE_CODE_LINK T = Block->Start; T; T = T->Next)
	{
		if (!(T->Flags & CODE_FLAG_IS_INST))
			continue;

		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)
			continue;

		INT32 BranchDisplacement = XedDecodedInstGetBranchDisplacement(&T->XedInst);
		PNATIVE_CODE_LINK JmpPos = NcValidateJmp(T, BranchDisplacement);
		if (!JmpPos)
		{
			printf("Failed to validate jump. Type: %s, Displacement: %d\n", 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, Block));
			T->Label = CurrentLabelId;
			++CurrentLabelId;
		}
		Block->HasRelativeJumps = TRUE;
		T->Flags |= CODE_FLAG_IS_REL_JMP;
	}
	return TRUE;
}

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->Prev;
		}
		if (Delta != 0 || !T)
			return NULL;
		while (T && (T->Flags & CODE_FLAG_IS_LABEL))
			T = T->Next;
		return T;
	}
	return Jmp;
}

PNATIVE_CODE_LINK NcDeepCopy(PNATIVE_CODE_LINK Link)
{
	if (Link->Flags & CODE_FLAG_IS_LABEL)
	{
		return new NATIVE_CODE_LINK(Link->Label, NULL);
	}
	else
	{
		XED_ICLASS_ENUM IClass = XedDecodedInstGetIClass(&Link->XedInst);
		printf("Doing %s\n", XedIClassEnumToString(IClass));

		PNATIVE_CODE_LINK NewLink = new NATIVE_CODE_LINK(Link->Flags, Link->RawData, Link->RawDataSize);
		NewLink->Label = Link->Label;
		XED_ERROR_ENUM DecodeError = XedDecode(&NewLink->XedInst, Link->RawData, Link->RawDataSize);
		if (DecodeError != XED_ERROR_NONE)
		{
			printf("XedDecode failed in NcDeepCopy: %s\n", XedErrorEnumToString(DecodeError));
			delete NewLink;
			return NULL;
		}
		printf("succeeded\n");
		return NewLink;
	}
}

PNATIVE_CODE_BLOCK NcDeepCopyPartialBlock(PNATIVE_CODE_LINK Start, PNATIVE_CODE_LINK End)
{
	if (!Start || !End)
		return NULL;

	PNATIVE_CODE_BLOCK Block = new NATIVE_CODE_BLOCK;
	if (!Block)
		return NULL;

	for (PNATIVE_CODE_LINK CurLink = Start; CurLink != End->Next; CurLink = CurLink->Next)
	{
		PNATIVE_CODE_LINK Temp = NcDeepCopy(CurLink);
		if (!Temp)
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}

		if (Temp->Flags & CODE_FLAG_IS_REL_JMP)
			Block->HasRelativeJumps = TRUE;

		NcAppendToBlock(Block, Temp);
	}

	return Block;
}

PNATIVE_CODE_BLOCK NcDeepCopyBlock(PNATIVE_CODE_BLOCK Block)
{
	return NcDeepCopyPartialBlock(Block->Start, Block->End);
}

BOOL NcDisassemble(PNATIVE_CODE_BLOCK Block, PVOID Buffer, ULONG BufferSize)
{
	PUCHAR Buf = (PUCHAR)Buffer;
	ULONG Offset = 0;

	while (Offset < BufferSize)
	{
		PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK;
		Link->Flags = CODE_FLAG_IS_INST;
		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));
			NcDeleteBlock(Block);
			delete Link;
			return FALSE;
		}
		Link->RawDataSize = XedDecodedInstGetLength(&Link->XedInst);
		Link->RawData = new UCHAR[Link->RawDataSize];
		memcpy(Link->RawData, (Buf + Offset), Link->RawDataSize);

		NcAppendToBlock(Block, Link);

		Offset += Link->RawDataSize;
	}

	NcCreateLabels(Block);

	return TRUE;
}

PVOID NcAssemble(PNATIVE_CODE_BLOCK Block)
{
	return NULL;
}

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

VOID NcDebugPrint(PNATIVE_CODE_BLOCK Block)
{
	HANDLE ConsoleHandle = GetStdHandle(STD_OUTPUT_HANDLE);
	if (!ConsoleHandle)
		return;

	for (PNATIVE_CODE_LINK T = Block->Start; T; T = T->Next)
	{
		if (T->Flags & CODE_FLAG_IS_LABEL)
		{
			SetConsoleTextAttribute(ConsoleHandle, FOREGROUND_GREEN | FOREGROUND_RED);
			printf("Label: %u\n", T->Label);
		}
		else
		{
			XED_ICLASS_ENUM IClass = XedDecodedInstGetIClass(&T->XedInst);
			if (T->Flags & CODE_FLAG_IS_REL_JMP)
			{
				SetConsoleTextAttribute(ConsoleHandle, FOREGROUND_GREEN | FOREGROUND_RED);
				printf("%s: %u\n", XedIClassEnumToString(IClass), T->Label);
			}
			else
			{
				SetConsoleTextAttribute(ConsoleHandle, FOREGROUND_GREEN | FOREGROUND_BLUE);
				printf("%s\n", XedIClassEnumToString(IClass));
			}
		}
	}
}

VOID NcPrintBlockCode(PNATIVE_CODE_BLOCK Block)
{
	for (PNATIVE_CODE_LINK T = Block->Start; T; T = T->Next)
	{
		if (!(T->Flags & CODE_FLAG_IS_LABEL))
		{
			for (uint32_t i = 0; i < T->RawDataSize; i++)
			{
				std::cout << std::hex << std::setw(2) << std::setfill('0') << (int)T->RawData[i] << ' ';
			}
		}
	}
}