#include "Obfuscator.h"

void obf_one_time_please()
{
	xed_tables_init();
}

bool obf_init_from_buffer(pobfuscator_t obf, void* buffer, uint32_t 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;
		}

		uint32_t 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, int32_t group_size)
{
	uint32_t cur_group_id = 0;
	uint32_t cur_offset = 0;

	if (group_size < 32)
		group_size = 32;

	//assign instructions to groups
	for (pcode_link_t t = obf->code_start->next; t; t = t->next)
	{
		if (!(t->flags & CLFLAG_IS_LABEL))
		{
			if (!(t->flags & CLFLAG_IS_GAGET))
			{
				if (cur_offset + t->raw_data_size > group_size)
				{
					++cur_group_id;
					cur_offset = 0;
				}
			}
			cur_offset += t->raw_data_size;
			t->group = cur_group_id;
		}
	}

	//assign labels to their proper groups
	for (pcode_link_t t = obf->code_start->next; t; t = t->next)
	{
		if (t->flags & CLFLAG_IS_LABEL)
		{
			pcode_link_t t2 = t;
			while (t2 && (t2->flags & CLFLAG_IS_LABEL)) { t2 = t2->next; }
			if (t2) t->group = t2->group;
		}
	}

	//create group descriptors
	obf->groups.clear();
	pcode_link_t start = obf->code_start->next;
	cur_offset = 0;
	pcode_link_t prev_meme = nullptr;
	for (pcode_link_t t = obf->code_start->next; t; t = t->next)
	{
		if (start->group != t->group)
		{
			obf->groups.emplace_back(0, start, t->prev, cur_offset);
			cur_offset = 0;
			start = t;
		}
		cur_offset += t->raw_data_size;
		prev_meme = t;
	}
	if (!prev_meme)
		return false;
	obf->groups.emplace_back(0, start, prev_meme, cur_offset);
	
	//append jumps to next group onto end
	for (uint32_t i = 0; i < obf->groups.size() - 1; i++)
	{
		pcode_group_t cur_group = &obf->groups[i];
		pcode_group_t next_group = &obf->groups[i+1];

		//add jump gadget to end of current group
		pcode_link_t gadget = new code_link_t;
		gadget->flags = CLFLAG_IS_GAGET;
		gadget->label_name = "";
		gadget->raw_data = new unsigned char[6];
		gadget->raw_data_size = 6;
		gadget->group = i;
		unsigned char gadget_data[] = { 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00 };
		memcpy(gadget->raw_data, gadget_data, 6);

		pcode_link_t abs_addr = new code_link_t;
		abs_addr->flags = (CLFLAG_IS_GAGET | CLFLAG_IS_ABS_ADDR);
		abs_addr->label_name = std::string("Group") + std::to_string(i + 1);
		abs_addr->raw_data = new unsigned char[8];
		abs_addr->raw_data_size = 8;
		abs_addr->group = i;

		cur_group->size_in_bytes += END_OF_GROUP_GAGT_SIZE;

		pcode_link_t real_next = cur_group->end->next;
		cur_group->end->next = gadget;
		gadget->next = abs_addr;
		abs_addr->next = real_next;

		real_next->prev = abs_addr;
		abs_addr->prev = gadget;
		gadget->prev = cur_group->end;

		cur_group->end = abs_addr;

		//add label to beginning of next group
		pcode_link_t next_group_label = new code_link_t;
		next_group_label->flags = CLFLAG_IS_LABEL;
		next_group_label->label_name = std::string("Group") + std::to_string(i + 1);
		next_group_label->group = i + 1;

		pcode_link_t real_prev = next_group->start->prev;
		next_group->start->prev = next_group_label;
		real_prev->next = next_group_label;
		next_group_label->next = next_group->start;
		next_group_label->prev = real_prev;
		next_group->start = next_group_label;
	}

	return true;
}

void obf_replace_rel_jmps(pobfuscator_t obf)

{	//	original_jump   -------------------------.
	//	jmp 0x0E(0xEB, 0x0E)  ------------------ | -----.
	//	jmp qword ptr[rip]		<----------------'	    |
	//	address here(8 bytes)							|
	//	not taken branch code	<-----------------------'

	for (pcode_link_t t = obf->code_start->next; t;)
	{
		pcode_link_t real_next = t->next;
		if (t->flags & CLFLAG_IS_REL_JUMP)
		{
			uint32_t jmp_delta_width = xed_decoded_inst_get_branch_displacement_width(&t->instruction);
			uint32_t opcode_size = t->raw_data_size - 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 gadget = new code_link_t;
			gadget->flags = CLFLAG_IS_GAGET;
			gadget->label_name = "";
			gadget->raw_data = new unsigned char[8];
			gadget->raw_data_size = 8;
			unsigned char gadget_data[] = { 0xEB, 0x0E, 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00 };
			memcpy(gadget->raw_data, gadget_data, 8);

			pcode_link_t abs_addr = new code_link_t;
			abs_addr->flags = (CLFLAG_IS_GAGET | CLFLAG_IS_ABS_ADDR);
			abs_addr->label_name = t->label_name;
			t->label_name = "";
			abs_addr->raw_data = new unsigned char[8];
			abs_addr->raw_data_size = 8;

			t->next = gadget;
			gadget->next = abs_addr;
			abs_addr->next = real_next;

			if (real_next) real_next->prev = abs_addr;
			abs_addr->prev = gadget;
			gadget->prev = t;

		}
		t = real_next;
	}
}

bool obf_resolve_abs_addresses(pobfuscator_t obf)
{
	for (pcode_link_t t = obf->code_start->next; t; t = t->next)
	{
		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;

		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;
				}
				if (!(temp2->flags & CLFLAG_IS_LABEL))
				{
					addr += temp2->raw_data_size;
				}
			}
			*(uint64_t*)((unsigned char*)t->raw_data) = addr;
		}
	}
	return true;
}

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)
		{
			int32_t 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)
{
	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;
}

bool obf_allocate_group_buffers(pobfuscator_t obf, FnAllocateMem alloc_mem)
{
	for (code_group_t& group : obf->groups)
	{
		group.base_address = (uint64_t)alloc_mem(group.size_in_bytes);
		if (!group.base_address)
			return false;
	}
	return true;
}

bool obf_copy_groups_into_buffers(pobfuscator_t obf, FnMemCopy mem_copy)
{
	for (uint32_t i = 0; i < obf->groups.size(); i++)
	{
		pcode_group_t group = &obf->groups[i];
		uint64_t cur_addr = group->base_address;
		for (pcode_link_t t = group->start; t && t->group == i; t = t->next)
		{
			if (!(t->flags & CLFLAG_IS_LABEL))
			{
				if (!cur_addr)
					return false;
				mem_copy((void*)cur_addr, t->raw_data, t->raw_data_size);
				cur_addr += t->raw_data_size;
			}
		}
		printf("finished copying group %u\n", i);
	}
	return true;
}

#include <Windows.h>
void obf_dbg_print_code(pobfuscator_t obf)
{
	HANDLE StdHandle = GetStdHandle(STD_OUTPUT_HANDLE);
	if (!StdHandle)
		return;

	for (pcode_link_t t = obf->code_start->next; t; t = t->next)
	{
		if (t->group % 2)
			SetConsoleTextAttribute(StdHandle, 12);
		else
			SetConsoleTextAttribute(StdHandle, 14);

		if (t->flags & CLFLAG_IS_REL_JUMP)
		{
			printf("\tRel jmp to:  %s\t\t", t->label_name.data());
		}
		else if (t->flags & CLFLAG_IS_LABEL)
		{
			//SetConsoleTextAttribute(StdHandle, 13);
			printf("Label:  %s   \n", t->label_name.data());
		}
		else if (t->flags & CLFLAG_IS_ABS_ADDR)
		{
			printf("\tAbs jmp to:  %s\t", t->label_name.data());
		}
		else
		{
			printf("\tInstruction:   \t\t");
		}
		
		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, uint32_t size)
{
	unsigned char* b = (unsigned char*)arr;
	for (uint32_t i = 0; i < size; i++)
	{
		std::cout << std::hex << std::setw(2) << std::setfill('0') << (int)b[i] << ' ';
	}
	std::cout << '\n';
	return;
}