#include "shadow_ssdt.h"
#include "kernel_function.h"
#include "kernel_api.h"
#include "utils.h"

PULONG W32pServiceTable = nullptr;


namespace ShadowSSDT
{
	ULONG64 get_sssdt_func_addr(ULONG id) {
		LONG dwtmp = 0;
		PULONG ServiceTableBase = NULL;
		ServiceTableBase = W32pServiceTable;
		dwtmp = ServiceTableBase[id];
		dwtmp = dwtmp >> 4;
		return (LONGLONG)dwtmp + (ULONGLONG)ServiceTableBase;
	}

	auto get_sssdt_extport_func_win10(const char* func_name)-> uintptr_t
	{
		uintptr_t fun_addrss = 0;

		auto process = kernel_function::FindProcess("explorer.exe");
		if (!process)
		{
			DbgPrintEx(77, 0, "[+]get explorer fail\n");
			return 0;
		}

		const auto win32k_base = reinterpret_cast<PVOID>(kernel_function::GetKernelModule("win32k.sys"));
		if (!win32k_base)
		{
			ObDereferenceObject(process);
			DbgPrintEx(77, 0, "[+]get win32k_base fail\n");
			return 0;
		}

		KAPC_STATE apc_state{ 0 };
		KeStackAttachProcess(process, &apc_state);
		W32pServiceTable = (PULONG)kernel_api::rtlfindexportedroutinebyname(win32k_base, "W32pServiceTable");
		if (!W32pServiceTable)
		{
			KeUnstackDetachProcess(&apc_state);
			ObDereferenceObject(process);
			DbgPrintEx(77, 0, "[+] getW32pServiceTable fail\n");
			return 0;
		}


		auto pdos_header = reinterpret_cast<PIMAGE_DOS_HEADER>(win32k_base);
		auto pNt_header = reinterpret_cast<PIMAGE_NT_HEADERS>((char*)win32k_base + pdos_header->e_lfanew);


		//±éÀúµ¼³ö±í
		const auto export_table = pNt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
		const auto export_data = reinterpret_cast<PIMAGE_EXPORT_DIRECTORY>((char*)win32k_base + export_table.VirtualAddress);

		const auto func_name_tab = reinterpret_cast<ULONG*>((char*)win32k_base + export_data->AddressOfNames);
		const auto func_index_tab = reinterpret_cast<SHORT*>((char*)win32k_base + export_data->AddressOfNameOrdinals);
		const auto func_addr_tab = reinterpret_cast<ULONG*>((char*)win32k_base + export_data->AddressOfFunctions);


		for (size_t i = 0; i < export_data->NumberOfNames; i++)
		{
			const char* func_names = reinterpret_cast<const char*>((char*)win32k_base + func_name_tab[i]);

			if (strstr(func_names, "__win32kstub_"))
			{
				const auto lpFunc_addr = reinterpret_cast<PUCHAR>((char*)win32k_base + func_addr_tab[func_index_tab[i]]);
				char* FunctionName = strstr(func_names, "Nt");
				if (!strcmp(FunctionName, func_name))
				{
					auto func_index = *(ULONG*)((uintptr_t)lpFunc_addr + 1);
					fun_addrss = get_sssdt_func_addr(func_index);
					break;
				}
			}
		}

		KeUnstackDetachProcess(&apc_state);
		ObDereferenceObject(process);
		return fun_addrss;
	}

	//auto get_sssdt_extport_func_win7(const char* func_name)
	//{
	//	UNREFERENCED_PARAMETER(func_name);

	//	auto process = kernel_function::FindProcess("explorer.exe");
	//	if (!process)
	//	{
	//		DbgPrintEx(77, 0, "[+]get explorer fail\n");
	//		return 0;
	//	}

	//	KAPC_STATE apc_state{ 0 };
	//	KeStackAttachProcess(process, &apc_state);

	//	do
	//	{
	//		auto dwPid = PsGetProcessId(process);

	//		UNICODE_STRING module_name{ 0 };
	//		RtlInitUnicodeString(&module_name, L"user32.dll");

	//		unsigned char* moudle_buff = nullptr;
	//		auto status = kernel_function::GetProcessModule(dwPid, module_name, &moudle_buff);
	//		if (!NT_SUCCESS(status))
	//			break;

	//		//char* FunctionName = strstr(func_name, "Nt");

	//		//DbgPrintEx(77,0,"")

	//		/*auto func_addr = utils::get_hmodule_expotr_func(moudle_buff, FunctionName);
	//		if (!func_addr)
	//			break;

	//		auto func_index = *(UINT32*)(func_addr + 4);

	//		DbgPrintEx(77, 0, "%p | %d\n", func_addr, func_index);*/



	//	} while (0);


	//	KeUnstackDetachProcess(&apc_state);
	//	ObDereferenceObject(process);
	//}


	auto fn_get_shadowssdt_extport_func(const char* func_name) -> uintptr_t
	{
		auto os = utils::GetVersion();
		if (os.dwBuildNumber == 7600 || os.dwBuildNumber == 7601)
		{
			return NULL;
		}
		else
		{
			return get_sssdt_extport_func_win10(func_name);
		}

		
	}

}