feat: 修改完测试

This commit is contained in:
381848900@qq.com 2024-12-19 15:45:49 +08:00
parent 0bbe0c4210
commit 578510c820
4 changed files with 114 additions and 56 deletions

View File

@ -10,7 +10,7 @@ use std::{
};
pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
/// Get a reference to a type at a given offset.
/// 获取一个给定偏移量的类型的引用
fn get_ref<T>(&self, offset: usize) -> Result<&T, PEParseError> {
let len = std::mem::size_of::<T>();
if offset + len > self.len() {
@ -20,7 +20,7 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
Ok(unsafe { &*ptr })
}
// 判断是64位还是32位
/// 判断是64位还是32位
fn is_64_bit(&self) -> Result<bool, PEParseError> {
// 先以32位加载可选头通过可选头的Magic字段判断
let optional_header: &ImageOptionalHeader32 = self.get_optional_header_32()?;
@ -28,6 +28,11 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
Ok(is_64_bit)
}
/// 通过RVA获取偏移
/// 如果在文件PE中获取的结果是FOA。
/// 如果是在内存PE中获取的结果是VA。
fn rva_to_offset(&self, rva: u32) -> Result<usize, PEParseError>;
/// 获取数据目录
fn get_data_directories(&self) -> Result<&[ImageDataDirectory], PEParseError> {
// 1. 获取数据目录偏移
@ -109,6 +114,14 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
Err(PEParseError::RvaToFoaError(foa))
}
/// RVA转换为VA
fn rva_to_va(&self, rva: u32) -> Result<usize, PEParseError> {
let image_base = self.get_image_base()?;
Ok(image_base + rva as usize)
}
fn get_image_base(&self) -> Result<usize, PEParseError>;
/// 将size与节对齐值进行对齐返回对齐后的值
/// size: 需要对齐的值
fn align_size_with_section_alignment(&self, size: u32) -> Result<u32, PEParseError> {
@ -295,8 +308,8 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
/// 以传入的RVA为开始偏移从中解析导出表
fn parse_export_table(&self, rva: u32) -> Result<&ImageExportDirectoryTable, PEParseError> {
let foa = self.rva_to_foa(rva)?;
let export_table = self.get_ref::<ImageExportDirectoryTable>(foa as usize)?;
let offset = self.rva_to_offset(rva)?;
let export_table = self.get_ref::<ImageExportDirectoryTable>(offset)?;
Ok(export_table)
}
@ -312,22 +325,22 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
let base = export_table.base; // 序号基数
let number_of_fun = export_table.number_of_functions;
let number_of_names = export_table.number_of_names;
// 名称指针表的FOA
let names_table_foa = self.rva_to_foa(export_table.address_of_names)?;
// 名称指针表的偏移
let names_table_offset = self.rva_to_offset(export_table.address_of_names)?;
let get_function_rva = |ordinal: u32| -> Result<u32, PEParseError> {
let functions_table_foa = self.rva_to_foa(export_table.address_of_functions)?;
let functions_table_offset = self.rva_to_offset(export_table.address_of_functions)?;
let function_item_foa =
functions_table_foa + ordinal * std::mem::size_of::<u32>() as u32;
functions_table_offset + ordinal as usize * std::mem::size_of::<u32>();
Ok(self.get_ref::<u32>(function_item_foa as usize)?.clone())
};
let check_forwarder = |rva: u32| -> Result<Option<ExportFunEnum>, PEParseError> {
if rva >= export_table_rva && rva < export_table_rva + export_table_size {
let forwarder_foa = self.rva_to_foa(rva)?;
let forwarder_offset = self.rva_to_offset(rva)?;
let forwarder = unsafe {
std::ffi::CStr::from_ptr(
self.get_ref::<i8>(forwarder_foa as usize)? as *const i8
self.get_ref::<i8>(forwarder_offset as usize)? as *const i8
)
.to_str()
.unwrap()
@ -340,19 +353,19 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
match proc_name {
ProcName::Named(proc_name_str) => {
for index in 0..number_of_names {
let name_item_foa = names_table_foa + index * std::mem::size_of::<u32>() as u32;
let name_item_foa = names_table_offset + index as usize * std::mem::size_of::<u32>() as usize;
let name_rva = self.get_ref::<u32>(name_item_foa as usize)?.clone();
let name_foa = self.rva_to_foa(name_rva)?;
let name_offset = self.rva_to_offset(name_rva)?;
let name = unsafe {
std::ffi::CStr::from_ptr(self.get_ref::<i8>(name_foa as usize)? as *const i8)
std::ffi::CStr::from_ptr(self.get_ref::<i8>(name_offset as usize)? as *const i8)
.to_str()
.unwrap()
};
if name == proc_name_str {
let ordinals_table_foa =
self.rva_to_foa(export_table.address_of_name_ordinals)?;
let ordinals_table_offset =
self.rva_to_offset(export_table.address_of_name_ordinals)?;
let ordinal_item_foa =
ordinals_table_foa + index * std::mem::size_of::<u16>() as u32;
ordinals_table_offset + index as usize * std::mem::size_of::<u16>();
let ordinal = self.get_ref::<u16>(ordinal_item_foa as usize)?.clone();
let rva = get_function_rva(ordinal as u32)?;
return check_forwarder(rva);
@ -373,19 +386,19 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
/// 以传入的RVA为开始偏移从中解析导入表
fn parse_import_tables(&self, rva: u32) -> Result<&[ImageImportDescriptor], PEParseError> {
let foa = self.rva_to_foa(rva)?;
let mut current_foa = foa;
let offset = self.rva_to_offset(rva)?;
let mut current_offset = offset;
let mut numbers = 0;
loop {
let import_table = self.get_ref::<ImageImportDescriptor>(current_foa as usize)?;
let import_table = self.get_ref::<ImageImportDescriptor>(current_offset as usize)?;
if import_table.original_first_thunk == 0 {
break;
}
numbers += 1;
current_foa += std::mem::size_of::<ImageImportDescriptor>() as u32;
current_offset += std::mem::size_of::<ImageImportDescriptor>();
}
let result_tables = unsafe {
let ptr = self.as_ptr().wrapping_offset(foa 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
};
@ -638,3 +651,31 @@ pub trait MutablePE: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {
Ok(())
}
}
pub trait FilePE: ReadOnlyPE {}
impl<T> ReadOnlyPE for T
where
T: FilePE,
{
fn rva_to_offset(&self, rva: u32) -> Result<usize, PEParseError> {
let result = self.rva_to_foa(rva)?;
Ok(result as usize)
}
fn get_image_base(&self) -> Result<usize, PEParseError> {
let nt_header = self.get_nt_header()?;
let image_base = match nt_header {
ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.image_base as usize,
ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.image_base as usize,
};
Ok(image_base)
}
}
// 默认实现MutablePE
// impl<T> MutablePE for T where T: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {}
pub trait FileMutPE: FilePE + MutablePE {}
impl<T> FileMutPE for T where T: FilePE + MutablePE {}

View File

@ -1,17 +1,17 @@
pub mod error;
pub mod header;
pub mod pe;
pub mod memory_pe;
mod constants;
pub mod error;
pub mod file_pe;
pub mod header;
pub mod memory_pe;
mod utils;
#[cfg(feature = "memmap2_impl")]
mod third_party {
use memmap2::{MmapMut, Mmap};
use crate::pe::{MutablePE, ReadOnlyPE};
use crate::file_pe::{FilePE, MutablePE};
use memmap2::{Mmap, MmapMut};
// 为文件映射实现PE结构
impl ReadOnlyPE for Mmap {}
impl FilePE for Mmap {}
impl ReadOnlyPE for MmapMut {}
// 为可变文件映射实现可操作PE结构
impl FilePE for MmapMut {}
impl MutablePE for MmapMut {}
}

View File

@ -2,21 +2,30 @@ use std::ops::{Deref, DerefMut};
use crate::{align_to, error::PEParseError};
use super::pe::{MutablePE, ReadOnlyPE};
use super::file_pe::{MutablePE, ReadOnlyPE};
use crate::constants::*;
pub struct MemoryPE {
pub struct MemoryPEStruct {
len: usize,
buf: *const u8, // 缓冲区的指针 指向内存中的数据
}
pub struct MemoryPEMut {
pub struct MemoryPEMutStruct {
len: usize,
buf: *mut u8, // 缓冲区的指针 指向内存中的数据
}
// DerefMut<Target = [u8]> + AsMut<[u8]>
impl MutablePE for MemoryPEMut{}
impl DerefMut for MemoryPEMut{
impl ReadOnlyPE for MemoryPEMutStruct{
fn rva_to_offset(&self, rva: u32) -> Result<usize, PEParseError> {
// 返回ImageBase + rva
let result = self.rva_to_va(rva)?;
Ok(result as usize)
}
fn get_image_base(&self) -> Result<usize, PEParseError> {
Ok(self.buf as usize)
}
}
impl DerefMut for MemoryPEMutStruct{
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe {
std::slice::from_raw_parts_mut(self.buf, self.len)
@ -24,7 +33,7 @@ impl DerefMut for MemoryPEMut{
}
}
impl AsMut<[u8]> for MemoryPEMut{
impl AsMut<[u8]> for MemoryPEMutStruct{
fn as_mut(&mut self) -> &mut [u8] {
unsafe {
std::slice::from_raw_parts_mut(self.buf, self.len)
@ -32,17 +41,7 @@ impl AsMut<[u8]> for MemoryPEMut{
}
}
impl Deref for MemoryPEMut{
type Target = [u8];
fn deref(&self) -> &Self::Target {
unsafe {
std::slice::from_raw_parts(self.buf, self.len)
}
}
}
impl ReadOnlyPE for MemoryPEMut{}
impl AsRef<[u8]> for MemoryPEMut{
impl AsRef<[u8]> for MemoryPEMutStruct{
fn as_ref(&self) -> &[u8] {
unsafe {
std::slice::from_raw_parts(self.buf, self.len)
@ -50,9 +49,27 @@ impl AsRef<[u8]> for MemoryPEMut{
}
}
impl ReadOnlyPE for MemoryPE{}
impl Deref for MemoryPEMutStruct{
type Target = [u8];
fn deref(&self) -> &Self::Target {
unsafe {
std::slice::from_raw_parts(self.buf, self.len)
}
}
}
impl Deref for MemoryPE{
impl ReadOnlyPE for MemoryPEStruct{
fn rva_to_offset(&self, rva: u32) -> Result<usize, PEParseError> {
let result = self.rva_to_va(rva)?;
Ok(result as usize)
}
fn get_image_base(&self) -> Result<usize, PEParseError> {
Ok(self.buf as usize)
}
}
impl Deref for MemoryPEStruct{
type Target = [u8];
fn deref(&self) -> &Self::Target {
@ -63,7 +80,7 @@ impl Deref for MemoryPE{
}
impl AsRef<[u8]> for MemoryPE{
impl AsRef<[u8]> for MemoryPEStruct{
fn as_ref(&self) -> &[u8] {
unsafe {
std::slice::from_raw_parts(self.buf, self.len)
@ -72,7 +89,7 @@ impl AsRef<[u8]> for MemoryPE{
}
/// 需要实现一个从内存中解析PE的方法
pub fn parse_pe_from_memory(buf: *const u8) -> Result<MemoryPE, PEParseError> {
pub fn parse_pe_from_memory(buf: *const u8) -> Result<MemoryPEStruct, PEParseError> {
// 1. 检查是否为PE文件
// 取出头两个字节
let dos_magic = unsafe {
@ -95,14 +112,14 @@ pub fn parse_pe_from_memory(buf: *const u8) -> Result<MemoryPE, PEParseError> {
let image_size = unsafe { *(buf.add(pe_offset as usize + IMAGE_SIZE_OFFSET) as *const u32) };
// 计算对齐后的内存大小
let aligned_image_size = align_to!(image_size, section_alignment);
Ok(MemoryPE {
Ok(MemoryPEStruct {
len: aligned_image_size as usize,
buf,
})
}
/// 还需要一个Mut的版本
pub fn parse_pe_from_memory_mut(buf: *mut u8) -> Result<MemoryPEMut, PEParseError> {
pub fn parse_pe_from_memory_mut(buf: *mut u8) -> Result<MemoryPEMutStruct, PEParseError> {
// 1. 检查是否为PE文件
// 取出头两个字节
let dos_magic = unsafe {
@ -125,7 +142,7 @@ pub fn parse_pe_from_memory_mut(buf: *mut u8) -> Result<MemoryPEMut, PEParseErro
let image_size = unsafe { *(buf.add(pe_offset as usize + IMAGE_SIZE_OFFSET) as *const u32) };
// 计算对齐后的内存大小
let aligned_image_size = align_to!(image_size, section_alignment);
Ok(MemoryPEMut {
Ok(MemoryPEMutStruct {
len: aligned_image_size as usize,
buf,
})

View File

@ -13,7 +13,7 @@ mod test {
use core::ffi;
// 1. 创建映射
use pe_parse::pe::ReadOnlyPE;
use pe_parse::file_pe::ReadOnlyPE;
let file = std::fs::OpenOptions::new()
.read(true)
.open("tests/test_dll.dll")
@ -63,7 +63,7 @@ mod test {
let find_function_str = ProcName::Named("test_def_add".to_string());
// let find_function_str = ProcName::Ordinal("div");
// 1. 还是先映射
use pe_parse::pe::ReadOnlyPE;
use pe_parse::file_pe::ReadOnlyPE;
let file = std::fs::OpenOptions::new()
.read(true)
.open("tests/test_dll.dll")