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