diff --git a/src-tauri/src/app_state.rs b/src-tauri/src/app_state.rs index acb4737..4f59334 100644 --- a/src-tauri/src/app_state.rs +++ b/src-tauri/src/app_state.rs @@ -1,8 +1,10 @@ -use std::io::Seek; -use std::mem::offset_of; +use std::io::{Seek, Write}; +use std::mem::{self, offset_of}; use crate::app_error::AppError; -use crate::pe_parse::header::ImageOptionalHeader; +use crate::pe_parse::header::{ImageOptionalHeader, ImageSectionHeader, SectionCharacteristics}; +use crate::pe_parse::pe::MutablePE; +use crate::services::file::expand_file_size; use crate::{pe_parse::pe::ReadOnlyPE, services::file}; use memmap2::{Mmap, MmapOptions}; @@ -62,11 +64,10 @@ impl AppState { Ok(()) } - /// 添加节 pub fn add_section( &mut self, section_name: &str, - section_size: usize, + section_size: u32, section_characteristics: u32, ) -> Result<(), AppError> { if section_name.len() > 8 { @@ -75,194 +76,49 @@ impl AppState { } let mut section_name_bytes = [0u8; 8]; section_name_bytes[..section_name.len()].copy_from_slice(section_name.as_bytes()); - - // 判断是否可以添加节 - // if(size_of_headers - 最后一个节表的偏移 - 40) >= 40，则可以添加节表。 - // 2. 获取size_of_headers let mmap = self.get_mmap_ref()?; - let origin_file_len = mmap.len(); - let optional_header = mmap.get_optional_header()?; - let optional_header_offset = mmap.get_optional_header_offset()?; - let file_header = mmap.get_file_header()?; - let file_header_offset = mmap.get_file_header_offset()?; - let (size_of_headers, section_alignment) = match optional_header { - ImageOptionalHeader::OptionalHeader32(image_optional_header32) => { - (image_optional_header32.size_of_headers, image_optional_header32.section_alignment) - } - ImageOptionalHeader::OptionalHeader64(image_optional_header64) => { - (image_optional_header64.size_of_headers, image_optional_header64.section_alignment) - } - }; - let section_table_offset = mmap.get_section_headers_offset()?; - // 获取节表的数量 - let number_of_sections = file_header.number_of_sections; - let section_table_size = 40 * number_of_sections as usize; - if size_of_headers - section_table_offset as u32 - (section_table_size as u32) < 40 { - // 2. 计算size_of_headers对齐之后的大小 例如:0x400 -> 0x1000 - let aligned_size_of_headers = - (size_of_headers + section_alignment - 1) & !(section_alignment - 1); - // 可以拓展的大小: 对齐之后的大小 - size_of_headers - let expand_size = aligned_size_of_headers - size_of_headers; - if expand_size < 40 { - // 如果可以拓展的大小小于40，则无法添加节表 + let file_path = self.file_path.as_ref().ok_or(AppError::NoFileOpened)?; + // 将section_size 进行文件对齐 + let file_alignment = mmap.align_size_with_file_alignment(section_size)?; + // 判断是否可以添加节 + // 1. 判断头部空余空间是否足够 + const SIZE_OF_SECTION_HEADER: u32 = mem::size_of::() as u32; + let header_empty_space = mmap.get_header_empty_space_size()?; + let mut rw_mmap_option = None; + + if header_empty_space < SIZE_OF_SECTION_HEADER { + // 头部空余空间不足 需要判断是否可以拓展 + // 获取最大可拓展的大小 + let max_expand_size = mmap.get_header_max_expand_size()?; + let file_alignment = mmap.align_size_with_file_alignment(SIZE_OF_SECTION_HEADER)?; + if file_alignment > max_expand_size { return Err(AppError::CannotModifyFile("无法拓展文件大小!".to_string())); } - // 3. 拓展文件大小 - // 1. 打开文件 - let mut file = std::fs::OpenOptions::new() - .read(true) - .write(true) - .open(self.file_path.as_ref().ok_or(AppError::NoFileOpened)?)?; - // 设置游标到文件末尾 - file.seek(std::io::SeekFrom::End(0))?; - // 2. 拓展文件大小 - file.set_len(file.metadata()?.len() + expand_size as u64)?; + // 否则可以通过拓展头部大小的方式进行节的添加 + // 1. 拓展文件大小 + expand_file_size(file_path, file_alignment)?; - // 2. 映射文件 - let mut rw_mmap = unsafe { - MmapOptions::new() - .map_mut(&file) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))? - }; - // 循环拷贝数据 把size_of_headers之后的数据往后移动expand_size - for i in (size_of_headers as usize..origin_file_len).rev() { - rw_mmap[i + expand_size as usize] = rw_mmap[i]; - } - - // 把[size_of_headers, size_of_headers + expand_size]的数据清零 - for i in size_of_headers as usize..(size_of_headers as usize + expand_size as usize - 1) - { - rw_mmap[i] = 0; - } - // 修整所有节的pointer_to_raw_data - for i in 0..number_of_sections { - let section_header_offset = section_table_offset + 40 * i as usize; - let pointer_to_raw_data = u32::from_le_bytes( - rw_mmap[section_header_offset + 20..section_header_offset + 24] - .try_into() - .unwrap(), - ); - rw_mmap[section_header_offset + 20..section_header_offset + 24] - .copy_from_slice(&(pointer_to_raw_data + expand_size as u32).to_le_bytes()); - } - - // TODO: 修改SizeOfImage - let origin_size_of_image = unsafe { - // 在0x38的位置 - u32::from_le_bytes( - rw_mmap[optional_header_offset + 0x38..optional_header_offset + 0x3C] - .try_into() - .unwrap(), - ) - }; - - let add_image_size = (expand_size + section_alignment - 1) & !(section_alignment - 1); - - rw_mmap[optional_header_offset + 0x38..optional_header_offset + 0x3C] - .copy_from_slice(&(origin_size_of_image + add_image_size as u32).to_le_bytes()); - - // 修改SizeOfHeaders,修改后的长度是：size_of_headers + expand_size例如：0x400 -> 0x1000 - let new_size_of_headers = size_of_headers + expand_size; - rw_mmap[optional_header_offset + 0x3C..optional_header_offset + 0x40] - .copy_from_slice(&new_size_of_headers.to_le_bytes()); - - rw_mmap.flush()?; - // 3. 重新映射文件 - self.mmap = Some(file::mmap_file(self.file_path.as_ref().unwrap())?); + // 2. 以可读可写的方式重新映射文件 + rw_mmap_option = Some(file::mmap_file_rw(file_path)?); + // 拓展头部大小 拓展一个节的大小 + rw_mmap_option.as_mut().unwrap().expand_headers(file_alignment)?; } - - // 可以添加节表 - let new_section_offset = self.get_mmap_ref()?.len(); - let new_section_header_offset = section_table_offset + section_table_size; - { - // 拓展文件大小 - let mut file = std::fs::OpenOptions::new() - .read(true) - .write(true) - .open(self.file_path.as_ref().ok_or(AppError::NoFileOpened)?)?; - // 设置游标到文件末尾 - file.seek(std::io::SeekFrom::End(0))?; - // 2. 拓展文件大小 - file.set_len(file.metadata()?.len() + section_size as u64)?; + // 2. 添加节 + // 先拓展文件大小 + expand_file_size(file_path, file_alignment)?; + // 如果rw_mmap_option为None，则需要重新映射文件 + if rw_mmap_option.is_none() { + rw_mmap_option = Some(file::mmap_file_rw(self.file_path.as_ref().unwrap())?); } - // 写入节表数据 - let mut rw_mmap = file::mmap_file_rw(self.file_path.as_ref().unwrap())?; - // 1. 写入节表数据 - // 1.1 写入节名 - rw_mmap[new_section_header_offset..new_section_header_offset + 8] - .copy_from_slice(§ion_name_bytes); + let rw_mmap = rw_mmap_option.as_mut().unwrap(); - // virtual_size需要对齐 - let mut virtual_size = - ((section_size as u32) + section_alignment - 1) & !(section_alignment - 1); - if virtual_size == 0 { - virtual_size = section_alignment; - } + let characteristics = SectionCharacteristics::from_bits(section_characteristics).unwrap(); - rw_mmap[new_section_header_offset + 8..new_section_header_offset + 12] - .copy_from_slice(&(virtual_size as u32).to_le_bytes()); + rw_mmap.add_section(§ion_name_bytes, section_size, characteristics)?; - // 写入节的VA - let last_section_header_offset = - section_table_offset + 40 * (number_of_sections - 1) as usize; - let last_section_va = u32::from_le_bytes( - rw_mmap[last_section_header_offset + 12..last_section_header_offset + 16] - .try_into() - .unwrap(), - ); - let last_sction_size = u32::from_le_bytes( - rw_mmap[last_section_header_offset + 16..last_section_header_offset + 20] - .try_into() - .unwrap(), - ); - let last_section_alignment_size = - (last_sction_size + section_alignment - 1) & !(section_alignment - 1); - - // 对齐后的VA, 对齐方式(last_section_va ) 和section_alignment进行对齐 - let aligned_va = ((last_section_va + last_section_alignment_size) + section_alignment - 1) - & !(section_alignment - 1); - - rw_mmap[new_section_header_offset + 12..new_section_header_offset + 16] - .copy_from_slice(&aligned_va.to_le_bytes()); - - // 写入新节的size_of_raw_data - rw_mmap[new_section_header_offset + 16..new_section_header_offset + 20] - .copy_from_slice(&(section_size as u32).to_le_bytes()); - - // 写入新节的PointerToRawData - rw_mmap[new_section_header_offset + 20..new_section_header_offset + 24] - .copy_from_slice(&(new_section_offset as u32).to_le_bytes()); - - // 1.4 写入节的特征 - rw_mmap[new_section_header_offset + 36..new_section_header_offset + 40] - .copy_from_slice(§ion_characteristics.to_le_bytes()); - - // 修改文件头的节表数量 - rw_mmap[file_header_offset + 2..file_header_offset + 4] - .copy_from_slice(&(number_of_sections + 1).to_le_bytes()); - - let origin_size_of_image = unsafe { - // 在0x38的位置 - u32::from_le_bytes( - rw_mmap[optional_header_offset + 0x38..optional_header_offset + 0x3C] - .try_into() - .unwrap(), - ) - }; - let mut add_image_size = - (section_size as u32 + section_alignment - 1) & !(section_alignment - 1); - - if add_image_size == 0 { - add_image_size = section_alignment; - } - - rw_mmap[optional_header_offset + 0x38..optional_header_offset + 0x3C] - .copy_from_slice(&(origin_size_of_image + add_image_size).to_le_bytes()); - - rw_mmap.flush()?; - - // 2. 重新映射文件 + // 3. 重新映射文件 self.mmap = Some(file::mmap_file(self.file_path.as_ref().unwrap())?); + Ok(()) } -} +} \ No newline at end of file diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index c80a719..d9d130b 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -204,11 +204,11 @@ pub fn command_write_data( } // 命令，添加节 -#[tauri::command] +#[tauri::command(async)] pub fn command_add_section( app_state: State<'_, Mutex>, section_name: String, - section_size: usize, + section_size: u32, section_characteristics: u32, ) -> Result<(), AppError> { let mut app_state = app_state.lock().unwrap(); @@ -216,43 +216,7 @@ pub fn command_add_section( Ok(()) } -// 测试一下expand_headers好使不好使 #[tauri::command] -pub fn command_expand_headers( - app_state: State<'_, Mutex>, - expand_size: u32, -) -> Result<(), AppError> { - let app_state = app_state.lock().unwrap(); - // app_state.expand_headers(expand_size)?; - let mmap = app_state.get_mmap_ref()?; - // 获取文件对齐后的大小 - let file_alignment = mmap.align_size_with_file_alignment(expand_size)?; - // 获取最大可以拓展的大小 - let max_expand_size = mmap.get_header_max_expand_size()?; - if expand_size > max_expand_size { - return Err(AppError::MutablePeError( - MutablePEError::CannotExpandHeader(expand_size, max_expand_size), - )); - } - - let file_path = app_state.file_path.as_ref().unwrap(); - let mut file = std::fs::OpenOptions::new() - .read(true) - .write(true) - .open(file_path)?; - // 文件指针移动到文件末尾 - file.seek(std::io::SeekFrom::End(0))?; - // 拓展文件大小 - file.set_len(file.metadata()?.len() + file_alignment as u64)?; - file.seek(std::io::SeekFrom::Start(0))?; - file.flush()?; - // 映射 - let mut mut_mmap = unsafe { - MmapOptions::new() - .map_mut(&file) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))? - }; - mut_mmap.expand_headers(expand_size)?; - +pub fn command_test() -> Result<(), AppError> { Ok(()) } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index f4966e7..83d82b7 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -26,7 +26,7 @@ pub fn run() { commands::command_get_pe_data_section_headers, commands::command_write_data, commands::command_add_section, - commands::command_expand_headers, + commands::command_test, ]) .setup(|app| { app.manage(Mutex::new(app_state::AppState::default())); diff --git a/src-tauri/src/pe_parse/error.rs b/src-tauri/src/pe_parse/error.rs index 1cf148a..ca8b1b9 100644 --- a/src-tauri/src/pe_parse/error.rs +++ b/src-tauri/src/pe_parse/error.rs @@ -26,4 +26,5 @@ pub enum MutablePEError { // 文件大小不足，无法扩容 #[error("文件大小不足，无法扩容! 需要的大小:{0:#X}, 当前的文件大小:{1:#X}")] CannotExpandFileSize(u32, u32), + } diff --git a/src-tauri/src/pe_parse/pe.rs b/src-tauri/src/pe_parse/pe.rs index 5c20224..470211a 100644 --- a/src-tauri/src/pe_parse/pe.rs +++ b/src-tauri/src/pe_parse/pe.rs @@ -177,12 +177,18 @@ pub trait ReadOnlyPE: Deref + Sized + AsRef<[u8]> { // 最大可扩容大小 = 节对齐大小 - 最后一个节区头最后一个字节的偏移 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), + 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 + 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 向下对齐 @@ -191,7 +197,21 @@ pub trait ReadOnlyPE: Deref + Sized + AsRef<[u8]> { 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]> { @@ -205,6 +225,12 @@ pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { 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. 获取节区数量 @@ -300,6 +326,97 @@ pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { 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(); + // TODO: VA计算有问题 + 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; + // 新节区的size_of_raw_data是section_size对齐到file_alignment的值 + // 如果section_size是0，那么size_of_raw_data也是0 + new_section.size_of_raw_data = file_alignment; + // 新节区的pointer_to_raw_data是上一个节区的pointer_to_raw_data+上一个节区的size_of_raw_data + new_section.pointer_to_raw_data = + last_section.pointer_to_raw_data + last_section.size_of_raw_data; + 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-tauri/src/services/file.rs b/src-tauri/src/services/file.rs index fd11c18..c19b6e9 100644 --- a/src-tauri/src/services/file.rs +++ b/src-tauri/src/services/file.rs @@ -1,3 +1,5 @@ +use std::io::{Seek, Write}; + use memmap2::*; /// 以只读的方式创建文件映射 @@ -19,3 +21,18 @@ pub fn mmap_file_rw(file_path: &str) -> Result { .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string())) } } + +// 拓展文件大小 +pub fn expand_file_size(file_path: &str, expand_size: u32) -> Result<(), std::io::Error> { + let mut file = std::fs::OpenOptions::new() + .read(true) + .write(true) + .open(file_path)?; + // 文件指针移动到文件末尾 + file.seek(std::io::SeekFrom::End(0))?; + // 拓展文件大小 + file.set_len(file.metadata()?.len() + expand_size as u64)?; + file.seek(std::io::SeekFrom::Start(0))?; + file.flush()?; + Ok(()) +} diff --git a/src/components/SectionHeaders/SectionHeaders.tsx b/src/components/SectionHeaders/SectionHeaders.tsx index ebe5717..a378568 100644 --- a/src/components/SectionHeaders/SectionHeaders.tsx +++ b/src/components/SectionHeaders/SectionHeaders.tsx @@ -12,6 +12,7 @@ import { Input, InputNumber, Select, + message, // message, } from "antd"; import { PlusOutlined, EditOutlined } from "@ant-design/icons"; @@ -183,27 +184,27 @@ const sectionCharacteristics = [ label: "数据对齐到 256 字节边界", }, { - value: 0x00A00000, + value: 0x00a00000, label: "数据对齐到 512 字节边界", }, { - value: 0x00B00000, + value: 0x00b00000, label: "数据对齐到 1024 字节边界", }, { - value: 0x00C00000, + value: 0x00c00000, label: "数据对齐到 2048 字节边界", }, { - value: 0x00D00000, + value: 0x00d00000, label: "数据对齐到 4096 字节边界", }, { - value: 0x00E00000, + value: 0x00e00000, label: "数据对齐到 8192 字节边界", }, { - value: 0x00F00000, + value: 0x00f00000, label: "数据对齐掩码", }, { @@ -236,7 +237,6 @@ const sectionCharacteristics = [ }, ]; - // 节区头大小 const sectionHeaderSize = 40; @@ -321,9 +321,11 @@ export default function SectionHeaders() { requestData(); setAddSectionLoading(false); setAddSectionVisible(false); + message.success("添加节成功"); }) .catch((_err) => { setAddSectionLoading(false); + message.error(`添加节失败:${_err}`); }); }; @@ -353,6 +355,7 @@ export default function SectionHeaders() { @@ -360,7 +363,9 @@ export default function SectionHeaders() {