From 302d8e2f72871a482d84500e10fcb1d4a22aca91 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Benjamin=20H=C3=B6glinger-Stelzer?= <nefarius@dhmx.at>
Date: Tue, 4 Jul 2023 17:17:57 +0200
Subject: [PATCH] Converted ZwQuerySystemInformation to dynamic call

---
 src/Domito.Internal.h | 49 ++++++++++++++++++++++---------------------
 src/Domito.cpp        | 14 ++++++++++---
 2 files changed, 36 insertions(+), 27 deletions(-)

diff --git a/src/Domito.Internal.h b/src/Domito.Internal.h
index 024aa89..1dccdc8 100644
--- a/src/Domito.Internal.h
+++ b/src/Domito.Internal.h
@@ -1,15 +1,15 @@
-/*  ___     _                     _   _                     
+/*  ___     _                     _   _
  * |_ _|_ _| |_ ___ _ _ _ _  __ _| | | |_ _  _ _ __  ___ ___
  *  | || ' \  _/ -_) '_| ' \/ _` | | |  _| || | '_ \/ -_|_-<
  * |___|_||_\__\___|_| |_||_\__,_|_|  \__|\_, | .__/\___/__/
- *                                        |__/|_|           
+ *                                        |__/|_|
  */
 
 #pragma once
 
-//
-// SDK/WDK
-// 
+ //
+ // SDK/WDK
+ // 
 #include <ntifs.h>
 #include <ntintsafe.h>
 #include <ntimage.h>
@@ -21,11 +21,11 @@
 #include "Domito.h"
 
 
-/*  _  _ _   ___  _ _       _        
- * | \| | |_|   \| | |  ___| |_ __   
- * | .` |  _| |) | | | / -_)  _/ _|_ 
+/*  _  _ _   ___  _ _       _
+ * | \| | |_|   \| | |  ___| |_ __
+ * | .` |  _| |) | | | / -_)  _/ _|_
  * |_|\_|\__|___/|_|_| \___|\__\__(_)
- *                                   
+ *
  */
 
  // Structure representing a loaded module
@@ -50,14 +50,6 @@ typedef struct _SYSTEM_MODULE_INFORMATION
     SYSTEM_MODULE_INFORMATION_ENTRY Module[ANYSIZE_ARRAY];
 } SYSTEM_MODULE_INFORMATION, * PSYSTEM_MODULE_INFORMATION;
 
-// Function prototype for ZwQuerySystemInformation
-EXTERN_C NTSYSAPI NTSTATUS NTAPI ZwQuerySystemInformation(
-    _In_ ULONG SystemInformationClass,
-    _Inout_ PVOID SystemInformation,
-    _In_ ULONG SystemInformationLength,
-    _Out_opt_ PULONG ReturnLength
-);
-
 typedef struct _LDR_DATA_TABLE_ENTRY
 {
     LIST_ENTRY64 InLoadOrderLinks;
@@ -97,6 +89,13 @@ typedef NTSTATUS(NTAPI* t_ZwQueryInformationProcess) (
     __out_opt PULONG ReturnLength
     );
 
+typedef NTSTATUS(NTAPI* t_ZwQuerySystemInformation)(
+    _In_ ULONG SystemInformationClass,
+    _Inout_ PVOID SystemInformation,
+    _In_ ULONG SystemInformationLength,
+    _Out_opt_ PULONG ReturnLength
+    );
+
 /*   ___
  *  / __|___ _ __  _ __  ___ _ _
  * | (__/ _ \ '  \| '  \/ _ \ ' \
@@ -109,22 +108,24 @@ typedef struct
     t_RtlImageDirectoryEntryToData RtlImageDirectoryEntryToData;
 
     t_ZwQueryInformationProcess ZwQueryInformationProcess;
-    
+
+    t_ZwQuerySystemInformation ZwQuerySystemInformation;
+
 } DOMITO_COMMON;
 
 extern DOMITO_COMMON G_Common;
 
 
-/*  __  __                          __  __                                       _   
- * |  \/  |___ _ __  ___ _ _ _  _  |  \/  |__ _ _ _  __ _ __ _ ___ _ __  ___ _ _| |_ 
+/*  __  __                          __  __                                       _
+ * |  \/  |___ _ __  ___ _ _ _  _  |  \/  |__ _ _ _  __ _ __ _ ___ _ __  ___ _ _| |_
  * | |\/| / -_) '  \/ _ \ '_| || | | |\/| / _` | ' \/ _` / _` / -_) '  \/ -_) ' \  _|
  * |_|  |_\___|_|_|_\___/_|  \_, | |_|  |_\__,_|_||_\__,_\__, \___|_|_|_\___|_||_\__|
- *                           |__/                        |___/                       
+ *                           |__/                        |___/
  */
 
-//
-// Default pool tag
-// 
+ //
+ // Default pool tag
+ // 
 #define DOMITO_POOL_TAG     'imoD'
 
 //
diff --git a/src/Domito.cpp b/src/Domito.cpp
index b79ec9e..73c43d3 100644
--- a/src/Domito.cpp
+++ b/src/Domito.cpp
@@ -12,6 +12,7 @@ DOMITO_COMMON G_Common = {};
 
 DECLARE_GLOBAL_CONST_UNICODE_STRING(G_QipRoutineName, L"ZwQueryInformationProcess");
 DECLARE_GLOBAL_CONST_UNICODE_STRING(G_IdetdRoutineName, L"RtlImageDirectoryEntryToData");
+DECLARE_GLOBAL_CONST_UNICODE_STRING(G_QsiRoutineName, L"ZwQuerySystemInformation");
 
 #ifndef LOG
 #define LOG(Format, ...) DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Domito][" __FUNCTION__ "] " Format " \n", __VA_ARGS__)
@@ -32,6 +33,8 @@ DomitoInit()
         (t_ZwQueryInformationProcess)MmGetSystemRoutineAddress((PUNICODE_STRING)&G_QipRoutineName);
     G_Common.RtlImageDirectoryEntryToData =
         (t_RtlImageDirectoryEntryToData)MmGetSystemRoutineAddress((PUNICODE_STRING)&G_IdetdRoutineName);
+    G_Common.ZwQuerySystemInformation =
+        (t_ZwQuerySystemInformation)MmGetSystemRoutineAddress((PUNICODE_STRING)&G_QsiRoutineName);
 
     return STATUS_SUCCESS; // TODO: unused currently
 }
@@ -56,10 +59,15 @@ DomitoFindModuleBaseAddress(
     ULONG bufferSize = 0;
     PSYSTEM_MODULE_INFORMATION moduleInfo = NULL;
 
+    if (!G_Common.ZwQuerySystemInformation)
+    {
+        return STATUS_NOT_IMPLEMENTED;
+    }
+
     const ULONG SystemModuleInformation = 11;
 
     // Query the required buffer size for module information
-    NTSTATUS status = ZwQuerySystemInformation(
+    NTSTATUS status = G_Common.ZwQuerySystemInformation(
         SystemModuleInformation,
         &bufferSize,
         0,
@@ -82,7 +90,7 @@ DomitoFindModuleBaseAddress(
     }
 
     // Retrieve the module information
-    status = ZwQuerySystemInformation(
+    status = G_Common.ZwQuerySystemInformation(
         SystemModuleInformation,
         moduleInfo,
         bufferSize,
@@ -133,7 +141,7 @@ DomitoFindExportedFunctionAddress(
 {
     NTSTATUS status = STATUS_NOT_FOUND;
     ULONG exportSize;
-    
+
     if (G_Common.RtlImageDirectoryEntryToData == NULL)
     {
         return STATUS_NOT_IMPLEMENTED;