#include "create_thread_callback.h"
#include "protect_filter.h"
#include "kernel_function.h"
#include "utils.h"
#include "kernel_api.h"
#include "ssdt.h"

#pragma warning(disable : 4309 4838)


EXTERN_C NTKERNELAPI NTSTATUS ZwOpenThread(
	_Out_ PHANDLE ThreadHandle,
	_In_ ACCESS_MASK DesiredAccess,
	_In_ POBJECT_ATTRIBUTES ObjectAttributes,
	_In_ PCLIENT_ID ClientId
	);


namespace thread_notify_routine
{

	// 打开线程句柄并终止
	NTSTATUS TerminateThreadByHandle(
			HANDLE ProcessId,
			HANDLE ThreadId)
	{
		NTSTATUS status;
		HANDLE threadHandle = NULL;
		OBJECT_ATTRIBUTES objAttr;
		CLIENT_ID clientId;

		clientId.UniqueProcess = ProcessId;
		clientId.UniqueThread = ThreadId;

		InitializeObjectAttributes(&objAttr, NULL, 0, NULL, NULL);

		// 打开线程句柄
		status = ZwOpenThread(
			&threadHandle,
			THREAD_TERMINATE,
			&objAttr,
			&clientId);

		if (!NT_SUCCESS(status)) {
			DbgPrintEx(77, 0, "[AnitCheat][%s] Failed to open thread %d: 0x%X\n",__FUNCTION__,
				(ULONG)(ULONG_PTR)ThreadId, status);
			return status;
		}

		// 终止线程
		status = SSDT::ZwTerminateThread(threadHandle, STATUS_THREAD_IS_TERMINATING);

		if (NT_SUCCESS(status)) {
			/*DbgPrintEx(77, 0, "[ThreadKiller] Thread %d terminated successfully\n",
				(ULONG)(ULONG_PTR)ThreadId);*/
		}
		else {
			DbgPrintEx(77,0,"[AnitCheat][%s] Failed to terminate thread %d: 0x%X\n", __FUNCTION__,
				(ULONG)(ULONG_PTR)ThreadId, status);
		}

		ZwClose(threadHandle);
		return status;
	}

	//保护的进程ID
	//auto protect_id = ULongToHandle(8976);
	void thread_notify_routine(HANDLE ProcessId, HANDLE ThreadId, BOOLEAN Create)
	{
		if (Create)
		{

			PEPROCESS process{ 0 };
			auto status = PsLookupProcessByProcessId(ProcessId, &process);
			if (!NT_SUCCESS(status))
				goto end;

			/*PEPROCESS protect_process{ 0 };
			status = PsLookupProcessByProcessId(protect_id, &protect_process);
			if (!NT_SUCCESS(status))
				goto end;*/

			if(!protect_filter::is_protect_pid(0, process))
				goto end;

			/*if (protect_process != process)
				goto end;*/

			auto cutor_process_id = PsGetCurrentProcessId();
			if (HandleToLong(cutor_process_id) == 4)
				goto end;

			if (ProcessId != cutor_process_id)
			{
				//DbgBreakPoint();

				DbgPrintEx(77, 0, "[!] 恶意远程线程检测！\n");

				DbgPrintEx(77, 0, "[!] 目标进程ID: %d\n", ProcessId);

				DbgPrintEx(77, 0, "[!] 创建者进程ID: %d\n", cutor_process_id);

				DbgPrintEx(77, 0, "[!] 被创建线程ID: %d\n", ThreadId);


				TerminateThreadByHandle(ProcessId, ThreadId);
	
			}
		end:	
			ObDereferenceObject(process);
		}
	}

	uintptr_t trampoline_ptr = NULL;

	auto create_thread_routine() -> NTSTATUS
	{
		trampoline_ptr = utils::get_beep_trampoline_ptr();
		if (trampoline_ptr)
		{
			char jmprax[] = { 0x48, 0xB8, 0xA0, 0x9A, 0xD5, 0x3E, 0xFE, 0x7F, 0x00, 0x00, 0xFF, 0xE0 };
			//mov rax,thread_notify_routine
			//jmp rax
			*(uintptr_t*)&jmprax[2] = (uintptr_t)thread_notify_routine;
			kernel_function::MmMDLPagesCopy((PVOID)trampoline_ptr, &jmprax, sizeof(jmprax), 0x0000002, KernelMode);


			return PsSetCreateThreadNotifyRoutine((PCREATE_PROCESS_NOTIFY_ROUTINE)trampoline_ptr);
		}

		return STATUS_INSUFFICIENT_RESOURCES;
	}

	auto unload_thread_routine() -> NTSTATUS
	{
		auto status = PsRemoveCreateThreadNotifyRoutine((PCREATE_PROCESS_NOTIFY_ROUTINE)trampoline_ptr);
		if (!NT_SUCCESS(status))
			return status;

		char nullcode[12]{ 0 };
		status = kernel_function::MmMDLPagesCopy((PVOID)trampoline_ptr, &nullcode, sizeof(nullcode), 0x0000002, KernelMode);
		if (!NT_SUCCESS(status))
			return status;

		return STATUS_SUCCESS;
	}
}