From 0f75801f8b642c17b56b70fe5c8bbf6dc2562bfe Mon Sep 17 00:00:00 2001 From: "381848900@qq.com" Date: Wed, 18 Dec 2024 08:32:13 +0800 Subject: [PATCH] first commit --- .gitignore | 7 + Cargo.lock | 112 +++++++++++ Cargo.toml | 14 ++ README.md | 2 + src/error.rs | 36 ++++ src/header.rs | 453 ++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 16 ++ src/pe.rs | 540 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/types.rs | 18 ++ 9 files changed, 1198 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 README.md create mode 100644 src/error.rs create mode 100644 src/header.rs create mode 100644 src/lib.rs create mode 100644 src/pe.rs create mode 100644 src/types.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b21bd68 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +# Generated by Cargo +# will have compiled files and executables +/target/ + +# Generated by Tauri +# will have schema files for capabilities auto-completion +/gen/schemas diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..cc30ac0 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,112 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "bitflags" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" +dependencies = [ + "serde", +] + +[[package]] +name = "libc" +version = "0.2.168" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aaeb2981e0606ca11d79718f8bb01164f1d6ed75080182d3abf017e6d244b6d" + +[[package]] +name = "memmap2" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f" +dependencies = [ + "libc", +] + +[[package]] +name = "pe_parse" +version = "0.1.0" +dependencies = [ + "bitflags", + "memmap2", + "serde", + "thiserror", +] + +[[package]] +name = "proc-macro2" +version = "1.0.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "serde" +version = "1.0.216" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b9781016e935a97e8beecf0c933758c97a5520d32930e460142b4cd80c6338e" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.216" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46f859dbbf73865c6627ed570e78961cd3ac92407a2d117204c49232485da55e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "syn" +version = "2.0.90" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93605438cbd668185516ab499d589afb7ee1859ea3d5fc8f6b0755e1c7443767" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d8749b4531af2117677a5fcd12b1348a3fe2b81e36e61ffeac5c4aa3273e36" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-ident" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..d7af2f4 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "pe_parse" +version = "0.1.0" +edition = "2021" + +[dependencies] +bitflags = {version = "2.6.0", features = ["serde"] } +memmap2 = "0.9.5" +serde = {version = "1", features = ["derive"] } +thiserror = "2.0.4" + +[features] +default = [] # 默认不启用任何特性 +memmap2_impl = [] # 定义一个 feature，并启用第三方库 \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..a2cd360 --- /dev/null +++ b/README.md @@ -0,0 +1,2 @@ +# `pe_parse` +用来解析PE \ No newline at end of file diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..0cd3319 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,36 @@ +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum PEParseError { + #[error("无效的DOSMagic值")] + InvalidDOSMagic, + #[error("Invalid NT header signature")] + InvalidNTSignature, + #[error("Invalid optional header magic number")] + InvalidOptionalMagic, + #[error("Invalid optional header size")] + InvalidOptionalSize, + #[error("解析超出了文件范围")] + OutOfBounds, + #[error("错误的RVA: {0}")] + RvaToFoaError(u32), + #[error("RVA对应的是一个空节区数据！RVA: {0}")] + RvaToEmptySection(u32), + #[error("错误的数据目录索引")] + InvalidDataDirectoryIndex, +} +/// PE操作的错误 +#[derive(Error, Debug)] +pub enum MutablePEError { + // 兼容PE解析错误 + #[error(transparent)] + PEParseError(#[from] PEParseError), + // 头部映射无法扩容，超出最大可扩容范围 + #[error("头部映射无法扩容, 需要扩容大小:{0:#X},超过了最大可扩容范围:{1:#X}")] + CannotExpandHeader(u32, u32), + + // 文件大小不足，无法扩容 + #[error("文件大小不足，无法扩容! 需要的大小:{0:#X}, 当前的文件大小:{1:#X}")] + CannotExpandFileSize(u32, u32), + +} diff --git a/src/header.rs b/src/header.rs new file mode 100644 index 0000000..6a945ee --- /dev/null +++ b/src/header.rs @@ -0,0 +1,453 @@ +use super::types::*; +use bitflags::bitflags; +use serde::Serialize; + +#[repr(C)] +#[derive(Debug, Clone, Copy, Serialize)] +pub struct ImageDosHeader { + pub e_magic: u16, // Magic number 固定值 0x5A4D + pub e_cblp: u16, + pub e_cp: u16, + pub e_crlc: u16, + pub e_cparhdr: u16, + pub e_minalloc: u16, + pub e_maxalloc: u16, + pub e_ss: u16, + pub e_sp: u16, + pub e_csum: u16, + pub e_ip: u16, + pub e_cs: u16, + pub e_lfarlc: u16, + pub e_ovno: u16, + pub e_res: [u16; 4], + pub e_oemid: u16, + pub e_oeminfo: u16, + pub e_res2: [u16; 10], + pub e_lfanew: Offset, // File address of new exe header nt头的偏移 +} + +#[repr(C)] +#[derive(Debug, Clone, Copy, Serialize)] +pub struct ImageNTHeader32 { + pub signature: u32, + pub file_header: ImageFileHeader, + pub optional_header: ImageOptionalHeader32, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy, Serialize)] +pub struct ImageNTHeader64 { + pub signature: u32, + pub file_header: ImageFileHeader, + pub optional_header: ImageOptionalHeader64, +} + +#[repr(C)] +#[derive(Serialize, Clone, Copy)] +#[serde(untagged)] +pub enum ImageNTHeader { + NTHeader32(ImageNTHeader32), + NTHeader64(ImageNTHeader64), +} + +#[repr(C)] +#[derive(Serialize)] +#[serde(untagged)] +pub enum ImageNTHeaderMut<'a> { + NTHeader32(&'a mut ImageNTHeader32), + NTHeader64(&'a mut ImageNTHeader64), +} + +impl<'a> ImageNTHeaderMut<'a> { + pub fn get_optional_header_mut(&mut self) -> ImageOptionalHeaderMut { + match self { + ImageNTHeaderMut::NTHeader32(header) => { + ImageOptionalHeaderMut::OptionalHeader32(&mut header.optional_header) + } + ImageNTHeaderMut::NTHeader64(header) => { + ImageOptionalHeaderMut::OptionalHeader64(&mut header.optional_header) + } + } + } +} + +#[repr(C)] +#[derive(Debug, Clone, Copy, Serialize)] +pub struct ImageFileHeader { + pub machine: u16, + pub number_of_sections: u16, + pub time_date_stamp: u32, + pub pointer_to_symbol_table: Offset, + pub number_of_symbols: u32, + pub size_of_optional_header: u16, + pub characteristics: FileCharacteristics, +} +impl Serialize for FileCharacteristics { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + // 直接返回bitflags的整数值 + serializer.serialize_u16(self.bits()) + } +} + +impl Serialize for DLLCharacteristics { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + serializer.serialize_u16(self.bits()) + } +} + +impl Serialize for SectionCharacteristics { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + serializer.serialize_u32(self.bits()) + } +} + +bitflags! { + #[repr(C)] + #[derive(Debug, Clone, Copy)] + pub struct FileCharacteristics: u16 { + const RELOCS_STRIPPED = 0x0001; + const EXECUTABLE_IMAGE = 0x0002; + const LINE_NUMS_STRIPPED = 0x0004; + const LOCAL_SYMS_STRIPPED = 0x0008; + const AGGRESSIVE_WS_TRIM = 0x0010; + const LARGE_ADDRESS_AWARE = 0x0020; + const BYTES_REVERSED_LO = 0x0080; + const MACHINE_32BIT = 0x0100; + const DEBUG_STRIPPED = 0x0200; + const REMOVABLE_RUN_FROM_SWAP = 0x0400; + const NET_RUN_FROM_SWAP = 0x0800; + const SYSTEM = 0x1000; + const DLL = 0x2000; + const UP_SYSTEM_ONLY = 0x4000; + const BYTES_REVERSED_HI = 0x8000; + } + #[repr(C)] + #[derive(Debug, Clone, Copy)] + pub struct DLLCharacteristics: u16 { + const RESERVED1 = 0x0001; + const RESERVED2 = 0x0002; + const RESERVED4 = 0x0004; + const RESERVED8 = 0x0008; + const HIGH_ENTROPY_VA = 0x0020; + const DYNAMIC_BASE = 0x0040; + const FORCE_INTEGRITY = 0x0080; + const NX_COMPAT = 0x0100; + const NO_ISOLATION = 0x0200; + const NO_SEH = 0x0400; + const NO_BIND = 0x0800; + const APPCONTAINER = 0x1000; + const WDM_DRIVER = 0x2000; + const GUARD_CF = 0x4000; + const TERMINAL_SERVER_AWARE = 0x8000; + } + #[repr(C)] + #[derive(Debug, Clone, Copy)] + pub struct SectionCharacteristics: u32 { + /// Reserved for future use. + const TYPE_REG = 0x00000000; + /// Reserved for future use. + const TYPE_DSECT = 0x00000001; + /// Reserved for future use. + const TYPE_NOLOAD = 0x00000002; + /// Reserved for future use. + const TYPE_GROUP = 0x00000004; + /// The section should not be padded to the next boundary. + /// This flag is obsolete and is replaced by IMAGE_SCN_ALIGN_1BYTES. + /// This is valid only for object files. + const TYPE_NO_PAD = 0x00000008; + /// Reserved for future use. + const TYPE_COPY = 0x00000010; + /// The section contains executable code. + const CNT_CODE = 0x00000020; + /// The section contains initialized data. + const CNT_INITIALIZED_DATA = 0x00000040; + /// The section contains uninitialized data. + const CNT_UNINITIALIZED_DATA = 0x00000080; + /// Reserved for future use. + const LNK_OTHER = 0x00000100; + /// The section contains comments or other information. + /// The .drectve section has this type. This is valid for object files only. + const LNK_INFO = 0x00000200; + /// Reserved for future use. + const TYPE_OVER = 0x00000400; + /// The section will not become part of the image. + /// This is valid only for object files. + const LNK_REMOVE = 0x00000800; + /// The section contains COMDAT data. This is valid only for object files. + const LNK_COMDAT = 0x00001000; + /// Unknown/Reserved. + const RESERVED = 0x00002000; + /// Unknown flag. + const MEM_PROTECTED = 0x00004000; + /// Unknown flag. + const NO_DEFER_SPEC_EXC = 0x00004000; + /// The section contains data referenced through the global pointer (GP). + const GPREL = 0x00008000; + /// Reserved for future use. + const MEM_FARDATA = 0x00008000; + /// Reserved for future use. + const MEM_SYSHEAP = 0x00010000; + /// Reserved for future use. + const MEM_PURGEABLE = 0x00020000; + /// Reserved for future use. + const MEM_16BIT = 0x00020000; + /// Reserved for future use. + const MEM_LOCKED = 0x00040000; + /// Reserved for future use. + const MEM_PRELOAD = 0x00080000; + /// Align data on a 1-byte boundary. Valid only for object files. + const ALIGN_1BYTES = 0x00100000; + /// Align data on a 2-byte boundary. Valid only for object files. + const ALIGN_2BYTES = 0x00200000; + /// Align data on a 4-byte boundary. Valid only for object files. + const ALIGN_4BYTES = 0x00300000; + /// Align data on an 8-byte boundary. Valid only for object files. + const ALIGN_8BYTES = 0x00400000; + /// Align data on a 16-byte boundary. Valid only for object files. + const ALIGN_16BYTES = 0x00500000; + /// Align data on a 32-byte boundary. Valid only for object files. + const ALIGN_32BYTES = 0x00600000; + /// Align data on a 64-byte boundary. Valid only for object files. + const ALIGN_64BYTES = 0x00700000; + /// Align data on a 128-byte boundary. Valid only for object files. + const ALIGN_128BYTES = 0x00800000; + /// Align data on a 256-byte boundary. Valid only for object files. + const ALIGN_256BYTES = 0x00900000; + /// Align data on a 512-byte boundary. Valid only for object files. + const ALIGN_512BYTES = 0x00A00000; + /// Align data on a 1024-byte boundary. Valid only for object files. + const ALIGN_1024BYTES = 0x00B00000; + /// Align data on a 2048-byte boundary. Valid only for object files. + const ALIGN_2048BYTES = 0x00C00000; + /// Align data on a 4096-byte boundary. Valid only for object files. + const ALIGN_4096BYTES = 0x00D00000; + /// Align data on an 8192-byte boundary. Valid only for object files. + const ALIGN_8192BYTES = 0x00E00000; + /// Mask for alignment. + const ALIGN_MASK = 0x00F00000; + /// The section contains extended relocations. + const LNK_NRELOC_OVFL = 0x01000000; + /// The section can be discarded as needed. + const MEM_DISCARDABLE = 0x02000000; + /// The section cannot be cached. + const MEM_NOT_CACHED = 0x04000000; + /// The section is not pageable. + const MEM_NOT_PAGED = 0x08000000; + /// The section can be shared in memory. + const MEM_SHARED = 0x10000000; + /// The section can be executed as code. + const MEM_EXECUTE = 0x20000000; + /// The section can be read. + const MEM_READ = 0x40000000; + /// The section can be written to. + const MEM_WRITE = 0x80000000; + } +} +#[repr(C)] +#[derive(Debug, Clone, Copy, Serialize)] +pub struct ImageOptionalHeader32 { + pub magic: u16, + pub major_linker_version: u8, + pub minor_linker_version: u8, + pub size_of_code: u32, + pub size_of_initialized_data: u32, + pub size_of_uninitialized_data: u32, + pub address_of_entry_point: u32, + pub base_of_code: u32, + pub base_of_data: u32, + pub image_base: u32, + pub section_alignment: u32, + pub file_alignment: u32, + pub major_operating_system_version: u16, + pub minor_operating_system_version: u16, + pub major_image_version: u16, + pub minor_image_version: u16, + pub major_subsystem_version: u16, + pub minor_subsystem_version: u16, + pub win32_version_value: u32, + pub size_of_image: u32, + pub size_of_headers: u32, + pub checksum: u32, + pub subsystem: u16, + pub dll_characteristics: DLLCharacteristics, + pub size_of_stack_reserve: u32, + pub size_of_stack_commit: u32, + pub size_of_heap_reserve: u32, + pub size_of_heap_commit: u32, + pub loader_flags: u32, + pub number_of_rva_and_sizes: u32, +} + +#[repr(C)] +#[derive(Copy, Clone, Debug, Serialize)] +pub struct ImageOptionalHeader64 { + pub magic: u16, + pub major_linker_version: u8, + pub minor_linker_version: u8, + pub size_of_code: u32, + pub size_of_initialized_data: u32, + pub size_of_uninitialized_data: u32, + pub address_of_entry_point: u32, + pub base_of_code: u32, + pub image_base: u64, + pub section_alignment: u32, + pub file_alignment: u32, + pub major_operating_system_version: u16, + pub minor_operating_system_version: u16, + pub major_image_version: u16, + pub minor_image_version: u16, + pub major_subsystem_version: u16, + pub minor_subsystem_version: u16, + pub win32_version_value: u32, + pub size_of_image: u32, + pub size_of_headers: u32, + pub checksum: u32, + pub subsystem: u16, + pub dll_characteristics: DLLCharacteristics, + pub size_of_stack_reserve: u64, + pub size_of_stack_commit: u64, + pub size_of_heap_reserve: u64, + pub size_of_heap_commit: u64, + pub loader_flags: u32, + pub number_of_rva_and_sizes: u32, +} +#[repr(C)] +#[derive(Serialize, Clone, Copy)] +#[serde(untagged)] +pub enum ImageOptionalHeader { + OptionalHeader32(ImageOptionalHeader32), + OptionalHeader64(ImageOptionalHeader64), +} + +#[repr(C)] +#[derive(Copy, Clone, Eq, PartialEq, Default, Debug)] +pub struct ImageDataDirectory { + pub virtual_address: u32, + pub size: u32, +} + +#[derive(Copy, Clone, Eq, PartialEq, Debug)] +pub enum ImageDirectoryEntry { + Export = 0, + Import = 1, + Resource = 2, + Exception = 3, + Security = 4, + BaseReloc = 5, + Debug = 6, + Architecture = 7, + GlobalPTR = 8, + TLS = 9, + LoadConfig = 10, + BoundImport = 11, + IAT = 12, + DelayImport = 13, + COMDescriptor = 14, + Reserved = 15, +} + +#[repr(C)] +#[derive(Serialize)] +#[serde(untagged)] +pub enum ImageOptionalHeaderMut<'a> { + OptionalHeader32(&'a mut ImageOptionalHeader32), + OptionalHeader64(&'a mut ImageOptionalHeader64), +} + +impl ImageOptionalHeader { + pub fn get_size_of_headers(&self) -> u32 { + match self { + ImageOptionalHeader::OptionalHeader32(header) => header.size_of_headers, + ImageOptionalHeader::OptionalHeader64(header) => header.size_of_headers, + } + } + + pub fn get_size_of_image(&self) -> u32 { + match self { + ImageOptionalHeader::OptionalHeader32(header) => header.size_of_image, + ImageOptionalHeader::OptionalHeader64(header) => header.size_of_image, + } + } + pub fn get_size_of_optional_header(&self) -> u32 { + match self { + ImageOptionalHeader::OptionalHeader32(_) => { + size_of::() as u32 + } + ImageOptionalHeader::OptionalHeader64(_) => { + size_of::() as u32 + } + } + } + /// 获取数据目录数量 + pub fn get_data_directory_count(&self) -> u32 { + // 如果数据目录数量大于16，则返回16 + match self { + ImageOptionalHeader::OptionalHeader32(header) => header.number_of_rva_and_sizes, + ImageOptionalHeader::OptionalHeader64(header) => header.number_of_rva_and_sizes, + } + .min(16) + } +} + +#[repr(C)] +#[derive(Debug, Clone, Copy, Serialize)] +pub struct ImageSectionHeader { + pub name: [u8; 8], + pub virtual_size: u32, + pub virtual_address: u32, + pub size_of_raw_data: u32, + pub pointer_to_raw_data: u32, + pub pointer_to_relocations: u32, + pub pointer_to_linenumbers: u32, + pub number_of_relocations: u16, + pub number_of_linenumbers: u16, + pub characteristics: SectionCharacteristics, +} + +#[repr(C)] +#[derive(Copy, Clone, Eq, PartialEq, Debug, Default)] +pub struct ImageImportDescriptor { + pub original_first_thunk: u32, + pub time_date_stamp: u32, + pub forwarder_chain: u32, + pub name: u32, + pub first_thunk: u32, +} + +bitflags! { + /// 导入查找表 32位 + /// [Import Lookup Table](https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#import-lookup-table) + #[derive(Debug, Clone, Copy, Serialize)] + pub struct ImportLookupTable32: u32 { + /// 第31位为1表示导入函数通过序号导入 + const IMPORT_BY_ORDINAL = 0x80000000; + /// 15-0 位表示导入函数的序号 + const IMPORT_BY_ORDINAL_MASK = 0x0000FFFF; + /// 30-0 位表示导入函数的RVA + const IMPORT_BY_RVA_MASK = 0x7FFFFFFF; + } + + /// 导入查找表 64位 + /// [Import Lookup Table](https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#import-lookup-table) + #[derive(Debug, Clone, Copy, Serialize)] + pub struct ImportLookupTable64: u64 { + /// 第63位为1表示导入函数通过序号导入 + const IMPORT_BY_ORDINAL = 0x8000000000000000; + /// 15-0 位表示导入函数的序号 + const IMPORT_BY_ORDINAL_MASK = 0x0000FFFF; + /// 30-0 位表示导入函数的RVA + const IMPORT_BY_RVA_MASK = 0x7FFFFFFFFFFFFFFF; + /// 62-31 must be zero. + const RESERVED = 0x7FFFFFFF00000000; + } + +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..5793e6e --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,16 @@ +pub mod error; +pub mod header; +pub mod pe; +pub mod types; + +#[cfg(feature = "memmap2_impl")] +mod third_party { + use memmap2::{MmapMut, Mmap}; + use crate::pe::{MutablePE, ReadOnlyPE}; + // 为文件映射实现PE结构 + impl ReadOnlyPE for Mmap {} + + impl ReadOnlyPE for MmapMut {} + // 为可变文件映射实现可操作PE结构 + impl MutablePE for MmapMut {} +} diff --git a/src/pe.rs b/src/pe.rs new file mode 100644 index 0000000..9a8fbff --- /dev/null +++ b/src/pe.rs @@ -0,0 +1,540 @@ +use super::{ + error::{MutablePEError, PEParseError}, + header::*, +}; +use std::{ + mem::{self}, + ops::{Deref, DerefMut}, +}; + +pub trait ReadOnlyPE: Deref + Sized + AsRef<[u8]> { + /// Get a reference to a type at a given offset. + fn get_ref(&self, offset: usize) -> Result<&T, PEParseError> { + let len = std::mem::size_of::(); + if offset + len > self.len() { + return Err(PEParseError::OutOfBounds); + } + let ptr = self.as_ptr().wrapping_offset(offset as isize) as *const T; + Ok(unsafe { &*ptr }) + } + + // 判断是64位还是32位 + fn is_64_bit(&self) -> Result { + // 先以32位加载可选头，通过可选头的Magic字段判断 + let optional_header: &ImageOptionalHeader32 = self.get_optional_header_32()?; + let is_64_bit = optional_header.magic == 0x20b; + Ok(is_64_bit) + } + + /// 获取数据目录 + fn get_data_directories(&self) -> Result<&[ImageDataDirectory], PEParseError> { + // 1. 获取数据目录偏移 + + let data_directories_offset = self.get_data_directories_offset()?; + // 2. 解析数据目录 + let data_directories = unsafe { + let ptr = self + .as_ptr() + .wrapping_offset(data_directories_offset as isize) + as *const ImageDataDirectory; + let data_directories_count = self.get_optional_header()?.get_data_directory_count(); + let result = std::slice::from_raw_parts(ptr, data_directories_count as usize); + result + }; + Ok(data_directories) + } + + /// 通过枚举获取某项数据目录的成员 + /// data_directory: 数据目录枚举 + fn get_data_directory( + &self, + data_directory: ImageDirectoryEntry, + ) -> Result<&ImageDataDirectory, PEParseError> { + let data_directories = self.get_data_directories()?; + let data_directory = data_directories + .get(data_directory as usize) + .ok_or(PEParseError::InvalidDataDirectoryIndex)?; + Ok(data_directory) + } + + /// 获取数据目录的偏移 + fn get_data_directories_offset(&self) -> Result { + // 1. 获取可选头和可选头的偏移 + let optional_header: ImageOptionalHeader = self.get_optional_header()?; + let optional_header_offset = self.get_optional_header_offset()? as u32; + // 2. 数据目录的偏移=可选头的偏移+可选头的大小 + let data_directories_offset = + optional_header_offset + optional_header.get_size_of_optional_header(); + Ok(data_directories_offset) + } + + /// 将RVA转换为FOA + /// rva: 需要转换的RVA + fn rva_to_foa(&self, rva: u32) -> Result { + let sections = self.get_section_headers()?; + for section in sections.iter() { + let section_start = section.virtual_address; + let section_end = + section_start + self.align_size_with_file_alignment(section.size_of_raw_data)?; + if rva >= section_start && rva < section_end { + if section.size_of_raw_data == 0 { + // 如果节区的SizeOfRawData是0，那么这个节区是没有数据的 映射到这个节区的RVA一定是无效的 + return Err(PEParseError::RvaToEmptySection(rva)); + } + let foa = rva - section_start + section.pointer_to_raw_data; + return Ok(foa); + } + } + Err(PEParseError::RvaToFoaError(rva)) + } + + /// 将FOA转换为RVA + /// foa: 需要转换的FOA + fn foa_to_rva(&self, foa: u32) -> Result { + let sections = self.get_section_headers()?; + for section in sections.iter() { + // 如果section的PointerToRawData是0，那么这个节区是没有数据的,直接跳过 + if section.pointer_to_raw_data == 0 { + continue; + } + let section_start = section.pointer_to_raw_data; + let section_end = section_start + section.size_of_raw_data; + if foa >= section_start && foa < section_end { + let rva = foa - section_start + section.virtual_address; + return Ok(rva); + } + } + Err(PEParseError::RvaToFoaError(foa)) + } + + /// 将size与节对齐值进行对齐，返回对齐后的值 + /// size: 需要对齐的值 + fn align_size_with_section_alignment(&self, size: u32) -> Result { + let nt_header = self.get_nt_header()?; + let section_alignment = match nt_header { + ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.section_alignment, + ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.section_alignment, + }; + let aligned_size = (size + section_alignment - 1) & !(section_alignment - 1); + Ok(aligned_size) + } + + /// 将size与文件对齐值进行对齐，返回对齐后的值 + /// size: 需要对齐的值 + fn align_size_with_file_alignment(&self, size: u32) -> Result { + let nt_header = self.get_nt_header()?; + let file_alignment = match nt_header { + ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.file_alignment, + ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.file_alignment, + }; + let aligned_size = (size + file_alignment - 1) & !(file_alignment - 1); + Ok(aligned_size) + } + + /// Get the DOS header without verifying its contents. + fn get_dos_header(&self) -> Result<&ImageDosHeader, PEParseError> { + let result = self.get_ref::(0)?; + Ok(result) + } + + /// 作为32位的NT头解析 + fn get_nt_header32(&self) -> Result<&ImageNTHeader32, PEParseError> { + let nt_offset = self.get_nt_headers_offset()?; + let result = self.get_ref::(nt_offset)?; + Ok(result) + } + /// 作为64位的NT头解析 + fn get_nt_header64(&self) -> Result<&ImageNTHeader64, PEParseError> { + let nt_offset = self.get_nt_headers_offset()?; + let result = self.get_ref::(nt_offset)?; + Ok(result) + } + /// 动态的解析NT头 无论是32位还是64位 + fn get_nt_header(&self) -> Result { + let is_64_bit = self.is_64_bit()?; + let result = match is_64_bit { + true => ImageNTHeader::NTHeader64(*self.get_nt_header64()?), + false => ImageNTHeader::NTHeader32(*self.get_nt_header32()?), + }; + Ok(result) + } + + // 获取nt头的偏移 + fn get_nt_headers_offset(&self) -> Result { + let dos_header = self.get_dos_header()?; + let nt_offset = dos_header.e_lfanew.0 as usize; + Ok(nt_offset) + } + + // 获取文件头的偏移 + fn get_file_header_offset(&self) -> Result { + // 1. 获取nt头偏移 + let nt_offset = self.get_nt_headers_offset()?; + // 2. 计算文件头偏移 + let file_header_offset = nt_offset + std::mem::size_of::(); + Ok(file_header_offset) + } + + // 获取文件头数据 + fn get_file_header(&self) -> Result<&ImageFileHeader, PEParseError> { + let file_header_offset = self.get_file_header_offset()?; + let result = self.get_ref::(file_header_offset)?; + Ok(result) + } + + // 获取节区数量 + fn get_number_of_sections(&self) -> Result { + let file_header = self.get_file_header()?; + Ok(file_header.number_of_sections as usize) + } + + // 获取可选头大小 + fn get_size_of_optional_header(&self) -> Result { + let file_header = self.get_file_header()?; + Ok(file_header.size_of_optional_header as usize) + } + /// 作为32位的可选头解析 + fn get_optional_header_32(&self) -> Result<&ImageOptionalHeader32, PEParseError> { + let file_header_offset = self.get_file_header_offset()?; + let optional_header_offset = file_header_offset + std::mem::size_of::(); + let result = self.get_ref::(optional_header_offset)?; + Ok(result) + } + + // 作为64位的可选头解析 + fn get_optional_header_64(&self) -> Result<&ImageOptionalHeader64, PEParseError> { + let file_header_offset = self.get_file_header_offset()?; + let optional_header_offset = file_header_offset + std::mem::size_of::(); + let result = self.get_ref::(optional_header_offset)?; + Ok(result) + } + /// 动态的解析可选头 无论是32位还是64位 + fn get_optional_header(&self) -> Result { + let is_64_bit = self.is_64_bit()?; + let result = match is_64_bit { + true => ImageOptionalHeader::OptionalHeader64(*self.get_optional_header_64()?), + false => ImageOptionalHeader::OptionalHeader32(*self.get_optional_header_32()?), + }; + Ok(result) + } + /// 获取可选头的偏移 + fn get_optional_header_offset(&self) -> Result { + let file_header_offset = self.get_file_header_offset()?; + let optional_header_offset = file_header_offset + std::mem::size_of::(); + Ok(optional_header_offset) + } + + /// 获取节区头的偏移 + fn get_section_headers_offset(&self) -> Result { + // 节区头偏移在可选头之后，可选头大小是可变的，所以需要计算 + let optional_header_size = self.get_size_of_optional_header()?; + let file_header_offset = self.get_file_header_offset()?; + let section_header_offset = + file_header_offset + std::mem::size_of::() + optional_header_size; + Ok(section_header_offset) + } + + // 获取所有节区头数据 + fn get_section_headers(&self) -> Result<&[ImageSectionHeader], PEParseError> { + // 1. 获取节区数量 + let number_of_sections = self.get_number_of_sections()?; + // 2. 获取节区头偏移 + let section_header_offset = self.get_section_headers_offset()?; + // 3. 获取节区头数据 + unsafe { + let ptr = self + .as_ptr() + .wrapping_offset(section_header_offset as isize) + as *const ImageSectionHeader; + let result = std::slice::from_raw_parts(ptr, number_of_sections); + Ok(result) + } + } + + /// 获取头映射最大可扩容大小 + fn get_header_max_expand_size(&self) -> Result { + // 最大可扩容大小 = 节对齐大小 - 最后一个节区头最后一个字节的偏移 + let nt_header: ImageNTHeader = self.get_nt_header()?; + let (section_alignment, file_alignment) = match nt_header { + ImageNTHeader::NTHeader32(nt_header) => ( + nt_header.optional_header.section_alignment, + nt_header.optional_header.file_alignment, + ), + ImageNTHeader::NTHeader64(nt_header) => ( + nt_header.optional_header.section_alignment, + nt_header.optional_header.file_alignment, + ), + }; + let section_header_offset = self.get_section_headers_offset()?; + let section_headers = self.get_number_of_sections()? as u32; + let last_section_header_offset = section_header_offset + + (section_headers as usize * std::mem::size_of::()); + let mut max_size_of_headers = section_alignment - last_section_header_offset as u32; + // max_size_of_headers 向下对齐 + max_size_of_headers &= !(file_alignment - 1); + + Ok(max_size_of_headers) + } + + /// 获取头部空余空间大小 + /// 头部空余空间大小 = 可选头的SizeOfHeaders - 最后一个节区头最后一个字节的偏移 + fn get_header_empty_space_size(&self) -> Result { + let nt_header: ImageNTHeader = self.get_nt_header()?; + let size_of_headers = match nt_header { + ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.size_of_headers, + ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.size_of_headers, + }; + let section_header_offset = self.get_section_headers_offset()?; + let section_headers = self.get_number_of_sections()? as u32; + let last_section_header_offset = section_header_offset + + (section_headers as usize * std::mem::size_of::()); + let empty_space_size = size_of_headers - last_section_header_offset as u32; + Ok(empty_space_size) + } +} +/// 可修改的PE trait +pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { + /// 实现一个get_mut + fn get_mut(&mut self, offset: usize) -> Result<&mut T, PEParseError> { + let len = std::mem::size_of::(); + if offset + len > self.len() { + return Err(PEParseError::OutOfBounds); + } + let ptr = self.as_mut_ptr().wrapping_offset(offset as isize) as *mut T; + Ok(unsafe { &mut *ptr }) + } + + /// 获取可变的文件头数据 + fn get_file_header_mut(&mut self) -> Result<&mut ImageFileHeader, PEParseError> { + let file_header_offset = self.get_file_header_offset()?; + let result = self.get_mut::(file_header_offset)?; + Ok(result) + } + /// 获取可变的节区头数据 + fn get_section_headers_mut(&mut self) -> Result<&mut [ImageSectionHeader], PEParseError> { + // 1. 获取节区数量 + let number_of_sections = self.get_number_of_sections()?; + // 2. 获取节区头偏移 + let section_header_offset = self.get_section_headers_offset()?; + // 3. 获取节区头数据 + unsafe { + let ptr = self + .as_mut_ptr() + .wrapping_offset(section_header_offset as isize) + as *mut ImageSectionHeader; + let result = std::slice::from_raw_parts_mut(ptr, number_of_sections); + Ok(result) + } + } + /// 获取可变的NT头数据64位 + fn get_nt_header64_mut(&mut self) -> Result<&mut ImageNTHeader64, PEParseError> { + let nt_offset = self.get_nt_headers_offset()?; + let result = self.get_mut::(nt_offset)?; + Ok(result) + } + + /// 获取可变的NT头数据32位 + fn get_nt_header32_mut(&mut self) -> Result<&mut ImageNTHeader32, PEParseError> { + let nt_offset = self.get_nt_headers_offset()?; + let result = self.get_mut::(nt_offset)?; + Ok(result) + } + + /// 获取可变的NT头数据 + fn get_nt_header_mut(&mut self) -> Result { + let is_64_bit = self.is_64_bit()?; + let result = match is_64_bit { + true => ImageNTHeaderMut::NTHeader64(self.get_nt_header64_mut()?), + false => ImageNTHeaderMut::NTHeader32(self.get_nt_header32_mut()?), + }; + Ok(result) + } + + /// 获取可变的数据目录引用 + fn get_data_directories_mut(&mut self) -> Result<&mut [ImageDataDirectory], PEParseError> { + // 1. 获取数据目录偏移 + let data_directories_offset = self.get_data_directories_offset()?; + // 2. 解析数据目录 + let data_directories = unsafe { + let ptr = self + .as_mut_ptr() + .wrapping_offset(data_directories_offset as isize) + as *mut ImageDataDirectory; + let data_directories_count = self.get_optional_header()?.get_data_directory_count(); + let result = std::slice::from_raw_parts_mut(ptr, data_directories_count as usize); + result + }; + Ok(data_directories) + } + + /// 通过枚举获取某项数据目录的可变引用 + /// data_directory: 数据目录枚举 + fn get_data_directory_mut( + &mut self, + data_directory: ImageDirectoryEntry, + ) -> Result<&mut ImageDataDirectory, PEParseError> { + let data_directories = self.get_data_directories_mut()?; + let data_directory = data_directories + .get_mut(data_directory as usize) + .ok_or(PEParseError::InvalidDataDirectoryIndex)?; + Ok(data_directory) + } + + /// 扩大头映射大小，调用者必须要保证当挪动所有节区数据时，不会超出文件范围 + /// 也就是，调用者必须要在扩大头映射大小之前，先扩大文件大小 + /// add_size: 需要扩大的大小，不需要进行对齐，会自动对齐 + fn expand_headers(&mut self, add_size: u32) -> Result<(), MutablePEError> { + // aligned_size: 需要增加的大小，对齐后的值 + let aligned_file_size = self.align_size_with_file_alignment(add_size)?; + let max_expand_size = self.get_header_max_expand_size()?; + let optional_header = self.get_optional_header()?; + let origin_size_for_headers = optional_header.get_size_of_headers(); + // 如果扩大后超过了最大可扩容大小，就不允许扩容 + if aligned_file_size > max_expand_size { + return Err(MutablePEError::CannotExpandHeader( + aligned_file_size, + origin_size_for_headers, + )); + } + // 检查一下把，万一调用者没有扩大文件大小 + let sections = self.get_section_headers()?; + let self_len = self.len(); + // 需要反向遍历，找到最后一个size_of_raw_data不为0的节区 + // 然后计算这个节区的PointerToRawData+（SizeOfRawData对齐到FileAlignment的值） + let mut last_section_end = 0; + for section in sections.iter().rev() { + if section.size_of_raw_data != 0 { + last_section_end = section.pointer_to_raw_data + + self.align_size_with_file_alignment(section.size_of_raw_data)?; + if last_section_end as usize + aligned_file_size as usize > self_len { + return Err(MutablePEError::CannotExpandFileSize( + last_section_end as u32 + aligned_file_size, + self_len as u32, + )); + } + break; + } + } + let mut nt_header = self.get_nt_header_mut()?; + let op = nt_header.get_optional_header_mut(); + // 更新可选头的SizeOfHeaders和SizeOfImage + match op { + ImageOptionalHeaderMut::OptionalHeader32(op) => { + op.size_of_headers += aligned_file_size; + } + ImageOptionalHeaderMut::OptionalHeader64(op) => { + op.size_of_headers += aligned_file_size; + } + } + let mut_section = self.get_section_headers_mut()?; + let start_offset = mut_section.first().unwrap().pointer_to_raw_data; + let end_offset = last_section_end; + // 更新节区头的PointerToRawData + for section in mut_section.iter_mut() { + section.pointer_to_raw_data += aligned_file_size; + } + { + let self_data = self.as_mut(); + // 把[start_offset, end_offset)的数据往后挪动aligned_file_size + self_data.copy_within( + start_offset as usize..end_offset as usize, + (start_offset + aligned_file_size) as usize, + ); + // 把[start_offset, start_offset+aligned_file_size)的数据填充为0 + self_data[start_offset as usize..(start_offset + aligned_file_size) as usize].fill(0); + } + Ok(()) + } + + /// 添加一个节区 + /// section_name: 节区名 + /// section_size: 节区大小 + /// section_characteristics: 节区特征 + fn add_section( + &mut self, + section_name: &[u8; 8], + section_size: u32, + section_characteristics: SectionCharacteristics, + ) -> Result<(), MutablePEError> { + // 1. 判断是否有空余空间 + let empty_space_size = self.get_header_empty_space_size()?; + const SIZE_OF_SECTION_HEADER: u32 = mem::size_of::() as u32; + // section_size与文件对齐后的大小 + let file_alignment = self.align_size_with_file_alignment(section_size)?; + + if empty_space_size < SIZE_OF_SECTION_HEADER { + // 判断是否能够扩容 + let max_expand_size = self.get_header_max_expand_size()?; + if max_expand_size < file_alignment { + return Err(MutablePEError::CannotExpandHeader( + file_alignment, + max_expand_size, + )); + } + self.expand_headers(file_alignment)?; + } + // 增加节区头 + let sections_offset = self.get_section_headers_offset()?; + let number_of_sections = self.get_number_of_sections()?; + let new_section_offset = + sections_offset + (number_of_sections * SIZE_OF_SECTION_HEADER as usize) as usize; + let new_section = unsafe { + let ptr = self + .as_mut_ptr() + .wrapping_offset(new_section_offset as isize) + as *mut ImageSectionHeader; + &mut *ptr + }; + new_section.name = *section_name; + // 新节区的virtual_size是section_size对齐到section_alignment的值 + let mut section_alignment = self.align_size_with_section_alignment(section_size)?; + if section_size == 0 { + section_alignment = self.align_size_with_section_alignment(1)?; + } + new_section.virtual_size = section_alignment; + // 如果section_size是0，那么virtual_size必须得是最小的section_alignment + // 新节区的virtual_address是上一个节区的virtual_address+(pointer_to_raw_data 对齐到 section_alignment的值) + let last_section = self.get_section_headers()?.last().unwrap(); + let mut last_section_size_of_raw_data_aligned = + self.align_size_with_section_alignment(last_section.size_of_raw_data)?; + if last_section_size_of_raw_data_aligned == 0 { + last_section_size_of_raw_data_aligned = self.align_size_with_section_alignment(1)?; + } + new_section.virtual_address = + last_section.virtual_address + last_section_size_of_raw_data_aligned; + new_section.size_of_raw_data = file_alignment; + // 如果上一个节区的size_of_raw_data是0，那么上一个节区的pointer_to_raw_data是可以乱写的 + // 正确的做法应该是，找到上一个size_of_raw_data不为0的节区，然后计算 + let sections = self.get_section_headers()?; + // 反向遍历 + for section in sections.iter().rev() { + if section.size_of_raw_data != 0 { + new_section.pointer_to_raw_data = section.pointer_to_raw_data + + self.align_size_with_file_alignment(section.size_of_raw_data)?; + break; + } + } + new_section.characteristics = section_characteristics; + // 其他字段最好填充为0 + new_section.number_of_linenumbers = 0; + new_section.number_of_relocations = 0; + new_section.pointer_to_linenumbers = 0; + new_section.pointer_to_relocations = 0; + // 更新文件头的NumberOfSections + let file_header = self.get_file_header_mut()?; + file_header.number_of_sections += 1; + + // 更新可选头的SizeOfImage + let mut nt_header = self.get_nt_header_mut()?; + let op = nt_header.get_optional_header_mut(); + match op { + ImageOptionalHeaderMut::OptionalHeader32(op) => { + op.size_of_image += section_alignment; + } + ImageOptionalHeaderMut::OptionalHeader64(op) => { + op.size_of_image += section_alignment; + } + } + Ok(()) + } +} diff --git a/src/types.rs b/src/types.rs new file mode 100644 index 0000000..c526003 --- /dev/null +++ b/src/types.rs @@ -0,0 +1,18 @@ +use std::fmt::{Formatter, UpperHex}; + +use serde::Serialize; + + +#[repr(C)] +#[derive(Debug, Clone, Copy, Serialize)] +pub struct Offset(pub u32); + +impl UpperHex for Offset { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +#[repr(C)] +#[derive(Debug, Clone, Copy, Serialize)] +pub struct RVA(pub u32);