Merge pull request #86 from orange-cpp/feature/improved_pe_scanner_interface

Feature/improved pe scanner interface
This commit is contained in:
2025-10-12 20:20:28 +03:00
committed by GitHub
2 changed files with 70 additions and 68 deletions

View File

@@ -7,7 +7,6 @@
#include <filesystem>
#include <optional>
#include <string_view>
#include <vector>
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<std::byte> data;
};
public:
[[nodiscard]]
static std::optional<std::uintptr_t> scan_for_pattern_in_loaded_module(const std::string_view& module_name,
@@ -35,9 +26,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

View File

@@ -11,6 +11,9 @@
#include <Windows.h>
#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<NtArchitecture::x64_bit>, ImageNtHeaders<NtArchitecture::x32_bit>>;
} // namespace
// Internal PE scanner functions
namespace
{
[[nodiscard]]
std::optional<NtHeaderVariant> 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<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 +337,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