1use std::collections::HashMap;
2
3use capstone::{
4 arch::{self, BuildsCapstone},
5 Capstone, CsResult,
6};
7use hexpatch_keystone::{Arch, Error, Keystone, 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(architecture: &Architecture) -> Result<Keystone, Error> {
199 match architecture {
200 Architecture::Aarch64 => Keystone::new(Arch::ARM64, Mode::LITTLE_ENDIAN),
201 Architecture::Aarch64_Ilp32 => Keystone::new(Arch::ARM64, Mode::LITTLE_ENDIAN),
202 Architecture::Arm => Keystone::new(Arch::ARM, Mode::ARM),
203 Architecture::I386 => Keystone::new(Arch::X86, Mode::MODE_32),
204 Architecture::X86_64 => Keystone::new(Arch::X86, Mode::MODE_64),
205 Architecture::X86_64_X32 => Keystone::new(Arch::X86, Mode::MODE_32),
206 Architecture::Hexagon => Keystone::new(Arch::HEXAGON, Mode::MODE_32),
207 Architecture::Mips => Keystone::new(Arch::MIPS, Mode::MIPS32),
208 Architecture::Mips64 => Keystone::new(Arch::MIPS, Mode::MIPS64),
209 Architecture::PowerPc => Keystone::new(Arch::PPC, Mode::PPC32),
210 Architecture::PowerPc64 => Keystone::new(Arch::PPC, Mode::PPC64),
211 Architecture::S390x => Keystone::new(Arch::SYSTEMZ, Mode::MODE_32),
212 Architecture::Sparc64 => Keystone::new(Arch::SPARC, Mode::SPARC64),
213 _ => Keystone::new(Arch::X86, Mode::MODE_64),
214 }
215 }
216
217 pub fn get_decoder(&self) -> CsResult<Capstone> {
218 let ret = match self {
219 Header::GenericHeader(header) => Self::get_decoder_for_arch(&header.architecture),
220 Header::CustomHeader(header) => Self::get_decoder_for_arch(&header.architecture),
221 Header::None => Capstone::new()
222 .x86()
223 .mode(capstone::arch::x86::ArchMode::Mode64)
224 .build(),
225 };
226 ret.map(|mut cs| {
227 cs.set_skipdata(true).expect(&t!("errors.set_skipdata"));
228 cs
229 })
230 }
231
232 pub fn get_encoder(&self) -> Result<Keystone, Error> {
233 match self {
234 Header::GenericHeader(header) => Self::get_encoder_for_arch(&header.architecture),
235 Header::CustomHeader(header) => Self::get_encoder_for_arch(&header.architecture),
236 Header::None => Keystone::new(Arch::X86, Mode::MODE_64),
237 }
238 }
239}
240
241impl UserData for Header {
242 fn add_fields<'lua, F: mlua::UserDataFields<Self>>(fields: &mut F) {
243 fields.add_field_method_get("bitness", |_, this| Ok(this.bitness()));
244 fields.add_field_method_get("entry_point", |_, this| Ok(this.entry_point()));
245 fields.add_field_method_get("architecture", |_, this| {
246 Ok(format!("{:?}", this.architecture()))
247 });
248 fields.add_field_method_get("sections", |_, this| Ok(this.get_sections()));
249 fields.add_field_method_get("text_section", |_, this| Ok(this.get_text_section()));
250 fields.add_field_method_get("symbols", |_, this| {
251 Ok(this
252 .get_symbols()
253 .map(|x| x.values().cloned().collect::<Vec<_>>()))
254 });
255 }
256
257 fn add_methods<'lua, M: mlua::UserDataMethods<Self>>(methods: &mut M) {
258 methods.add_method("symbol_to_address", |_, this, symbol: String| {
259 Ok(this.symbol_to_address(&symbol))
260 });
261 methods.add_method(
262 "virtual_to_physical_address",
263 |_, this, virtual_address: u64| Ok(this.virtual_to_physical_address(virtual_address)),
264 );
265 }
266}
267
268#[cfg(test)]
269mod test {
270 use crate::headers::generic::FileType;
271
272 use super::*;
273 #[test]
274 fn test_parse_elf() {
275 let data = include_bytes!("../../test/elf.bin");
276 let header = Header::parse_header(data, "./elf.bin", &FileSystem::new_local(".").unwrap());
277 if let Header::GenericHeader(header) = &header {
278 assert_eq!(header.file_type, FileType::Elf64);
279 } else {
280 panic!("Failed to parse ELF header.");
281 }
282 assert_eq!(header.architecture(), Architecture::X86_64);
283 assert_eq!(header.bitness(), 64);
284 assert_eq!(header.endianness(), Endianness::Little);
285 }
286
287 #[test]
288 fn test_parse_pe() {
289 let data = include_bytes!("../../test/pe.bin");
290 let header = Header::parse_header(data, "./pe.bin", &FileSystem::new_local(".").unwrap());
291 if let Header::GenericHeader(header) = &header {
292 assert_eq!(header.file_type, FileType::Pe64);
293 } else {
294 panic!("Failed to parse PE header.");
295 }
296 assert_eq!(header.architecture(), Architecture::X86_64);
297 assert_eq!(header.bitness(), 64);
298 assert_eq!(header.endianness(), Endianness::Little);
299 }
300
301 #[test]
302 fn test_parse_macho() {
303 let data = include_bytes!("../../test/macho.bin");
304 let header =
305 Header::parse_header(data, "./macho.bin", &FileSystem::new_local(".").unwrap());
306 if let Header::GenericHeader(header) = &header {
307 assert_eq!(header.file_type, FileType::MachO64);
308 } else {
309 panic!("Failed to parse Mach-O header.");
310 }
311 assert_eq!(header.architecture(), Architecture::X86_64);
312 assert_eq!(header.bitness(), 64);
313 assert_eq!(header.endianness(), Endianness::Little);
314 }
315}