#include "Obfuscator.h"


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

VOID ObfGenerateOpaqueBranches(POPBR_SETS Obf, PNATIVE_CODE_BLOCK Block, ULONG Depth)
{
	if (Depth > Obf->MaxDepth)
		return;

	ULONG InstructionCount = NcCountInstructions(Block, TRUE);
	if (InstructionCount > Obf->MinBranchSize)
	{
		ULONG TargetCount = (ULONG)((FLOAT)InstructionCount / Obf->Divisor);
		ULONG CurrentCount = 0;
		PNATIVE_CODE_LINK NewBlockStart = 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 (++CurrentCount >= TargetCount)
			{
				if (Depth >= Obf->MinDepthForBranch && (rand() % 100) <= Obf->ChanceForBranch && CurrentCount >= Obf->MinBranchSize)
				{
					NATIVE_CODE_BLOCK NotTaken, Taken;
					ObfCreateOpaqueBranches(NewBlockStart, T, &NotTaken, &Taken);
					ObfGenerateOpaqueBranches(Obf, &NotTaken, Depth + 1);
					ObfGenerateOpaqueBranches(Obf, &Taken, Depth + 1);
					ObfCombineOpaqueBranches(&NotTaken, &Taken, NcGenUnusedLabelId(Obf->ParentBlock), NcGenUnusedLabelId(Obf->ParentBlock));
					ObfInsertOpaqueBranchBlock(NewBlockStart, T, &NotTaken);
					T = NotTaken.End;
				}
				else
				{
					NATIVE_CODE_BLOCK TempBlock;
					if (NcDeepCopyPartialBlock(NewBlockStart, T, &TempBlock))
					{
						ObfGenerateOpaqueBranches(Obf, &TempBlock, Depth + 1);
						ObfInsertOpaqueBranchBlock(NewBlockStart, T, &TempBlock);
					}
					T = TempBlock.End;
				}
				NewBlockStart = T->Next;
				CurrentCount = 0;
			}
			T = T->Next;
		}
		if (NewBlockStart) //Deal with remaining instructions in the block
		{
			if (Depth >= Obf->MinDepthForBranch && (rand() % 100) <= Obf->ChanceForBranch && CurrentCount >= Obf->MinBranchSize)
			{
				NATIVE_CODE_BLOCK NotTaken, Taken;
				ObfCreateOpaqueBranches(NewBlockStart, Block->End, &NotTaken, &Taken);
				ObfGenerateOpaqueBranches(Obf, &NotTaken, Depth + 1);
				ObfGenerateOpaqueBranches(Obf, &Taken, Depth + 1);
				ObfCombineOpaqueBranches(&NotTaken, &Taken, NcGenUnusedLabelId(Obf->ParentBlock), NcGenUnusedLabelId(Obf->ParentBlock));
				ObfInsertOpaqueBranchBlock(NewBlockStart, Block->End, &NotTaken);
			}
			else
			{
				NATIVE_CODE_BLOCK TempBlock;
				if (NcDeepCopyPartialBlock(NewBlockStart, Block->End, &TempBlock))
				{
					ObfGenerateOpaqueBranches(Obf, &TempBlock, Depth + 1);
					ObfInsertOpaqueBranchBlock(NewBlockStart, Block->End, &TempBlock);
				}
			}
		}
	}
}

VOID ObfMutateInstructions(PINSTMUT_SETS Obf, PNATIVE_CODE_BLOCK Block)
{
	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 ((rand() % 100) <= Obf->MutateChance)
		{
			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;
		}

		T = RealNext;
	}
}