admin/omath
1
0
Fork 0
mirror of https://github.com/orange-cpp/omath.git synced 2026-02-12 22:53:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactors PE scanner to eliminate redundant section extraction

Browse Source 
Simplifies the PE pattern scanner by removing the redundant `extract_section_from_pe_file` function.

The extracted data is now returned directly from a new `ExtractedSection` struct within the main scanning function, streamlining the code and reducing duplication. This improves readability and maintainability of the PE scanner.
This commit is contained in:
Orange
2025-10-12 19:53:36 +03:00
parent 5c30f57c4c
commit cd0a864e7c
2 changed files with 63 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
include/omath/utility/pe_pattern_scan.hpp
View File

@@ -7,7 +7,6 @@
#include <filesystem>
#include <optional>
#include <string_view>
#include <vector>
namespace omath
{
struct PeSectionScanResult
@@ -18,13 +17,6 @@ namespace omath
};
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]]
@@ -35,9 +27,5 @@ namespace omath
static std::optional<PeSectionScanResult>
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<Section> extract_section_from_pe_file(const std::filesystem::path& path_to_file,
const std::string_view& section_name);
};
} // namespace omath

118
source/utility/pe_pattern_scan.cpp
View File

@@ -211,6 +211,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<std::byte> data;
};
[[nodiscard]]
std::optional<ExtractedSection> 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<char*>(&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, &section_name](auto& concrete_headers) -> std::optional<ExtractedSection>
{
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<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 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 +330,4 @@ namespace omath
.raw_base_addr = pe_section->raw_base_addr,
.target_offset = offset};
}
std::optional<PePatternScanner::Section>
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<char*>(&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, &section_name](auto& concrete_headers) -> std::optional<Section>
{
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<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.value());
}
} // namespace omath