#include "mouse_keybord_hook.h"
#include "shadow_ssdt.h"
#include "kernel_function.h"
#include "utils.h"
#include "inline_hooks.h"
#include "mouse_key_win10.h"
#include "mouse_key_win7.h"



namespace mouse_keybord_hook
{
	unsigned char* _ntSendInput;

	auto get_sendinput_hook_addr() -> unsigned char*
	{
		auto os = utils::GetVersion();
		if (os.dwBuildNumber == 7600 || os.dwBuildNumber == 7601)
		{
			return NULL;
		}
		else if (os.dwBuildNumber >= 26100)
		{
			auto gSessionGlobalSlots = kernel_function::GetgSessionGlobalSlots();

			auto W32GetSessionState = *(ULONG64*)(*(ULONG64*)((*(ULONG64*)gSessionGlobalSlots)) + 0x88);
			W32GetSessionState = *(ULONG64*)(W32GetSessionState + 0x150);
			auto NtUserSendInput = *(unsigned char**)(W32GetSessionState + 0xB18);

			return NtUserSendInput;
		}
		else
		{
			auto function = reinterpret_cast<unsigned char*>(ShadowSSDT::fn_get_shadowssdt_extport_func("NtUserSendInput"));
			if (!function)
				return NULL;

			unsigned int break_byte_count = 0;
			unsigned int index = 0;
			while (break_byte_count < 2)
			{
				if (function[index++] == 0x48)
					break_byte_count++;
			}

			return *(unsigned char**)((*(ULONG*)&function[index + 2]) + &function[index + 6]);
		}
	}

	auto install_mouse_keybord_hook() -> BOOL
	{
		auto _eprocess_process = kernel_function::FindProcess("winlogon.exe");
		if (!_eprocess_process)
			return FALSE;

		auto apc_state = kernel_function::ke_stack_attch_process(_eprocess_process);

		_ntSendInput = get_sendinput_hook_addr();

		auto os = utils::GetVersion();
		if (os.dwBuildNumber == 7600 || os.dwBuildNumber == 7601)
		{

		}
		else
		{
			inline_hooks_manager::fn_get_instance()->inline_install_hook(_ntSendInput, mouse_key_win10::hkNtUserSendInput, (void**)&origon_NtUserSendInput_win10);
		}

		kernel_function::ke_unstack_detach_process(apc_state);
		ObDereferenceObject(_eprocess_process);
		return 0;
	}

	auto remove_mouse_keybord_hook() -> BOOL
	{
		auto _eprocess_process = kernel_function::FindProcess("winlogon.exe");
		if (!_eprocess_process)
			return FALSE;

		auto apc_state = kernel_function::ke_stack_attch_process(_eprocess_process);

		inline_hooks_manager::fn_get_instance()->inline_remov_hook((void**)_ntSendInput);

		kernel_function::ke_unstack_detach_process(apc_state);
		ObDereferenceObject(_eprocess_process);
		return 0;
	}

}