From 378c52eb47a15810ecf61c6bd5de4a0800e8f089 Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sun, 17 Apr 2022 11:51:31 +0200
Subject: [PATCH] More stuff

---
 src/driver/ept.cpp      | 12 +++++++-----
 src/driver/irp.cpp      | 19 ++++++++++++++++---
 src/runner/main.cpp     | 13 +++++++++----
 src/shared/irp_data.hpp |  1 +
 4 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/driver/ept.cpp b/src/driver/ept.cpp
index ff31dc5..ae67b69 100644
--- a/src/driver/ept.cpp
+++ b/src/driver/ept.cpp
@@ -425,6 +425,9 @@ namespace vmx
 		{
 			if (hook->target_page->flags == hook->original_entry.flags)
 			{
+				const auto* data_source = translation_hint ? &translation_hint->page[0] : virtual_target;
+				memcpy(&hook->fake_page[0], data_source, PAGE_SIZE);
+
 				hook->target_page->flags = hook->readwrite_entry.flags;
 			}
 
@@ -441,7 +444,6 @@ namespace vmx
 		this->split_large_page(physical_address);
 
 		const auto* data_source = translation_hint ? &translation_hint->page[0] : virtual_target;
-
 		memcpy(&hook->fake_page[0], data_source, PAGE_SIZE);
 		hook->physical_base_address = physical_base_address;
 
@@ -516,7 +518,7 @@ namespace vmx
 		auto current_destination = reinterpret_cast<uint64_t>(destination);
 		auto current_length = length;
 
-		 ept_translation_hint* current_hints = nullptr;
+		ept_translation_hint* current_hints = nullptr;
 
 		auto destructor = utils::finally([&current_hints]()
 		{
@@ -531,7 +533,7 @@ namespace vmx
 			const auto data_to_write = min(page_remaining, current_length);
 
 			auto* new_hint = memory::allocate_non_paged_object<ept_translation_hint>();
-			if(!new_hint)
+			if (!new_hint)
 			{
 				throw std::runtime_error("Failed to allocate hint");
 			}
@@ -541,12 +543,12 @@ namespace vmx
 			current_hints->virtual_base_address = aligned_destination;
 			current_hints->physical_base_address = memory::get_physical_address(aligned_destination);
 
-			if(!current_hints->physical_base_address)
+			if (!current_hints->physical_base_address)
 			{
 				throw std::runtime_error("Failed to resolve physical address");
 			}
 
-			memcpy(&current_hints->page[0], aligned_destination, PAGE_SIZE);			
+			memcpy(&current_hints->page[0], aligned_destination, PAGE_SIZE);
 
 			current_length -= data_to_write;
 			current_destination += data_to_write;
diff --git a/src/driver/irp.cpp b/src/driver/irp.cpp
index 572e93b..f301e70 100644
--- a/src/driver/irp.cpp
+++ b/src/driver/irp.cpp
@@ -41,7 +41,7 @@ namespace
 	void apply_hook(const hook_request* request)
 	{
 		auto* buffer = new uint8_t[request->source_data_size];
-		if(!buffer)
+		if (!buffer)
 		{
 			throw std::runtime_error("Failed to copy buffer");
 		}
@@ -80,17 +80,27 @@ namespace
 
 		t.join();
 
-		if(!translation_hints)
+		if (!translation_hints)
 		{
 			debug_log("Failed to generate tranlsation hints");
 			return;
 		}
 
-		hypervisor::get_instance()->install_ept_hook(request->target_address, buffer, request->source_data_size, translation_hints);
+		hypervisor::get_instance()->install_ept_hook(request->target_address, buffer, request->source_data_size,
+		                                             translation_hints);
 
 		debug_log("Done1\n");
 	}
 
+	void unhook()
+	{
+		const auto instance = hypervisor::get_instance();
+		if(instance)
+		{
+			instance->disable_all_ept_hooks();
+		}
+	}
+
 	_Function_class_(DRIVER_DISPATCH) NTSTATUS io_ctl_handler(
 		PDEVICE_OBJECT /*device_object*/, const PIRP irp)
 	{
@@ -113,6 +123,9 @@ namespace
 			case HOOK_DRV_IOCTL:
 				apply_hook(static_cast<hook_request*>(irp_sp->Parameters.DeviceIoControl.Type3InputBuffer));
 				break;
+			case UNHOOK_DRV_IOCTL:
+				unhook();
+				break;
 			default:
 				debug_log("Invalid IOCTL Code: 0x%X\n", ioctr_code);
 				irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
diff --git a/src/runner/main.cpp b/src/runner/main.cpp
index 47db47d..5aafaad 100644
--- a/src/runner/main.cpp
+++ b/src/runner/main.cpp
@@ -67,19 +67,24 @@ void unsafe_main(const int /*argc*/, char* /*argv*/[])
 
 	hook_request hook_request{};
 	hook_request.process_id = _pid; //GetCurrentProcessId();
-	hook_request.target_address = (void*)0x465FF7;//0x14007DCF7;
+	hook_request.target_address = (void*)0x41297A;//0x14007DCF7;
 
-	uint8_t buffer[1];
-	buffer[0] = 0xEB;
+	uint8_t buffer[] = {0x90, 0x90};
 
 	hook_request.source_data = buffer;
-	hook_request.source_data_size = 1;
+	hook_request.source_data_size = sizeof(buffer);
 
 	input.assign(reinterpret_cast<uint8_t*>(&hook_request),
 	             reinterpret_cast<uint8_t*>(&hook_request) + sizeof(hook_request));
 
 	(void)driver_device.send(HOOK_DRV_IOCTL, input);
 
+	printf("Press any key to disable all hooks!\n");
+	_getch();
+
+	input.resize(0);
+	(void)driver_device.send(UNHOOK_DRV_IOCTL, input);
+
 	printf("Press any key to exit!\n");
 	_getch();
 }
diff --git a/src/shared/irp_data.hpp b/src/shared/irp_data.hpp
index 782f0db..0227c21 100644
--- a/src/shared/irp_data.hpp
+++ b/src/shared/irp_data.hpp
@@ -2,6 +2,7 @@
 
 #define HELLO_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_NEITHER, FILE_ANY_ACCESS)
 #define HOOK_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_NEITHER, FILE_ANY_ACCESS)
+#define UNHOOK_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x802, METHOD_NEITHER, FILE_ANY_ACCESS)
 
 static_assert(sizeof(void*) == 8);