#include "OpaqueBranching.h"

XED_ICLASS_ENUM ObfGetRandomJccClass()
{
	switch (rand() % 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;
}

BOOL ObfInsertOpaqueBranchBlock(PNATIVE_CODE_LINK Start, PNATIVE_CODE_LINK End, PNATIVE_CODE_BLOCK OpaqueBranchBlock)
{
	OpaqueBranchBlock->Start->Prev = Start->Prev;
	OpaqueBranchBlock->End->Next = End->Next;

	if (Start->Prev)
		Start->Prev->Next = OpaqueBranchBlock->Start;
	if (End->Next)
		End->Next->Prev = OpaqueBranchBlock->End;

	if (Start->Block->Start == Start)
		Start->Block->Start = OpaqueBranchBlock->Start;

	if (Start->Block->End == End)
		Start->Block->End = OpaqueBranchBlock->End;

	//Update group for the current isntructions
	for (PNATIVE_CODE_LINK T = OpaqueBranchBlock->Start; T && T != OpaqueBranchBlock->End->Next; T = T->Next)
		T->Block = Start->Block;

	PNATIVE_CODE_LINK EndBlock = End->Next;
	for (PNATIVE_CODE_LINK T = Start; T && T != EndBlock;)
	{
		PNATIVE_CODE_LINK RealNext = T->Next;
		delete T;
		T = RealNext;
	}
	return TRUE;
}