#include "Obfuscator.h"

BOOL JitEmitRipRelativeMovD(PNATIVE_CODE_BLOCK Block, INT32 RipDelta, PUCHAR Data)
{
	UCHAR RawData[] = { 0xC7, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST | CODE_FLAG_DO_NOT_DIVIDE, RawData, sizeof(RawData));
	*(PINT32)&Link->RawData[2] = RipDelta;
	RtlCopyMemory(&Link->RawData[6], Data, 4);
	XedDecode(&Link->XedInstruction, Link->RawData, Link->RawDataSize);
	NcAppendToBlock(Block, Link);
	return TRUE;
}

BOOL JitEmitRipRelativeMovW(PNATIVE_CODE_BLOCK Block, INT32 RipDelta, PUCHAR Data)
{
	UCHAR RawData[] = { 0x66, 0xC7, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST | CODE_FLAG_DO_NOT_DIVIDE, RawData, sizeof(RawData));
	*(PINT32)&Link->RawData[3] = RipDelta;
	RtlCopyMemory(&Link->RawData[7], Data, 2);
	XedDecode(&Link->XedInstruction, Link->RawData, Link->RawDataSize);
	NcAppendToBlock(Block, Link);
	return TRUE;
}

BOOL JitEmitRipRelativeMovB(PNATIVE_CODE_BLOCK Block, INT32 RipDelta, PUCHAR Data)
{
	UCHAR RawData[] = { 0xC6, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00 };

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST | CODE_FLAG_DO_NOT_DIVIDE, RawData, sizeof(RawData));
	*(PINT32)&Link->RawData[2] = RipDelta;
	Link->RawData[6] = *Data;
	XedDecode(&Link->XedInstruction, Link->RawData, Link->RawDataSize);
	NcAppendToBlock(Block, Link);
	return TRUE;
}

PNATIVE_CODE_BLOCK JitEmitPreRipMov(PNATIVE_CODE_LINK Link, INT32 Delta)
{
	ULONG FourByte = Link->RawDataSize / 4;
	ULONG TwoByte = (Link->RawDataSize - (FourByte * 4)) / 2;
	ULONG OneByte = (Link->RawDataSize - (FourByte * 4) - (TwoByte * 2));

	PNATIVE_CODE_BLOCK Block = new NATIVE_CODE_BLOCK;

	Block->Start = Block->End = new NATIVE_CODE_LINK;
	PUCHAR DataOffset = Link->RawData;
	ULONG Count = FourByte;
	while (Count)
	{
		//Account for remaining MOVs
		INT32 RipDelta = (((Count - 1) * DWORD_MOV_INST_LENGTH) + (TwoByte * WORD_MOV_INST_LENGTH) + (OneByte * BYTE_MOV_INST_LENGTH));
		//Account for already MOVd instructions
		RipDelta += ((FourByte - Count) * 4);
		RipDelta += Delta;
		//Add the actual instruction
		if (!JitEmitRipRelativeMovD(Block, RipDelta, DataOffset))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
		DataOffset += 4;
		--Count;
	}

	if (TwoByte)
	{
		INT32 RipDelta = (OneByte * BYTE_MOV_INST_LENGTH);
		RipDelta += (FourByte * 4);
		RipDelta += Delta;
		if (!JitEmitRipRelativeMovW(Block, RipDelta, DataOffset))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
		DataOffset += 2;
	}

	if (OneByte)
	{
		INT32 RipDelta = 0;
		RipDelta += (FourByte * 4) + (TwoByte * 2);
		RipDelta += Delta;
		if (!JitEmitRipRelativeMovB(Block, RipDelta, DataOffset))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

	PNATIVE_CODE_LINK StartLink = Block->Start;
	Block->Start = Block->Start->Next;
	if (Block->Start)
		Block->Start->Prev = NULL;
	delete StartLink;

	return Block;
}

PNATIVE_CODE_BLOCK JitEmitPostRipMov(PNATIVE_CODE_LINK Link, INT32 Delta)
{
	ULONG FourByte = Link->RawDataSize / 4;
	ULONG TwoByte = (Link->RawDataSize - (FourByte * 4)) / 2;
	ULONG OneByte = (Link->RawDataSize - (FourByte * 4) - (TwoByte * 2));

	PNATIVE_CODE_BLOCK Block = new NATIVE_CODE_BLOCK;

	Block->Start = Block->End = new NATIVE_CODE_LINK;
	ULONG ZeroValue = 0;
	ULONG Count = FourByte;
	while (Count)
	{
		INT32 RipDelta = Link->RawDataSize - ((FourByte - Count) * 4);
		RipDelta += (FourByte - (Count - 1)) * DWORD_MOV_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += Delta;
		ZeroValue = rand();
		if (!JitEmitRipRelativeMovD(Block, RipDelta, (PUCHAR)&ZeroValue))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
		--Count;
	}

	if (TwoByte)
	{
		INT32 RipDelta = Link->RawDataSize - (FourByte * 4);
		RipDelta += (FourByte * DWORD_MOV_INST_LENGTH);
		RipDelta += WORD_MOV_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += Delta;
		ZeroValue = rand();
		if (!JitEmitRipRelativeMovW(Block, RipDelta, (PUCHAR)&ZeroValue))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

	if (OneByte)
	{
		INT32 RipDelta = Link->RawDataSize - (FourByte * 4) - (TwoByte * 2);
		RipDelta += (FourByte * DWORD_MOV_INST_LENGTH);
		RipDelta += (TwoByte * WORD_MOV_INST_LENGTH);
		RipDelta += BYTE_MOV_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += Delta;
		ZeroValue = rand();
		if (!JitEmitRipRelativeMovB(Block, RipDelta, (PUCHAR)&ZeroValue))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

	PNATIVE_CODE_LINK StartLink = Block->Start;
	Block->Start = Block->Start->Next;
	if (Block->Start)
		Block->Start->Prev = NULL;
	delete StartLink;

	return Block;
}