Added Pattern Scannermain

Signed-off-by: Toni Uhlig <matzeton@googlemail.com>

3 files changed, 180 insertions, 0 deletions

diff --git a/driver-kmem.cpp b/driver-kmem.cpp
index 60f7374..3b8fafe 100644
--- a/driver-kmem.cpp
+++ b/driver-kmem.cpp
@@ -17,6 +17,53 @@ NTSTATUS DriverEntry(_In_ struct _DRIVER_OBJECT *DriverObject,
   UNREFERENCED_PARAMETER(DriverObject);
   UNREFERENCED_PARAMETER(RegistryPath);
 
+  const eastl::array<uint8_t, 10> buffer = {0x41, 0xDE, 0xAD, 0xC0, 0xDE,
+                                            0xCA, 0xFE, 0xCA, 0xFE, 0x41};
+  const eastl::array<uint8_t, 2> pattern_00 = {0xCA, 0xFE};
+  const eastl::array<uint8_t, 2> pattern_01 = {0xFE, 0x41};
+  const eastl::array<uint8_t, 1> pattern_02 = {0x41};
+  const eastl::array<uint8_t, 10> pattern_03 = {0x41, 0x00, 0x00, 0x00, 0x00,
+                                                0x00, 0x00, 0x00, 0x00, 0x41};
+  eastl::vector<size_t> found_offsets;
+
+  auto found = PatternScanner::SearchWithMask(
+      buffer.data(), eastl::size(buffer), pattern_00.data(),
+      eastl::size(pattern_00), "xx", found_offsets);
+  if (!found) {
+    DbgPrint("%s\n", "First pattern not found!");
+    return STATUS_UNSUCCESSFUL;
+  }
+  found = PatternScanner::SearchWithMask(
+      buffer.data(), eastl::size(buffer), pattern_01.data(),
+      eastl::size(pattern_01), "xx", found_offsets);
+  if (!found) {
+    DbgPrint("%s\n", "Second pattern not found!");
+    return STATUS_UNSUCCESSFUL;
+  }
+  found = PatternScanner::SearchWithMask(
+      buffer.data(), eastl::size(buffer), pattern_02.data(),
+      eastl::size(pattern_02), "x", found_offsets);
+  if (!found) {
+    DbgPrint("%s\n", "Third pattern not found!");
+    return STATUS_UNSUCCESSFUL;
+  }
+  found = PatternScanner::SearchWithMask(buffer, pattern_03, "x????????x",
+                                         found_offsets);
+  if (!found) {
+    DbgPrint("%s\n", "Fourth pattern not found!");
+    return STATUS_UNSUCCESSFUL;
+  }
+  found = PatternScanner::SearchWithMask(
+      buffer, {0xDE, 0xAD, 0x00, 0x00, 0x00, 0x00, 0xCA, 0xFE}, "xx????xx",
+      found_offsets);
+  if (!found) {
+    DbgPrint("%s\n", "Fifth pattern not found!");
+    return STATUS_UNSUCCESSFUL;
+  }
+  for (const auto offset : found_offsets) {
+    DbgPrint("Offset: %zu\n", offset);
+  }
+
   DbgPrint("%s\n", "Starting thread..");
   auto args = eastl::make_shared<ThreadArgs>();
   thread.Start(
@@ -156,9 +203,24 @@ NTSTATUS DriverEntry(_In_ struct _DRIVER_OBJECT *DriverObject,
             DbgPrint("%s\n", page.toString().c_str());
           }
 
+          PatternScanner::ProcessModule scanner(pep, obj, {0x4D, 0x5A, 0x90},
+                                                "xxx");
+          eastl::vector<size_t> results;
+          auto found = scanner.Scan(L"Explorer.EXE", results);
+          if (!found)
+            DbgPrint("%s\n", "PatternScanner::ProcessModule was unsuccessful");
+          if (results.size() != 1)
+            DbgPrint(
+                "PatternScanner::ProcessModule was unsuccessful: %zu results\n",
+                results.size());
+          else
+            DbgPrint("PatternScanner::ProcessModule found address for 'MZ\\x90': %p\n", results[0]);
+
           ::CloseProcess(&pep, &obj);
         }
 
+        DbgPrint("%s\n", "Done.");
+
         return STATUS_SUCCESS;
       },
       args);
diff --git a/memory.cpp b/memory.cpp
index eb13135..af58279 100644
--- a/memory.cpp
+++ b/memory.cpp
@@ -1,4 +1,5 @@
 #include <EASTL/finally.h>
+#include <EASTL/shared_ptr.h>
 #include <eastl_compat.hpp>
 #include <except.h>
 
@@ -601,3 +602,73 @@ NTSTATUS WriteVirtualMemory(_In_ PEPROCESS pep, _In_ const UCHAR *sourceAddress,
 
   return status;
 }
+
+bool PatternScanner::SearchWithMask(const uint8_t *buffer, size_t buffer_size,
+                                    const uint8_t *pattern, size_t pattern_size,
+                                    const eastl::string_view &mask,
+                                    eastl::vector<size_t> &results) {
+  bool found_some = false;
+
+  if (mask.length() != pattern_size)
+    return false;
+
+  for (size_t i = 0; i < buffer_size - pattern_size + 1; ++i) {
+    bool match = true;
+    for (size_t j = 0; j < pattern_size; ++j) {
+      if (mask[j] == 'x' && buffer[i + j] != pattern[j]) {
+        match = false;
+        break;
+      }
+    }
+    if (match) {
+      found_some = true;
+      results.push_back(i);
+    }
+  }
+
+  return found_some;
+}
+
+bool PatternScanner::ProcessModule::Scan(const eastl::wstring_view &module_name,
+                                         eastl::vector<size_t> &results,
+                                         size_t max_results) {
+  if (eastl::size(m_pattern) > 64)
+    return false;
+
+  const auto mods = ::GetModules(m_pep, FALSE);
+  const auto mod = eastl::find_if(mods.begin(), mods.end(),
+                                  [&module_name](const Module &mod) {
+                                    return mod.BaseDllName == module_name;
+                                  });
+  if (mod == mods.end())
+    return false;
+
+  auto copy_buffer = eastl::make_shared<eastl::array<uint8_t, 4096 * 8>>();
+  const auto pages = ::GetPages(m_obj, m_max_pages);
+  for (const auto page : pages) {
+    if (page.BaseAddress < mod->DllBase ||
+        page.BaseAddress + page.RegionSize > mod->DllBase + mod->SizeOfImage)
+      continue;
+    SIZE_T offset = 0;
+    while (offset < page.RegionSize - eastl::size(m_pattern) + 1) {
+      SIZE_T in_out_size =
+          eastl::min(eastl::size(*copy_buffer), page.RegionSize - offset);
+      const auto ret =
+          ::ReadVirtualMemory(m_pep, page.BaseAddress + offset,
+                              eastl::begin(*copy_buffer), &in_out_size);
+      if (!NT_SUCCESS(ret) || in_out_size == 0)
+        break;
+
+      if (SearchWithMask(eastl::begin(*copy_buffer), in_out_size,
+                         m_pattern.begin(), eastl::size(m_pattern), m_mask,
+                         results)) {
+        if (results.size() >= max_results)
+          return false;
+      }
+
+      offset += in_out_size - eastl::size(m_pattern) + 1;
+    }
+  }
+
+  return true;
+}
diff --git a/memory.hpp b/memory.hpp
index c8db72c..0e86c00 100644
--- a/memory.hpp
+++ b/memory.hpp
@@ -1,6 +1,8 @@
 #ifndef MEMORY_H
 #define MEMORY_H 1
 
+#include <EASTL/array.h>
+#include <EASTL/initializer_list.h>
 #include <EASTL/string.h>
 #include <EASTL/vector.h>
 #include <cstdint>
@@ -175,4 +177,49 @@ private:
   NTSTATUS m_last_error;
   SIZE_T m_last_size;
 };
+
+namespace PatternScanner {
+bool SearchWithMask(const uint8_t *buffer, size_t buffer_size,
+                    const uint8_t *pattern, size_t pattern_size,
+                    const eastl::string_view &mask,
+                    eastl::vector<size_t> &results);
+
+template <size_t PM, size_t N>
+SearchWithMask(const eastl::array<uint8_t, N> &buffer,
+               const eastl::array<uint8_t, PM> &pattern,
+               const eastl::string_view &mask, eastl::vector<size_t> &results) {
+  return SearchWithMask(buffer.data(), eastl::size(buffer), pattern.data(),
+                        eastl::size(pattern), mask, results);
+}
+
+template <size_t N>
+SearchWithMask(const eastl::array<uint8_t, N> &buffer,
+               const std::initializer_list<uint8_t> &pattern,
+               const eastl::string_view &mask, eastl::vector<size_t> &results) {
+  return SearchWithMask(buffer.data(), eastl::size(buffer), pattern.begin(),
+                        eastl::size(pattern), mask, results);
+}
+
+class ProcessModule {
+public:
+  ProcessModule(_In_ PEPROCESS pep, _In_ HANDLE obj,
+                const std::initializer_list<uint8_t> &pattern,
+                const eastl::string_view &mask)
+      : m_max_pages(8192), m_pep(pep), m_obj(obj), m_pattern(pattern),
+        m_mask(mask), m_offset(0) {}
+  ProcessModule(const ProcessModule &) = delete;
+
+  void SetMaxPages(SIZE_T new_max_pages) { m_max_pages = new_max_pages; }
+  bool Scan(const eastl::wstring_view &module_name,
+            eastl::vector<size_t> &results, size_t max_results = 128);
+
+private:
+  SIZE_T m_max_pages;
+  PEPROCESS m_pep;
+  HANDLE m_obj;
+  const std::initializer_list<uint8_t> &m_pattern;
+  const eastl::string_view &m_mask;
+  size_t m_offset;
+};
+} // namespace PatternScanner
 #endif