From 8a04910bd43dcce066f2300be07bc17380d891f8 Mon Sep 17 00:00:00 2001
From: CJ <chris@apolyta.com>
Date: Wed, 19 Oct 2022 03:05:46 -0700
Subject: [PATCH] fixed some stuff

---
 kutils.hpp | 46 +++++++++++++++++++++++++++++++++++-----------
 1 file changed, 35 insertions(+), 11 deletions(-)

diff --git a/kutils.hpp b/kutils.hpp
index 3951b06..3eb94c7 100644
--- a/kutils.hpp
+++ b/kutils.hpp
@@ -163,6 +163,22 @@ typedef struct _IMAGE_DATA_DIRECTORY {
   unsigned Size;
 } IMAGE_DATA_DIRECTORY, *PIMAGE_DATA_DIRECTORY;
 
+typedef struct _IMAGE_SECTION_HEADER {
+  unsigned char Name[8];
+  union {
+    unsigned PhysicalAddress;
+    unsigned VirtualSize;
+  } Misc;
+  unsigned VirtualAddress;
+  unsigned SizeOfRawData;
+  unsigned PointerToRawData;
+  unsigned PointerToRelocations;
+  unsigned PointerToLinenumbers;
+  unsigned short NumberOfRelocations;
+  unsigned short NumberOfLinenumbers;
+  unsigned Characteristics;
+} IMAGE_SECTION_HEADER, *PIMAGE_SECTION_HEADER;
+
 typedef struct _IMAGE_OPTIONAL_HEADER64 {
   short Magic;
   unsigned char MajorLinkerVersion;
@@ -514,34 +530,41 @@ extern "C" NTSTATUS ZwQuerySystemInformation(
     PULONG ReturnLength);
 
 namespace Driver {
-FORCEINLINE PVOID GetDriverExportByHash(_In_ PVOID lpDriverBase,
+/// <summary>
+/// walk export directory of a module given its base and the hash of the export
+/// string.
+/// </summary>
+/// <param name="lpDriverBase"></param>
+/// <param name="nStrHash"></param>
+/// <returns></returns>
+FORCEINLINE PVOID GetDriverExportByHash(_In_ PVOID ModuleBase,
                                         _In_ ULONG nStrHash) {
-  PIMAGE_DOS_HEADER lpDosHeader = (PIMAGE_DOS_HEADER)lpDriverBase;
+  PIMAGE_DOS_HEADER lpDosHeader = (PIMAGE_DOS_HEADER)ModuleBase;
   PIMAGE_NT_HEADERS64 lpNtHeader =
-      (PIMAGE_NT_HEADERS64)(lpDosHeader->e_lfanew + (ULONG64)lpDriverBase);
+      (PIMAGE_NT_HEADERS64)(lpDosHeader->e_lfanew + (ULONG64)ModuleBase);
 
   PIMAGE_EXPORT_DIRECTORY lpExportDir =
-      (PIMAGE_EXPORT_DIRECTORY)((ULONG64)lpDriverBase +
+      (PIMAGE_EXPORT_DIRECTORY)((ULONG64)ModuleBase +
                                 lpNtHeader->OptionalHeader
                                     .DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]
                                     .VirtualAddress);
 
   ULONG32* lpNameArr =
-      (ULONG32*)(lpExportDir->AddressOfNames + (ULONG64)lpDriverBase);
+      (ULONG32*)(lpExportDir->AddressOfNames + (ULONG64)ModuleBase);
 
   ULONG32* lpFuncs =
-      (ULONG32*)(lpExportDir->AddressOfFunctions + (ULONG64)lpDriverBase);
+      (ULONG32*)(lpExportDir->AddressOfFunctions + (ULONG64)ModuleBase);
 
   USHORT* lpOrdinals =
-      (USHORT*)(lpExportDir->AddressOfNameOrdinals + (ULONG64)lpDriverBase);
+      (USHORT*)(lpExportDir->AddressOfNameOrdinals + (ULONG64)ModuleBase);
 
   for (auto nIdx = 0u; nIdx < lpExportDir->NumberOfFunctions; ++nIdx) {
     if (!lpNameArr[nIdx] || !lpOrdinals[nIdx])
       continue;
 
-    if (hashstr::hash((PCHAR)((ULONG64)lpDriverBase + lpNameArr[nIdx])) ==
+    if (hashstr::hash((PCHAR)((ULONG64)ModuleBase + lpNameArr[nIdx])) ==
         nStrHash)
-      return (PVOID)((ULONG64)lpDriverBase + lpFuncs[lpOrdinals[nIdx]]);
+      return (PVOID)((ULONG64)ModuleBase + lpFuncs[lpOrdinals[nIdx]]);
   }
   return NULL;
 }
@@ -575,7 +598,8 @@ FORCEINLINE PVOID GetKernelBase() {
       if (lpNtHeaders->Signature != PE_HEADER_MAGIC)
         continue;
 
-      if (lpNtHeaders->OptionalHeader.SizeOfImage < 0x1000000)
+      if (!GetDriverExportByHash((PVOID)nPage,
+                                 HSTRING("ExAllocatePoolWithTag")))
         continue;
 
       return (PVOID)nPage;
@@ -863,4 +887,4 @@ inline auto Scan(void* base, u32 size, const char* pattern, const char* mask)
   return nullptr;
 }
 }  // namespace Signature
-}  // namespace KUtils
+}  // namespace KUtils
\ No newline at end of file