#include "ssdt.h"
#include "kernel_function.h"
#include "utils.h"


namespace SSDT
{
	auto get_service_descriptortable() -> PSERVICE_DESCRIPTOR_TABLE
	{
		auto ntoskrnl = kernel_function::GetKernelModule("ntoskrnl.exe");

		auto KiSystemServiceRepeat = utils::FindSectionsCode(ntoskrnl, "\x4C\x8D\x15\x00\x00\x00\x00\x4C\x8D\x1D\x00\x00\x00\x00\xF7",
			"xxx????xxx????x", ".text");

		return (PSERVICE_DESCRIPTOR_TABLE)(*(UINT32*)(KiSystemServiceRepeat + 3) + KiSystemServiceRepeat + 7);

	}

	auto get_ssdt_function_aadr(ULONG32 number) -> uintptr_t
	{
		auto KeServiceDescriptorTable = get_service_descriptortable();
		if (!KeServiceDescriptorTable)
			return NULL;

		auto table_base = KeServiceDescriptorTable->ServiceTableBase;
		auto offset = (*(ULONG*)((ULONG64)table_base + number * 4)) >> 4;
		return ((ULONG64)table_base + offset);
	}


	auto get_func_number(unsigned char* function) -> ULONG32
	{
		unsigned int count = 0;
		unsigned int index = 0;
		while (count < 1)
		{
			if (function[index++] == 0xB8) {
				count++;
				break;
			}
		}

		return *(ULONG32*)&function[index++];
	}

	
	auto ZwTerminateThread(HANDLE ThreadHandle, NTSTATUS ExitStatus) -> NTSTATUS
	{
		typedef NTSTATUS(__fastcall* fnZwTerminateThread)(HANDLE ThreadHandle, NTSTATUS ExitStatus);
		static fnZwTerminateThread _ZwTerminateThread;
		if (_ZwTerminateThread == nullptr)
		{
			auto func = utils::kernel_get_expotr_func(L"ZwQuerySection") + 0x40;
			auto func_index = get_func_number(func);

			_ZwTerminateThread = (fnZwTerminateThread)get_ssdt_function_aadr(func_index);
		}


		return _ZwTerminateThread(ThreadHandle, ExitStatus);

	}


}