Improved heap search for bigger NT heaps

Author: leonav-unizar

CMakeLists.txt

@@ -28,6 +28,7 @@ endif()
 ## Ease for includes
 add_library(common_includes INTERFACE)
 add_library(interrogate_include INTERFACE)
+target_include_directories(common_includes INTERFACE ${PROJECT_SOURCE_DIR}/include/ntdll)
 target_include_directories(common_includes INTERFACE ${PROJECT_SOURCE_DIR}/include/keyreaper)
 target_include_directories(interrogate_include INTERFACE ${PROJECT_SOURCE_DIR}/include/interrogate)
 

include/keyreaper/key_scanner.h

@@ -38,7 +38,7 @@ class ScannerFacade {
 
   // Query
   bool IsProcessAlive() const;
-  std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> DoScan();
+  std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> DoScan(bool extended_search_enabled);
   std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> GetKeys();  
   error_handling::ProgramResult ExportKeysToJSON(std::string output_json);
   error_handling::ProgramResult ExportKeysToBinary();

include/keyreaper/process_capturer.h

@@ -5,6 +5,8 @@
 #include <vector>
 #include <tlhelp32.h>
 
+#include "ntdll.h"
+#include "winntheap.h"
 #include "injection/custom_ipc.h"
 #include "program_result.h"
 
@@ -14,9 +16,15 @@ typedef _Return_type_success_(return >= 0) LONG NTSTATUS;
 
 typedef NTSTATUS(NTAPI *pNtSuspendProcess)(
   HANDLE ProcessHandle
-); // Undocumented NTDLL function
+);
 }
 
+typedef NTSTATUS (NTAPI* pNtQueryVirtualMemory)(
+  HANDLE, PVOID, int, PVOID, SIZE_T, PSIZE_T);
+
+typedef NTSTATUS (NTAPI* pNtQueryInformationProcess)(
+  HANDLE, int, PVOID, ULONG, PULONG);
+
 namespace process_manipulation {
 
 // from x64dbg
@@ -150,7 +158,7 @@ class ProcessCapturer {
   error_handling::ProgramResult GetMemoryChunk(LPCVOID start, SIZE_T size, BYTE* buffer, SIZE_T* bytes_read);
   void WriteBufferToFile(unsigned char* buffer, SIZE_T size, std::string file_name);
 
-  error_handling::ProgramResult EnumerateHeaps(std::vector<HeapInformation>* heaps);
+  error_handling::ProgramResult EnumerateHeaps(std::vector<HeapInformation>& heaps, bool extended_search = false);
 
   /**
    * From the information about the heap, copies the whole heap to a buffer.
@@ -226,6 +234,13 @@ class ProcessCapturer {
    */
   error_handling::ProgramResult GetKeyBlobFromRemote(HCRYPTKEY key_handle, DWORD blob_type, std::vector<BYTE>& key_blob);
 
+  /**
+   * Returns a copy of the Process Environment Block of the target process
+   * @param peb A pointer to a PEB structure
+   */
+  error_handling::ProgramResult CopyPEB(void* peb);
+  void* GetPEBLocation();
+
  private:
   /**
    * Function for initializing dynamically imported functions, such as `NtSuspendProcess`.
@@ -240,6 +255,17 @@ class ProcessCapturer {
   error_handling::ProgramResult StartControllerServerOnProcess();
   error_handling::ProgramResult StopControllerServerOnProcess(bool terminate = false);
 
+    /**
+   * Looks for heaps by enumerating the memory regions of the process and matching 
+   * for the HEAP_ENTRY structure
+   * @param heaps (OUT) vector where the heap region base and size will be placed
+   */
+  error_handling::ProgramResult ExtendedHeapSearch(std::vector<HeapInformation>& heaps);
+  error_handling::ProgramResult SimpleHeapSearch(std::vector<HeapInformation>& heaps);
+
+  HANDLE proc_handle_;
+  pNtQueryInformationProcess fnNtQueryInformationProcess_;
+  pNtQueryVirtualMemory fnNtQueryVirtualMemory_;
   static nt_suspend::pNtSuspendProcess fNtPauseProcess;
   custom_ipc::CustomClient injection_client_;
   DWORD pid_;