#include "RipMovInst.h"


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

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK;
	Link->Flags = CODE_FLAG_IS_INST;
	Link->RawData = new UCHAR[DWORD_MOV_INST_LENGTH];
	Link->RawDataSize = DWORD_MOV_INST_LENGTH;
	memcpy(Link->RawData, Bytes, DWORD_MOV_INST_LENGTH);

	*(PINT32)&Link->RawData[DWORD_MOV_INST_RIP_OFFSET] = RipDelta;
	memcpy(&Link->RawData[DWORD_MOV_INST_MOV_OFFSET], Data, 4);

	XED_ERROR_ENUM DecodeError = XedDecode(&Link->XedInst, Link->RawData, Link->RawDataSize);
	if (DecodeError != XED_ERROR_NONE)
	{
		printf("XedDecode failed with error %s\n", XedErrorEnumToString(DecodeError));
		delete Link;
		return FALSE;
	}
	Link->Block = Block;
	Link->Prev = Block->End;
	Block->End->Next = Link;
	Block->End = Link;
	return TRUE;
}

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

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK;
	Link->Flags = CODE_FLAG_IS_INST;
	Link->RawData = new UCHAR[WORD_MOV_INST_LENGTH];
	Link->RawDataSize = WORD_MOV_INST_LENGTH;
	memcpy(Link->RawData, Bytes, WORD_MOV_INST_LENGTH);

	*(PINT32)&Link->RawData[WORD_MOV_INST_RIP_OFFSET] = RipDelta;
	memcpy(&Link->RawData[WORD_MOV_INST_MOV_OFFSET], Data, 2);

	XED_ERROR_ENUM DecodeError = XedDecode(&Link->XedInst, Link->RawData, Link->RawDataSize);
	if (DecodeError != XED_ERROR_NONE)
	{
		printf("XedDecode failed with error %s\n", XedErrorEnumToString(DecodeError));
		delete Link;
		return FALSE;
	}
	Link->Block = Block;
	Link->Prev = Block->End;
	Block->End->Next = Link;
	Block->End = Link;
	return TRUE;
}

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

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK;
	Link->Flags = CODE_FLAG_IS_INST;
	Link->RawData = new UCHAR[BYTE_MOV_INST_LENGTH];
	Link->RawDataSize = BYTE_MOV_INST_LENGTH;
	memcpy(Link->RawData, Bytes, BYTE_MOV_INST_LENGTH);

	*(PINT32)&Link->RawData[BYTE_MOV_INST_RIP_OFFSET] = RipDelta;
	Link->RawData[BYTE_MOV_INST_MOV_OFFSET] = *Data;

	XED_ERROR_ENUM DecodeError = XedDecode(&Link->XedInst, Link->RawData, Link->RawDataSize);
	if (DecodeError != XED_ERROR_NONE)
	{
		printf("XedDecode failed with error %s\n", XedErrorEnumToString(DecodeError));
		delete Link;
		return FALSE;
	}
	Link->Block = Block;
	Link->Prev = Block->End;
	Block->End->Next = Link;
	Block->End = Link;
	return TRUE;
}

PNATIVE_CODE_BLOCK ObfEmitPreMovForInst(PNATIVE_CODE_LINK Link)
{
	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);
		//Add the actual instruction
		if (!ObfEmitRipRelativeMovD(Block, RipDelta, DataOffset))
		{
			NcDelete(Block);
			delete Block;
			return NULL;
		}
		DataOffset += 4;
		--Count;
	}

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

	if (OneByte)
	{
		INT32 RipDelta = 0;
		RipDelta += (FourByte * 4) + (TwoByte * 2);
		if (!ObfEmitRipRelativeMovB(Block, RipDelta, DataOffset))
		{
			NcDelete(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 ObfEmitPostMovForInst(PNATIVE_CODE_LINK Link)
{
	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);
		if (!ObfEmitRipRelativeMovD(Block, RipDelta, (PUCHAR)&ZeroValue))
		{
			NcDelete(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);

		if (!ObfEmitRipRelativeMovW(Block, RipDelta, (PUCHAR)&ZeroValue))
		{
			NcDelete(Block);
			delete Block;
			return NULL;
		}
	}

	if (OneByte)
	{
		INT32 RipDelta = Link->RawDataSize - (FourByte * 4) - (TwoByte * 2);
		RipDelta += (FourByte * DWORD_MOV_INST_LENGTH);
		RipDelta += WORD_MOV_INST_LENGTH;
		RipDelta += BYTE_MOV_INST_LENGTH;
		RipDelta *= (-1);
		if (!ObfEmitRipRelativeMovB(Block, RipDelta, (PUCHAR)&ZeroValue))
		{
			NcDelete(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;
}