#include "threads_util.h"

#include <peconv.h>
#include <tlhelp32.h>
#include "../utils/ntddk.h"
#include "custom_buffer.h"

#ifdef _DEBUG
#include <iostream>
#endif


namespace pesieve {
	namespace util {

		// Thread info structures:
		typedef struct _THREAD_LAST_SYSCALL_INFORMATION
		{
			PVOID FirstArgument;
			USHORT SystemCallNumber;
		} THREAD_LAST_SYSCALL_INFORMATION, * PTHREAD_LAST_SYSCALL_INFORMATION;


		bool query_thread_details(IN DWORD tid, OUT pesieve::util::thread_info& info)
		{
			static auto mod = GetModuleHandleA("ntdll.dll");
			if (!mod) return false;

			static auto pNtQueryInformationThread = reinterpret_cast<decltype(&NtQueryInformationThread)>(GetProcAddress(mod, "NtQueryInformationThread"));
			if (!pNtQueryInformationThread)  return false;

			const DWORD thAccess = THREAD_QUERY_INFORMATION | THREAD_GET_CONTEXT;
			HANDLE hThread = OpenThread(thAccess, 0, tid);
			if (!hThread) {
				hThread = OpenThread(THREAD_QUERY_INFORMATION, 0, tid);
				if (!hThread) return false;
			}
			bool isOk = false;
			ULONG returnedLen = 0;
			LPVOID startAddr = 0;
			NTSTATUS status = 0;
			status = pNtQueryInformationThread(hThread, ThreadQuerySetWin32StartAddress, &startAddr, sizeof(LPVOID), &returnedLen);
			if (status == 0 && returnedLen == sizeof(startAddr)) {
				info.start_addr = (ULONGLONG)startAddr;
				isOk = true;
			}
			returnedLen = 0;
			THREAD_LAST_SYSCALL_INFORMATION syscallInfo = { 0 };
			status = pNtQueryInformationThread(hThread, ThreadLastSystemCall, &syscallInfo, sizeof(syscallInfo), &returnedLen);
			if (status == 0 && returnedLen >= sizeof(syscallInfo)) {
				info.last_syscall = syscallInfo.SystemCallNumber;
				isOk = true;
			}
			CloseHandle(hThread);
			info.is_filled = isOk;
			return isOk;
		}

	}; // namespace util
}; // namespace pesieve


size_t pesieve::util::query_threads_details(IN OUT std::map<DWORD, pesieve::util::thread_info>& threads_info)
{
	size_t queried_count = 0;
	for (auto itr = threads_info.begin(); itr != threads_info.end(); ++itr) {
		pesieve::util::thread_info& info = itr->second;
		if (query_thread_details(info.tid, info)) queried_count++;
	}
	return queried_count;
}

bool pesieve::util::fetch_threads_info(IN DWORD pid, OUT std::map<DWORD, thread_info>& threads_info)
{
	AutoBuffer bBuf;
	NTSTATUS status = STATUS_UNSUCCESSFUL;

	while (status != STATUS_SUCCESS) {
		ULONG ret_len = 0;
		status = NtQuerySystemInformation(SystemProcessInformation, bBuf.buf, bBuf.buf_size, &ret_len);
		if (status == STATUS_INFO_LENGTH_MISMATCH) {
			if (!bBuf.alloc(ret_len)) {
				return false;
			}
			continue; // try again
		}
		break; //other error, or success
	};

	if (status != STATUS_SUCCESS) {
		return false;
	}

	bool pid_found = false;
	SYSTEM_PROCESS_INFORMATION* info = (SYSTEM_PROCESS_INFORMATION*)bBuf.buf;
	while (info) {
		if (info->UniqueProcessId == pid) {
			pid_found = true;
			break;
		}
		if (!info->NextEntryOffset) {
			break;
		}
		size_t record_size = info->NextEntryOffset;
		if (record_size < sizeof(SYSTEM_PROCESS_INFORMATION)) {
			// Record size smaller than expected, probably it is an old system that doesn not support the new version of this API
#ifdef _DEBUG
			std::cout << "The new version of SYSTEM_PROCESS_INFORMATION is not supported!\n";
#endif
			break;
		}
		info = (SYSTEM_PROCESS_INFORMATION*)((ULONG_PTR)info + info->NextEntryOffset);
		if (!peconv::validate_ptr(bBuf.buf, bBuf.buf_size, info, sizeof(SYSTEM_PROCESS_INFORMATION))) {
			break;
		}
	}
	if (!pid_found) {
		return false;
	}
	const size_t thread_count = info->NumberOfThreads;
	if (!peconv::validate_ptr(bBuf.buf, bBuf.buf_size, info->Threads, thread_count * sizeof(info->Threads[0]))) {
		return false;
	}
	bool fetched = false;
	for (size_t i = 0; i < thread_count; i++) {
		const DWORD tid = MASK_TO_DWORD((ULONGLONG)info->Threads[i].ClientId.UniqueThread);
		auto itr = threads_info.find(tid);
		if (itr == threads_info.end()) {
			threads_info[tid] = thread_info(tid);
		}
		fetched = true;
		thread_info &threadi = threads_info[tid];
		threadi.is_extended = true;
		threadi.ext.sys_start_addr = (ULONG_PTR)info->Threads[i].StartAddress;
		threadi.ext.state = info->Threads[i].ThreadState;
		threadi.ext.wait_reason = info->Threads[i].WaitReason;
		threadi.ext.wait_time  = info->Threads[i].WaitTime;
		threadi.is_filled = true;
	}
	return fetched;
}

bool pesieve::util::fetch_threads_by_snapshot(IN DWORD pid, OUT std::map<DWORD, thread_info>& threads_info)
{
	HANDLE hThreadSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	if (hThreadSnapShot == INVALID_HANDLE_VALUE) {
#ifdef _DEBUG
		const DWORD err = GetLastError();
		std::cerr << "[-] Could not create threads snapshot. Error: " << std::dec << err << std::endl;
#endif
		return false;
	}
	THREADENTRY32 th32 = { 0 };
	th32.dwSize = sizeof(THREADENTRY32);

	//check all threads in the process:
	if (!Thread32First(hThreadSnapShot, &th32)) {
		CloseHandle(hThreadSnapShot);
#ifdef _DEBUG
		std::cerr << "[-] Could not enumerate thread. Error: " << GetLastError() << std::endl;
#endif
		return false;
	}

	bool fetched = false;
	do {
		if (th32.th32OwnerProcessID != pid) {
			continue;
		}
		
		fetched = true;

		const DWORD tid = th32.th32ThreadID;
		auto itr = threads_info.find(tid);
		if (itr == threads_info.end()) {
			threads_info[tid] = thread_info(tid);
		}
	} while (Thread32Next(hThreadSnapShot, &th32));

	CloseHandle(hThreadSnapShot);
	return fetched;
}