#include "hmdm_ctx.h"
#include "vdm_ctx/vdm_ctx.hpp"

int __cdecl main(int argc, char** argv)
{
	if (argc < 2)
	{
		std::printf("> please provide a path to a driver...\n");
		return -1;
	}

	const auto [drv_handle, drv_key, drv_status] = vdm::load_drv();
	std::printf("> drv handle: 0x%x, drv key: %s, load status: 0x%x\n",
		drv_handle, drv_key.c_str(), drv_status);

	if (drv_status != STATUS_SUCCESS || drv_handle == INVALID_HANDLE_VALUE)
	{
		std::printf("> failed to load driver... reason -> 0x%x\n", drv_status);
		return -1;
	}

	// read physical memory using the driver...
	vdm::read_phys_t _read_phys =
		[&](void* addr, void* buffer, std::size_t size) -> bool
	{
		return vdm::read_phys(addr, buffer, size);
	};

	// write physical memory using the driver...
	vdm::write_phys_t _write_phys =
		[&](void* addr, void* buffer, std::size_t size) -> bool
	{
		return vdm::write_phys(addr, buffer, size);
	};

	// use VDM to syscall into ExAllocatePool...
	vdm::vdm_ctx vdm(_read_phys, _write_phys);
	drv::kalloc_t _kalloc = [&](std::size_t size) -> void*
	{
		using ex_alloc_pool_t = 
			void* (*)(std::uint32_t, std::uint32_t);

		static const auto ex_alloc_pool =
			reinterpret_cast<void*>(
				utils::kmodule::get_export(
					"ntoskrnl.exe", "ExAllocatePool"));

		return vdm.syscall<ex_alloc_pool_t>(ex_alloc_pool, NULL, size);
	};

	// use VDM to syscall into memcpy exported by ntoskrnl.exe...
	drv::kmemcpy_t _kmemcpy = 
		[&](void* dest, const void* src, std::size_t size) -> void*
	{
		static const auto kmemcpy = 
			reinterpret_cast<void*>(
				utils::kmodule::get_export(
					"ntoskrnl.exe", "memcpy"));

		return vdm.syscall<decltype(&memcpy)>(kmemcpy, dest, src, size);
	};

	drv::drv_buffer_t drv_buffer;
	utils::open_binary_file(argv[1], drv_buffer);
	drv::hmdm_ctx drv_mapper({ _kalloc, _kmemcpy });

	const auto [drv_base, drv_entry] = drv_mapper.map_module(drv_buffer);
	std::printf("> driver base -> 0x%p, driver entry -> 0x%p\n", drv_base, drv_entry);

	if (!drv_base || !drv_entry)
	{
		std::printf("> failed to map driver...\n");
		return -1;
	}
	
	// call driver entry... its up to you to do this using whatever method...
	// with VDM you can syscall into it... with msrexec you will use msrexec::exec...
	const auto entry_result = vdm.syscall<NTSTATUS(*)(std::uintptr_t)>(
		reinterpret_cast<void*>(drv_entry), drv_base);

	std::printf("> entry result -> 0x%p\n", entry_result);
	

	const auto unload_status = vdm::unload_drv(drv_handle, drv_key);
	if (unload_status != STATUS_SUCCESS)
	{
		std::printf("> failed to unload driver... reason -> 0x%x\n", unload_status);
		return -1;
	}
}