feat: 增加节区、关闭文件

src-tauri/src/app_error.rs

@@ -25,6 +25,10 @@ pub enum AppError {
     /// 无法修改文件
     #[error("无法修改文件!: {0}")]
     CannotModifyFile(String),
+
+    /// 节名过长
+    #[error("节名过长!")]
+    SectionNameTooLong,
 }
 
 impl serde::Serialize for AppError {

src-tauri/src/app_state.rs

@@ -1,6 +1,9 @@
-use memmap2::Mmap;
+use std::io::Seek;
+use std::option;
+
 use crate::app_error::AppError;
 use crate::{pe_parse::pe::ReadOnlyPE, services::file};
+use memmap2::{Mmap, MmapOptions};
 
 /// 应用的状态结构体
 #[derive(Default)]
@@ -10,7 +13,7 @@ pub struct AppState {
     /// 文件Mmap的只读内存映射
     mmap: Option<Mmap>,
     /// 文件是否是64位
-    pub is_64_bit: Option<bool>
+    pub is_64_bit: Option<bool>,
 }
 
 impl AppState {
@@ -24,6 +27,14 @@ impl AppState {
         Ok(())
     }
 
+    /// close_file 关闭文件
+    pub fn close_file(&mut self) -> Result<(), AppError> {
+        self.file_path = None;
+        self.mmap = None;
+        self.is_64_bit = None;
+        Ok(())
+    }
+
     /// 获取文件内存映射的引用
     pub fn get_mmap_ref(&self) -> Result<&Mmap, AppError> {
         self.mmap.as_ref().ok_or(AppError::NoFileOpened)
@@ -50,4 +61,216 @@ impl AppState {
         Ok(())
     }
 
-}
+    /// 添加节
+    pub fn add_section(
+        &mut self,
+        section_name: &str,
+        section_size: usize,
+        section_characteristics: u32,
+    ) -> Result<(), AppError> {
+        // 1. section_name不能超过8个字符
+        if section_name.len() > 8 {
+            return Err(AppError::SectionNameTooLong);
+        }
+        // as_bytes 没有补0
+        let mut section_name_bytes = [0u8; 8];
+        // 循环拷贝section_name的数据
+        for i in 0..section_name.len() {
+            if i < section_name.len() {
+                section_name_bytes[i] = section_name.as_bytes()[i];
+            }
+        }
+
+        // 判断是否可以添加节
+        // 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 file_header_offset = mmap.get_file_header_offset()?;
+        let optional_header_offset = mmap.get_optional_header_offset()?;
+        let file_header = mmap.get_file_header()?;
+        let size_of_headers = match optional_header {
+            crate::pe_parse::header::ImageOptionalHeader::OptionalHeader32(
+                image_optional_header32,
+            ) => image_optional_header32.size_of_headers,
+            crate::pe_parse::header::ImageOptionalHeader::OptionalHeader64(
+                image_optional_header64,
+            ) => image_optional_header64.size_of_headers,
+        };
+        let section_table_offset = mmap.get_section_headers_offset()?;
+        // 获取节表的数量
+        let number_of_sections = file_header.number_of_sections;
+        // 1. 获取偏移
+        let section_alignment = match optional_header {
+            crate::pe_parse::header::ImageOptionalHeader::OptionalHeader32(
+                image_optional_header32,
+            ) => image_optional_header32.section_alignment,
+            crate::pe_parse::header::ImageOptionalHeader::OptionalHeader64(
+                image_optional_header64,
+            ) => image_optional_header64.section_alignment,
+        };
+
+        // 3. 判断是否可以添加节表
+        let section_table_size = 40 * number_of_sections as usize;
+        // TODO: 这里有问题
+        // 要判断第一个节表的pointer_to_raw_data - size_of_headers是否大于40
+        let first_section_header_offset = section_table_offset;
+        let first_pointer_to_raw_data = u32::from_le_bytes(
+            mmap[first_section_header_offset + 20..first_section_header_offset + 24]
+                .try_into()
+                .unwrap(),
+        );
+
+        if first_pointer_to_raw_data - size_of_headers < 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，则无法添加节表
+                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)?;
+
+            // 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());
+
+            rw_mmap.flush()?;
+            // 3. 重新映射文件
+            self.mmap = Some(file::mmap_file(self.file_path.as_ref().unwrap())?);
+        }
+        // 4. 添加节表
+        // 新的节区文件偏移是文件末尾
+        // TODO: 如果增加的大小为0是否有问题?
+        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)?;
+        }
+        // 写入节表数据
+        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(&section_name_bytes);
+        // 写入节的VA
+        // TODO: 新节的VA应该是最后一个节的VA + section_size对齐后的大小
+        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(),
+        );
+        // 对齐后的VA, 对齐方式(last_section_va + section_size) 和section_alignment进行对齐
+        let aligned_va = (last_section_va + section_size as u32 + 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(&section_characteristics.to_le_bytes());
+
+        // 修改文件头的节表数量
+        rw_mmap[file_header_offset + 2..file_header_offset + 4]
+            .copy_from_slice(&(number_of_sections + 1).to_le_bytes());
+
+        // 修改size_of_headers的大小
+        let new_size_of_headers = size_of_headers + 40;
+        rw_mmap[optional_header_offset + 0x3C..optional_header_offset + 0x40]
+            .copy_from_slice(&new_size_of_headers.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(),
+            )
+        };
+        // TODO: 增加的大小应该是section_size对齐后的大小
+        let add_image_size =
+            (section_size as u32 + 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).to_le_bytes());
+
+        rw_mmap.flush()?;
+
+        // 2. 重新映射文件
+        self.mmap = Some(file::mmap_file(self.file_path.as_ref().unwrap())?);
+        Ok(())
+    }
+}

src-tauri/src/commands.rs

@@ -43,6 +43,13 @@ pub fn command_open_file(
     Ok(())
 }
 
+// 命令，关闭文件
+#[tauri::command]
+pub fn command_close_file(app_state: State<'_, Mutex<AppState>>) -> Result<(), AppError> {
+    app_state.lock().unwrap().close_file()?;
+    Ok(())
+}
+
 #[tauri::command(async)]
 pub fn command_get_file_pe_node_tree_data(
     app_state: State<'_, Mutex<AppState>>,
@@ -180,4 +187,18 @@ pub fn command_write_data(
     let mut app_state = app_state.lock().unwrap();
     app_state.write_data(offset, &data)?;
     Ok(())
+}
+
+
+// 命令，添加节
+#[tauri::command]
+pub fn command_add_section(
+    app_state: State<'_, Mutex<AppState>>,
+    section_name: String,
+    section_size: usize,
+    section_characteristics: u32,
+) -> Result<(), AppError> {
+    let mut app_state = app_state.lock().unwrap();
+    app_state.add_section(&section_name, section_size, section_characteristics)?;
+    Ok(())
 }

src-tauri/src/lib.rs

@@ -17,6 +17,7 @@ pub fn run() {
         .invoke_handler(tauri::generate_handler![
             commands::set_complete,
             commands::command_open_file,
+            commands::command_close_file,
             commands::command_get_file_pe_node_tree_data,
             commands::command_get_pe_data_dos_header,
             commands::command_get_pe_data_nt_header,
@@ -24,6 +25,7 @@ pub fn run() {
             commands::command_get_pe_data_optional_header,
             commands::command_get_pe_data_section_headers,
             commands::command_write_data,
+            commands::command_add_section,
         ])
         .setup(|app| {
             app.manage(Mutex::new(app_state::AppState::default()));

src-tauri/src/pe_parse/pe.rs

@@ -103,14 +103,14 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized {
         };
         Ok(result)
     }
-
+    /// 获取可选头的偏移
     fn get_optional_header_offset(&self) -> Result<usize, PEParseError> {
         let file_header_offset = self.get_file_header_offset()?;
         let optional_header_offset = file_header_offset + std::mem::size_of::<ImageFileHeader>();
         Ok(optional_header_offset)
     }
 
-    // 获取节区头的偏移
+    /// 获取节区头的偏移
     fn get_section_headers_offset(&self) -> Result<usize, PEParseError> {
         // 节区头偏移在可选头之后，可选头大小是可变的，所以需要计算
         let optional_header_size = self.get_size_of_optional_header()?;

src-tauri/src/services/file.rs

@@ -18,4 +18,4 @@ pub fn mmap_file_rw(file_path: &str) -> Result<MmapMut, std::io::Error> {
             .map_mut(&file)
             .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))
     }
-}
+}