1use crate::{MapleError, Result};
2use std::mem;
3
4#[repr(C)]
5pub struct ImageDosHeader {
6 pub e_magic: u16,
7 pub e_cblp: u16,
8 pub e_cp: u16,
9 pub e_crlc: u16,
10 pub e_cparhdr: u16,
11 pub e_minalloc: u16,
12 pub e_maxalloc: u16,
13 pub e_ss: u16,
14 pub e_sp: u16,
15 pub e_csum: u16,
16 pub e_ip: u16,
17 pub e_cs: u16,
18 pub e_lfarlc: u16,
19 pub e_ovno: u16,
20 pub e_res: [u16; 4],
21 pub e_oemid: u16,
22 pub e_oeminfo: u16,
23 pub e_res2: [u16; 10],
24 pub e_lfanew: u32,
25}
26
27#[repr(C)]
28pub struct ImageFileHeader {
29 pub machine: u16,
30 pub number_of_sections: u16,
31 pub time_date_stamp: u32,
32 pub pointer_to_symbol_table: u32,
33 pub number_of_symbols: u32,
34 pub size_of_optional_header: u16,
35 pub characteristics: u16,
36}
37
38#[repr(C)]
39pub struct ImageOptionalHeader32 {
40 pub magic: u16,
41 pub major_linker_version: u8,
42 pub minor_linker_version: u8,
43 pub size_of_code: u32,
44 pub size_of_initialized_data: u32,
45 pub size_of_uninitialized_data: u32,
46 pub address_of_entry_point: u32,
47 pub base_of_code: u32,
48 pub base_of_data: u32,
49 pub image_base: u32,
50 pub section_alignment: u32,
51 pub file_alignment: u32,
52 pub major_operating_system_version: u16,
53 pub minor_operating_system_version: u16,
54 pub major_image_version: u16,
55 pub minor_image_version: u16,
56 pub major_subsystem_version: u16,
57 pub minor_subsystem_version: u16,
58 pub win32_version_value: u32,
59 pub size_of_image: u32,
60 pub size_of_headers: u32,
61 pub checksum: u32,
62 pub subsystem: u16,
63 pub dll_characteristics: u16,
64 pub size_of_stack_reserve: u32,
65 pub size_of_stack_commit: u32,
66 pub size_of_heap_reserve: u32,
67 pub size_of_heap_commit: u32,
68 pub loader_flags: u32,
69 pub number_of_rva_and_sizes: u32,
70}
71
72#[repr(C)]
73pub struct ImageOptionalHeader64 {
74 pub magic: u16,
75 pub major_linker_version: u8,
76 pub minor_linker_version: u8,
77 pub size_of_code: u32,
78 pub size_of_initialized_data: u32,
79 pub size_of_uninitialized_data: u32,
80 pub address_of_entry_point: u32,
81 pub base_of_code: u32,
82 pub image_base: u64,
83 pub section_alignment: u32,
84 pub file_alignment: u32,
85 pub major_operating_system_version: u16,
86 pub minor_operating_system_version: u16,
87 pub major_image_version: u16,
88 pub minor_image_version: u16,
89 pub major_subsystem_version: u16,
90 pub minor_subsystem_version: u16,
91 pub win32_version_value: u32,
92 pub size_of_image: u32,
93 pub size_of_headers: u32,
94 pub checksum: u32,
95 pub subsystem: u16,
96 pub dll_characteristics: u16,
97 pub size_of_stack_reserve: u64,
98 pub size_of_stack_commit: u64,
99 pub size_of_heap_reserve: u64,
100 pub size_of_heap_commit: u64,
101 pub loader_flags: u32,
102 pub number_of_rva_and_sizes: u32,
103}
104
105#[repr(C)]
106pub struct ImageDataDirectory {
107 pub virtual_address: u32,
108 pub size: u32,
109}
110
111#[repr(C)]
112pub struct ImageSectionHeader {
113 pub name: [u8; 8],
114 pub virtual_size: u32,
115 pub virtual_address: u32,
116 pub size_of_raw_data: u32,
117 pub pointer_to_raw_data: u32,
118 pub pointer_to_relocations: u32,
119 pub pointer_to_line_numbers: u32,
120 pub number_of_relocations: u16,
121 pub number_of_line_numbers: u16,
122 pub characteristics: u32,
123}
124
125#[repr(C)]
126pub struct ImageImportDescriptor {
127 pub original_first_thunk: u32,
128 pub time_date_stamp: u32,
129 pub forwarder_chain: u32,
130 pub name: u32,
131 pub first_thunk: u32,
132}
133
134#[repr(C)]
135pub struct ImageImportByName {
136 pub hint: u16,
137 pub name: [u8; 1],
138}
139
140#[repr(C)]
141pub struct ImageBaseRelocation {
142 pub virtual_address: u32,
143 pub size_of_block: u32,
144}
145
146#[repr(C)]
147pub struct ImageExportDirectory {
148 pub characteristics: u32,
149 pub time_date_stamp: u32,
150 pub major_version: u16,
151 pub minor_version: u16,
152 pub name: u32,
153 pub base: u32,
154 pub number_of_functions: u32,
155 pub number_of_names: u32,
156 pub address_of_functions: u32,
157 pub address_of_names: u32,
158 pub address_of_name_ordinals: u32,
159}
160
161#[repr(C)]
162pub struct ImageTlsDirectory32 {
163 pub start_address_of_raw_data: u32,
164 pub end_address_of_raw_data: u32,
165 pub address_of_index: u32,
166 pub address_of_callbacks: u32,
167 pub size_of_zero_fill: u32,
168 pub characteristics: u32,
169}
170
171#[repr(C)]
172pub struct ImageTlsDirectory64 {
173 pub start_address_of_raw_data: u64,
174 pub end_address_of_raw_data: u64,
175 pub address_of_index: u64,
176 pub address_of_callbacks: u64,
177 pub size_of_zero_fill: u32,
178 pub characteristics: u32,
179}
180
181#[repr(C)]
182pub struct ImageResourceDirectory {
183 pub characteristics: u32,
184 pub time_date_stamp: u32,
185 pub major_version: u16,
186 pub minor_version: u16,
187 pub number_of_name_entries: u16,
188 pub number_of_id_entries: u16,
189}
190
191#[repr(C)]
192pub struct ImageResourceDirectoryEntry {
193 pub name_or_id: u32,
194 pub offset_to_data_or_directory: u32,
195}
196
197#[repr(C)]
198pub struct ImageResourceDataEntry {
199 pub offset_to_data: u32,
200 pub size: u32,
201 pub code_page: u32,
202 pub reserved: u32,
203}
204
205pub const IMAGE_DOS_SIGNATURE: u16 = 0x5A4D;
206pub const IMAGE_NT_SIGNATURE: u32 = 0x00004550;
207pub const IMAGE_NT_OPTIONAL_HDR32_MAGIC: u16 = 0x010b;
208pub const IMAGE_NT_OPTIONAL_HDR64_MAGIC: u16 = 0x020b;
209
210pub const IMAGE_DIRECTORY_ENTRY_EXPORT: usize = 0;
211pub const IMAGE_DIRECTORY_ENTRY_IMPORT: usize = 1;
212pub const IMAGE_DIRECTORY_ENTRY_RESOURCE: usize = 2;
213pub const IMAGE_DIRECTORY_ENTRY_BASERELOC: usize = 5;
214pub const IMAGE_DIRECTORY_ENTRY_TLS: usize = 9;
215
216pub const IMAGE_REL_BASED_ABSOLUTE: u16 = 0;
217pub const IMAGE_REL_BASED_HIGH: u16 = 1;
218pub const IMAGE_REL_BASED_LOW: u16 = 2;
219pub const IMAGE_REL_BASED_HIGHLOW: u16 = 3;
220pub const IMAGE_REL_BASED_HIGHADJ: u16 = 4;
221pub const IMAGE_REL_BASED_DIR64: u16 = 10;
222
223pub const IMAGE_ORDINAL_FLAG32: u32 = 0x80000000;
224pub const IMAGE_ORDINAL_FLAG64: u64 = 0x8000000000000000;
225
226#[derive(Clone, Copy, Debug)]
227pub enum PEArchitecture {
228 PE32,
229 PE32Plus,
230}
231
232pub enum OptionalHeader<'a> {
233 PE32(&'a ImageOptionalHeader32),
234 PE32Plus(&'a ImageOptionalHeader64),
235}
236
237pub struct PEParser<'a> {
238 data: &'a [u8],
239 _dos_header: &'a ImageDosHeader,
240 _nt_headers_offset: usize,
241 file_header: &'a ImageFileHeader,
242 pub optional_header: &'a ImageOptionalHeader64,
243 data_directories: &'a [ImageDataDirectory],
244 sections: Vec<&'a ImageSectionHeader>,
245}
246
247impl<'a> PEParser<'a> {
248 pub fn new(data: &'a [u8]) -> Result<Self> {
249 if data.len() < mem::size_of::<ImageDosHeader>() {
250 return Err(MapleError::InvalidPEFormat(
251 "File too small for DOS header".to_string(),
252 ));
253 }
254
255 let dos_header = unsafe { &*(data.as_ptr() as *const ImageDosHeader) };
256
257 if dos_header.e_magic != IMAGE_DOS_SIGNATURE {
258 return Err(MapleError::InvalidPEFormat(
259 "Invalid DOS signature".to_string(),
260 ));
261 }
262
263 let nt_headers_offset = dos_header.e_lfanew as usize;
264 if nt_headers_offset + 4 > data.len() {
265 return Err(MapleError::InvalidPEFormat(
266 "Invalid NT headers offset".to_string(),
267 ));
268 }
269
270 let nt_signature = u32::from_le_bytes([
271 data[nt_headers_offset],
272 data[nt_headers_offset + 1],
273 data[nt_headers_offset + 2],
274 data[nt_headers_offset + 3],
275 ]);
276
277 if nt_signature != IMAGE_NT_SIGNATURE {
278 return Err(MapleError::InvalidPEFormat(
279 "Invalid NT signature".to_string(),
280 ));
281 }
282
283 let file_header_offset = nt_headers_offset + 4;
284 if file_header_offset + mem::size_of::<ImageFileHeader>() > data.len() {
285 return Err(MapleError::InvalidPEFormat(
286 "Invalid file header".to_string(),
287 ));
288 }
289
290 let file_header =
291 unsafe { &*(data[file_header_offset..].as_ptr() as *const ImageFileHeader) };
292
293 let optional_header_offset = file_header_offset + mem::size_of::<ImageFileHeader>();
294 if optional_header_offset + mem::size_of::<ImageOptionalHeader64>() > data.len() {
295 return Err(MapleError::InvalidPEFormat(
296 "Invalid optional header".to_string(),
297 ));
298 }
299
300 let optional_header =
301 unsafe { &*(data[optional_header_offset..].as_ptr() as *const ImageOptionalHeader64) };
302
303 if optional_header.magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC {
304 return Err(MapleError::InvalidPEFormat(
305 "Not a 64-bit executable".to_string(),
306 ));
307 }
308
309 let data_directories_offset =
310 optional_header_offset + mem::size_of::<ImageOptionalHeader64>();
311 let data_directories = unsafe {
312 std::slice::from_raw_parts(
313 data[data_directories_offset..].as_ptr() as *const ImageDataDirectory,
314 optional_header.number_of_rva_and_sizes as usize,
315 )
316 };
317
318 let sections_offset = optional_header_offset + file_header.size_of_optional_header as usize;
319 let mut sections = Vec::new();
320
321 for i in 0..file_header.number_of_sections {
322 let section_offset =
323 sections_offset + i as usize * mem::size_of::<ImageSectionHeader>();
324 if section_offset + mem::size_of::<ImageSectionHeader>() > data.len() {
325 return Err(MapleError::InvalidPEFormat(
326 "Invalid section header".to_string(),
327 ));
328 }
329
330 let section =
331 unsafe { &*(data[section_offset..].as_ptr() as *const ImageSectionHeader) };
332 sections.push(section);
333 }
334
335 Ok(PEParser {
336 data,
337 _dos_header: dos_header,
338 _nt_headers_offset: nt_headers_offset,
339 file_header,
340 optional_header,
341 data_directories,
342 sections,
343 })
344 }
345
346 pub fn image_base(&self) -> u64 {
347 self.optional_header.image_base
348 }
349
350 pub fn size_of_image(&self) -> u32 {
351 self.optional_header.size_of_image
352 }
353
354 pub fn entry_point(&self) -> u32 {
355 self.optional_header.address_of_entry_point
356 }
357
358 pub fn sections(&self) -> &[&ImageSectionHeader] {
359 &self.sections
360 }
361
362 pub fn data_directory(&self, index: usize) -> Option<&ImageDataDirectory> {
363 self.data_directories.get(index)
364 }
365
366 pub fn rva_to_offset(&self, rva: u32) -> Option<usize> {
367 for section in &self.sections {
368 if rva >= section.virtual_address
369 && rva < section.virtual_address + section.virtual_size
370 {
371 let offset = rva - section.virtual_address;
372 return Some(section.pointer_to_raw_data as usize + offset as usize);
373 }
374 }
375 None
376 }
377
378 pub fn get_data_at_rva(&self, rva: u32, size: usize) -> Option<&[u8]> {
379 let offset = self.rva_to_offset(rva)?;
380 if offset + size <= self.data.len() {
381 Some(&self.data[offset..offset + size])
382 } else {
383 None
384 }
385 }
386
387 pub fn get_import_directory(&self) -> Option<&ImageDataDirectory> {
388 self.data_directory(IMAGE_DIRECTORY_ENTRY_IMPORT)
389 }
390
391 pub fn get_export_directory(&self) -> Option<&ImageDataDirectory> {
392 self.data_directory(IMAGE_DIRECTORY_ENTRY_EXPORT)
393 }
394
395 pub fn get_base_relocation_directory(&self) -> Option<&ImageDataDirectory> {
396 self.data_directory(IMAGE_DIRECTORY_ENTRY_BASERELOC)
397 }
398
399 pub fn get_tls_directory(&self) -> Option<&ImageDataDirectory> {
400 self.data_directory(IMAGE_DIRECTORY_ENTRY_TLS)
401 }
402
403 pub fn get_resource_directory(&self) -> Option<&ImageDataDirectory> {
404 self.data_directory(IMAGE_DIRECTORY_ENTRY_RESOURCE)
405 }
406
407 pub fn is_dll(&self) -> bool {
408 const IMAGE_FILE_DLL: u16 = 0x2000;
409 (self.file_header.characteristics & IMAGE_FILE_DLL) != 0
410 }
411}