#include "Obfuscator.h"

//snake case is honestly so disgusting
void obf_one_time_please()
{
	xed_tables_init();
}

bool obf_init_from_buffer(pobfuscator_t obf, void* buffer, int buffer_size)
{
	obf->current_label_id = 0;
	obf->machine_mode = XED_MACHINE_MODE_LONG_64;
	obf->addr_width = XED_ADDRESS_WIDTH_64b;

	unsigned long long off = 0;

	obf->code_start = new code_link_t;
	obf->code_end = obf->code_start;
	obf->code_start->flags = 0;
	obf->code_start->group = 0;
	obf->code_start->label_name = "omegalawl";
	obf->code_start->prev = obf->code_start->next = nullptr;

	while (off < buffer_size)
	{
		pcode_link_t link = new code_link_t;
		link->flags = 0;
		link->group = 0;

		xed_decoded_inst_zero(&link->instruction);
		xed_decoded_inst_set_mode(&link->instruction, obf->machine_mode, obf->addr_width);
		xed_error_enum_t err = xed_decode(&link->instruction, (unsigned char*)((unsigned char*)buffer + off), 15);
		if (err != XED_ERROR_NONE)
		{
			printf("Failed decoding instruction at %llu with error \"%s\"(%d)", off, xed_error_enum_t2str(err), err);
			return false;
		}

		unsigned int inst_len = xed_decoded_inst_get_length(&link->instruction);
		link->raw_data_size = inst_len;

		link->raw_data = (unsigned char*)malloc(inst_len);
		if (!link->raw_data)
		{
			printf("outta memory son.\n");
			return false;
		}

		memcpy(link->raw_data, ((unsigned char*)buffer + off), inst_len);

		//filter out 8 byte wide jumps cuz they aint relative dawg and wont be showin up in my shellcod
		xed_category_enum_t cat = xed_decoded_inst_get_category(&link->instruction);
		if (cat == XED_CATEGORY_COND_BR || cat == XED_CATEGORY_UNCOND_BR)
		{
			unsigned int disp_width = xed_decoded_inst_get_branch_displacement_width(&link->instruction);
			if (disp_width != 8)
			{
				link->flags |= CLFLAG_IS_REL_JUMP;

				//int jump_delta = xed_decoded_inst_get_branch_displacement(&link->instruction);
				//printf("Jump delta is %d\n", jump_delta);
			}
		}

		link->prev = obf->code_end;
		link->next = nullptr;
		obf->code_end->next = link;
		obf->code_end = link;

		off += inst_len;
	}
}

bool obf_create_groups(pobfuscator_t obf, int group_size)
{

	obf->groups.clear();

	if (group_size < 24)
		return false;

	int cur_group_id = 0, cur_size_in_bytes = 0;
	pcode_link_t start = obf->code_start->next;
	for (pcode_link_t t = obf->code_start->next; t; t = t->next)
	{
		if (!(t->flags & CLFLAG_IS_GAGET))
		{
			if (cur_size_in_bytes + t->raw_data_size + 16 > group_size)
			{
				
				pcode_link_t push_rax = new code_link_t;
				push_rax->flags = 0;
				push_rax->label_name = "";
				push_rax->raw_data = new unsigned char[1];
				push_rax->raw_data_size = 1;
				push_rax->group = cur_group_id;
				*(unsigned char*)push_rax->raw_data = 0x50;
				push_rax->label_name = "";

				pcode_link_t mov_address = new code_link_t;
				mov_address->flags = CLFLAG_IS_GROUP_JMP;
				mov_address->label_name = t->label_name;
				mov_address->raw_data = new unsigned char[10];
				mov_address->raw_data_size = 10;
				mov_address->group = cur_group_id;
				unsigned char mov_address_data[] = { 0x48, 0xB8, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x0F };
				memcpy(mov_address->raw_data, mov_address_data, 10);

				pcode_link_t xchg_rax_rsp = new code_link_t;
				xchg_rax_rsp->flags = 0;
				xchg_rax_rsp->label_name = "";
				xchg_rax_rsp->raw_data = new unsigned char[4];
				xchg_rax_rsp->raw_data_size = 4;
				xchg_rax_rsp->group = cur_group_id;
				unsigned char xchg_rax_rsp_data[] = { 0x48, 0x87, 0x04, 0x24 };
				memcpy(xchg_rax_rsp->raw_data, xchg_rax_rsp_data, 4);

				pcode_link_t ret = new code_link_t;
				ret->flags = 0;
				ret->label_name = "";
				ret->raw_data = new unsigned char[1];
				ret->raw_data_size = 1;
				ret->group = cur_group_id;
				*(unsigned char*)ret->raw_data = 0xC3;

				t->prev->next = push_rax;
				push_rax->next = mov_address;
				mov_address->next = xchg_rax_rsp;
				xchg_rax_rsp->next = ret;
				ret->next = t;

				ret->prev = xchg_rax_rsp;
				xchg_rax_rsp->prev = mov_address;
				mov_address->prev = push_rax;
				push_rax->prev = t->prev;
				t->prev = ret;

				printf("creating group %d\n", cur_group_id);
				obf->groups.emplace_back();
				obf->groups.back().size_in_bytes = cur_size_in_bytes + 16;
				obf->groups.back().start = start;
				cur_size_in_bytes = 0;
				cur_group_id++;
				start = t;
			}
		}

		cur_size_in_bytes += t->raw_data_size;
		t->group = cur_group_id;
	}

	obf->groups.emplace_back();
	obf->groups.back().size_in_bytes = cur_size_in_bytes + 16;
	obf->groups.back().start = start;

	return true;
}

void obf_replace_rel_jmps(pobfuscator_t obf)

{	//	original_jump   -------------------------.
	//	jmp 0x10			0xEB, 0x10			 |
	//	push rax			0x50,			<----'
	//	mov rax,abs_address 0x48, 0xB8, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x0F, 
	//	xchg rax,[rsp]		0x48, 0x87, 0x04, 0x24, 
	//	ret					0xC3 

	for (pcode_link_t t = obf->code_start->next; t;)
	{
		if (t->flags & CLFLAG_IS_REL_JUMP)
		{
			pcode_link_t real_next = t->next;
			unsigned int inst_len = xed_decoded_inst_get_length(&t->instruction);
			unsigned int jmp_delta_width = xed_decoded_inst_get_branch_displacement_width(&t->instruction);
			unsigned int opcode_size = inst_len - jmp_delta_width;

			switch (jmp_delta_width)
			{
			case 1:
				*(char*)((unsigned char*)t->raw_data + opcode_size) = (char)2; break;
			case 2:
				*(short*)((unsigned char*)t->raw_data + opcode_size) = (short)2; break;
			case 4:
				*(int*)((unsigned char*)t->raw_data + opcode_size) = (int)2; break;
			}
			t->flags = CLFLAG_IS_GAGET;

			pcode_link_t jmp_around_gagt = new code_link_t;
			jmp_around_gagt->flags = CLFLAG_IS_GAGET;
			jmp_around_gagt->label_name = "";
			jmp_around_gagt->raw_data = new unsigned char[2];
			jmp_around_gagt->raw_data_size = 2;
			unsigned char jmp_around_gagt_data[] = { 0xEB, 0x10 };
			memcpy(jmp_around_gagt->raw_data, jmp_around_gagt_data, 2);


			pcode_link_t push_rax = new code_link_t;
			push_rax->flags = CLFLAG_IS_GAGET;
			push_rax->label_name = "";
			push_rax->raw_data = new unsigned char[1];
			push_rax->raw_data_size = 1;
			*(unsigned char*)push_rax->raw_data = 0x50;
			push_rax->label_name = "";


			pcode_link_t mov_address = new code_link_t;
			mov_address->flags = (CLFLAG_IS_ABS_ADDR | CLFLAG_IS_GAGET);
			mov_address->label_name = t->label_name;
			mov_address->raw_data = new unsigned char[10];
			mov_address->raw_data_size = 10;
			unsigned char mov_address_data[] = { 0x48, 0xB8, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x0F };
			memcpy(mov_address->raw_data, mov_address_data, 10);

			pcode_link_t xchg_rax_rsp = new code_link_t;
			xchg_rax_rsp->flags = CLFLAG_IS_GAGET;
			xchg_rax_rsp->label_name = "";
			xchg_rax_rsp->raw_data = new unsigned char[4];
			xchg_rax_rsp->raw_data_size = 4;
			unsigned char xchg_rax_rsp_data[] = { 0x48, 0x87, 0x04, 0x24 };
			memcpy(xchg_rax_rsp->raw_data, xchg_rax_rsp_data, 4);

			pcode_link_t ret = new code_link_t;
			ret->flags = CLFLAG_IS_GAGET;
			ret->label_name = "";
			ret->raw_data = new unsigned char[1];
			ret->raw_data_size = 1;
			*(unsigned char*)ret->raw_data = 0xC3;

			t->next = jmp_around_gagt;
			jmp_around_gagt->next = push_rax;
			push_rax->next = mov_address;
			mov_address->next = xchg_rax_rsp;
			xchg_rax_rsp->next = ret;
			ret->next = real_next;

			real_next->prev = ret;
			ret->prev = xchg_rax_rsp;
			xchg_rax_rsp->prev = mov_address;
			mov_address->prev = push_rax;
			push_rax->prev = jmp_around_gagt;
			jmp_around_gagt->prev = t;


			t = real_next;
			continue;
		}
		t = t->next;
	}
}

bool obf_replace_abs_jmps(pobfuscator_t obf)
{
	for (pcode_link_t t = obf->code_start->next; t;)
	{
		if (t->flags & CLFLAG_IS_ABS_ADDR)
		{
			pcode_link_t temp;
			for (temp = t; temp && temp->prev; temp = temp->prev)
			{
				if (temp->flags & CLFLAG_IS_LABEL && temp->label_name == t->label_name)
				{
					goto have_label_found;
				}
			}
			for (temp = t; temp && temp->next; temp = temp->next)
			{
				if (temp->flags & CLFLAG_IS_LABEL && temp->label_name == t->label_name)
				{
					goto have_label_found;
				}
			}
			return false;
			continue;

		have_label_found:
			uint64_t addr = obf->groups[temp->group].base_address;
			for (pcode_link_t temp2 = obf->groups[temp->group].start; temp2 && temp2->group == temp->group; temp2 = temp2->next)
			{
				if (temp2 == temp)
				{
					break;
				}
				addr += temp->raw_data_size;
			}
			*(uint64_t*)((unsigned char*)t->raw_data + 2) = addr;
		}
	}
	return true;
}

size_t obf_get_group_size(pobfuscator_t obf, int group_id)
{
	if (group_id < obf->groups.size())
		return obf->groups[group_id].size_in_bytes;
}

void obf_copy_group_to_buffer(pobfuscator_t obf, void* buffer, int group_id)
{

}

bool obf_gen_all_labels(pobfuscator_t obf)
{
	for (pcode_link_t t = obf->code_start->next; t; t = t->next)
	{
		if (t->flags & CLFLAG_IS_REL_JUMP)
		{
			int jump_delta = xed_decoded_inst_get_branch_displacement(&t->instruction);
			if (!obf_gen_label(obf, t, jump_delta))
				return false;
		}
	}
	return true;
}

bool obf_gen_label(pobfuscator_t obf, pcode_link_t jmp, int32_t delta)
{
	obf->current_label_id++;
	pcode_link_t temp;
	//when going positive, counting starts at NEXT instruction(excluding size of jmp)
	//when negative, counting INCLUDES the size of the jmp instructrion
	if (delta > 0)
	{
		temp = jmp->next;
		while (delta && temp)
		{
			delta -= temp->instruction._decoded_length;
			//if (delta == 0) break;
			temp = temp->next;
		}
		if (temp && temp->flags & CLFLAG_IS_LABEL)
		{
			jmp->label_name = temp->label_name;
			return true;
		}
	}
	else if (delta < 0)
	{
		temp = jmp;
		while (temp)
		{
			delta += temp->instruction._decoded_length;
			if (delta == 0) break;
			temp = temp->prev;
		}

		if (temp && temp->prev && (temp->prev->flags & CLFLAG_IS_LABEL))
		{
			jmp->label_name = temp->prev->label_name;
			return true;
		}
	}
	else return false;

	if (!temp)
		return false;

	//couldnt find label, adding new one
	pcode_link_t new_label = new code_link_t;
	new_label->flags = CLFLAG_IS_LABEL;
	new_label->label_name = std::to_string(obf->current_label_id);
	jmp->label_name = new_label->label_name;

	new_label->next = temp;
	new_label->prev = temp->prev;
	if (temp->prev)
		temp->prev->next = new_label;
	temp->prev = new_label;

	return true;
}

void obf_dbg_print_code(pobfuscator_t obf)
{
	for (pcode_link_t t = obf->code_start->next; t; t = t->next)
	{
		if (!(t->flags & CLFLAG_IS_LABEL))
		{
			obf_print_byte_array(t->raw_data, t->raw_data_size);
		}
		/*if (t->flags & CLFLAG_IS_REL_JUMP)
		{
			printf("\tJump to: %s\n", t->label_name.data());
		}
		else if (t->flags & CLFLAG_IS_LABEL)
		{
			printf("Label: %s\n", t->label_name.data());
		}
		else
		{
			printf("\tRegular Instruction.\n");
		}*/
	}
}

void obf_dbg_print_group(pobfuscator_t obf, int group_id)
{
	if (group_id >= obf->groups.size())
		return;

	for (pcode_link_t t = obf->groups[group_id].start; t && t->group == group_id; t = t->next)
	{
		if (!(t->flags & CLFLAG_IS_LABEL))
		{
			obf_print_byte_array(t->raw_data, t->raw_data_size);
		}
	}
}

#include <iostream>
#include <iomanip>
void obf_print_byte_array(void* arr, unsigned int size)
{
	unsigned char* b = (unsigned char*)arr;
	for (int i = 0; i < size; i++)
	{
		std::cout << std::hex << std::setw(2) << std::setfill('0') << (int)b[i] << ' ';
	}
	std::cout << '\n';
	return;
}