#include "inline_hooks.h"
#include "kernel_function.h"
#include "hde/hde64.h"

#pragma warning(disable : 4309 4838)

inline_hooks_manager* inline_hooks_manager::instance;

inline_hooks_manager* inline_hooks_manager::fn_get_instance()
{
	if (instance == nullptr)
	{
		instance = (inline_hooks_manager*)ExAllocatePoolWithTag(NonPagedPool, sizeof(inline_hooks_manager), 'hook');

		instance->m_cur_hook_count = 0;
		instance->m_trampo_line = (unsigned char*)ExAllocatePoolWithTag(NonPagedPool, PAGE_SIZE, 'line');
		RtlSecureZeroMemory(instance->m_Info, sizeof(HOOK_INFO) * HOOK_MAX_COUNT);
		RtlSecureZeroMemory(instance->m_trampo_line, PAGE_SIZE);
	}

	return instance;
}

auto inline_hooks_manager::inline_install_hook(void* target_func, void* my_function, void** origon_func) -> BOOL
{
	if (m_cur_hook_count > HOOK_MAX_COUNT)
		return FALSE;

	auto ori_func_addr = reinterpret_cast<unsigned char*>(target_func);

	hde64s hde{ 0 };
	uint64_t break_byte_count = NULL;
	while (break_byte_count < 14)
	{
		hde64_disasm(ori_func_addr + break_byte_count, &hde);
		break_byte_count += hde.len;
	}

	//保存原来的hook信息方便进行恢复
	m_Info[m_cur_hook_count].origon_func = target_func;
	memcpy(m_Info[m_cur_hook_count].destroy_code, target_func, 14);
	m_cur_hook_count++;

	//开始hook
	//2D963190000 - FF25 00000000 A09AA40CFC7F0000 - jmp USER32.GetWindow
	char jmpcode[14] = { 0xFF,0x25,0,0,0,0,0xA0,0x9A,0xA4,0x0C,0xFC,0x7F,0x00,0x00 };
	*(uint64_t*)&jmpcode[6] = (uint64_t)my_function;

	//蹦床构建
	*origon_func = bulid_trampo_line(ori_func_addr, break_byte_count);

	auto status = kernel_function::MmMDLPagesCopy(ori_func_addr, &jmpcode, sizeof(jmpcode), 0x0000002, KernelMode);
	if (!NT_SUCCESS(status))
		return FALSE;

	return TRUE;
}
 
auto inline_hooks_manager::inline_remov_hook(void** origon_func) -> BOOL
{
	for (size_t i = 0; i < HOOK_MAX_COUNT; i++)
	{
		if (m_cur_hook_count <= i)
			break;

		if (m_Info[i].origon_func)
		{
			if (m_Info[i].origon_func == origon_func)
			{
				kernel_function::MmMDLPagesCopy(m_Info[i].origon_func,
					&m_Info[i].destroy_code, 14, 0x0000002, KernelMode);
			}

			if (origon_func == nullptr) {
				kernel_function::MmMDLPagesCopy(m_Info[i].origon_func,
					&m_Info[i].destroy_code, 14, 0x0000002, KernelMode);
			}
		}
	}
	return TRUE;
}

auto inline_hooks_manager::bulid_trampo_line(void* target_func, UINT64 break_byte_count) -> void*
{
	//push 低32位
	//mov[rsp+4],高32位
	//ret

	/*138DADF0000 - 68 A09AA40C - push 0CA49AA0
	138DADF0005 - C7 44 24 04 FC7F0000 - mov[rsp + 04], 00007FFC
	138DADF000D - C3 - ret*/

	char push_mov_ret[] =
	{   0x68, 0xA0, 0x9A, 0xA4, 0x0C,
	    0xC7, 0x44, 0x24, 0x04, 0xFC, 0x7F, 0x00, 0x00,
	    0xC3
	};

	*(uint32_t*)&push_mov_ret[1] = (((uint64_t)target_func + break_byte_count) & 0xffffffff);
	*(uint32_t*)&push_mov_ret[9] = ((((uint64_t)target_func + break_byte_count) >> 32) & 0xffffffff);

	memcpy(instance->m_trampo_line, target_func, break_byte_count);
	memcpy(instance->m_trampo_line + break_byte_count, &push_mov_ret, sizeof(push_mov_ret));

	auto ret = reinterpret_cast<void*>(instance->m_trampo_line);

	instance->m_trampo_line += break_byte_count + sizeof(push_mov_ret);

	return ret;
}