#include <ntifs.h>
#include "../SymbolicAccess/Utils/Log.h"
#include "hypervisor_gateway.h"
#include "ASM/AsmCallset.h"
#include "vmcall_reason.h"
#include "ntapi.h"

#define IOCTL_POOL_MANAGER_ALLOCATE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x900, METHOD_BUFFERED, FILE_SPECIAL_ACCESS)

namespace hvgt
{
	void broadcast_vmoff(KDPC* Dpc, PVOID DeferredContext, PVOID SystemArgument1, PVOID SystemArgument2)
	{
		UNREFERENCED_PARAMETER(Dpc);
		UNREFERENCED_PARAMETER(DeferredContext);

		__vm_call(VMCALL_VMXOFF, 0, 0, 0);
		KeSignalCallDpcSynchronize(SystemArgument2);
		KeSignalCallDpcDone(SystemArgument1);
	}

	void broadcast_invept_all_contexts(KDPC* Dpc, PVOID DeferredContext, PVOID SystemArgument1, PVOID SystemArgument2)
	{
		UNREFERENCED_PARAMETER(Dpc);
		UNREFERENCED_PARAMETER(DeferredContext);

		__vm_call(VMCALL_INVEPT_CONTEXT, true, 0, 0);
		KeSignalCallDpcSynchronize(SystemArgument2);
		KeSignalCallDpcDone(SystemArgument1);
	}

	void broadcast_invept_single_context(KDPC* Dpc, PVOID DeferredContext, PVOID SystemArgument1, PVOID SystemArgument2)
	{
		UNREFERENCED_PARAMETER(Dpc);
		UNREFERENCED_PARAMETER(DeferredContext);

		__vm_call(VMCALL_INVEPT_CONTEXT, false, 0, 0);
		KeSignalCallDpcSynchronize(SystemArgument2);
		KeSignalCallDpcDone(SystemArgument1);
	}

	/// <summary>
	/// Turn off virtual machine
	/// </summary>
	void vmoff()
	{
		KeGenericCallDpc(broadcast_vmoff, NULL);
	}

	/// <summary>
	/// Invalidates mappings in the translation lookaside buffers (TLBs) 
	/// and paging-structure caches that were derived from extended page tables (EPT)
	/// </summary>
	/// <param name="invept_all"> If true invalidates all contexts otherway invalidate only single context (currently hv doesn't use more than 1 context)</param>
	void invept(bool invept_all)
	{
		if (invept_all == true) KeGenericCallDpc(broadcast_invept_all_contexts, NULL);
		else KeGenericCallDpc(broadcast_invept_single_context, NULL);
	}

	/// <summary>
	/// Set/Unset presence of hypervisor
	/// </summary>
	/// <param name="value"> If false, hypervisor is not visible via cpuid interface, If true, it become visible</param>
	void hypervisor_visible(bool value)
	{
		if (value == true)
			__vm_call(VMCALL_UNHIDE_HV_PRESENCE, 0, 0, 0);
		else
			__vm_call(VMCALL_HIDE_HV_PRESENCE, 0, 0, 0);
	}

	/// <summary>
	/// Unhook all functions and invalidate tlb
	/// 卸载全部hook 并刷新tlb
	/// </summary>
	/// <returns> status </returns>
	bool ept_unhook()
	{
		bool status = __vm_call(VMCALL_EPT_UNHOOK_FUNCTION, true, 0, 0);
		invept(false);
		return status;
	}

	/// <summary>
	/// Unhook single function and invalidate tlb
	/// 卸载指定函数 并刷新tlb
	/// </summary>
	/// <param name="function_address"></param>
	/// <returns> status </returns>
	bool ept_unhook(void* function_address)
	{
		bool status = __vm_call(VMCALL_EPT_UNHOOK_FUNCTION, false, (unsigned __int64)function_address, 0);
		invept(false);
		return status;
	}

	//读取隐形软件断点
	bool get_hide_software_breakpoint(void* target_address, void* buffer, unsigned __int64 buffer_size)
	{
		bool status = __vm_call(VMCALL_READ_SOFTWARE_BREAKPOINT, (unsigned __int64)target_address, (unsigned __int64)buffer, buffer_size);
		invept(false);

		return status;
	}

	//设置隐形软件断点
	bool set_hide_software_breakpoint(void* target_address, void* buffer, unsigned __int64 buffer_size)
	{
		bool status = __vm_call(VMCALL_HIDE_SOFTWARE_BREAKPOINT, (unsigned __int64)target_address, (unsigned __int64)buffer, buffer_size);
		invept(false);

		return status;
	}

	/// <summary>
	/// Hook function via ept and invalidates mappings
	/// </summary>
	/// <param name="target_address">Address of function which we want to hook</param>
	/// <param name="hook_function">Address of function which is used to call original function</param>
	/// <param name="origin_function">Address of function which is used to call original function</param>
	/// <returns> status </returns>
	bool hook_function(void* target_address, void* proxy_function, void** origin_function)
	{
		bool status = __vm_call_ex(VMCALL_EPT_HOOK_FUNCTION, (unsigned __int64)target_address, (unsigned __int64)proxy_function, 0, (unsigned __int64)origin_function, 0, 0, 0, 0, 0);
		invept(false);

		return status;
	}

	/// <summary>
	/// Hook function via ept and invalidates mappings
	/// </summary>
	/// <param name="target_address">Address of function which we want to hook</param>
	/// <param name="hook_function">Address of function which is used to call original function</param>
	/// <param name="trampoline_address">Address of some memory which isn't used with size at least 13 and withing 2GB range of target function
	/// Use only if you can function you want to hook use relative offeset in first 13 bytes of it. For example if you want hook NtYieldExecution which
	/// size is 15 bytes you have to find a codecave witihn ntoskrnl.exe image with size atleast 13 bytes and pass it there</param>
	/// <param name="origin_function">Address of function which is used to call original function</param>
	/// <returns> status </returns>
	bool hook_function(void* target_address, void* proxy_function, void* trampoline_address, void** origin_function)
	{
		bool status = __vm_call_ex(VMCALL_EPT_HOOK_FUNCTION, (unsigned __int64)target_address, (unsigned __int64)proxy_function, (unsigned __int64)trampoline_address, (unsigned __int64)origin_function, 0, 0, 0, 0, 0);
		invept(false);

		return status;
	}

	/// <summary>
	/// Check if we can communicate with hypervisor
	/// </summary>
	/// <returns> status </returns>
	bool test_vmcall()
	{
		return __vm_call(VMCALL_TEST, 0, 0, 0);
	}

	//测试VT驱动是否加载
	NTSTATUS TestVMM()
	{
		__try
		{
			if (test_vmcall() == FALSE)  //尝试是否能够成功执行vmx指令
			{
				outLog("vt驱动没有安装!!!\n");
				return STATUS_UNSUCCESSFUL;
			}
		}
		__except (EXCEPTION_EXECUTE_HANDLER)
		{
			outLog("不支持vmx指令!!!\n");
			return STATUS_UNSUCCESSFUL;
		}
		outLog("VT驱动已安装成功!!!\n");
		return STATUS_SUCCESS;
	}

	/// <summary>
	/// Send irp with information to allocate memory
	/// </summary>
	/// <returns> status </returns>
	bool perform_memory_allocation()
	{
		PDEVICE_OBJECT airhv_device_object;
		KEVENT event;
		PIRP irp;
		IO_STATUS_BLOCK io_status = { 0 };
		UNICODE_STRING airhv_name;
		PFILE_OBJECT file_object;

		RtlInitUnicodeString(&airhv_name, L"\\Device\\airhv");

		NTSTATUS status = IoGetDeviceObjectPointer(&airhv_name, 0, &file_object, &airhv_device_object);

		ObReferenceObjectByPointer(airhv_device_object, FILE_ALL_ACCESS, NULL, KernelMode);

		// We don't need this so we instantly dereference file object
		ObDereferenceObject(file_object);

		if (NT_SUCCESS(status) == false)
		{
			LogError("Couldn't get hypervisor device object pointer");
			return false;
		}

		KeInitializeEvent(&event, NotificationEvent, 0);
		irp = IoBuildDeviceIoControlRequest(IOCTL_POOL_MANAGER_ALLOCATE, airhv_device_object, 0, 0, 0, 0, 0, &event, &io_status);

		if (irp == NULL)
		{
			LogError("Couldn't create Irp");
			ObDereferenceObject(airhv_device_object);
			return false;
		}

		else
		{
			status = IofCallDriver(airhv_device_object, irp);

			if (status == STATUS_PENDING)
				KeWaitForSingleObject(&event, Executive, KernelMode, 0, 0);

			ObDereferenceObject(airhv_device_object);
			return true;
		}
	}
}