#include "create_thread_callback.h"
#include "protect_filter.h"
#include "kernel_function.h"

#pragma warning(disable : 4309)
#pragma warning(disable : 4838)

namespace thread_notify_routine
{
	//保护的进程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)
			{
				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);
			}

			

		end:	
			ObDereferenceObject(process);
		}
	}

	auto get_beep_trampoline_ptr() -> uintptr_t
	{
		uintptr_t trampoline_ptr = NULL;

		auto beep = kernel_function::GetKernelModule("beep.sys");
		auto nt_header = reinterpret_cast<PIMAGE_NT_HEADERS>(beep + reinterpret_cast<PIMAGE_DOS_HEADER>(beep)->e_lfanew);
		auto file_header = &nt_header->FileHeader;

		auto section_header = IMAGE_FIRST_SECTION(nt_header);
		for (size_t i = 0; i < file_header->NumberOfSections; i++)
		{
			if (!strcmp((char*)section_header[i].Name, ".text"))
			{
				trampoline_ptr = ((beep + section_header[i].VirtualAddress) + section_header[i].Misc.VirtualSize);
			}
		}
		return trampoline_ptr;
	}

	auto create_thread_routine() -> NTSTATUS
	{
		auto trampoline_ptr = 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 trampoline_ptr = get_beep_trampoline_ptr();
		return PsRemoveCreateThreadNotifyRoutine((PCREATE_PROCESS_NOTIFY_ROUTINE)trampoline_ptr);
	}


	
}