wraith/structures/pe/
nt_headers.rs

1//! NT headers (IMAGE_NT_HEADERS)
2
3use super::data_directory::DataDirectory;
4
5pub const NT_SIGNATURE: u32 = 0x00004550; // "PE\0\0"
6pub const PE32_MAGIC: u16 = 0x10b;
7pub const PE32PLUS_MAGIC: u16 = 0x20b;
8
9#[repr(C)]
10#[derive(Debug, Clone, Copy)]
11pub struct FileHeader {
12    pub machine: u16,
13    pub number_of_sections: u16,
14    pub time_date_stamp: u32,
15    pub pointer_to_symbol_table: u32,
16    pub number_of_symbols: u32,
17    pub size_of_optional_header: u16,
18    pub characteristics: u16,
19}
20
21#[repr(C)]
22#[derive(Debug, Clone, Copy)]
23pub struct OptionalHeader32 {
24    pub magic: u16,
25    pub major_linker_version: u8,
26    pub minor_linker_version: u8,
27    pub size_of_code: u32,
28    pub size_of_initialized_data: u32,
29    pub size_of_uninitialized_data: u32,
30    pub address_of_entry_point: u32,
31    pub base_of_code: u32,
32    pub base_of_data: u32,
33    pub image_base: u32,
34    pub section_alignment: u32,
35    pub file_alignment: u32,
36    pub major_operating_system_version: u16,
37    pub minor_operating_system_version: u16,
38    pub major_image_version: u16,
39    pub minor_image_version: u16,
40    pub major_subsystem_version: u16,
41    pub minor_subsystem_version: u16,
42    pub win32_version_value: u32,
43    pub size_of_image: u32,
44    pub size_of_headers: u32,
45    pub check_sum: u32,
46    pub subsystem: u16,
47    pub dll_characteristics: u16,
48    pub size_of_stack_reserve: u32,
49    pub size_of_stack_commit: u32,
50    pub size_of_heap_reserve: u32,
51    pub size_of_heap_commit: u32,
52    pub loader_flags: u32,
53    pub number_of_rva_and_sizes: u32,
54    pub data_directory: [DataDirectory; 16],
55}
56
57#[repr(C)]
58#[derive(Debug, Clone, Copy)]
59pub struct OptionalHeader64 {
60    pub magic: u16,
61    pub major_linker_version: u8,
62    pub minor_linker_version: u8,
63    pub size_of_code: u32,
64    pub size_of_initialized_data: u32,
65    pub size_of_uninitialized_data: u32,
66    pub address_of_entry_point: u32,
67    pub base_of_code: u32,
68    pub image_base: u64,
69    pub section_alignment: u32,
70    pub file_alignment: u32,
71    pub major_operating_system_version: u16,
72    pub minor_operating_system_version: u16,
73    pub major_image_version: u16,
74    pub minor_image_version: u16,
75    pub major_subsystem_version: u16,
76    pub minor_subsystem_version: u16,
77    pub win32_version_value: u32,
78    pub size_of_image: u32,
79    pub size_of_headers: u32,
80    pub check_sum: u32,
81    pub subsystem: u16,
82    pub dll_characteristics: u16,
83    pub size_of_stack_reserve: u64,
84    pub size_of_stack_commit: u64,
85    pub size_of_heap_reserve: u64,
86    pub size_of_heap_commit: u64,
87    pub loader_flags: u32,
88    pub number_of_rva_and_sizes: u32,
89    pub data_directory: [DataDirectory; 16],
90}
91
92#[repr(C)]
93#[derive(Debug, Clone, Copy)]
94pub struct NtHeaders32 {
95    pub signature: u32,
96    pub file_header: FileHeader,
97    pub optional_header: OptionalHeader32,
98}
99
100#[repr(C)]
101#[derive(Debug, Clone, Copy)]
102pub struct NtHeaders64 {
103    pub signature: u32,
104    pub file_header: FileHeader,
105    pub optional_header: OptionalHeader64,
106}
107
108/// architecture-independent NT headers access
109pub enum NtHeaders {
110    Headers32(NtHeaders32),
111    Headers64(NtHeaders64),
112}
113
114impl NtHeaders {
115    /// get number of sections
116    pub fn number_of_sections(&self) -> u16 {
117        match self {
118            Self::Headers32(h) => h.file_header.number_of_sections,
119            Self::Headers64(h) => h.file_header.number_of_sections,
120        }
121    }
122
123    /// get entry point RVA
124    pub fn entry_point(&self) -> u32 {
125        match self {
126            Self::Headers32(h) => h.optional_header.address_of_entry_point,
127            Self::Headers64(h) => h.optional_header.address_of_entry_point,
128        }
129    }
130
131    /// get image base
132    pub fn image_base(&self) -> u64 {
133        match self {
134            Self::Headers32(h) => h.optional_header.image_base as u64,
135            Self::Headers64(h) => h.optional_header.image_base,
136        }
137    }
138
139    /// get size of image
140    pub fn size_of_image(&self) -> u32 {
141        match self {
142            Self::Headers32(h) => h.optional_header.size_of_image,
143            Self::Headers64(h) => h.optional_header.size_of_image,
144        }
145    }
146
147    /// get section alignment
148    pub fn section_alignment(&self) -> u32 {
149        match self {
150            Self::Headers32(h) => h.optional_header.section_alignment,
151            Self::Headers64(h) => h.optional_header.section_alignment,
152        }
153    }
154
155    /// get data directory by index
156    pub fn data_directory(&self, index: usize) -> Option<&DataDirectory> {
157        match self {
158            Self::Headers32(h) => h.optional_header.data_directory.get(index),
159            Self::Headers64(h) => h.optional_header.data_directory.get(index),
160        }
161    }
162
163    /// check if 64-bit
164    pub fn is_64bit(&self) -> bool {
165        matches!(self, Self::Headers64(_))
166    }
167}