diff --git a/include/omath/utility/pe_pattern_scan.hpp b/include/omath/utility/pe_pattern_scan.hpp index 8d582744..a5de4205 100644 --- a/include/omath/utility/pe_pattern_scan.hpp +++ b/include/omath/utility/pe_pattern_scan.hpp @@ -7,7 +7,6 @@ #include #include #include -#include namespace omath { struct PeSectionScanResult @@ -18,14 +17,6 @@ namespace omath }; class PePatternScanner final { - private: - struct Section - { - std::uint64_t virtual_base_addr; - std::uint64_t raw_base_addr; - std::vector data; - }; - public: [[nodiscard]] static std::optional scan_for_pattern_in_loaded_module(const std::string_view& module_name, @@ -35,9 +26,5 @@ namespace omath static std::optional scan_for_pattern_in_file(const std::filesystem::path& path_to_file, const std::string_view& pattern, const std::string_view& target_section_name = ".text"); - - [[nodiscard]] - static std::optional
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 f7f8d024..90c47e5a 100644 --- a/source/utility/pe_pattern_scan.cpp +++ b/source/utility/pe_pattern_scan.cpp @@ -11,6 +11,9 @@ #include #endif +// Internal PE shit defines +// Big thx for linuxpe sources as ref +// Link: https://github.com/can1357/linux-pe namespace { constexpr std::uint16_t opt_hdr32_magic = 0x010B; @@ -178,7 +181,11 @@ namespace using NtHeaderVariant = std::variant, ImageNtHeaders>; +} // namespace +// Internal PE scanner functions +namespace +{ [[nodiscard]] std::optional get_nt_header_from_file(std::fstream& file, const DosHeader& dos_header) { @@ -211,6 +218,69 @@ namespace constexpr std::uint32_t nt_hdr_magic = 0x4550; return std::visit([](const auto& header) -> bool { return header.signature != nt_hdr_magic; }, variant); } + + struct ExtractedSection + { + std::uint64_t virtual_base_addr; + std::uint64_t raw_base_addr; + std::vector data; + }; + + [[nodiscard]] + std::optional extract_section_from_pe_file(const std::filesystem::path& path_to_file, + const std::string_view& section_name) + { + std::fstream file(path_to_file, std::ios::binary | std::ios::in); + + if (!file.is_open()) [[unlikely]] + return std::nullopt; + + DosHeader dos_header{}; + file.read(reinterpret_cast(&dos_header), sizeof(dos_header)); + + if (invalid_dos_header_file(dos_header)) [[unlikely]] + return std::nullopt; + + const auto nt_headers = get_nt_header_from_file(file, dos_header); + + if (!nt_headers) + return std::nullopt; + + if (invalid_nt_header_file(nt_headers.value())) [[unlikely]] + return std::nullopt; + + return std::visit( + [&file, &dos_header, §ion_name](auto& concrete_headers) -> std::optional + { + constexpr std::size_t size_of_signature = sizeof(concrete_headers.signature); + 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(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(¤t_section), sizeof(current_section)); + + if (std::string_view(current_section.name) != section_name) + continue; + + std::vector section_data(current_section.size_raw_data); + + file.seekg(current_section.ptr_raw_data, std::ios::beg); + file.read(reinterpret_cast(section_data.data()), + static_cast(section_data.size())); + return ExtractedSection{.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.value()); + } } // namespace namespace omath @@ -267,59 +337,4 @@ namespace omath .raw_base_addr = pe_section->raw_base_addr, .target_offset = offset}; } - std::optional - PePatternScanner::extract_section_from_pe_file(const std::filesystem::path& path_to_file, - const std::string_view& section_name) - { - std::fstream file(path_to_file, std::ios::binary | std::ios::in); - - if (!file.is_open()) [[unlikely]] - return std::nullopt; - - DosHeader dos_header{}; - file.read(reinterpret_cast(&dos_header), sizeof(dos_header)); - - if (invalid_dos_header_file(dos_header)) [[unlikely]] - return std::nullopt; - - const auto nt_headers = get_nt_header_from_file(file, dos_header); - - if (!nt_headers) - return std::nullopt; - - if (invalid_nt_header_file(nt_headers.value())) [[unlikely]] - return std::nullopt; - - return std::visit( - [&file, &dos_header, §ion_name](auto& concrete_headers) -> std::optional
- { - constexpr std::size_t size_of_signature = sizeof(concrete_headers.signature); - 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(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(¤t_section), sizeof(current_section)); - - if (std::string_view(current_section.name) != section_name) - continue; - - std::vector section_data(current_section.size_raw_data); - - file.seekg(current_section.ptr_raw_data, std::ios::beg); - file.read(reinterpret_cast(section_data.data()), - static_cast(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.value()); - } } // namespace omath \ No newline at end of file