288 lines
12 KiB
Rust
288 lines
12 KiB
Rust
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)]
|
||
pub struct AppState {
|
||
/// 文件路径
|
||
pub file_path: Option<String>,
|
||
/// 文件Mmap的只读内存映射
|
||
mmap: Option<Mmap>,
|
||
/// 文件是否是64位
|
||
pub is_64_bit: Option<bool>,
|
||
}
|
||
|
||
impl AppState {
|
||
/// 初始化打开文件
|
||
pub fn init_open_file(&mut self, file_path: &str) -> Result<(), AppError> {
|
||
self.file_path = Some(file_path.to_string());
|
||
self.mmap = Some(file::mmap_file(file_path)?);
|
||
// 读取PE格式来判断是否是64位
|
||
let mmap: &Mmap = self.mmap.as_ref().unwrap();
|
||
self.is_64_bit = Some(mmap.is_64_bit()?);
|
||
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)
|
||
}
|
||
|
||
/// 设置文件内存映射
|
||
pub fn set_mmap(&mut self, mmap: Mmap) {
|
||
self.mmap = Some(mmap);
|
||
}
|
||
|
||
/// 在指定偏移处写入数据，此举会同步修改文件内容
|
||
pub fn write_data(&mut self, offset: usize, data: &[u8]) -> Result<(), AppError> {
|
||
// 1. 重新以可读可写打开文件
|
||
let file_path = self.file_path.as_ref().ok_or(AppError::NoFileOpened)?;
|
||
let mut rw_mmap = file::mmap_file_rw(file_path)?;
|
||
// 2. 向内存映射中写入数据
|
||
rw_mmap[offset..(offset + data.len())].copy_from_slice(data);
|
||
rw_mmap.flush()?;
|
||
|
||
// 3. 重新打开文件
|
||
self.mmap = Some(file::mmap_file(file_path)?);
|
||
self.is_64_bit = Some(self.get_mmap_ref()?.is_64_bit()?);
|
||
|
||
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;
|
||
|
||
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，则无法添加节表
|
||
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. 添加节表
|
||
// 新的节区文件偏移是文件末尾
|
||
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(§ion_name_bytes);
|
||
|
||
// virtual_size需要对齐
|
||
let mut virtual_size = ((section_size as u32) + section_alignment - 1) & !(section_alignment - 1);
|
||
if virtual_size == 0 {
|
||
virtual_size = section_alignment;
|
||
}
|
||
|
||
rw_mmap[new_section_header_offset + 8..new_section_header_offset + 12]
|
||
.copy_from_slice(&(virtual_size as u32).to_le_bytes());
|
||
|
||
|
||
// 写入节的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());
|
||
|
||
// 修改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());
|
||
|
||
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. 重新映射文件
|
||
self.mmap = Some(file::mmap_file(self.file_path.as_ref().unwrap())?);
|
||
Ok(())
|
||
}
|
||
}
|