#include "hooks.h"

namespace hooks
{
	void* get_addr_hook(void* base_addr, const char* func_name)
	{
		DBG_PRINT("");
		DBG_PRINT("=============== %s ==============", __FUNCTION__);
		DBG_PRINT("func_name: %s", func_name);
		if (!strcmp(func_name, "MmMapIoSpace"))
			return &map_io_space;
		else if (!strcmp(func_name, "MmUnmapIoSpace"))
			return &unmap_io_space;
		else if (!strcmp(func_name, "MmGetPhysicalAddress"))
			return &get_phys_addr;
		else if (!strcmp(func_name, "IoCreateSymbolicLink"))
			return &create_sym_link;
		else if (!strcmp(func_name, "IoCreateDevice"))
			return &create_device;
		return driver_util::get_kmode_export("ntoskrnl.exe", func_name);
	}

	PHYSICAL_ADDRESS get_phys_addr(
		void* base_addr
	)
	{
		DBG_PRINT("");
		DBG_PRINT("=============== %s ==============", __FUNCTION__);
		DBG_PRINT("getting physical address of: 0x%p", base_addr);
		DBG_PRINT("base_addr value: 0x%p", *(ULONG64*)base_addr);
		auto result = MmGetPhysicalAddress(base_addr);
		DBG_PRINT("physical address: 0x%p", result);
		return result;
	}

	void* map_io_space(
		PHYSICAL_ADDRESS    phys_addr,
		SIZE_T              size,
		MEMORY_CACHING_TYPE cache_type
	)
	{
		DBG_PRINT("");
		DBG_PRINT("=============== %s ==============", __FUNCTION__);
		DBG_PRINT("mapping physical memory 0x%p of size 0x%x", phys_addr, size);
		mapped_io_space_addr = MmMapIoSpace(phys_addr, size, cache_type);
		DBG_PRINT("mapped io space 0x%p, value: 0x%p", mapped_io_space_addr, *(ULONG64*)mapped_io_space_addr);
		last_mapped_ptr = (void*)(*(ULONG64*)mapped_io_space_addr);
		return mapped_io_space_addr;
	}

	void* allocate_pool(
		POOL_TYPE pool_type,
		SIZE_T    size
	)
	{
		DBG_PRINT("");
		DBG_PRINT("=============== %s ==============", __FUNCTION__);
		DBG_PRINT("allocating size: 0x%p", size);

		auto lp_pool = ExAllocatePool(pool_type, size);
		DBG_PRINT("pool allocated at 0x%p", lp_pool);
		return lp_pool;
	}

	void unmap_io_space(
		void*  base_addr,
		SIZE_T size
	)
	{
		DBG_PRINT("");
		DBG_PRINT("=============== %s ==============", __FUNCTION__);
		DBG_PRINT("value of 0x%p is 0x%p", base_addr, *(ULONG64*)base_addr);
		MmUnmapIoSpace(base_addr, size);
	}

	NTSTATUS device_control(
		PDEVICE_OBJECT  device_obj,
		PIRP  irp
	)
	{
		auto irp_stack = IoGetCurrentIrpStackLocation(irp);
		auto ioctl_code = irp_stack->Parameters.Read.ByteOffset.LowPart;
		auto master_irp = irp->AssociatedIrp.MasterIrp;

		DBG_PRINT("");
		DBG_PRINT("=============== %s ==============", __FUNCTION__);
		DBG_PRINT("ioctl called with ioctl code: 0x%p", ioctl_code);
		DBG_PRINT("ioctl called from: 0x%x", PsGetCurrentProcessId());

		if (ioctl_code == WRITE_MEMORY_IOCTL)
		{
			DBG_PRINT("master_irp->MdlAddress: 0x%p", master_irp->MdlAddress);
			DBG_PRINT("master_irp->Type: 0x%p", (__m128i*) master_irp->Type);
			DBG_PRINT("master_irp->Flags: 0x%p", master_irp->Flags);
			*(ULONG64*)mapped_io_space_addr = (ULONG64)&allocate_pool;
		}
		return reinterpret_cast<decltype(&device_control)>(orig_irp_ioctl)(device_obj, irp);
	}

	NTSTATUS create_sym_link(
		PUNICODE_STRING sym_link,
		PUNICODE_STRING device_name
	)
	{
		DBG_PRINT("");
		DBG_PRINT("=============== %s ==============", __FUNCTION__);
		DBG_PRINT("creating symbolic link of device %ws to %ws", device_name->Buffer, sym_link->Buffer);
		DBG_PRINT("IRP_MJ_DEVICE_CONTROL: 0x%p", unfair_drv_obj->MajorFunction[IRP_MJ_DEVICE_CONTROL]);

		//
		// right before we return from driver entry we want to hook the ioctl.
		//
		orig_irp_ioctl = unfair_drv_obj->MajorFunction[IRP_MJ_DEVICE_CONTROL];
		unfair_drv_obj->MajorFunction[IRP_MJ_DEVICE_CONTROL] = &device_control;

		//
		// swap unload pointers
		//
		orig_driver_unload = unfair_drv_obj->DriverUnload;
		unfair_drv_obj->DriverUnload = &driver_unload;
		return IoCreateSymbolicLink(sym_link, device_name);
	}

	NTSTATUS create_device(
		PDRIVER_OBJECT  drv_obj,
		ULONG           drv_ext,
		PUNICODE_STRING device_name,
		DEVICE_TYPE     device_type,
		ULONG           device_char,
		BOOLEAN         exclusive,
		PDEVICE_OBJECT* device_obj
	)
	{
		unfair_drv_obj = drv_obj;
		return IoCreateDevice(
			drv_obj,
			drv_ext,
			device_name,
			device_type,
			device_char,
			exclusive,
			device_obj
		);
	}

	void driver_unload(
		PDRIVER_OBJECT drv_obj
	)
	{
		DBG_PRINT("unfairgame driver unloading....");
		return reinterpret_cast<decltype(&driver_unload)>(orig_driver_unload)(drv_obj);
	}
}