pe_assembler/helpers/pe_reader/
mod.rs

1use crate::types::{
2    CoffHeader, DosHeader, ExportTable, ImportTable, NtHeader, OptionalHeader, PeHeader, PeInfo, PeProgram, PeSection,
3    SectionHeader,
4};
5use byteorder::{LittleEndian, ReadBytesExt};
6use gaia_types::{helpers::Architecture, GaiaError};
7use std::io::{Read, Seek, SeekFrom};
8
9/// PE 文件读取器的通用 trait
10pub trait PeReader<R: Read + Seek> {
11    /// 获取二进制读取器的可变引用
12    fn get_viewer(&mut self) -> &mut R;
13
14    /// 获取诊断信息的可变引用
15    fn add_diagnostics(&mut self, error: impl Into<GaiaError>);
16
17    fn get_position(&mut self) -> Result<u64, GaiaError>
18    where
19        R: Seek,
20    {
21        let position = self.get_viewer().stream_position()?;
22        Ok(position)
23    }
24
25    fn set_position(&mut self, offset: u64) -> Result<u64, GaiaError>
26    where
27        R: Seek,
28    {
29        let position = self.get_viewer().seek(SeekFrom::Start(offset))?;
30        Ok(position)
31    }
32
33    /// 获取缓存的节头信息
34    fn get_section_headers(&mut self) -> Result<&[SectionHeader], GaiaError>;
35
36    /// 读取 PE 头部信息（通用实现）
37    fn get_pe_header(&mut self) -> Result<&PeHeader, GaiaError>;
38
39    /// 将 RVA 转换为文件偏移（通用实现）
40    fn rva_to_file_offset(&self, rva: u32, sections: &[PeSection]) -> Result<u32, GaiaError> {
41        for section in sections {
42            if rva >= section.virtual_address && rva < section.virtual_address + section.virtual_size {
43                let offset_in_section = rva - section.virtual_address;
44                return Ok(section.pointer_to_raw_data + offset_in_section);
45            }
46        }
47        Err(GaiaError::invalid_data(&format!("无法将 RVA 0x{:08X} 转换为文件偏移", rva)))
48    }
49
50    /// 解析导入表（通用实现）
51    fn parse_import_table(&mut self, header: &PeHeader, sections: &[PeSection]) -> Result<ImportTable, GaiaError> {
52        // 检查数据目录表是否包含导入表信息
53        if header.optional_header.data_directories.len() < 2 {
54            return Ok(ImportTable::new());
55        }
56
57        let import_dir = &header.optional_header.data_directories[1]; // 导入表是第2个数据目录
58        if import_dir.virtual_address == 0 || import_dir.size == 0 {
59            return Ok(ImportTable::new());
60        }
61
62        // 将 RVA 转换为文件偏移
63        let file_offset = self.rva_to_file_offset(import_dir.virtual_address, sections)?;
64
65        // 保存当前位置
66        let current_pos = self.get_position()?;
67
68        // 定位到导入表
69        self.set_position(file_offset as u64)?;
70
71        let mut import_table = ImportTable::new();
72
73        // 读取导入描述符
74        loop {
75            let import_lookup_table = self.get_viewer().read_u32::<LittleEndian>()?;
76            let time_date_stamp = self.get_viewer().read_u32::<LittleEndian>()?;
77            let forwarder_chain = self.get_viewer().read_u32::<LittleEndian>()?;
78            let name_rva = self.get_viewer().read_u32::<LittleEndian>()?;
79            let import_address_table = self.get_viewer().read_u32::<LittleEndian>()?;
80
81            // 如果所有字段都为0，表示导入表结束
82            if import_lookup_table == 0
83                && time_date_stamp == 0
84                && forwarder_chain == 0
85                && name_rva == 0
86                && import_address_table == 0
87            {
88                break;
89            }
90
91            let mut dll_name = String::new();
92            let mut functions = Vec::new();
93
94            // 读取 DLL 名称
95            if name_rva != 0 {
96                let name_offset = self.rva_to_file_offset(name_rva, sections)?;
97                let saved_pos = self.get_position()?;
98                self.set_position(name_offset as u64)?;
99
100                let mut name_bytes = Vec::new();
101                loop {
102                    let byte = self.get_viewer().read_u8()?;
103                    if byte == 0 {
104                        break;
105                    }
106                    name_bytes.push(byte);
107                }
108                dll_name = String::from_utf8_lossy(&name_bytes).to_string();
109                self.set_position(saved_pos)?;
110            }
111
112            // 读取函数名称（从导入查找表）
113            if import_lookup_table != 0 {
114                let lookup_offset = self.rva_to_file_offset(import_lookup_table, sections)?;
115                let saved_pos = self.get_position()?;
116                self.set_position(lookup_offset as u64)?;
117
118                loop {
119                    let entry = if header.optional_header.magic == 0x20b {
120                        // PE32+
121                        self.get_viewer().read_u64::<LittleEndian>()?
122                    }
123                    else {
124                        // PE32
125                        self.get_viewer().read_u32::<LittleEndian>()? as u64
126                    };
127
128                    if entry == 0 {
129                        break;
130                    }
131
132                    // 检查是否是按名称导入（最高位为0）
133                    let is_ordinal = if header.optional_header.magic == 0x20b {
134                        (entry & 0x8000000000000000) != 0
135                    }
136                    else {
137                        (entry & 0x80000000) != 0
138                    };
139
140                    if !is_ordinal {
141                        let hint_name_rva =
142                            entry & if header.optional_header.magic == 0x20b { 0x7FFFFFFFFFFFFFFF } else { 0x7FFFFFFF };
143                        let hint_name_offset = self.rva_to_file_offset(hint_name_rva as u32, sections)?;
144                        let func_pos = self.get_position()?;
145                        self.set_position(hint_name_offset as u64)?;
146
147                        // 跳过 hint（2字节）
148                        self.get_viewer().read_u16::<LittleEndian>()?;
149
150                        // 读取函数名
151                        let mut func_name_bytes = Vec::new();
152                        loop {
153                            let byte = self.get_viewer().read_u8()?;
154                            if byte == 0 {
155                                break;
156                            }
157                            func_name_bytes.push(byte);
158                        }
159                        let func_name = String::from_utf8_lossy(&func_name_bytes).to_string();
160                        functions.push(func_name);
161
162                        self.set_position(func_pos)?;
163                    }
164                    else {
165                        // 按序号导入
166                        let ordinal = entry & 0xFFFF;
167                        functions.push(format!("Ordinal_{}", ordinal));
168                    }
169                }
170
171                self.set_position(saved_pos)?;
172            }
173
174            // 添加导入条目
175            if !dll_name.is_empty() {
176                use crate::types::tables::ImportEntry;
177                let entry = ImportEntry { dll_name, functions };
178                import_table.entries.push(entry);
179            }
180        }
181
182        // 恢复位置
183        self.set_position(current_pos)?;
184
185        Ok(import_table)
186    }
187
188    /// 解析导出表（通用实现）
189    fn parse_export_table(&mut self, header: &PeHeader, sections: &[PeSection]) -> Result<ExportTable, GaiaError> {
190        // 检查数据目录表是否包含导出表信息
191        if header.optional_header.data_directories.is_empty() {
192            return Ok(ExportTable { name: String::new(), functions: Vec::new() });
193        }
194
195        let export_dir = &header.optional_header.data_directories[0]; // 导出表是第1个数据目录
196        if export_dir.virtual_address == 0 || export_dir.size == 0 {
197            return Ok(ExportTable { name: String::new(), functions: Vec::new() });
198        }
199
200        // 将 RVA 转换为文件偏移
201        let file_offset = self.rva_to_file_offset(export_dir.virtual_address, sections)?;
202
203        // 保存当前位置
204        let current_pos = self.get_position()?;
205
206        // 定位到导出表
207        self.set_position(file_offset as u64)?;
208
209        // 读取导出目录表
210        let _export_flags = self.get_viewer().read_u32::<LittleEndian>()?;
211        let _time_date_stamp = self.get_viewer().read_u32::<LittleEndian>()?;
212        let _major_version = self.get_viewer().read_u16::<LittleEndian>()?;
213        let _minor_version = self.get_viewer().read_u16::<LittleEndian>()?;
214        let name_rva = self.get_viewer().read_u32::<LittleEndian>()?;
215        let _ordinal_base = self.get_viewer().read_u32::<LittleEndian>()?;
216        let _number_of_functions = self.get_viewer().read_u32::<LittleEndian>()?;
217        let number_of_names = self.get_viewer().read_u32::<LittleEndian>()?;
218        let _address_of_functions = self.get_viewer().read_u32::<LittleEndian>()?;
219        let address_of_names = self.get_viewer().read_u32::<LittleEndian>()?;
220        let _address_of_name_ordinals = self.get_viewer().read_u32::<LittleEndian>()?;
221
222        // 读取模块名称
223        let mut name = String::new();
224        if name_rva != 0 {
225            let name_offset = self.rva_to_file_offset(name_rva, sections)?;
226            let saved_pos = self.get_position()?;
227            self.set_position(name_offset as u64)?;
228
229            let mut name_bytes = Vec::new();
230            loop {
231                let byte = self.get_viewer().read_u8()?;
232                if byte == 0 {
233                    break;
234                }
235                name_bytes.push(byte);
236            }
237            name = String::from_utf8_lossy(&name_bytes).to_string();
238            self.set_position(saved_pos)?;
239        }
240
241        // 读取函数名称
242        let mut functions = Vec::new();
243        if address_of_names != 0 && number_of_names > 0 {
244            let names_offset = self.rva_to_file_offset(address_of_names, sections)?;
245            let saved_pos = self.get_position()?;
246            self.set_position(names_offset as u64)?;
247
248            for _ in 0..number_of_names {
249                let name_rva = self.get_viewer().read_u32::<LittleEndian>()?;
250                if name_rva != 0 {
251                    let func_name_offset = self.rva_to_file_offset(name_rva, sections)?;
252                    let func_pos = self.get_position()?;
253                    self.set_position(func_name_offset as u64)?;
254
255                    let mut func_name_bytes = Vec::new();
256                    loop {
257                        let byte = self.get_viewer().read_u8()?;
258                        if byte == 0 {
259                            break;
260                        }
261                        func_name_bytes.push(byte);
262                    }
263                    let func_name = String::from_utf8_lossy(&func_name_bytes).to_string();
264                    functions.push(func_name);
265
266                    self.set_position(func_pos)?;
267                }
268            }
269
270            self.set_position(saved_pos)?;
271        }
272
273        // 恢复位置
274        self.set_position(current_pos)?;
275
276        Ok(ExportTable { name, functions })
277    }
278
279    /// 创建 PE 信息视图（通用实现）
280    fn create_pe_info(&mut self) -> Result<PeInfo, GaiaError> {
281        let header = self.get_pe_header()?.clone();
282
283        // 根据机器类型确定架构
284        let target_arch = match header.coff_header.machine {
285            0x014c => Architecture::X86,
286            0x8664 => Architecture::X86_64,
287            0x01c0 => Architecture::ARM32,
288            0xaa64 => Architecture::ARM64,
289            unknown => {
290                tracing::warn!("未知的机器类型: {:04x}", unknown);
291                Architecture::Unknown
292            }
293        };
294
295        // 获取当前文件大小
296        let current_pos = self.get_position()?;
297        self.get_viewer().seek(std::io::SeekFrom::End(0))?;
298        let file_size = self.get_position()?;
299        self.set_position(current_pos)?;
300
301        Ok(PeInfo {
302            target_arch,
303            subsystem: header.optional_header.subsystem,
304            entry_point: header.optional_header.address_of_entry_point,
305            image_base: header.optional_header.image_base,
306            section_count: header.coff_header.number_of_sections,
307            file_size,
308        })
309    }
310    /// 强制读取完整的 [PeProgram]，并缓存结果
311    fn get_program(&mut self) -> Result<&PeProgram, GaiaError>;
312}
313
314/// 解析 PE 头部（通用实现）
315pub fn read_pe_head<R: Read + Seek>(reader: &mut impl PeReader<R>) -> Result<PeHeader, GaiaError> {
316    // 保存当前位置
317    let original_pos = reader.get_position()?;
318
319    // 重置到文件开头
320    reader.set_position(0)?;
321
322    // 读取 DOS 头
323    let dos_header = DosHeader::read(reader.get_viewer())?;
324
325    // 验证 DOS 签名 (MZ)
326    if dos_header.e_magic != 0x5A4D {
327        let error = GaiaError::invalid_data("无效的 DOS 签名 (MZ)");
328        reader.add_diagnostics(error);
329    }
330
331    // 跳转到 NT 头位置
332    reader.set_position(dos_header.e_lfanew as u64)?;
333
334    // 读取 NT 头
335    let nt_header = NtHeader::read(reader.get_viewer())?;
336
337    // 验证 PE 签名 (PE\0\0)
338    if nt_header.signature != 0x00004550 {
339        let error = GaiaError::invalid_data("无效的 PE 签名 (PE)");
340        reader.add_diagnostics(error);
341    }
342
343    // 读取 COFF 头
344    let coff_header = CoffHeader::read(reader.get_viewer())?;
345
346    // 验证 COFF 头中的节数量
347    if coff_header.number_of_sections == 0 {
348        let error = GaiaError::invalid_data("PE 文件必须至少有一个节");
349        reader.add_diagnostics(error);
350    }
351
352    // 读取可选头
353    let optional_header = OptionalHeader::read(reader.get_viewer())?;
354
355    // 验证可选头的魔数
356    match optional_header.magic {
357        0x10b => {} // PE32
358        0x20b => {} // PE32+
359        _ => {
360            let error = GaiaError::invalid_data("无效的可选头魔数");
361            reader.add_diagnostics(error);
362            return Err(GaiaError::invalid_data("无效的可选头魔数"));
363        }
364    }
365
366    // 恢复原始位置
367    reader.set_position(original_pos)?;
368
369    Ok(PeHeader { dos_header, nt_header, coff_header, optional_header })
370}
371
372/// 读取节头信息（通用实现）
373pub fn read_pe_section_headers<R: Read + Seek>(reader: &mut impl PeReader<R>) -> Result<Vec<SectionHeader>, GaiaError> {
374    // 先读取主头部
375    let header = reader.get_pe_header()?.clone();
376    let original_pos = reader.get_position()?;
377
378    // 读取节头部
379    let mut section_headers = Vec::with_capacity(header.coff_header.number_of_sections as usize);
380
381    // 定位到节头部位置
382    let section_header_offset = header.dos_header.e_lfanew as u64
383        + 4 // PE signature
384        + std::mem::size_of::<CoffHeader>() as u64
385        + header.coff_header.size_of_optional_header as u64;
386
387    reader.set_position(section_header_offset)?;
388
389    for _ in 0..header.coff_header.number_of_sections {
390        let section_header = SectionHeader::read(reader.get_viewer())?;
391        section_headers.push(section_header);
392    }
393
394    // 恢复原始位置
395    reader.set_position(original_pos)?;
396
397    Ok(section_headers)
398}
399
400/// 从节头读取节数据（通用实现）
401pub fn read_section_from_header<R: Read + Seek>(
402    reader: &mut impl PeReader<R>,
403    header: &SectionHeader,
404) -> Result<PeSection, GaiaError> {
405    let name = String::from_utf8_lossy(&header.name).trim_end_matches('\0').to_string();
406
407    let mut data = Vec::new();
408    if header.size_of_raw_data > 0 && header.pointer_to_raw_data > 0 {
409        let original_pos = reader.get_position()?;
410        reader.set_position(header.pointer_to_raw_data as u64)?;
411        data.resize(header.size_of_raw_data as usize, 0);
412        reader.get_viewer().read_exact(&mut data)?;
413        reader.set_position(original_pos)?;
414    }
415
416    Ok(PeSection {
417        name,
418        virtual_size: header.virtual_size,
419        virtual_address: header.virtual_address,
420        size_of_raw_data: header.size_of_raw_data,
421        pointer_to_raw_data: header.pointer_to_raw_data,
422        pointer_to_relocations: header.pointer_to_relocations,
423        pointer_to_line_numbers: header.pointer_to_line_numbers,
424        number_of_relocations: header.number_of_relocations,
425        number_of_line_numbers: header.number_of_line_numbers,
426        characteristics: header.characteristics,
427        data,
428    })
429}
430
431pub fn read_pe_program<R: Read + Seek>(reader: &mut impl PeReader<R>) -> Result<PeProgram, GaiaError> {
432    let header = reader.get_pe_header()?.clone();
433    let section_headers = reader.get_section_headers()?.to_vec();
434
435    // 读取节数据
436    let mut sections = Vec::new();
437    for section_header in section_headers {
438        let section = read_section_from_header(reader, &section_header)?;
439        sections.push(section);
440    }
441
442    // 解析导入表
443    let imports = reader.parse_import_table(&header, &sections)?;
444
445    // 解析导出表（EXE 文件通常没有导出表）
446    let exports = reader.parse_export_table(&header, &sections)?;
447    Ok(PeProgram { header, sections, imports, exports })
448}
pe_assembler/helpers/pe_reader/mod.rs

pe_assembler/helpers/pe_reader/
mod.rs
