#include "Jit.h"
#include "Flags.h"
#include "RipXorInst.h"
#include "RipAndInst.h"
#include "RipOrInst.h"
#include "RipMovInst.h"


BOOL JitMutateInstForXor(PNATIVE_CODE_LINK Link, PUCHAR ToMutate, PJIT_BITWISE_DATA JitData)
{
	ULONG FourByte = Link->RawDataSize / 4;
	ULONG TwoByte = (Link->RawDataSize - (FourByte * 4)) / 2;
	ULONG OneByte = (Link->RawDataSize - (FourByte * 4) - (TwoByte * 2));

	while (FourByte)
	{
		*(PULONG)ToMutate ^= JitData->Data[2 - FourByte];
		ToMutate += 4;
		FourByte--;
	}

	if (TwoByte)
	{
		*(PUSHORT)ToMutate ^= (USHORT)JitData->Data[3];
		ToMutate += 2;
	}

	if (OneByte)
		*(PUCHAR)ToMutate ^= (UCHAR)JitData->Data[3];

	return TRUE;
}

VOID JitMutateInstForOr(PNATIVE_CODE_LINK Link, PJIT_BITWISE_DATA JitData)
{

}

VOID JitMutateInstForAnd(PNATIVE_CODE_LINK Link, PJIT_BITWISE_DATA JitData)
{

}

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

INLINE BOOL JitiEmitWrapperD(ULONG OpType, PNATIVE_CODE_BLOCK Block, INT32 RipDelta, ULONG Value)
{
	switch (OpType)
	{
	case JIT_BITWISE_XOR: return JitEmitRipRelativeXorD(Block, RipDelta, Value);
	case JIT_BITWISE_AND: return JitEmitRipRelativeAndD(Block, RipDelta, Value);
	case JIT_BITWISE_OR: return JitEmitRipRelativeOrD(Block, RipDelta, Value);
	}
}
INLINE BOOL JitiEmitWrapperW(ULONG OpType, PNATIVE_CODE_BLOCK Block, INT32 RipDelta, ULONG Value)
{
	switch (OpType)
	{
	case JIT_BITWISE_XOR: return JitEmitRipRelativeXorW(Block, RipDelta, Value);
	case JIT_BITWISE_AND: return JitEmitRipRelativeAndW(Block, RipDelta, Value);
	case JIT_BITWISE_OR: return JitEmitRipRelativeOrW(Block, RipDelta, Value);
	}
}
INLINE BOOL JitiEmitWrapperB(ULONG OpType, PNATIVE_CODE_BLOCK Block, INT32 RipDelta, ULONG Value)
{
	switch (OpType)
	{
	case JIT_BITWISE_XOR: return JitEmitRipRelativeXorB(Block, RipDelta, Value);
	case JIT_BITWISE_AND: return JitEmitRipRelativeAndB(Block, RipDelta, Value);
	case JIT_BITWISE_OR: return JitEmitRipRelativeOrB(Block, RipDelta, Value);
	}
}

PNATIVE_CODE_BLOCK JitEmitPreRipBitwiseOp(PNATIVE_CODE_LINK Link, PJIT_BITWISE_DATA JitData, ULONG OpType, BOOL SaveFlags, 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;
	if (!Block)
		return NULL;

	if (SaveFlags)
	{
		PNATIVE_CODE_LINK PushF = FlgEmitPushfqInst();
		PushF->Flags |= CODE_FLAG_DO_NOT_DIVIDE;
		NcAppendToBlock(Block, PushF);
	}

	ULONG Count = FourByte;
	while (Count)
	{
		INT32 RipDelta = (((Count - 1) * DWORD_RIP_INST_LENGTH) + (TwoByte * WORD_RIP_INST_LENGTH) + (OneByte * BYTE_RIP_INST_LENGTH));
		if (SaveFlags)
			RipDelta += 1;
		RipDelta += ((FourByte - Count) * 4);
		RipDelta += Delta;
		if (!JitiEmitWrapperD(OpType, Block, RipDelta, JitData->Data[FourByte - Count]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
		--Count;
	}

	if (TwoByte)
	{
		INT32 RipDelta = (OneByte * BYTE_RIP_INST_LENGTH);
		if (SaveFlags)
			RipDelta += 1;
		RipDelta += (FourByte * 4);
		RipDelta += Delta;
		if (!JitiEmitWrapperW(OpType, Block, RipDelta, JitData->Data[3]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

	if (OneByte)
	{
		INT32 RipDelta = 0;
		if (SaveFlags)
			RipDelta += 1;
		RipDelta += (FourByte * 4) + (TwoByte * 2);
		RipDelta += Delta;
		if (!JitiEmitWrapperB(OpType, Block, RipDelta, JitData->Data[4]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

	if (SaveFlags)
	{
		PNATIVE_CODE_LINK PopF = FlgEmitPopfqInst();
		PopF->Flags |= CODE_FLAG_DO_NOT_DIVIDE;
		NcAppendToBlock(Block, PopF);
	}

	return Block;
}

PNATIVE_CODE_BLOCK JitEmitPostRipBitwiseOp(PNATIVE_CODE_LINK Link, PJIT_BITWISE_DATA JitData, ULONG OpType, BOOL SaveFlags, 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;
	if (!Block)
		return NULL;

	if (SaveFlags)
	{
		PNATIVE_CODE_LINK PushF = FlgEmitPushfqInst();
		PushF->Flags |= CODE_FLAG_DO_NOT_DIVIDE;
		NcAppendToBlock(Block, PushF);
	}

	ULONG Count = FourByte;
	while (Count)
	{
		INT32 RipDelta = Link->RawDataSize - ((FourByte - Count) * 4);
		if (SaveFlags)
			RipDelta += 1;
		RipDelta += (FourByte - (Count - 1)) * DWORD_RIP_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += Delta;
		if (!JitiEmitWrapperD(OpType, Block, RipDelta, JitData->Data[FourByte - Count]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
		--Count;
	}

	if (TwoByte)
	{
		INT32 RipDelta = Link->RawDataSize - (FourByte * 4);
		if (SaveFlags)
			RipDelta += 1;
		RipDelta += (FourByte * DWORD_RIP_INST_LENGTH);
		RipDelta += WORD_RIP_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += Delta;
		if (!JitiEmitWrapperW(OpType, Block, RipDelta, JitData->Data[3]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

	if (OneByte)
	{
		INT32 RipDelta = Link->RawDataSize - (FourByte * 4) - (TwoByte * 2);
		if (SaveFlags)
			RipDelta += 1;
		RipDelta += (FourByte * DWORD_RIP_INST_LENGTH);
		RipDelta += WORD_RIP_INST_LENGTH;
		RipDelta += BYTE_RIP_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += Delta;
		if (!JitiEmitWrapperB(OpType, Block, RipDelta, JitData->Data[4]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

	if (SaveFlags)
	{
		PNATIVE_CODE_LINK PopF = FlgEmitPopfqInst();
		PopF->Flags |= CODE_FLAG_DO_NOT_DIVIDE;
		NcAppendToBlock(Block, PopF);
	}

	return Block;
}