#include "OpaqueBranching.h"

XED_ICLASS_ENUM ObfGetRandomJccClass()
{
	switch (rand() % 15)
	{
	case 0: return XED_ICLASS_JL;
	case 1: return XED_ICLASS_JLE;
	case 2:	return XED_ICLASS_JNB;
	case 3: return XED_ICLASS_JNBE;
	case 4: return XED_ICLASS_JNL;
	case 5: return XED_ICLASS_JNLE;
	case 6: return XED_ICLASS_JNO;
	case 7: return XED_ICLASS_JNP;
	case 8: return XED_ICLASS_JNS;
	case 9: return XED_ICLASS_JNZ;
	case 10: return XED_ICLASS_JO;
	case 11: return XED_ICLASS_JP;
	case 12: return XED_ICLASS_JRCXZ;
	case 13: return XED_ICLASS_JS;
	case 14: return XED_ICLASS_JZ;
	}
	return XED_ICLASS_JLE;
}

PNATIVE_CODE_LINK ObfGenRandomJcc(ULONG LabelId, ULONG DisplacementWidth)
{
	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;

	XedInst1(&EncoderInstruction, MachineState, ObfGetRandomJccClass(), DisplacementWidth, XedRelBr(0, DisplacementWidth));

	XedEncoderRequestZeroSetMode(&EncoderRequest, &MachineState);
	if (!XedConvertToEncoderRequest(&EncoderRequest, &EncoderInstruction))
		return NULL;

	if (XED_ERROR_NONE != XedEncode(&EncoderRequest, EncodeBuffer, 15, &ReturnedSize))
		return NULL;

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST, EncodeBuffer, ReturnedSize);
	if (XED_ERROR_NONE != XedDecode(&Link->XedInstruction, Link->RawData, Link->RawDataSize))
	{
		delete Link;
			return NULL;
	}
	Link->Label = LabelId;
	Link->Flags = (CODE_FLAG_IS_INST | CODE_FLAG_IS_REL_JMP);

	return Link;
}

PNATIVE_CODE_LINK ObfGenJmpToLabel(ULONG LabelId, ULONG DisplacementWidth)
{
	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;

	XedInst1(&EncoderInstruction, MachineState, XED_ICLASS_JMP, DisplacementWidth, XedRelBr(0, DisplacementWidth));

	XedEncoderRequestZeroSetMode(&EncoderRequest, &MachineState);
	if (!XedConvertToEncoderRequest(&EncoderRequest, &EncoderInstruction))
		return NULL;

	if (XED_ERROR_NONE != XedEncode(&EncoderRequest, EncodeBuffer, 15, &ReturnedSize))
		return NULL;

	PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK(CODE_FLAG_IS_INST, EncodeBuffer, ReturnedSize);
	if (XED_ERROR_NONE != XedDecode(&Link->XedInstruction, Link->RawData, Link->RawDataSize))
	{
		delete Link;
		return NULL;
	}
	Link->Label = LabelId;
	Link->Flags = (CODE_FLAG_IS_INST | CODE_FLAG_IS_REL_JMP);

	return Link;
}

PNATIVE_CODE_BLOCK ObfGenOpaqueBranch(PNATIVE_CODE_LINK Start, PNATIVE_CODE_LINK End)
{
	if (!Start || !End || !Start->Block || Start->Block != End->Block)
		return NULL;

	PNATIVE_CODE_BLOCK NotTaken = NcDeepCopyPartialBlock(Start, End);
	if (!NotTaken)
	{
		return NULL;
	}
	PNATIVE_CODE_BLOCK Taken = NcDeepCopyPartialBlock(Start, End);
	if (!Taken)
	{
		NcDeleteBlock(NotTaken);
		delete NotTaken;
		return NULL;
	}

	ULONG JccLabel = NcGenUnusedLabelId(Start->Block);
	ULONG JmpLabel = NcGenUnusedLabelId(Start->Block);
	Start->Block->LabelIds.push_back(JccLabel);
	Start->Block->LabelIds.push_back(JmpLabel);

	PNATIVE_CODE_LINK Jcc = ObfGenRandomJcc(JccLabel);
	if (!Jcc)
	{
		NcDeleteBlock(Taken);
		delete Taken;
		NcDeleteBlock(NotTaken);
		delete NotTaken;
		return NULL;
	}
	PNATIVE_CODE_LINK Jmp = ObfGenJmpToLabel(JmpLabel);
	if (!Jmp)
	{
		delete Jcc;
		NcDeleteBlock(Taken);
		delete Taken;
		NcDeleteBlock(NotTaken);
		delete NotTaken;
		return NULL;
	}

	NcPrependToBlock(NotTaken, Jcc);
	NcAppendToBlock(NotTaken, Jmp);

	NcPrependToBlock(Taken, new NATIVE_CODE_LINK(JccLabel, Taken));
	NcAppendToBlock(Taken, new NATIVE_CODE_LINK(JmpLabel, Taken));

	NcInsertBlockAfter(NotTaken->End, Taken, FALSE);
	NotTaken->End = Taken->End;

	delete Taken;
	return NotTaken;
}