#include "RipXorInst.h"


BOOL ObfEmitPushfqInst(PNATIVE_CODE_BLOCK Block)
{
	UCHAR RawData[] = { 0x9C };
	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST, RawData, 1);
	XedDecode(&Link->XedInstruction, Link->RawData, 1);
	NcAppendToBlock(Block, Link);
	return TRUE;
}

BOOL ObfEmitPopfqInst(PNATIVE_CODE_BLOCK Block)
{
	UCHAR RawData[] = { 0x9D };
	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST, RawData, 1);
	XedDecode(&Link->XedInstruction, Link->RawData, 1);
	NcAppendToBlock(Block, Link);
	return TRUE;
}

BOOL ObfEmitRipRelativeXorD(PNATIVE_CODE_BLOCK Block, INT32 RipDelta, ULONG Value)
{
	UCHAR RawData[] = { 0x81, 0x35, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST, RawData, sizeof(RawData));
	*(PINT32)&Link->RawData[DWORD_XOR_INST_RIP_OFFSET] = RipDelta;
	*(PULONG)&Link->RawData[DWORD_XOR_INST_XOR_OFFSET] = Value;
	XedDecode(&Link->XedInstruction, Link->RawData, Link->RawDataSize);
	NcAppendToBlock(Block, Link);
	return TRUE;
}

BOOL ObfEmitRipRelativeXorW(PNATIVE_CODE_BLOCK Block, INT32 RipDelta, ULONG Value)
{
	UCHAR RawData[] = { 0x66, 0x81, 0x35, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST, RawData, sizeof(RawData));
	*(PINT32)&Link->RawData[WORD_XOR_INST_RIP_OFFSET] = RipDelta;
	*(PUSHORT)&Link->RawData[WORD_XOR_INST_XOR_OFFSET] = (USHORT)Value;
	XedDecode(&Link->XedInstruction, Link->RawData, Link->RawDataSize);
	NcAppendToBlock(Block, Link);
	return TRUE;
}

BOOL ObfEmitRipRelativeXorB(PNATIVE_CODE_BLOCK Block, INT32 RipDelta, ULONG Value)
{
	UCHAR RawData[] = { 0x80, 0x35, 0x00, 0x00, 0x00, 0x00, 0x00 };

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST, RawData, sizeof(RawData));
	*(PINT32)&Link->RawData[BYTE_XOR_INST_RIP_OFFSET] = RipDelta;
	*(PUCHAR)&Link->RawData[BYTE_XOR_INST_XOR_OFFSET] = (UCHAR)Value;
	XedDecode(&Link->XedInstruction, Link->RawData, Link->RawDataSize);
	NcAppendToBlock(Block, Link);
	return TRUE;
}

VOID ObfXorInstBytes(PNATIVE_CODE_LINK Link, PXOR_INST_DATA XorData)
{
	ULONG FourByte = Link->RawDataSize / 4;
	ULONG TwoByte = (Link->RawDataSize - (FourByte * 4)) / 2;
	ULONG OneByte = (Link->RawDataSize - (FourByte * 4) - (TwoByte * 2));

	PUCHAR Buffer = Link->RawData;
	while (FourByte)
	{
		*(PULONG)Buffer ^= XorData->Data[2 - FourByte];
		Buffer += 4;
		FourByte--;
	}

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

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

}

BOOL ObfiCheckXorFlagCollisions(CONST XED_FLAG_SET* FlagsRead, XED_FLAG_SET Ledger)
{
	return ((FlagsRead->s.zf && FlagsRead->s.zf == Ledger.s.zf) ||
		(FlagsRead->s.sf && FlagsRead->s.sf == Ledger.s.sf) ||
		(FlagsRead->s.pf && FlagsRead->s.pf == Ledger.s.pf) ||
		(FlagsRead->s.of && FlagsRead->s.of == Ledger.s.of) ||
		(FlagsRead->s.cf && FlagsRead->s.cf == Ledger.s.cf) ||
		(FlagsRead->s.af && FlagsRead->s.af == Ledger.s.af)
		);
}

VOID ObfiUpdateXorLedger(CONST XED_FLAG_SET* FlagsWritten, XED_FLAG_SET* Ledger)
{
	if (FlagsWritten->s.zf)
		Ledger->s.zf = FALSE;
	if (FlagsWritten->s.sf)
		Ledger->s.sf = FALSE;
	if (FlagsWritten->s.pf)
		Ledger->s.pf = FALSE;
	if (FlagsWritten->s.of)
		Ledger->s.of = FALSE;
	if (FlagsWritten->s.cf)
		Ledger->s.cf = FALSE;
	if (FlagsWritten->s.af)
		Ledger->s.af = FALSE;
}

BOOL ObfDoesInstWriteToAllXorFlags(PNATIVE_CODE_LINK Link)
{
	CONST XED_SIMPLE_FLAG* SimpleFlags = XedDecodedInstGetRflagsInfo(&Link->XedInstruction);
	CONST XED_FLAG_SET* FlagsWritten = XedSimpleFlagGetWrittenFlagSet(SimpleFlags);
	CONST XED_FLAG_SET* FlagsUndefined = XedSimpleFlagGetUndefinedFlagSet(SimpleFlags);

	return (FlagsWritten->s.zf &&
		FlagsWritten->s.sf &&
		FlagsWritten->s.pf &&
		FlagsWritten->s.of &&
		FlagsWritten->s.cf &&
		FlagsUndefined->s.af
		);
}

BOOL ObfAreXorFlagsClobberedBeforeUse(PNATIVE_CODE_LINK Link)
{
	XED_FLAG_SET Ledger;
	Ledger.s.zf = TRUE;
	Ledger.s.sf = TRUE;
	Ledger.s.pf = TRUE;
	Ledger.s.of = TRUE;
	Ledger.s.cf = TRUE;
	Ledger.s.af = TRUE;

	for (PNATIVE_CODE_LINK T = Link->Next; T; T = T->Next)
	{
		if (T->Flags & CODE_FLAG_IS_LABEL)
			continue;

		CONST XED_SIMPLE_FLAG* SimpleFlags = XedDecodedInstGetRflagsInfo(&T->XedInstruction);
		CONST XED_FLAG_SET* FlagsRead = XedSimpleFlagGetReadFlagSet(SimpleFlags);
		CONST XED_FLAG_SET* FlagsWritten = XedSimpleFlagGetWrittenFlagSet(SimpleFlags);

		if (ObfiCheckXorFlagCollisions(FlagsRead, Ledger))
			return FALSE;

		ObfiUpdateXorLedger(FlagsWritten, &Ledger);

		if (Ledger.flat == 0)
			return TRUE;
	}
	return FALSE;
}

PNATIVE_CODE_BLOCK ObfEmitPreXorForInst(PNATIVE_CODE_LINK Link, PXOR_INST_DATA XorData, BOOL SaveFlags, INT32 DeltaToInst)
{
	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 && !ObfEmitPushfqInst(Block))
	{
		NcDeleteBlock(Block);
		delete Block;
		return NULL;
	}
		
	ULONG Count = FourByte;
	while (Count)
	{
		//Account for remaining XORs
		INT32 RipDelta = (((Count - 1) * DWORD_XOR_INST_LENGTH) + (TwoByte * WORD_XOR_INST_LENGTH) + (OneByte * BYTE_XOR_INST_LENGTH));
		//Account for POPFQ
		if (SaveFlags)
			RipDelta += 1;
		//Account for already XORd instructions
		RipDelta += ((FourByte - Count) * 4);
		RipDelta += DeltaToInst;
		//Add the actual instruction
		if (!ObfEmitRipRelativeXorD(Block, RipDelta, XorData->Data[FourByte-Count]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
		--Count;
	}

	if (TwoByte)
	{
		INT32 RipDelta = (OneByte * BYTE_XOR_INST_LENGTH);
		if (SaveFlags)
			RipDelta += 1;
		RipDelta += (FourByte * 4);
		RipDelta += DeltaToInst;
		if (!ObfEmitRipRelativeXorW(Block, RipDelta, XorData->Data[3]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

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

	if (SaveFlags && !ObfEmitPopfqInst(Block))
	{
		NcDeleteBlock(Block);
		delete Block;
		return NULL;
	}

	return Block;
}

PNATIVE_CODE_BLOCK ObfEmitPostXorForInst(PNATIVE_CODE_LINK Link, PXOR_INST_DATA XorData, BOOL SaveFlags, INT32 DeltaToInst)
{
	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 && !ObfEmitPushfqInst(Block))
	{
		NcDeleteBlock(Block);
		delete Block;
		return NULL;
	}

	ULONG Count = FourByte;
	while (Count)
	{
		INT32 RipDelta = Link->RawDataSize - ((FourByte - Count) * 4);
		if (SaveFlags)
			RipDelta += 1;
		RipDelta += (FourByte - (Count - 1)) * DWORD_XOR_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += DeltaToInst;
		if (!ObfEmitRipRelativeXorD(Block, RipDelta, XorData->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_XOR_INST_LENGTH);
		RipDelta += WORD_XOR_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += DeltaToInst;
		if (!ObfEmitRipRelativeXorW(Block, RipDelta, XorData->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_XOR_INST_LENGTH);
		RipDelta += WORD_XOR_INST_LENGTH;
		RipDelta += BYTE_XOR_INST_LENGTH;
		RipDelta *= (-1);
		RipDelta += DeltaToInst;
		if (!ObfEmitRipRelativeXorB(Block, RipDelta, XorData->Data[4]))
		{
			NcDeleteBlock(Block);
			delete Block;
			return NULL;
		}
	}

	if (SaveFlags && !ObfEmitPopfqInst(Block))
	{
		NcDeleteBlock(Block);
		delete Block;
		return NULL;
	}

	return Block;
}