From 2ff291b255b51f410cabcbfaab9adea0ff49a0df Mon Sep 17 00:00:00 2001
From: Orange <orange-cpp@yandex.ru>
Date: Thu, 9 Oct 2025 19:58:19 +0300
Subject: [PATCH] Supports broader architectures in PE scanner

Updates the PE scanner implementation to support both 32-bit and 64-bit architectures.
Leverages `std::variant` and a generic `ImageNtHeaders` to abstract architecture-specific details.
Simplifies the logic for retrieving section data, generalizing the process for improved maintainability.
---
 .idea/editor.xml                             |   4 +-
 include/omath/system/pe/file_header.hpp      | 111 ++++++++----------
 include/omath/system/pe/image_nt_headers.hpp |  23 ++++
 include/omath/system/pe/optional_header.hpp  | 115 ++++++++++++++++++-
 include/omath/system/pe/section_header.hpp   |  70 ++++-------
 include/omath/system/pe_headers_x64.hpp      |  32 ------
 include/omath/utility/pe_pattern_scan.hpp    |  18 ++-
 source/utility/pe_pattern_scan.cpp           |  93 ++++++++++-----
 tests/general/unit_test_pattern_scan.cpp     |   9 +-
 9 files changed, 288 insertions(+), 187 deletions(-)
 create mode 100644 include/omath/system/pe/image_nt_headers.hpp
 delete mode 100644 include/omath/system/pe_headers_x64.hpp

diff --git a/.idea/editor.xml b/.idea/editor.xml
index fde5348..373c50f 100644
--- a/.idea/editor.xml
+++ b/.idea/editor.xml
@@ -201,7 +201,7 @@
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppStaticDataMemberInUnnamedStruct/@EntryIndexedValue" value="WARNING" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppStaticSpecifierOnAnonymousNamespaceMember/@EntryIndexedValue" value="SUGGESTION" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppStringLiteralToCharPointerConversion/@EntryIndexedValue" value="WARNING" type="string" />
-    <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppTabsAreDisallowed/@EntryIndexedValue" value="WARNING" type="string" />
+    <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppTabsAreDisallowed/@EntryIndexedValue" value="DO_NOT_SHOW" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppTemplateArgumentsCanBeDeduced/@EntryIndexedValue" value="HINT" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppTemplateParameterNeverUsed/@EntryIndexedValue" value="HINT" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppTemplateParameterShadowing/@EntryIndexedValue" value="WARNING" type="string" />
@@ -215,7 +215,7 @@
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnmatchedPragmaEndRegionDirective/@EntryIndexedValue" value="WARNING" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnmatchedPragmaRegionDirective/@EntryIndexedValue" value="WARNING" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnnamedNamespaceInHeaderFile/@EntryIndexedValue" value="WARNING" type="string" />
-    <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnnecessaryWhitespace/@EntryIndexedValue" value="WARNING" type="string" />
+    <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnnecessaryWhitespace/@EntryIndexedValue" value="DO_NOT_SHOW" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnsignedZeroComparison/@EntryIndexedValue" value="WARNING" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnusedIncludeDirective/@EntryIndexedValue" value="WARNING" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUseAlgorithmWithCount/@EntryIndexedValue" value="SUGGESTION" type="string" />
diff --git a/include/omath/system/pe/file_header.hpp b/include/omath/system/pe/file_header.hpp
index a07f965..f120fa6 100644
--- a/include/omath/system/pe/file_header.hpp
+++ b/include/omath/system/pe/file_header.hpp
@@ -5,73 +5,52 @@
 #pragma once
 #include <cstdint>
 
-enum class MachineId : std::uint16_t
+namespace omath::system::pe
 {
-    UNKNOWN = 0x0000,
-    TARGET_HOST = 0x0001, // Useful for indicating we want to interact with the host and not a WoW guest.
-    I386 = 0x014C, // Intel 386.
-    R3000 = 0x0162, // MIPS little-endian, 0x160 big-endian
-    R4000 = 0x0166, // MIPS little-endian
-    R10000 = 0x0168, // MIPS little-endian
-    WCEMIPSV2 = 0x0169, // MIPS little-endian WCE v2
-    ALPHA = 0x0184, // Alpha_AXP
-    SH3 = 0x01A2, // SH3 little-endian
-    SH3DSP = 0x01A3,
-    SH3E = 0x01A4, // SH3E little-endian
-    SH4 = 0x01A6, // SH4 little-endian
-    SH5 = 0x01A8, // SH5
-    ARM = 0x01C0, // ARM Little-Endian
-    THUMB = 0x01C2, // ARM Thumb/Thumb-2 Little-Endian
-    ARMNT = 0x01C4, // ARM Thumb-2 Little-Endian
-    AM33 = 0x01D3,
-    POWERPC = 0x01F0, // IBM PowerPC Little-Endian
-    POWERPCP = 0x01F1,
-    IA64 = 0x0200, // Intel 64
-    MIPS16 = 0x0266, // MIPS
-    ALPHA64 = 0x0284, // ALPHA64
-    MIPSFPU = 0x0366, // MIPS
-    MIPSFPU16 = 0x0466, // MIPS
-    AXP64 = 0x0284,
-    TRICORE = 0x0520, // Infineon
-    CEF = 0x0CEF,
-    EBC = 0x0EBC, // EFI Byte Code
-    AMD64 = 0x8664, // AMD64 (K8)
-    M32R = 0x9041, // M32R little-endian
-    ARM64 = 0xAA64, // ARM64 Little-Endian
-    CEE = 0xC0EE,
-};
-
-union FileCharacteristics
-{
-    std::uint16_t flags;
-    struct
+    enum class MachineId : std::uint16_t
     {
-        std::uint16_t relocs_stripped : 1; // Relocation info stripped from file.
-        std::uint16_t executable : 1; // File is executable  (i.e. no unresolved external references).
-        std::uint16_t lines_stripped : 1; // Line nunbers stripped from file.
-        std::uint16_t local_symbols_stripped : 1; // Local symbols stripped from file.
-        std::uint16_t aggressive_ws_trim : 1; // Aggressively trim working set
-        std::uint16_t large_address_aware : 1; // App can handle >2gb addresses
-        std::uint16_t _pad0 : 1;
-        std::uint16_t bytes_reversed_lo : 1; // Bytes of machine word are reversed.
-        std::uint16_t machine_32 : 1; // 32 bit word machine.
-        std::uint16_t debug_stripped : 1; // Debugging info stripped from file in .DBG file
-        std::uint16_t runnable_from_swap : 1; // If Image is on removable media, copy and run from the swap file.
-        std::uint16_t net_run_from_swap : 1; // If Image is on Net, copy and run from the swap file.
-        std::uint16_t system_file : 1; // System File.
-        std::uint16_t dll_file : 1; // File is a DLL.
-        std::uint16_t up_system_only : 1; // File should only be run on a UP machine
-        std::uint16_t bytes_reversed_hi : 1; // Bytes of machine word are reversed.
+        UNKNOWN = 0x0000,
+        TARGET_HOST = 0x0001, // Useful for indicating we want to interact with the host and not a WoW guest.
+        I386 = 0x014C, // Intel 386.
+        R3000 = 0x0162, // MIPS little-endian, 0x160 big-endian
+        R4000 = 0x0166, // MIPS little-endian
+        R10000 = 0x0168, // MIPS little-endian
+        WCEMIPSV2 = 0x0169, // MIPS little-endian WCE v2
+        ALPHA = 0x0184, // Alpha_AXP
+        SH3 = 0x01A2, // SH3 little-endian
+        SH3DSP = 0x01A3,
+        SH3E = 0x01A4, // SH3E little-endian
+        SH4 = 0x01A6, // SH4 little-endian
+        SH5 = 0x01A8, // SH5
+        ARM = 0x01C0, // ARM Little-Endian
+        THUMB = 0x01C2, // ARM Thumb/Thumb-2 Little-Endian
+        ARMNT = 0x01C4, // ARM Thumb-2 Little-Endian
+        AM33 = 0x01D3,
+        POWERPC = 0x01F0, // IBM PowerPC Little-Endian
+        POWERPCP = 0x01F1,
+        IA64 = 0x0200, // Intel 64
+        MIPS16 = 0x0266, // MIPS
+        ALPHA64 = 0x0284, // ALPHA64
+        MIPSFPU = 0x0366, // MIPS
+        MIPSFPU16 = 0x0466, // MIPS
+        AXP64 = 0x0284,
+        TRICORE = 0x0520, // Infineon
+        CEF = 0x0CEF,
+        EBC = 0x0EBC, // EFI Byte Code
+        AMD64 = 0x8664, // AMD64 (K8)
+        M32R = 0x9041, // M32R little-endian
+        ARM64 = 0xAA64, // ARM64 Little-Endian
+        CEE = 0xC0EE,
     };
-};
 
-struct FileHeader
-{
-    MachineId machine;
-    uint16_t num_sections;
-    uint32_t timedate_stamp;
-    uint32_t ptr_symbols;
-    uint32_t num_symbols;
-    uint16_t size_optional_header;
-    FileCharacteristics characteristics;
-};
\ No newline at end of file
+    struct FileHeader
+    {
+        MachineId machine;
+        uint16_t num_sections;
+        uint32_t timedate_stamp;
+        uint32_t ptr_symbols;
+        uint32_t num_symbols;
+        uint16_t size_optional_header;
+        std::uint16_t characteristics;
+    };
+}
\ No newline at end of file
diff --git a/include/omath/system/pe/image_nt_headers.hpp b/include/omath/system/pe/image_nt_headers.hpp
new file mode 100644
index 0000000..35c6430
--- /dev/null
+++ b/include/omath/system/pe/image_nt_headers.hpp
@@ -0,0 +1,23 @@
+//
+// Created by Vlad on 10/9/2025.
+//
+
+#pragma once
+#include "file_header.hpp"
+#include "optional_header.hpp"
+
+namespace omath::system::pe
+{
+    enum class NtArchitecture
+    {
+        x32_bit,
+        x64_bit,
+    };
+    template<NtArchitecture architecture>
+    struct ImageNtHeaders
+    {
+        std::uint32_t signature;
+        FileHeader file_header;
+        OptionalHeader<architecture == NtArchitecture::x64_bit> optional_header;
+    };
+} // namespace omath::system::pe
\ No newline at end of file
diff --git a/include/omath/system/pe/optional_header.hpp b/include/omath/system/pe/optional_header.hpp
index 9a6f9f5..afe0b23 100644
--- a/include/omath/system/pe/optional_header.hpp
+++ b/include/omath/system/pe/optional_header.hpp
@@ -1,5 +1,118 @@
 //
 // Created by Vlad on 10/8/2025.
 //
-
 #pragma once
+#include <cstdint>
+
+namespace omath::system::pe
+{
+    static constexpr std::uint16_t opt_hdr32_magic = 0x010B;
+    static constexpr std::uint16_t opt_hdr64_magic = 0x020B;
+    enum class SubsystemId : uint16_t
+    {
+        unknown = 0x0000, // Unknown subsystem.
+        native = 0x0001, // Image doesn't require a subsystem.
+        windows_gui = 0x0002, // Image runs in the Windows GUI subsystem.
+        windows_cui = 0x0003, // Image runs in the Windows character subsystem
+        os2_cui = 0x0005, // image runs in the OS/2 character subsystem.
+        posix_cui = 0x0007, // image runs in the Posix character subsystem.
+        native_windows = 0x0008, // image is a native Win9x driver.
+        windows_ce_gui = 0x0009, // Image runs in the Windows CE subsystem.
+        efi_application = 0x000A, //
+        efi_boot_service_driver = 0x000B, //
+        efi_runtime_driver = 0x000C, //
+        efi_rom = 0x000D,
+        xbox = 0x000E,
+        windows_boot_application = 0x0010,
+        xbox_code_catalog = 0x0011,
+    };
+
+    struct DataDirectory
+    {
+        uint32_t rva;
+        uint32_t size;
+    };
+    struct OptionalHeaderX64
+    {
+        // Standard fields.
+        uint16_t magic;
+        std::uint16_t linker_version;
+
+        uint32_t size_code;
+        uint32_t size_init_data;
+        uint32_t size_uninit_data;
+
+        uint32_t entry_point;
+        uint32_t base_of_code;
+
+        // NT additional fields.
+        uint64_t image_base;
+        uint32_t section_alignment;
+        uint32_t file_alignment;
+
+        std::uint32_t os_version;
+        std::uint32_t img_version;
+        std::uint32_t subsystem_version;
+        uint32_t win32_version_value;
+
+        uint32_t size_image;
+        uint32_t size_headers;
+
+        uint32_t checksum;
+        SubsystemId subsystem;
+        SubsystemId characteristics;
+
+        uint64_t size_stack_reserve;
+        uint64_t size_stack_commit;
+        uint64_t size_heap_reserve;
+        uint64_t size_heap_commit;
+
+        uint32_t ldr_flags;
+
+        uint32_t num_data_directories;
+        DataDirectory data_directories[16];
+    };
+    struct OptionalHeaderX86
+    {
+        // Standard fields.
+        uint16_t magic;
+        uint16_t linker_version;
+
+        uint32_t size_code;
+        uint32_t size_init_data;
+        uint32_t size_uninit_data;
+
+        uint32_t entry_point;
+        uint32_t base_of_code;
+        uint32_t base_of_data;
+
+        // NT additional fields.
+        uint32_t image_base;
+        uint32_t section_alignment;
+        uint32_t file_alignment;
+
+        std::uint32_t os_version;
+        std::uint32_t img_version;
+        std::uint32_t subsystem_version;
+        uint32_t win32_version_value;
+
+        uint32_t size_image;
+        uint32_t size_headers;
+
+        uint32_t checksum;
+        SubsystemId subsystem;
+        std::uint16_t characteristics;
+
+        uint32_t size_stack_reserve;
+        uint32_t size_stack_commit;
+        uint32_t size_heap_reserve;
+        uint32_t size_heap_commit;
+
+        uint32_t ldr_flags;
+
+        uint32_t num_data_directories;
+        DataDirectory data_directories[16];
+    };
+    template<bool x64 = true>
+    using OptionalHeader = std::conditional_t<x64, OptionalHeaderX64, OptionalHeaderX86>;
+} // namespace omath::system::pe
\ No newline at end of file
diff --git a/include/omath/system/pe/section_header.hpp b/include/omath/system/pe/section_header.hpp
index 4147932..868eb24 100644
--- a/include/omath/system/pe/section_header.hpp
+++ b/include/omath/system/pe/section_header.hpp
@@ -4,59 +4,29 @@
 
 #pragma once
 #include <cstdint>
-union section_characteristics_t
-{
-    std::uint32_t flags;
-    struct
-    {
-        std::uint32_t _pad0 : 5;
-        std::uint32_t cnt_code : 1; // Section contains code.
-        std::uint32_t cnt_init_data : 1; // Section contains initialized data.
-        std::uint32_t cnt_uninit_data : 1; // Section contains uninitialized data.
-        std::uint32_t _pad1 : 1;
-        std::uint32_t lnk_info : 1; // Section contains comments or some other type of information.
-        std::uint32_t _pad2 : 1;
-        std::uint32_t lnk_remove : 1; // Section contents will not become part of image.
-        std::uint32_t lnk_comdat : 1; // Section contents comdat.
-        std::uint32_t _pad3 : 1;
-        std::uint32_t no_defer_spec_exc : 1; // Reset speculative exceptions handling bits in the TLB entries for this
-                                             // section.
-        std::uint32_t mem_far : 1;
-        std::uint32_t _pad4 : 1;
-        std::uint32_t mem_purgeable : 1;
-        std::uint32_t mem_locked : 1;
-        std::uint32_t mem_preload : 1;
-        std::uint32_t alignment : 4; // Alignment calculated as: n ? 1 << ( n - 1 ) : 16
-        std::uint32_t lnk_nreloc_ovfl : 1; // Section contains extended relocations.
-        std::uint32_t mem_discardable : 1; // Section can be discarded.
-        std::uint32_t mem_not_cached : 1; // Section is not cachable.
-        std::uint32_t mem_not_paged : 1; // Section is not pageable.
-        std::uint32_t mem_shared : 1; // Section is shareable.
-        std::uint32_t mem_execute : 1; // Section is executable.
-        std::uint32_t mem_read : 1; // Section is readable.
-        std::uint32_t mem_write : 1; // Section is writeable.
-    };
-};
 
 // Section header
 //
-struct section_header_t
+namespace omath::system::pe
 {
-    char name[8];
-    union
+    struct SectionHeader
     {
-        std::uint32_t physical_address;
-        std::uint32_t virtual_size;
+        char name[8];
+        union
+        {
+            std::uint32_t physical_address;
+            std::uint32_t virtual_size;
+        };
+        std::uint32_t virtual_address;
+
+        std::uint32_t size_raw_data;
+        std::uint32_t ptr_raw_data;
+
+        std::uint32_t ptr_relocs;
+        std::uint32_t ptr_line_numbers;
+        std::uint32_t  num_relocs;
+        std::uint32_t  num_line_numbers;
+
+        std::uint32_t characteristics;
     };
-    std::uint32_t virtual_address;
-
-    std::uint32_t size_raw_data;
-    std::uint32_t ptr_raw_data;
-
-    std::uint32_t ptr_relocs;
-    std::uint32_t ptr_line_numbers;
-    uint16_t num_relocs;
-    uint16_t num_line_numbers;
-
-    section_characteristics_t characteristics;
-};
\ No newline at end of file
+}
\ No newline at end of file
diff --git a/include/omath/system/pe_headers_x64.hpp b/include/omath/system/pe_headers_x64.hpp
deleted file mode 100644
index 2656ab2..0000000
--- a/include/omath/system/pe_headers_x64.hpp
+++ /dev/null
@@ -1,32 +0,0 @@
-//
-// Created by Vlad on 10/6/2025.
-//
-
-#pragma once
-#include <cstdint>
-namespace omath::system
-{
-    struct ImageDosHeader
-    {
-        uint16_t e_magic; // Magic number
-        uint16_t e_cblp; // Bytes on last page of file
-        uint16_t e_cp; // Pages in file
-        uint16_t e_crlc; // Relocations
-        uint16_t e_cparhdr; // Size of header in paragraphs
-        uint16_t e_minalloc; // Minimum extra paragraphs needed
-        uint16_t e_maxalloc; // Maximum extra paragraphs needed
-        uint16_t e_ss; // Initial (relative) SS value
-        uint16_t e_sp; // Initial SP value
-        uint16_t e_csum; // Checksum
-        uint16_t e_ip; // Initial IP value
-        uint16_t e_cs; // Initial (relative) CS value
-        uint16_t e_lfarlc; // File address of relocation table
-        uint16_t e_ovno; // Overlay number
-        uint16_t e_res[4]; // Reserved words
-        uint16_t e_oemid; // OEM identifier (for e_oeminfo)
-        uint16_t e_oeminfo; // OEM information; e_oemid specific
-        uint16_t e_res2[10]; // Reserved words
-        int32_t e_lfanew; // File address of new exe header
-    };
-
-} // namespace omath::system
\ No newline at end of file
diff --git a/include/omath/utility/pe_pattern_scan.hpp b/include/omath/utility/pe_pattern_scan.hpp
index bb83bb6..ef06739 100644
--- a/include/omath/utility/pe_pattern_scan.hpp
+++ b/include/omath/utility/pe_pattern_scan.hpp
@@ -8,22 +8,34 @@
 #include <optional>
 #include <string_view>
 #include <vector>
-
 namespace omath
 {
+    struct PeSectionScanResult
+    {
+        std::uint64_t virtual_base_addr;
+        std::uint64_t raw_base_addr;
+        std::ptrdiff_t target_offset;
+    };
     class PePatternScanner final
     {
+    private:
+        struct Section
+        {
+            std::uint64_t virtual_base_addr;
+            std::uint64_t raw_base_addr;
+            std::vector<std::byte> data;
+        };
     public:
         [[nodiscard]]
         static std::optional<std::uintptr_t> scan_for_pattern_in_loaded_module(const std::string_view& module_name,
                                                                                const std::string_view& pattern);
 
         [[nodiscard]]
-        static std::optional<std::uintptr_t> scan_for_pattern_in_file(const std::filesystem::path& path_to_file,
+        static std::optional<PeSectionScanResult> scan_for_pattern_in_file(const std::filesystem::path& path_to_file,
                                                                       const std::string_view& pattern);
 
         [[nodiscard]]
-        static std::optional<std::vector<std::byte>>
+        static std::optional<Section>
         extract_section_from_pe_file(const std::filesystem::path& path_to_file, const std::string_view& section_name);
     };
 } // namespace omath
\ No newline at end of file
diff --git a/source/utility/pe_pattern_scan.cpp b/source/utility/pe_pattern_scan.cpp
index 7191418..a78dc26 100644
--- a/source/utility/pe_pattern_scan.cpp
+++ b/source/utility/pe_pattern_scan.cpp
@@ -2,12 +2,14 @@
 // Created by Vlad on 10/7/2025.
 //
 #include "omath/utility/pe_pattern_scan.hpp"
+#include "omath/system/pe/image_nt_headers.hpp"
 #include "omath/system/pe/section_header.hpp"
 #include "omath/utility/pattern_scan.hpp"
 #include <fstream>
 #include <omath/system/pe/dos_header.hpp>
 #include <span>
 #include <stdexcept>
+#include <variant>
 #ifdef _WIN32
 #include <Windows.h>
 #endif
@@ -44,32 +46,38 @@ namespace omath
         throw std::runtime_error("Pattern scan for loaded modules is only for windows platform");
 #endif
     }
-    std::optional<std::uintptr_t> PePatternScanner::scan_for_pattern_in_file(const std::filesystem::path& path_to_file,
-                                                                             const std::string_view& pattern)
+    std::optional<PeSectionScanResult>
+    PePatternScanner::scan_for_pattern_in_file(const std::filesystem::path& path_to_file,
+                                               const std::string_view& pattern)
     {
         const auto pe_section = extract_section_from_pe_file(path_to_file, ".text");
 
         if (!pe_section.has_value())
             return std::nullopt;
 
-        const auto scan_result = PatternScanner::scan_for_pattern(pe_section->cbegin(), pe_section->cend(), pattern);
+        const auto scan_result =
+                PatternScanner::scan_for_pattern(pe_section->data.cbegin(), pe_section->data.cend(), pattern);
 
-        if (scan_result == pe_section->cend())
+        if (scan_result == pe_section->data.cend())
             return std::nullopt;
+        const auto offset = std::distance(pe_section->data.begin(), scan_result);
 
-        return std::distance(pe_section->begin(), pe_section->end());
+        return PeSectionScanResult{.virtual_base_addr = pe_section->virtual_base_addr,
+                                   .raw_base_addr = pe_section->raw_base_addr,
+                                   .target_offset = offset};
     }
-    std::optional<std::vector<std::byte>>
+    std::optional<PePatternScanner::Section>
     PePatternScanner::extract_section_from_pe_file([[maybe_unused]] const std::filesystem::path& path_to_file,
                                                    [[maybe_unused]] const std::string_view& section_name)
     {
 #ifdef _WIN32
+        using namespace system::pe;
         std::fstream file(path_to_file, std::ios::binary | std::ios::in);
 
         if (!file.is_open()) [[unlikely]]
             return std::nullopt;
 
-        system::pe::DosHeader dos_header{};
+        DosHeader dos_header{};
         file.read(reinterpret_cast<char*>(&dos_header), sizeof(dos_header));
 
         if (dos_header.e_magic != 0x5A4D) [[unlikely]]
@@ -77,34 +85,57 @@ namespace omath
 
         file.seekg(dos_header.e_lfanew, std::ios::beg);
 
-        IMAGE_NT_HEADERS32 nt_headers{};
-        file.read(reinterpret_cast<char*>(&nt_headers), sizeof(nt_headers));
+        std::variant<ImageNtHeaders<NtArchitecture::x64_bit>, ImageNtHeaders<NtArchitecture::x32_bit>> nt_headers;
 
-        if (nt_headers.Signature != 0x00004550) [[unlikely]]
-            return std::nullopt;
-
-        constexpr size_t size_of_signature = 4;
-        const auto offset_to_segment_table = dos_header.e_lfanew + nt_headers.FileHeader.SizeOfOptionalHeader
-                                             + sizeof(IMAGE_FILE_HEADER) + size_of_signature;
-
-        file.seekg(static_cast<std::fstream::off_type>(offset_to_segment_table), std::ios::beg);
-
-        for (std::size_t i = 0; i < nt_headers.FileHeader.NumberOfSections; i++)
         {
-            section_header_t current_section{};
-            file.read(reinterpret_cast<char*>(&current_section), sizeof(IMAGE_SECTION_HEADER));
+            ImageNtHeaders<NtArchitecture::x32_bit> x86_headers;
+            file.seekg(dos_header.e_lfanew, std::ios::beg);
+            file.read(reinterpret_cast<char*>(&x86_headers), sizeof(x86_headers));
 
-            if (std::string_view(current_section.name) != section_name)
-                continue;
-
-            std::vector<std::byte> section_data(current_section.size_raw_data);
-
-            file.seekg(current_section.ptr_raw_data, std::ios::beg);
-            file.read(reinterpret_cast<char*>(section_data.data()), static_cast<std::streamsize>(section_data.size()));
-
-            return section_data;
+            if (x86_headers.optional_header.magic == opt_hdr32_magic)
+                nt_headers = x86_headers;
+            else
+            {
+                ImageNtHeaders<NtArchitecture::x64_bit> x64_headers;
+                file.seekg(dos_header.e_lfanew, std::ios::beg);
+                file.read(reinterpret_cast<char*>(&x64_headers), sizeof(x64_headers));
+                nt_headers = x64_headers;
+            }
         }
-        return std::nullopt;
+        return std::visit(
+                [&file, &dos_header, &section_name](auto& concrete_headers) -> std::optional<Section>
+                {
+                    if (concrete_headers.signature != 0x00004550) [[unlikely]]
+                        return std::nullopt;
+
+                    constexpr std::size_t size_of_signature = 4;
+                    const auto offset_to_segment_table = dos_header.e_lfanew
+                                                         + concrete_headers.file_header.size_optional_header
+                                                         + sizeof(FileHeader) + size_of_signature;
+
+                    file.seekg(static_cast<std::fstream::off_type>(offset_to_segment_table), std::ios::beg);
+
+                    for (std::size_t i = 0; i < concrete_headers.file_header.num_sections; i++)
+                    {
+                        SectionHeader current_section{};
+                        file.read(reinterpret_cast<char*>(&current_section), sizeof(current_section));
+
+                        if (std::string_view(current_section.name) != section_name)
+                            continue;
+
+                        std::vector<std::byte> section_data(current_section.size_raw_data);
+
+                        file.seekg(current_section.ptr_raw_data, std::ios::beg);
+                        file.read(reinterpret_cast<char*>(section_data.data()),
+                                  static_cast<std::streamsize>(section_data.size()));
+                        return Section{.virtual_base_addr = current_section.virtual_address
+                                                            + concrete_headers.optional_header.image_base,
+                                       .raw_base_addr = current_section.ptr_raw_data,
+                                       .data = std::move(section_data)};
+                    }
+                    return std::nullopt;
+                },
+                nt_headers);
 #else
         throw std::runtime_error("Pattern scan for loaded modules is only for windows platform");
 #endif
diff --git a/tests/general/unit_test_pattern_scan.cpp b/tests/general/unit_test_pattern_scan.cpp
index 59c873a..7f7b6fd 100644
--- a/tests/general/unit_test_pattern_scan.cpp
+++ b/tests/general/unit_test_pattern_scan.cpp
@@ -5,7 +5,7 @@
 #include "gtest/gtest.h"
 #include <omath/utility/pattern_scan.hpp>
 #include <source_location>
-
+#include <print>
 TEST(unit_test_pattern_scan, read_test)
 {
     const auto result = omath::PatternScanner::parse_pattern("FF ? ?? E9");
@@ -48,7 +48,12 @@ TEST(unit_test_pattern_scan, corner_case_3)
 
 TEST(unit_test_pattern_scan, corner_case_4)
 {
-    const auto result = omath::PatternScanner::parse_pattern("XZ");
+    const auto result = omath::PatternScanner::parse_pattern("X ? ?? E9 ");
+    const auto result2 = omath::PePatternScanner::scan_for_pattern_in_file(
+            std::filesystem::path{
+                    R"(C:\Users\Vlad\CLionProjects\aether\out\Release\aether.dll)"},
+            "48 89 5C 24 ? 57 48 83 EC ? 8B DA 48 8B F9 FF 15 ? ? ? ? 83 FB ? 75 ? B9 ? ? ? ? E8 ? ? ? ? 33 DB 48 85 C0 74 ? 48 8D 0D ? ? ? ? 48 89 38 48 89 48 ? EB");
 
+    std::println("{:x}", result2->virtual_base_addr + result2->target_offset);
     EXPECT_FALSE(result.has_value());
 }
\ No newline at end of file