#include "Obfuscator.h"

XED_ICLASS_ENUM ObfGetRandomJccClass()
{
	switch (RndGetRandomInt(0, 14))
	{
	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 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;
}

BOOL ObfCreateOpaqueBranches(PNATIVE_CODE_LINK Start, PNATIVE_CODE_LINK End, PNATIVE_CODE_BLOCK NotTaken, PNATIVE_CODE_BLOCK Taken)
{
	return (NcDeepCopyPartialBlock(Start, End, Taken) && !NcDeepCopyPartialBlock(Start, End, NotTaken));
}

BOOL ObfCombineOpaqueBranches(PNATIVE_CODE_BLOCK NotTaken, PNATIVE_CODE_BLOCK Taken, ULONG JccLabel, ULONG JmpLabel)
{
	PNATIVE_CODE_LINK Jcc = ObfGenRandomJcc(JccLabel);
	if (!Jcc)
		return FALSE;
	PNATIVE_CODE_LINK Jmp = ObfGenJmpToLabel(JmpLabel);
	if (!Jmp)
	{
		delete Jcc;
		return FALSE;
	}

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