#include "obfuscation.hpp"
using namespace asmjit;

namespace obfuscation
{
	obfuscate::obfuscate(instruction_info_t info)
		:
		instruction(info)
	{
		// idk what i did to zydis but + 1 fixed it?
		// this is brutal... TODO: figure out tf is going on here...
		switch (instruction.first.mnemonic + 1)
		{
		case ZYDIS_MNEMONIC_JLE:
		case ZYDIS_MNEMONIC_JB:
		case ZYDIS_MNEMONIC_JBE:
		case ZYDIS_MNEMONIC_JCXZ:
		case ZYDIS_MNEMONIC_JECXZ:
		case ZYDIS_MNEMONIC_JKNZD:
		case ZYDIS_MNEMONIC_JKZD:
		case ZYDIS_MNEMONIC_JL:
		case ZYDIS_MNEMONIC_JNB:
		case ZYDIS_MNEMONIC_JNBE:
		case ZYDIS_MNEMONIC_JNL:
		case ZYDIS_MNEMONIC_JNLE:
		case ZYDIS_MNEMONIC_JNO:
		case ZYDIS_MNEMONIC_JNP:
		case ZYDIS_MNEMONIC_JNS:
		case ZYDIS_MNEMONIC_JNZ:
		case ZYDIS_MNEMONIC_JO:
		case ZYDIS_MNEMONIC_JP:
		case ZYDIS_MNEMONIC_JRCXZ:
		case ZYDIS_MNEMONIC_JS:
		case ZYDIS_MNEMONIC_JZ:
		{
			const auto rva_fix_offset = instruction.first.length -
				(instruction.first.operands[0].size / 8);

			const auto rva_fix_addr =
				instruction.second.data() + rva_fix_offset;

			reloc_t inline_jmp_reloc
			{
				reloc_type::next_instruction_addr,
				JMP_RIP_ADDR_IDX
			};

			reloc_t inline_jmp_branch
			{
				reloc_type::jcc,
				JMP_RIP_ADDR_IDX,
				*reinterpret_cast<std::int32_t*>(rva_fix_addr)
			};

			std::printf("			> fixing JCC rva...\n");
			std::printf("				> new rva = 0x%x\n", JMP_RIP_SIZE);
			std::printf("				> old rva = 0x%x\n",
				*reinterpret_cast<std::int32_t*>(rva_fix_addr));

			// when you inherit obfuscate please be mindful of JCC rvas...
			*reinterpret_cast<std::int32_t*>(rva_fix_addr) = JMP_RIP_SIZE;

			gadget_stack.push_back({ instruction.second, {} });
			gadget_stack.push_back({ jmp_rip, inline_jmp_reloc });
			gadget_stack.push_back({ jmp_rip, inline_jmp_branch });
			break;
		}
		case ZYDIS_MNEMONIC_JMP:
		{
			const auto rva_fix_offset = instruction.first.length -
				(instruction.first.operands[0].size / 8);

			const auto rva_fix_addr =
				instruction.second.data() + rva_fix_offset;

			reloc_t inline_jmp_reloc
			{
				reloc_type::jcc,
				JMP_RIP_ADDR_IDX,
				*reinterpret_cast<std::int32_t*>(rva_fix_addr)
			};

			gadget_stack.push_back({ jmp_rip, inline_jmp_reloc });
			break;
		}
		case ZYDIS_MNEMONIC_RET:
		{
			gadget_stack.push_back({ instruction.second, {} });
			break;
		}
		default: // not a JCC, JMP, or RET...
		{
			reloc_t inline_jmp_reloc
			{
				reloc_type::next_instruction_addr,
				JMP_RIP_ADDR_IDX
			};

			gadget_stack.push_back({ instruction.second, {} });
			gadget_stack.push_back({ jmp_rip, inline_jmp_reloc });
			break;
		}
		}
	}

	auto obfuscate::get_size() const -> std::uint32_t
	{
		std::uint32_t result = 0u;
		for (auto& gadget : gadget_stack)
			result += gadget.first.size();

		return result;
	}

	auto obfuscate::get_instruc() const -> ZydisDecodedInstruction
	{
		return instruction.first;
	}

	auto obfuscate::get_gadget() const -> gadget_stack_t
	{
		return gadget_stack;
	}

	mutation::mutation(instruction_info_t info)
		: obfuscate(info)
	{
		JitRuntime rt;
		CodeHolder push_code, pop_code;

		push_code.init(rt.environment());
		pop_code.init(rt.environment());

		std::random_device rd; 
		std::mt19937 gen(rd());

		std::vector<std::uint8_t> stack_palindrome;
		std::uniform_int_distribution<> push_reg(0, 15), num_push(1, 10);
		x86::Assembler push_asm(&push_code), pop_asm(&pop_code);
		const auto push_pop_num = num_push(gen);
		std::vector<x86::Gpq> reg_nums;

		// generate random amount of pushes and subsequent pops...
		for (auto idx = 0u; idx < push_pop_num; ++idx)
		{
			const auto reg_num = push_reg(gen);
			const auto gp_reg = x86::Gpq(reg_num);

			reg_nums.push_back(gp_reg);
			push_asm.push(gp_reg);
		}

		stack_palindrome.insert
		(
			stack_palindrome.end(), 
			push_asm.bufferData(),
			push_asm.bufferData() + 
			push_asm.code()->codeSize()
		);

		std::reverse(reg_nums.begin(), reg_nums.end());
		for (auto idx = 0u; idx < push_pop_num; ++idx)
			pop_asm.pop(reg_nums[idx]);

		stack_palindrome.insert
		(
			stack_palindrome.end(),
			pop_asm.bufferData(),
			pop_asm.bufferData() + 
			pop_asm.code()->codeSize()
		);

		// insert the stack palandrome before the actual
		// instruction gets executed...
		gadget_stack.insert(
			gadget_stack.begin(), 
				{ std::move(stack_palindrome), {reloc_type::none, NULL, NULL } });
	}
}