#include "Obfuscator.h"


BOOL ObfiRandomizeInstruction(PNATIVE_CODE_LINK Link, PUCHAR ToMutate, PVOID Context)
{
	for (UINT i = 0; i < Link->RawDataSize; i++)
		ToMutate[i] = RndGetRandomInt(0, 255);
	return TRUE;
}

VOID ObfGenOpaqueBranch(PNATIVE_CODE_BLOCK Block, FLOAT BranchSizePercentage)
{
	UINT BlockSize = NcCountInstructions(Block, TRUE);
	UINT BranchSize = (FLOAT)BlockSize * BranchSizePercentage;
	UINT BranchStartPos = RndGetRandomInt(0, BlockSize - 1 - BranchSize);
	PNATIVE_CODE_LINK BranchStart = BranchStartPos ? NULL : Block->Start;
	for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next;)
	{
		if (T->Flags & CODE_FLAG_IS_LABEL || (T->Flags & CODE_FLAG_DO_NOT_DIVIDE && !(T->Flags & CODE_FLAG_GROUP_END)))
		{
			T = T->Next;
			continue;
		}

		if (BranchStartPos)
		{
			--BranchStartPos;
			if (BranchStartPos == 0)
				BranchStart = T;
			T = T->Next;
			continue;
		}

		if (BranchSize)
		{
			--BranchSize;
			T = T->Next;
			continue;
		}

		NATIVE_CODE_BLOCK NotTaken, Taken;
		ObfCreateOpaqueBranches(BranchStart, T, &NotTaken, &Taken);
		ObfCombineOpaqueBranches(&NotTaken, &Taken, NcGenUnusedLabelId(Block), NcGenUnusedLabelId(Block));
		NcInsertBlockStartToEnd(BranchStart, T, &NotTaken);
		return;
	}
}

UINT ObfMutateInstructions(PNATIVE_CODE_BLOCK Block, UINT MutateChance, BOOL MutateFirst, BOOL MutateOnlyFirst)
{
	UINT MutatedInstructions = 0;
	for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next;)
	{
		if ((T->Flags & CODE_FLAG_IS_LABEL) || (T->Flags & CODE_FLAG_DO_NOT_DIVIDE) || (T->Flags & CODE_FLAG_IS_REL_JMP))
		{
			T = T->Next;
			continue;
		}

		PNATIVE_CODE_LINK RealNext = T->Next;

		if (RndGetRandomInt(0, 100) <= MutateChance || MutateFirst)
		{
			MutateFirst = FALSE;

			PNATIVE_CODE_BLOCK PreOp = JitEmitPreRipMov(T);
			PNATIVE_CODE_BLOCK PostOp = JitEmitPostRipMov(T);

			PreOp->Start->Flags |= CODE_FLAG_GROUP_START;
			PostOp->End->Flags |= CODE_FLAG_GROUP_END;
			T->Flags |= CODE_FLAG_DO_NOT_DIVIDE;
			T->Flags |= CODE_FLAG_HAS_ASM_OP;
			T->AsmOperations.emplace_back((FN_INST_ASM_OP)ObfiRandomizeInstruction, (PVOID)NULL);

			NcInsertBlockBefore(T, PreOp, FALSE);
			NcInsertBlockAfter(T, PostOp, FALSE);

			if (T == Block->End)
				Block->End = PostOp->End;
			if (T == Block->Start)
				Block->Start = PreOp->Start;

			if (Block->Start == T)
				Block->Start = PreOp->Start;
			if (Block->End == T)
				Block->End = PostOp->End;

			delete PreOp;
			delete PostOp;

			MutatedInstructions++;

			if (MutateOnlyFirst)
				return MutatedInstructions;

		}

		T = RealNext;
	}
	return MutatedInstructions;
}