diff --git a/src/error.rs b/src/error.rs index 0af36a6..4bf5f79 100644 --- a/src/error.rs +++ b/src/error.rs @@ -20,6 +20,9 @@ pub enum PEParseError { InvalidDataDirectoryIndex(u16), #[error("错误的导出函数序号: {0}")] ExportOrdinalNotFound(u16), + + #[error("目标的缓冲区不够大")] + TargetBufferTooSmall } /// PE操作的错误 #[derive(Error, Debug)] diff --git a/src/file_pe.rs b/src/file_pe.rs index 1f057dd..945774e 100644 --- a/src/file_pe.rs +++ b/src/file_pe.rs @@ -122,6 +122,16 @@ pub trait ReadOnlyPE: Deref + Sized + AsRef<[u8]> { fn get_image_base(&self) -> Result; + /// 获取SizeOfImage大小 + fn get_size_of_image(&self) -> Result { + let nt_header = self.get_nt_header()?; + let size_of_image = match nt_header { + ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.size_of_image, + ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.size_of_image, + }; + Ok(size_of_image) + } + /// 将size与节对齐值进行对齐,返回对齐后的值 /// size: 需要对齐的值 fn align_size_with_section_alignment(&self, size: u32) -> Result { @@ -353,13 +363,16 @@ pub trait ReadOnlyPE: Deref + Sized + AsRef<[u8]> { match proc_name { ProcName::Named(proc_name_str) => { for index in 0..number_of_names { - let name_item_foa = names_table_offset + index as usize * std::mem::size_of::() as usize; + let name_item_foa = + names_table_offset + index as usize * std::mem::size_of::() as usize; let name_rva = self.get_ref::(name_item_foa as usize)?.clone(); let name_offset = self.rva_to_offset(name_rva)?; let name = unsafe { - std::ffi::CStr::from_ptr(self.get_ref::(name_offset as usize)? as *const i8) - .to_str() - .unwrap() + std::ffi::CStr::from_ptr( + self.get_ref::(name_offset as usize)? as *const i8 + ) + .to_str() + .unwrap() }; if name == proc_name_str { let ordinals_table_offset = @@ -398,12 +411,23 @@ pub trait ReadOnlyPE: Deref + Sized + AsRef<[u8]> { current_offset += std::mem::size_of::(); } let result_tables = unsafe { - let ptr = self.as_ptr().wrapping_offset(offset as isize) as *const ImageImportDescriptor; + let ptr = + self.as_ptr().wrapping_offset(offset as isize) as *const ImageImportDescriptor; let result = std::slice::from_raw_parts(ptr, numbers); result }; Ok(result_tables) } + + /// 获取映射的头部大小 + fn get_size_of_header(&self) -> Result { + let nt_header = self.get_nt_header()?; + let size_of_header = 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, + }; + Ok(size_of_header) + } } /// 可修改的PE trait pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { @@ -652,7 +676,39 @@ pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { } } -pub trait FilePE: ReadOnlyPE {} +pub trait FilePE: ReadOnlyPE { + /// 将PE文件拷贝到指定地址 + /// 仅拷贝节表数据 拷贝后的数据已经对齐 + fn copy_file_pe_to_memory(&self, target: &mut [u8]) -> Result<(), PEParseError> { + // 如果长度不够PE中的size_of_image对齐到sction_alignment的值,就返回错误 + let size_of_image = self.get_size_of_image()?; + let section_alignment = self.align_size_with_section_alignment(size_of_image)?; + if target.len() < section_alignment as usize { + return Err(PEParseError::TargetBufferTooSmall); + } + // 长度够了 现在开始拷贝 + // 1. 拷贝头部[0..size_of_headers) + let size_of_header = self.get_size_of_header()?; + target[..size_of_header as usize] + .copy_from_slice(&self.as_ref()[..size_of_header as usize]); + // 2. 遍历节表 + let sections = self.get_section_headers()?; + for item in sections.iter() { + // 如果节表的SizeOfRawData是0,那么这个节区是没有数据的,直接跳过 + if item.size_of_raw_data == 0 { + continue; + } + // 拷贝到target的VirtualAddress处 + let va = item.virtual_address as usize; + let copy_size = item.size_of_raw_data; + // copy_size = copy_size对齐到file_alignment的值 + let copy_size = self.align_size_with_file_alignment(copy_size)? as usize; + let foa = item.pointer_to_raw_data as usize; + target[va..va + copy_size].copy_from_slice(&self.as_ref()[foa..foa + copy_size]); + } + Ok(()) + } +} impl ReadOnlyPE for T where @@ -678,4 +734,4 @@ where pub trait FileMutPE: FilePE + MutablePE {} -impl FileMutPE for T where T: FilePE + MutablePE {} \ No newline at end of file +impl FileMutPE for T where T: FilePE + MutablePE {}