1use std::collections::HashMap;
2
3use capstone::{
4 arch::{self, BuildsCapstone},
5 Capstone, CsResult,
6};
7use keystone_engine::{Arch, Keystone, KeystoneError, Mode};
8use mlua::UserData;
9use object::{Architecture, Endianness};
10
11use crate::app::files::filesystem::FileSystem;
12
13use super::{
14 bitness::Bitness, custom_header::CustomHeader, generic::GenericHeader, section::Section,
15};
16
17#[derive(Debug, PartialEq, Eq, Clone, Default)]
18pub enum Header {
19 GenericHeader(GenericHeader),
20 CustomHeader(CustomHeader),
21 #[default]
22 None,
23}
24
25impl Header {
26 pub fn parse_header(bytes: &[u8], file_path: &str, filesystem: &FileSystem) -> Header {
27 let header = GenericHeader::parse_header(bytes, file_path, filesystem);
28 match header {
29 Some(header) => Header::GenericHeader(header),
30 None => Header::None,
31 }
32 }
33
34 pub fn bitness(&self) -> u32 {
35 match self {
36 Header::GenericHeader(header) => match header.bitness {
37 Bitness::Bit32 => 32,
38 Bitness::Bit64 => 64,
39 },
40 Header::CustomHeader(header) => match header.bitness {
41 Bitness::Bit32 => 32,
42 Bitness::Bit64 => 64,
43 },
44 Header::None => 64,
45 }
46 }
47
48 pub fn endianness(&self) -> Endianness {
49 match self {
50 Header::GenericHeader(header) => header.endianness,
51 Header::CustomHeader(header) => header.endianness,
52 Header::None => object::Endianness::Little,
53 }
54 }
55
56 pub fn entry_point(&self) -> u64 {
57 match self {
58 Header::GenericHeader(header) => header.entry,
59 Header::CustomHeader(header) => header.entry,
60 Header::None => 0,
61 }
62 }
63
64 pub fn architecture(&self) -> Architecture {
65 match self {
66 Header::GenericHeader(header) => header.architecture,
67 Header::CustomHeader(header) => header.architecture,
68 Header::None => Architecture::Unknown,
69 }
70 }
71
72 pub fn get_sections(&self) -> Vec<Section> {
73 match self {
74 Header::GenericHeader(header) => header.sections.clone(),
75 Header::CustomHeader(header) => header.sections.clone(),
76 Header::None => Vec::new(),
77 }
78 }
79
80 pub fn get_text_section(&self) -> Option<Section> {
81 match self {
82 Header::GenericHeader(header) => {
83 for section in &header.sections {
84 if section.name == ".text" || section.name == "__text" {
85 return Some(section.clone());
86 }
87 }
88 None
89 }
90 Header::CustomHeader(header) => {
91 for section in &header.sections {
92 if section.name == ".text" || section.name == "__text" {
93 return Some(section.clone());
94 }
95 }
96 None
97 }
98 Header::None => None,
99 }
100 }
101
102 pub fn get_symbols(&self) -> Option<&HashMap<u64, String>> {
103 match self {
104 Header::GenericHeader(header) => Some(&header.symbols),
105 Header::CustomHeader(header) => Some(&header.symbols),
106 Header::None => None,
107 }
108 }
109
110 pub fn symbol_to_address(&self, symbol: &str) -> Option<u64> {
111 match self {
112 Header::GenericHeader(header) => header.symbols_by_name.get(symbol).cloned(),
113 Header::CustomHeader(header) => header.symbols_by_name.get(symbol).cloned(),
114 Header::None => None,
115 }
116 }
117
118 pub fn virtual_to_physical_address(&self, virtual_address: u64) -> Option<u64> {
119 self.get_sections()
120 .iter()
121 .find(|x| {
122 virtual_address >= x.virtual_address && virtual_address < x.virtual_address + x.size
123 })
124 .map(|x| x.file_offset + virtual_address - x.virtual_address)
125 }
126
127 pub fn physical_to_virtual_address(&self, physical_address: u64) -> Option<u64> {
128 self.get_sections()
129 .iter()
130 .find(|x| {
131 physical_address >= x.file_offset && physical_address < x.file_offset + x.size
132 })
133 .map(|x| x.virtual_address + physical_address - x.file_offset)
134 }
135
136 pub(super) fn get_decoder_for_arch(architecture: &Architecture) -> CsResult<Capstone> {
137 match architecture {
138 Architecture::Aarch64 => Capstone::new()
139 .arm64()
140 .mode(arch::arm64::ArchMode::Arm)
141 .build(),
142 Architecture::Aarch64_Ilp32 => Capstone::new()
143 .arm64()
144 .mode(arch::arm64::ArchMode::Arm)
145 .build(),
146 Architecture::Arm => Capstone::new().arm().mode(arch::arm::ArchMode::Arm).build(),
147 Architecture::I386 => Capstone::new()
148 .x86()
149 .mode(arch::x86::ArchMode::Mode32)
150 .build(),
151 Architecture::X86_64 => Capstone::new()
152 .x86()
153 .mode(arch::x86::ArchMode::Mode64)
154 .build(),
155 Architecture::X86_64_X32 => Capstone::new()
156 .x86()
157 .mode(arch::x86::ArchMode::Mode64)
158 .build(),
159 Architecture::Mips => Capstone::new()
160 .mips()
161 .mode(arch::mips::ArchMode::Mips32)
162 .build(),
163 Architecture::Mips64 => Capstone::new()
164 .mips()
165 .mode(arch::mips::ArchMode::Mips64)
166 .build(),
167 Architecture::PowerPc => Capstone::new()
168 .ppc()
169 .mode(arch::ppc::ArchMode::Mode32)
170 .build(),
171 Architecture::PowerPc64 => Capstone::new()
172 .ppc()
173 .mode(arch::ppc::ArchMode::Mode64)
174 .build(),
175 Architecture::Riscv32 => Capstone::new()
176 .riscv()
177 .mode(arch::riscv::ArchMode::RiscV32)
178 .build(),
179 Architecture::Riscv64 => Capstone::new()
180 .riscv()
181 .mode(arch::riscv::ArchMode::RiscV64)
182 .build(),
183 Architecture::S390x => Capstone::new()
184 .sysz()
185 .mode(arch::sysz::ArchMode::Default)
186 .build(),
187 Architecture::Sparc64 => Capstone::new()
188 .sparc()
189 .mode(arch::sparc::ArchMode::V9)
190 .build(),
191 _ => Capstone::new()
192 .x86()
193 .mode(arch::x86::ArchMode::Mode64)
194 .build(),
195 }
196 }
197
198 pub(super) fn get_encoder_for_arch(
199 architecture: &Architecture,
200 ) -> Result<Keystone, KeystoneError> {
201 match architecture {
202 Architecture::Aarch64 => Keystone::new(Arch::ARM64, Mode::LITTLE_ENDIAN),
203 Architecture::Aarch64_Ilp32 => Keystone::new(Arch::ARM64, Mode::LITTLE_ENDIAN),
204 Architecture::Arm => Keystone::new(Arch::ARM, Mode::ARM),
205 Architecture::I386 => Keystone::new(Arch::X86, Mode::MODE_32),
206 Architecture::X86_64 => Keystone::new(Arch::X86, Mode::MODE_64),
207 Architecture::X86_64_X32 => Keystone::new(Arch::X86, Mode::MODE_32),
208 Architecture::Hexagon => Keystone::new(Arch::HEXAGON, Mode::MODE_32),
209 Architecture::Mips => Keystone::new(Arch::MIPS, Mode::MIPS32),
210 Architecture::Mips64 => Keystone::new(Arch::MIPS, Mode::MIPS64),
211 Architecture::PowerPc => Keystone::new(Arch::PPC, Mode::PPC32),
212 Architecture::PowerPc64 => Keystone::new(Arch::PPC, Mode::PPC64),
213 Architecture::S390x => Keystone::new(Arch::SYSTEMZ, Mode::MODE_32),
214 Architecture::Sparc64 => Keystone::new(Arch::SPARC, Mode::SPARC64),
215 _ => Keystone::new(Arch::X86, Mode::MODE_64),
216 }
217 }
218
219 pub fn get_decoder(&self) -> CsResult<Capstone> {
220 let ret = match self {
221 Header::GenericHeader(header) => Self::get_decoder_for_arch(&header.architecture),
222 Header::CustomHeader(header) => Self::get_decoder_for_arch(&header.architecture),
223 Header::None => Capstone::new()
224 .x86()
225 .mode(capstone::arch::x86::ArchMode::Mode64)
226 .build(),
227 };
228 ret.map(|mut cs| {
229 cs.set_skipdata(true).expect("Failed to set skipdata");
230 cs
231 })
232 }
233
234 pub fn get_encoder(&self) -> Result<Keystone, KeystoneError> {
235 match self {
236 Header::GenericHeader(header) => Self::get_encoder_for_arch(&header.architecture),
237 Header::CustomHeader(header) => Self::get_encoder_for_arch(&header.architecture),
238 Header::None => Keystone::new(Arch::X86, Mode::MODE_64),
239 }
240 }
241}
242
243impl UserData for Header {
244 fn add_fields<'lua, F: mlua::UserDataFields<Self>>(fields: &mut F) {
245 fields.add_field_method_get("bitness", |_, this| Ok(this.bitness()));
246 fields.add_field_method_get("entry_point", |_, this| Ok(this.entry_point()));
247 fields.add_field_method_get("architecture", |_, this| {
248 Ok(format!("{:?}", this.architecture()))
249 });
250 fields.add_field_method_get("sections", |_, this| Ok(this.get_sections()));
251 fields.add_field_method_get("text_section", |_, this| Ok(this.get_text_section()));
252 fields.add_field_method_get("symbols", |_, this| {
253 Ok(this
254 .get_symbols()
255 .map(|x| x.values().cloned().collect::<Vec<_>>()))
256 });
257 }
258
259 fn add_methods<'lua, M: mlua::UserDataMethods<Self>>(methods: &mut M) {
260 methods.add_method("symbol_to_address", |_, this, symbol: String| {
261 Ok(this.symbol_to_address(&symbol))
262 });
263 methods.add_method(
264 "virtual_to_physical_address",
265 |_, this, virtual_address: u64| Ok(this.virtual_to_physical_address(virtual_address)),
266 );
267 }
268}
269
270#[cfg(test)]
271mod test {
272 use crate::headers::generic::FileType;
273
274 use super::*;
275 #[test]
276 fn test_parse_elf() {
277 let data = include_bytes!("../../test/elf.bin");
278 let header = Header::parse_header(data, "./elf.bin", &FileSystem::new_local(".").unwrap());
279 if let Header::GenericHeader(header) = &header {
280 assert_eq!(header.file_type, FileType::Elf64);
281 } else {
282 panic!("Failed to parse ELF header.");
283 }
284 assert_eq!(header.architecture(), Architecture::X86_64);
285 assert_eq!(header.bitness(), 64);
286 assert_eq!(header.endianness(), Endianness::Little);
287 }
288
289 #[test]
290 fn test_parse_pe() {
291 let data = include_bytes!("../../test/pe.bin");
292 let header = Header::parse_header(data, "./pe.bin", &FileSystem::new_local(".").unwrap());
293 if let Header::GenericHeader(header) = &header {
294 assert_eq!(header.file_type, FileType::Pe64);
295 } else {
296 panic!("Failed to parse PE header.");
297 }
298 assert_eq!(header.architecture(), Architecture::X86_64);
299 assert_eq!(header.bitness(), 64);
300 assert_eq!(header.endianness(), Endianness::Little);
301 }
302
303 #[test]
304 fn test_parse_macho() {
305 let data = include_bytes!("../../test/macho.bin");
306 let header =
307 Header::parse_header(data, "./macho.bin", &FileSystem::new_local(".").unwrap());
308 if let Header::GenericHeader(header) = &header {
309 assert_eq!(header.file_type, FileType::MachO64);
310 } else {
311 panic!("Failed to parse Mach-O header.");
312 }
313 assert_eq!(header.architecture(), Architecture::X86_64);
314 assert_eq!(header.bitness(), 64);
315 assert_eq!(header.endianness(), Endianness::Little);
316 }
317}