hex_patch/app/plugins/
header_context.rs

1use std::collections::HashMap;
2
3use mlua::UserData;
4use object::{Architecture, Endianness};
5
6use crate::headers::{bitness::Bitness, custom_header::CustomHeader, section::Section};
7
8#[derive(Debug, Clone, Default)]
9pub struct HeaderContext {
10    pub bitness: Option<Bitness>,
11    pub endianness: Option<Endianness>,
12    pub entry: Option<u64>,
13    pub architecture: Option<Architecture>,
14    pub sections: Vec<Section>,
15    pub symbols: HashMap<u64, String>,
16}
17
18impl HeaderContext {
19    pub fn try_into_custom_header(self) -> Option<CustomHeader> {
20        if let (Some(bitness), Some(endianness), Some(entry), Some(architecture)) =
21            (self.bitness, self.endianness, self.entry, self.architecture)
22        {
23            let symbols_by_name = self.symbols.iter().map(|(k, v)| (v.clone(), *k)).collect();
24            Some(CustomHeader {
25                bitness,
26                entry,
27                endianness,
28                architecture,
29                sections: self.sections,
30                symbols: self.symbols,
31                symbols_by_name,
32            })
33        } else {
34            None
35        }
36    }
37}
38
39impl UserData for HeaderContext {
40    fn add_methods<'lua, M: mlua::UserDataMethods<Self>>(methods: &mut M) {
41        methods.add_method_mut("set_bitness", |_, this, bitness: u8| {
42            if this.bitness.is_some() {
43                Err(mlua::Error::external("bitness already set"))
44            } else {
45                match bitness {
46                    32 => {
47                        this.bitness = Some(Bitness::Bit32);
48                        Ok(())
49                    }
50                    64 => {
51                        this.bitness = Some(Bitness::Bit64);
52                        Ok(())
53                    }
54                    _ => Err(mlua::Error::external("invalid bitness")),
55                }
56            }
57        });
58
59        methods.add_method_mut("set_entry", |_, this, entry_point: u64| {
60            if this.entry.is_some() {
61                Err(mlua::Error::external("entry point already set"))
62            } else {
63                this.entry = Some(entry_point);
64                Ok(())
65            }
66        });
67
68        methods.add_method_mut("set_endianness", |_, this, endianness: String| {
69            if this.endianness.is_some() {
70                Err(mlua::Error::external("endianness already set"))
71            } else {
72                match endianness.as_str() {
73                    "little" => {
74                        this.endianness = Some(Endianness::Little);
75                        Ok(())
76                    }
77                    "big" => {
78                        this.endianness = Some(Endianness::Big);
79                        Ok(())
80                    }
81                    _ => Err(mlua::Error::external("invalid endianness")),
82                }
83            }
84        });
85
86        methods.add_method_mut("set_architecture", |_, this, architecture: String| {
87            if this.architecture.is_some() {
88                Err(mlua::Error::external("architecture already set"))
89            } else {
90                let architecture = match architecture.as_str() {
91                    "Unknown" => Architecture::Unknown,
92                    "Aarch64" => Architecture::Aarch64,
93                    "Aarch64_Ilp32" => Architecture::Aarch64_Ilp32,
94                    "Arm" => Architecture::Arm,
95                    "Avr" => Architecture::Avr,
96                    "Bpf" => Architecture::Bpf,
97                    "Csky" => Architecture::Csky,
98                    "I386" => Architecture::I386,
99                    "X86_64" => Architecture::X86_64,
100                    "X86_64_X32" => Architecture::X86_64_X32,
101                    "Hexagon" => Architecture::Hexagon,
102                    "LoongArch64" => Architecture::LoongArch64,
103                    "Mips" => Architecture::Mips,
104                    "Mips64" => Architecture::Mips64,
105                    "Msp430" => Architecture::Msp430,
106                    "PowerPc" => Architecture::PowerPc,
107                    "PowerPc64" => Architecture::PowerPc64,
108                    "Riscv32" => Architecture::Riscv32,
109                    "Riscv64" => Architecture::Riscv64,
110                    "S390x" => Architecture::S390x,
111                    "Sbf" => Architecture::Sbf,
112                    "Sharc" => Architecture::Sharc,
113                    "Sparc" => Architecture::Sparc,
114                    "Sparc32Plus" => Architecture::Sparc32Plus,
115                    "Sparc64" => Architecture::Sparc64,
116                    "Wasm32" => Architecture::Wasm32,
117                    "Wasm64" => Architecture::Wasm64,
118                    "Xtensa" => Architecture::Xtensa,
119                    _ => return Err(mlua::Error::external("invalid architecture")),
120                };
121                this.architecture = Some(architecture);
122                Ok(())
123            }
124        });
125
126        methods.add_method_mut(
127            "add_section",
128            |_, this, (name, virtual_address, file_offset, size): (String, u64, u64, u64)| {
129                this.sections.push(Section {
130                    name,
131                    virtual_address,
132                    file_offset,
133                    size,
134                });
135                Ok(())
136            },
137        );
138
139        methods.add_method_mut("add_symbol", |_, this, (address, name): (u64, String)| {
140            this.symbols.insert(address, name);
141            Ok(())
142        });
143    }
144}
hex_patch/app/plugins/header_context.rs

hex_patch/app/plugins/
header_context.rs
