#include "kdmapper.h"

BOOL kdmapper::MapDriver(HANDLE device_handle, uint64_t ker_base, char* map_base)
{
	const auto pdos_header = reinterpret_cast<PIMAGE_DOS_HEADER>(map_base);
	const auto pNT_header = reinterpret_cast<PIMAGE_NT_HEADERS>(map_base + pdos_header->e_lfanew);

	auto image_size = pNT_header->OptionalHeader.SizeOfImage;

	auto map_image_base = reinterpret_cast<char*>(VirtualAlloc(nullptr, image_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE));
	if (!map_image_base)
	{
		//printf("vritualalloc fail!!\n");
		return false;
	}

	//image_size = image_size - TotalVirtualHeaderSize;

	auto kernel_base = kd_driver::ExAllocatePool(device_handle, 0, image_size);
	if (!kernel_base)
	{
		//printf("[-] ExAllocatePool fail\n");
		return false;
	}

	//写入pe头
	memcpy(map_image_base, map_base, pNT_header->OptionalHeader.SizeOfHeaders);

	//节区拉伸
	const auto section_header = IMAGE_FIRST_SECTION(pNT_header);
	for (size_t i = 0; i < pNT_header->FileHeader.NumberOfSections; i++)
	{
		auto section_virtual = reinterpret_cast<char*>(map_image_base + section_header[i].VirtualAddress);
		memcpy(section_virtual, reinterpret_cast<char*>(map_base + section_header[i].PointerToRawData), section_header[i].SizeOfRawData);
	}


	//修复重定位表
	//auto PReloc = reinterpret_cast<PIMAGE_BASE_RELOCATION>(map_image_base + pNT_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);

	auto current_base_relocation = reinterpret_cast<PIMAGE_BASE_RELOCATION>(reinterpret_cast<ULONG64>(map_image_base) + pNT_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
	const auto reloc_end = reinterpret_cast<PIMAGE_BASE_RELOCATION>(reinterpret_cast<ULONG64>(current_base_relocation) + pNT_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size);

	while (current_base_relocation < reloc_end && current_base_relocation->SizeOfBlock)
	{
		WORD* NumOfBlock = (WORD*)(current_base_relocation + 1);

		for (size_t i = 0; i < (current_base_relocation->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / 2; i++)
		{
			if (NumOfBlock[i] >> 12 == IMAGE_REL_BASED_DIR64)
			{
				auto RelocAddr = reinterpret_cast<uint64_t*>((char*)map_image_base + (NumOfBlock[i] & 0xFFF) + current_base_relocation->VirtualAddress);
				*RelocAddr -= pNT_header->OptionalHeader.ImageBase;
				

				*RelocAddr += (DWORD64)kernel_base;
			}
		}

		current_base_relocation = reinterpret_cast<PIMAGE_BASE_RELOCATION>(reinterpret_cast<DWORD64>(current_base_relocation) + current_base_relocation->SizeOfBlock);
	}


	//修复加载配置
	auto load_config_directory = pNT_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG].VirtualAddress;
	if (load_config_directory)  //有些驱动并没有加载表 如果没有就不用修复了
	{
		auto load_config_struct = (PIMAGE_LOAD_CONFIG_DIRECTORY)((uintptr_t)map_image_base + load_config_directory);
		auto stack_cookie = load_config_struct->SecurityCookie;
		stack_cookie = stack_cookie - (uintptr_t)kernel_base + (uintptr_t)map_image_base; //since our local image is already relocated the base returned will be kernel address

		if (*(uintptr_t*)(stack_cookie) != 0x2B992DDFA232) 
			return false;
		auto new_cookie = 0x2B992DDFA232 ^ GetCurrentProcessId() ^ GetCurrentThreadId(); // here we don't really care about the value of stack cookie, it will still works and produce nice result
		if (new_cookie == 0x2B992DDFA232)
			new_cookie = 0x2B992DDFA233;

		*(uintptr_t*)(stack_cookie) = new_cookie; // the _security_cookie_complement will be init by the driver itself if they use crt
	}

	//修复导入表
	auto pimpor = reinterpret_cast<PIMAGE_IMPORT_DESCRIPTOR>(map_image_base + pNT_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
	while (pimpor->Name)
	{
		const auto Module = utils::GetKernelModule((char*)map_image_base + pimpor->Name);

		auto Int = reinterpret_cast<PIMAGE_THUNK_DATA>(map_image_base + pimpor->OriginalFirstThunk);
		auto Iat = reinterpret_cast<PIMAGE_THUNK_DATA>(map_image_base + pimpor->FirstThunk);

		while (Int->u1.AddressOfData && Iat->u1.AddressOfData)
		{
			if (!IMAGE_SNAP_BY_ORDINAL(Int->u1.AddressOfData))
			{
				auto Names = reinterpret_cast<PIMAGE_IMPORT_BY_NAME>(map_image_base + Int->u1.AddressOfData);
				Iat->u1.AddressOfData = kd_driver::get_kernel_module_export(device_handle, ker_base, Names->Name);
			}

			Int++; Iat++;
		}
		pimpor++;
	}


	kd_driver::write(device_handle, (uint64_t)kernel_base, (uint64_t*)map_image_base, image_size);

	auto entrypoint = kernel_base + pNT_header->OptionalHeader.AddressOfEntryPoint;
	NTSTATUS status = 0;
	kd_driver::func_call(device_handle, (void*)entrypoint, &status, 0, 0);


	//DWORD64 imagehead = 0;
	//kd_driver::write(device_handle, kernel_base, &imagehead, pNT_header->OptionalHeader.SizeOfHeaders);

	VirtualFree(map_image_base, 0, MEM_RELEASE);
	return true;
}