darklua_core/rules/rename_variables/
rename_processor.rs

1use crate::nodes::{Expression, Identifier, LocalFunctionStatement, TypeField};
2use crate::process::utils::{identifier_permutator, CharPermutator};
3use crate::process::{utils::KEYWORDS, NodeProcessor, Scope};
4
5use std::cmp::Ordering;
6use std::collections::{HashMap, HashSet};
7use std::iter::FromIterator;
8use std::mem;
9
10#[derive(Debug)]
11pub struct RenameProcessor {
12    real_to_obfuscated: Vec<HashMap<String, (String, bool)>>,
13    permutator: CharPermutator,
14    avoid_identifier: HashSet<String>,
15    reuse_identifiers: Vec<String>,
16    include_functions: bool,
17}
18
19impl RenameProcessor {
20    pub fn new<I: IntoIterator<Item = String>>(iter: I, include_functions: bool) -> Self {
21        let mut avoid_identifier = HashSet::from_iter(iter);
22        avoid_identifier.extend(KEYWORDS.iter().map(|s| (*s).to_owned()));
23
24        Self {
25            real_to_obfuscated: Vec::new(),
26            permutator: identifier_permutator(),
27            avoid_identifier,
28            reuse_identifiers: Vec::new(),
29            include_functions,
30        }
31    }
32
33    pub fn add(&mut self, real: String, obfuscated: String, reuse: bool) {
34        if let Some(dictionary) = self.real_to_obfuscated.last_mut() {
35            dictionary.insert(real, (obfuscated, reuse));
36        } else {
37            let mut dictionary = HashMap::new();
38            dictionary.insert(real, (obfuscated, reuse));
39            self.real_to_obfuscated.push(dictionary);
40        }
41    }
42
43    pub fn get_obfuscated_name(&self, real: &str) -> Option<&String> {
44        self.real_to_obfuscated
45            .iter()
46            .rev()
47            .find_map(|dictionary| dictionary.get(real).map(|(name, _)| name))
48    }
49
50    pub fn generate_identifier(&mut self) -> String {
51        if let Some(identifier) = self.reuse_identifiers.pop() {
52            identifier
53        } else {
54            let generated = self.permutator.next().unwrap();
55
56            if self.filter_identifier(&generated) {
57                generated
58            } else {
59                self.generate_identifier()
60            }
61        }
62    }
63
64    fn filter_identifier(&self, identifier: &str) -> bool {
65        !self.avoid_identifier.contains(identifier)
66            && !identifier.chars().next().unwrap().is_ascii_digit()
67    }
68
69    fn replace_identifier(&mut self, identifier: &mut String) {
70        let original = mem::take(identifier);
71        let obfuscated_name = self.generate_identifier();
72
73        identifier.push_str(&obfuscated_name);
74
75        self.add(original, obfuscated_name, true);
76    }
77}
78
79fn sort_char(a: char, b: char) -> Ordering {
80    if a == b {
81        Ordering::Equal
82    } else if a.is_ascii_digit() && b.is_ascii_digit()
83        || a.is_lowercase() && b.is_lowercase()
84        || a.is_uppercase() && b.is_uppercase()
85    {
86        a.cmp(&b)
87    } else {
88        match (a, b) {
89            (a, _) if a.is_ascii_digit() => Ordering::Greater,
90            (_, b) if b.is_ascii_digit() => Ordering::Less,
91            ('_', _) => Ordering::Greater,
92            (_, '_') => Ordering::Less,
93            (a, _) if a.is_lowercase() => Ordering::Less,
94            (_, b) if b.is_lowercase() => Ordering::Greater,
95            _ => Ordering::Equal,
96        }
97    }
98}
99
100fn sort_identifiers(a: &str, b: &str) -> Ordering {
101    let mut b_chars = b.chars();
102
103    for a_char in a.chars() {
104        if let Some(b_char) = b_chars.next() {
105            match sort_char(a_char, b_char) {
106                Ordering::Less => return Ordering::Less,
107                Ordering::Greater => return Ordering::Greater,
108                Ordering::Equal => {}
109            }
110        } else {
111            return Ordering::Greater;
112        }
113    }
114
115    if b_chars.next().is_some() {
116        Ordering::Less
117    } else {
118        Ordering::Equal
119    }
120}
121
122impl Scope for RenameProcessor {
123    fn push(&mut self) {
124        self.real_to_obfuscated.push(HashMap::new())
125    }
126
127    fn pop(&mut self) {
128        if let Some(dictionary) = self.real_to_obfuscated.pop() {
129            self.reuse_identifiers.extend(
130                dictionary
131                    .into_values()
132                    .filter_map(|(name, reuse)| reuse.then_some(name)),
133            );
134            self.reuse_identifiers
135                .sort_by(|a, b| sort_identifiers(a, b).reverse());
136        }
137    }
138
139    fn insert(&mut self, identifier: &mut String) {
140        self.replace_identifier(identifier);
141    }
142
143    fn insert_self(&mut self) {
144        self.add("self".to_owned(), "self".to_owned(), false);
145    }
146
147    fn insert_local(&mut self, identifier: &mut String, _value: Option<&mut Expression>) {
148        self.replace_identifier(identifier);
149    }
150
151    fn insert_local_function(&mut self, function: &mut LocalFunctionStatement) {
152        if self.include_functions {
153            self.replace_identifier(function.mutate_identifier().mutate_name());
154        } else {
155            let name = function.mutate_identifier().get_name();
156            self.add(name.clone(), name.to_owned(), false);
157        }
158    }
159}
160
161impl NodeProcessor for RenameProcessor {
162    fn process_variable_expression(&mut self, variable: &mut Identifier) {
163        if let Some(obfuscated_name) = self.get_obfuscated_name(variable.get_name()) {
164            variable.set_name(obfuscated_name);
165        }
166    }
167
168    fn process_type_field(&mut self, type_field: &mut TypeField) {
169        if let Some(obfuscated_name) =
170            self.get_obfuscated_name(type_field.get_namespace().get_name())
171        {
172            type_field.mutate_namespace().set_name(obfuscated_name);
173        }
174    }
175}
176
177#[cfg(test)]
178mod test {
179    use super::*;
180
181    fn new_scope() -> RenameProcessor {
182        RenameProcessor::new(Vec::new(), true)
183    }
184
185    #[test]
186    fn pop_root_should_not_panic() {
187        new_scope().pop();
188    }
189
190    #[test]
191    fn should_get_mapped_name_from_inserted_names() {
192        let mut scope = new_scope();
193        let real = "a".to_owned();
194        let obfuscated = "b".to_owned();
195
196        scope.add(real.clone(), obfuscated.clone(), true);
197
198        assert_eq!(&obfuscated, scope.get_obfuscated_name(&real).unwrap());
199    }
200
201    #[test]
202    fn mapped_name_should_not_exist_after_pop() {
203        let mut scope = new_scope();
204        let real = "a".to_owned();
205        let obfuscated = "def".to_owned();
206
207        scope.push();
208        scope.add(real.clone(), obfuscated, true);
209        scope.pop();
210
211        assert_eq!(None, scope.get_obfuscated_name(&real));
212    }
213
214    #[test]
215    fn remapped_name_should_exist_after_pop() {
216        let mut scope = new_scope();
217        let real = "a".to_owned();
218        let obfuscated = "b".to_owned();
219        let other_obfuscated = "c".to_owned();
220
221        scope.add(real.clone(), obfuscated.clone(), true);
222
223        scope.push();
224        scope.add(real.clone(), other_obfuscated, true);
225        scope.pop();
226
227        assert_eq!(&obfuscated, scope.get_obfuscated_name(&real).unwrap());
228    }
229
230    #[test]
231    fn sort_char_digits() {
232        assert_eq!(sort_char('0', '1'), Ordering::Less);
233        assert_eq!(sort_char('1', '2'), Ordering::Less);
234        assert_eq!(sort_char('4', '2'), Ordering::Greater);
235        assert_eq!(sort_char('5', '5'), Ordering::Equal);
236    }
237
238    #[test]
239    fn sort_char_lowercase_letters() {
240        assert_eq!(sort_char('a', 'f'), Ordering::Less);
241        assert_eq!(sort_char('y', 'i'), Ordering::Greater);
242        assert_eq!(sort_char('t', 't'), Ordering::Equal);
243    }
244
245    #[test]
246    fn sort_char_uppercase_letters() {
247        assert_eq!(sort_char('A', 'F'), Ordering::Less);
248        assert_eq!(sort_char('Y', 'I'), Ordering::Greater);
249        assert_eq!(sort_char('T', 'T'), Ordering::Equal);
250    }
251
252    #[test]
253    fn sort_char_underscore_is_less_than_digit() {
254        for digit in "0123456789".chars() {
255            assert_eq!(sort_char('_', digit), Ordering::Less);
256        }
257    }
258
259    #[test]
260    fn sort_char_underscore_is_greather_than_letters() {
261        assert_eq!(sort_char('_', 'a'), Ordering::Greater);
262        assert_eq!(sort_char('_', 'A'), Ordering::Greater);
263        assert_eq!(sort_char('a', '_'), Ordering::Less);
264        assert_eq!(sort_char('A', '_'), Ordering::Less);
265    }
266
267    #[test]
268    fn sort_char_lowercase_is_less_than_uppercase() {
269        assert_eq!(sort_char('a', 'A'), Ordering::Less);
270        assert_eq!(sort_char('A', 'a'), Ordering::Greater);
271        assert_eq!(sort_char('z', 'A'), Ordering::Less);
272        assert_eq!(sort_char('A', 'z'), Ordering::Greater);
273    }
274
275    #[test]
276    fn sort_identifiers_compare_chars() {
277        assert_eq!(sort_identifiers("foo", "foo"), Ordering::Equal);
278        assert_eq!(sort_identifiers("aA", "ab"), Ordering::Greater);
279        assert_eq!(sort_identifiers("foo1", "foo9"), Ordering::Less);
280    }
281
282    #[test]
283    fn sort_identifiers_shorter_is_less() {
284        assert_eq!(sort_identifiers("foo", "fooo"), Ordering::Less);
285    }
286
287    #[test]
288    fn sort_identifiers_longer_is_greather() {
289        assert_eq!(sort_identifiers("foo", "fo"), Ordering::Greater);
290    }
291}
darklua_core/rules/rename_variables/rename_processor.rs

darklua_core/rules/rename_variables/
rename_processor.rs
