Untitled
unknown
c_cpp
a year ago
7.4 kB
6
Indexable
Never
#include <cstdint> #include <memory> #include <pe/pe_parser.h> #define TRY(expr) \ ({ \ const auto TRY_expr = (expr); \ if (!TRY_expr.has_value()) { \ return std::unexpected(TRY_expr.error()); \ } \ TRY_expr.value(); \ }) #define TRY_OPT(expr, error) \ ({ \ const auto TRY_OPT_expr = (expr); \ if (!TRY_OPT_expr.has_value()) { \ return std::unexpected(error); \ } \ TRY_OPT_expr.value(); \ }) std::expected<void, ParserError> PEParser::parse_file_headers() { // Read the DOS header and verify that it's correct. m_fs.read(reinterpret_cast<char *>(&m_parsed->m_dos), sizeof(DosHeader)); if (m_parsed->m_dos.m_magic != kDosSignature) { return std::unexpected(ParserError::kInvalidDosSignature); } // Read the NT header. m_fs.seekg(m_parsed->m_dos.m_lfanew); m_fs.read(reinterpret_cast<char *>(&m_parsed->m_nt), sizeof(NtHeader)); if (m_parsed->m_nt.m_signature != kNtSignature) { return std::unexpected(ParserError::kInvalidNtSignature); } // Read the optional NT header (it's not optional at all, thanks Microsoft.) const auto old_pos{m_fs.tellg()}; BaseOptional base_opt{}; m_fs.read(reinterpret_cast<char *>(&base_opt), sizeof(BaseOptional)); m_fs.seekg(old_pos); // Verify that it's either a 32- or 64-bit PE file. if (base_opt.m_magic == kOptionalHeaderMagic32) { auto &new_opt{m_parsed->m_optional.emplace<Optional32>()}; m_fs.read(reinterpret_cast<char *>(&new_opt), sizeof(Optional32)); } else if (base_opt.m_magic == kOptionalHeaderMagic64) { auto &new_opt{m_parsed->m_optional.emplace<Optional64>()}; m_fs.read(reinterpret_cast<char *>(&new_opt), sizeof(Optional64)); } else { return std::unexpected(ParserError::kInvalidOptionalMagic); } // Read section headers... m_fs.seekg(old_pos + static_cast<ptrdiff_t>(m_parsed->m_nt.m_size_of_optional_header)); m_parsed->m_sections.resize(m_parsed->m_nt.m_number_of_sections); for (auto § : m_parsed->m_sections) { m_fs.read(reinterpret_cast<char *>(§.m_header), sizeof(SectionHeader)); const auto data_size{sect.m_header.m_size_of_raw_data}; if (data_size == 0) { continue; } // ... and their data too :) const auto old_pos{m_fs.tellg()}; sect.m_data.resize(data_size); m_fs.seekg(sect.m_header.m_ptr_to_raw_data); m_fs.read(reinterpret_cast<char *>(sect.m_data.data()), data_size); m_fs.seekg(old_pos); } return {}; } std::expected<void, ParserError> PEParser::parse_relocations() { // Get the base relocation directory, need to convert from RVA to file offset. const auto &reloc_dir{m_parsed->get_data_directory(kDirBaseReloc)}; m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(reloc_dir.m_rva), ParserError::kInvalidRva)); while (true) { // Keep reading relocation blocks until we hit one that's invalid. There is // always an invalid one at the end :^) BaseRelocation reloc{}; m_fs.read(reinterpret_cast<char *>(&reloc), sizeof(BaseRelocation)); if (reloc.m_virtual_address == 0 || reloc.m_size_of_block < sizeof(BaseRelocation)) { break; } // Read all base relocations in this relocation block. auto &reloc_block{m_parsed->m_relocs.emplace_back()}; reloc_block.m_rva = reloc.m_virtual_address; reloc_block.m_relocs.resize((reloc.m_size_of_block - sizeof(BaseRelocation)) / sizeof(uint16_t)); for (auto &base_reloc : reloc_block.m_relocs) { m_fs.read(reinterpret_cast<char *>(&base_reloc), sizeof(uint16_t)); } } return {}; } std::expected<void, ParserError> PEParser::parse_import_descriptors() { // Get the import directory, need to convert from RVA to file offset. const auto &import_dir{m_parsed->get_data_directory(kDirImport)}; m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(import_dir.m_rva), ParserError::kInvalidRva)); while (true) { ImportDescriptor import_desc{}; m_fs.read(reinterpret_cast<char *>(&import_desc), sizeof(ImportDescriptor)); const auto next_pos{m_fs.tellg()}; if (import_desc.m_name == 0) { break; } ParsedImportThunk thunk{}; // Read the import name, need to convert from RVA to file offset. This is null terminated, // so we need to use a stack buffer first. std::array<char, kMaxPathSize> module_name{}; m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(import_desc.m_name), ParserError::kInvalidRva)); m_fs.get(module_name.data(), module_name.size() - 1, '\0'); thunk.m_module.assign_range(module_name); // Read all import thunks and the imports descriptors they point to. m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(import_desc.m_first_thunk), ParserError::kInvalidRva)); if (m_parsed->get_optional_magic() == kOptionalHeaderMagic32) { TRY(fill_import_thunk32(thunk)); } else if (m_parsed->get_optional_magic() == kOptionalHeaderMagic64) { TRY(fill_import_thunk64(thunk)); } m_parsed->m_imports.emplace_back(thunk); m_fs.seekg(next_pos); } return {}; } std::expected<void, ParserError> PEParser::fill_import_thunk32(ParsedImportThunk &thunk) { while (true) { ThunkData32 thunk_data{}; m_fs.read(reinterpret_cast<char *>(&thunk_data), sizeof(ThunkData32)); const auto next_pos{m_fs.tellg()}; if (thunk_data.m_address_of_data == 0) { break; } ImportByName import_desc{}; m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(thunk_data.m_address_of_data), ParserError::kInvalidRva)); m_fs.read(reinterpret_cast<char *>(&import_desc), sizeof(uint16_t)); m_fs.get(import_desc.m_name.data(), import_desc.m_name.size() - 1, '\0'); ParsedImportByNameDescriptor parsed_desc{}; parsed_desc.m_name.assign_range(import_desc.m_name); parsed_desc.m_offset = next_pos - static_cast<std::streampos>(sizeof(ThunkData32)); thunk.m_imports.emplace_back(parsed_desc); m_fs.seekg(next_pos); } return {}; } std::expected<void, ParserError> PEParser::fill_import_thunk64(ParsedImportThunk &thunk) { while (true) { ThunkData64 thunk_data{}; m_fs.read(reinterpret_cast<char *>(&thunk_data), sizeof(ThunkData64)); const auto next_pos{m_fs.tellg()}; if (thunk_data.m_address_of_data == 0) { break; } ImportByName import_desc{}; m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(thunk_data.m_address_of_data), ParserError::kInvalidRva)); m_fs.read(reinterpret_cast<char *>(&import_desc), sizeof(uint16_t)); m_fs.get(import_desc.m_name.data(), import_desc.m_name.size() - 1, '\0'); ParsedImportByNameDescriptor parsed_desc{}; parsed_desc.m_name.assign_range(import_desc.m_name); parsed_desc.m_offset = next_pos - static_cast<std::streampos>(sizeof(ThunkData64)); thunk.m_imports.emplace_back(parsed_desc); m_fs.seekg(next_pos); } return {}; } std::expected<SharedPEImage, ParserError> PEParser::parse(const std::filesystem::path &file_path) { m_fs.open(file_path, std::ios::binary); if (!m_fs.is_open()) { return std::unexpected(ParserError::kInvalidFile); } m_parsed = std::make_shared<PEImage>(); TRY(parse_file_headers()); TRY(parse_import_descriptors()); TRY(parse_relocations()); return m_parsed; }