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

namespace utils
{
	auto GetVersion() -> OSVERSIONINFOW
	{
		OSVERSIONINFOW osver{ 0 };
		osver.dwOSVersionInfoSize = sizeof(OSVERSIONINFOW);
		RtlGetVersion(&osver);
		return osver;
	}

	auto patternCheck(const char* data, const char* pattern, const char* mask) -> bool
	{
		size_t len = strlen(mask);

		for (size_t i = 0; i < len; i++)
		{
			if (data[i] == pattern[i] || mask[i] == '?')
				continue;
			else
				return false;
		}

		return true;
	}
	auto patternFind(unsigned long long addr, unsigned long size, const char* pattern, const char* mask) -> ULONG64
	{
		size -= (unsigned long)strlen(mask);
		for (unsigned long i = 0; i < size; i++)
		{
			if (patternCheck((const char*)addr + i, pattern, mask))
				return addr + i;
		}
		return 0;
	}



	auto FindSectionsCode(unsigned long long Base, const char* pattern, const char* mask, const char* name) -> ULONG64
	{
		auto image_dos = (PIMAGE_DOS_HEADER)Base;
		auto ntheader = (PIMAGE_NT_HEADERS)(Base + image_dos->e_lfanew);

		auto section = IMAGE_FIRST_SECTION(ntheader);
		for (size_t i = 0; i < ntheader->FileHeader.NumberOfSections; i++)
		{
			PIMAGE_SECTION_HEADER p = &section[i];
			if (strstr((const char*)p->Name, name))
			{
				unsigned long long result = patternFind(Base + p->VirtualAddress, p->Misc.VirtualSize, pattern, mask);
				if (result) return result;

			}
		}
		return 0;
	}

	auto kernel_get_expotr_func(wchar_t* func_name)-> uintptr_t
	{
		UNICODE_STRING unicode_func_name{ 0 };
		RtlInitUnicodeString(&unicode_func_name, func_name);
		return (uintptr_t)MmGetSystemRoutineAddress(&unicode_func_name);
	}


	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 = (uintptr_t*)((beep + section_header[i].VirtualAddress) + section_header[i].Misc.VirtualSize);
			}
		}

		if (trampoline_ptr)
		{
			unsigned int count = NULL;
			unsigned int index = NULL;
			while (trampoline_ptr[index] != NULL)
			{
				if (trampoline_ptr[index] == NULL)
					count++;

				if (count > 0x08)
					break;

				index++;
			}
			return (uintptr_t)&trampoline_ptr[index];
		}

		return NULL;
	}


}