#include "NativeCode.h"

_NATIVE_CODE_LINK::_NATIVE_CODE_LINK()
{
	XedDecodedInstZero(&XedInstruction);
	XedDecodedInstSetMode(&XedInstruction, 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, BOOL Decode)
	: _NATIVE_CODE_LINK()
{
	Flags = F;
	RawDataSize = Rds;
	RawData = new UCHAR[Rds];
	if (Rd)
	{
		RtlCopyMemory(RawData, Rd, Rds);
		if (Decode)
			XedDecode(&XedInstruction, RawData, RawDataSize);
	}
}

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

_NATIVE_CODE_BLOCK::_NATIVE_CODE_BLOCK()
{
	Start = End = NULL;
	LabelIds.clear();
}

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 NcCalcBlockSize(PNATIVE_CODE_BLOCK Block)
{
	ULONG TotalSize = 0;
	for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
	{
		if (T->Flags & CODE_FLAG_IS_LABEL)
			continue;
		TotalSize += T->RawDataSize;
	}
	return TotalSize;
}

ULONG NcGenUnusedLabelId(PNATIVE_CODE_BLOCK Block)
{
	ULONG ReturnLabelId = rand();
	while (StdFind(Block->LabelIds.begin(), Block->LabelIds.end(), ReturnLabelId) != Block->LabelIds.end())
		ReturnLabelId = rand();
	Block->LabelIds.push_back(ReturnLabelId);
	return ReturnLabelId;
}

VOID NcChangeLabelId(PNATIVE_CODE_BLOCK Block, ULONG Original, ULONG New)
{
	for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; 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)
{
	for (PNATIVE_CODE_LINK T = Block2->Start; T && T != Block2->End->Next; T = T->Next)
	{
		if ((T->Flags & CODE_FLAG_IS_LABEL) && StdFind(Block1->LabelIds.begin(), Block1->LabelIds.end(), T->Label) != Block1->LabelIds.end())
		{
			NcChangeLabelId(Block2, T->Label, NcGenUnusedLabelId(Block1));
		}
	}
}

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

	if (FixLabels && Block->LabelIds.size() && Link->Block->LabelIds.size())
		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 != Block->End->Next; T = T->Next)
		T->Block = Link->Block;

	return TRUE;
}

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

	if (FixLabels && Block->LabelIds.size() && Link->Block->LabelIds.size())
		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 != Block->End->Next; 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->XedInstruction);
		if (Category != XED_CATEGORY_COND_BR && Category != XED_CATEGORY_UNCOND_BR)
			continue;

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

		CONST XED_INST* Inst = XedDecodedInstInst(&T->XedInstruction);
		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->XedInstruction);
		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));
			Block->LabelIds.push_back(CurrentLabelId);
			T->Label = CurrentLabelId;
			++CurrentLabelId;
		}
		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->XedInstruction);
			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->XedInstruction);
			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 NcDeepCopyLink(PNATIVE_CODE_LINK Link)
{
	if (Link->Flags & CODE_FLAG_IS_LABEL)
	{
		return new NATIVE_CODE_LINK(Link->Label, NULL);
	}
	else
	{		
		PNATIVE_CODE_LINK NewLink = new NATIVE_CODE_LINK(Link->Flags, Link->RawData, Link->RawDataSize);
		NewLink->Label = Link->Label;
		XED_ERROR_ENUM DecodeError = XedDecode(&NewLink->XedInstruction, Link->RawData, Link->RawDataSize);
		if (DecodeError != XED_ERROR_NONE)
		{
			printf("XedDecode failed in NcDeepCopyLink: %s %u\n", XedErrorEnumToString(DecodeError), Link->RawDataSize);
			delete NewLink;
			return NULL;
		}
		return NewLink;
	}
}

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

	Block->LabelIds.clear();
	Block->Start = Block->End = NULL;

	for (ULONG L : Start->Block->LabelIds)
		Block->LabelIds.push_back(L);

	for (PNATIVE_CODE_LINK CurLink = Start; CurLink && CurLink != End->Next; CurLink = CurLink->Next)
	{
		PNATIVE_CODE_LINK Temp = NcDeepCopyLink(CurLink);
		if (!Temp)
		{
			NcDeleteBlock(Block);
			return FALSE;
		}
		NcAppendToBlock(Block, Temp);
	}

	return TRUE;
}

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

BOOL NcGetDeltaToLabel(PNATIVE_CODE_LINK Link, PINT32 DeltaOut)
{
	INT32 Delta = 0;
	//First checking backwards because I feel like thats the direction most jmps are in
	for (PNATIVE_CODE_LINK T = Link; T; T = T->Prev)
	{
		if (T->Flags & CODE_FLAG_IS_LABEL)
		{
			if (T->Label == Link->Label)
			{
				*DeltaOut = Delta;
				return TRUE;
			}
			continue;
		}
		Delta -= T->RawDataSize;
	}

	//Now check forwards
	Delta = 0;
	for (PNATIVE_CODE_LINK T = Link->Next; T; T = T->Next)
	{
		if (T->Flags & CODE_FLAG_IS_LABEL)
		{
			if (T->Label == Link->Label)
			{
				*DeltaOut = Delta;
				return TRUE;
			}
			continue;
		}
		Delta += T->RawDataSize;
	}
	return FALSE;
}

BOOL NcFixRelJmps(PNATIVE_CODE_BLOCK Block)
{
	for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next;)
	{
		if (T->Flags & CODE_FLAG_IS_REL_JMP)
		{
			INT32 BranchDisp = 0;
			if (!NcGetDeltaToLabel(T, &BranchDisp))
				return FALSE;

			ULONG DispWidth = XedDecodedInstGetBranchDisplacementWidthBits(&T->XedInstruction);
			if (log2(abs(BranchDisp)) + 1 > DispWidth)
			{
				//duh oh
				if (DispWidth == 32)
					return FALSE;

				//Grow displacement width to required size
				DispWidth *= 2;

				//Check again
				if (log2(abs(BranchDisp)) + 1 > DispWidth)
				{
					if (DispWidth == 32)
						return FALSE;

					//Grow once more if not already at 32
					DispWidth *= 2;
				}

				//Encode new instruction 
				XED_STATE MachineState;
				MachineState.mmode = XED_MACHINE_MODE_LONG_64;
				MachineState.stack_addr_width = XED_ADDRESS_WIDTH_64b;
				XED_ENCODER_INSTRUCTION EncoderInstruction;
				XED_ENCODER_REQUEST EncoderRequest;
				UCHAR EncodeBuffer[15];
				UINT ReturnedSize;
				XED_ICLASS_ENUM IClass = XedDecodedInstGetIClass(&T->XedInstruction);

				//Do the encoding
				XedInst1(&EncoderInstruction, MachineState, IClass, DispWidth, XedRelBr(0, DispWidth));
				XedEncoderRequestZeroSetMode(&EncoderRequest, &MachineState);
				if (!XedConvertToEncoderRequest(&EncoderRequest, &EncoderInstruction))
					return FALSE;
				if (XED_ERROR_NONE != XedEncode(&EncoderRequest, EncodeBuffer, 15, &ReturnedSize))
					return FALSE;

				//fixup T->RawData
				delete[] T->RawData;
				T->RawDataSize = ReturnedSize;
				T->RawData = new UCHAR[ReturnedSize];
				RtlCopyMemory(T->RawData, EncodeBuffer, ReturnedSize);

				//Decode instruction so its proper and all that
				XedDecodedInstZeroSetMode(&T->XedInstruction, &MachineState);
				if (XED_ERROR_NONE != XedDecode(&T->XedInstruction, T->RawData, T->RawDataSize))
					return FALSE;

				//Go back to the start and loop through all labels again because now this instruction is larger :))))
				T = Block->Start;
				continue;
			}
			else
			{
				DispWidth = XedDecodedInstGetBranchDisplacementWidth(&T->XedInstruction);
				switch (DispWidth)
				{
				case 1: *(PINT8)&T->RawData[T->RawDataSize - DispWidth] = (INT8)BranchDisp; break;
				case 2: *(PINT16)&T->RawData[T->RawDataSize - DispWidth] = (INT16)BranchDisp; break;
				case 4: *(PINT32)&T->RawData[T->RawDataSize - DispWidth] = (INT32)BranchDisp; break;
				}
			}
		}

		T = T->Next;
	}
	return TRUE;
}

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->XedInstruction, (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->XedInstruction);
		Link->RawData = new UCHAR[Link->RawDataSize];
		RtlCopyMemory(Link->RawData, (Buf + Offset), Link->RawDataSize);

		NcAppendToBlock(Block, Link);

		Offset += Link->RawDataSize;
	}

	NcCreateLabels(Block);

	return TRUE;
}

PVOID NcAssemble(PNATIVE_CODE_BLOCK Block, PULONG OutSize)
{
	if (!NcFixRelJmps(Block))
		return NULL;

	*OutSize = NcCalcBlockSize(Block);

	PUCHAR Buffer = (PUCHAR)malloc(*OutSize);
	if (!Buffer)
		return NULL;

	PUCHAR BufferOffset = Buffer;

	for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
	{
		if (T->Flags & CODE_FLAG_IS_LABEL)
			continue;

		RtlCopyMemory(BufferOffset, T->RawData, T->RawDataSize);
		BufferOffset += T->RawDataSize;
	}

	return Buffer;
}

VOID NcDeleteBlock(PNATIVE_CODE_BLOCK Block)
{
	if (!Block->Start || !Block->End)
		return;

	PNATIVE_CODE_LINK BlockEnding = Block->End->Next;

	for (PNATIVE_CODE_LINK T = Block->Start; T && T != BlockEnding;)
	{
		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 != Block->End->Next; 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->XedInstruction);
			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 != Block->End->Next; 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] << ' ';
			}
		}
	}
}