darklua_core/rules/rename_variables/
rename_processor.rs

1use crate::nodes::{Expression, Identifier, LocalFunctionStatement};
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
169#[cfg(test)]
170mod test {
171    use super::*;
172
173    fn new_scope() -> RenameProcessor {
174        RenameProcessor::new(Vec::new(), true)
175    }
176
177    #[test]
178    fn pop_root_should_not_panic() {
179        new_scope().pop();
180    }
181
182    #[test]
183    fn should_get_mapped_name_from_inserted_names() {
184        let mut scope = new_scope();
185        let real = "a".to_owned();
186        let obfuscated = "b".to_owned();
187
188        scope.add(real.clone(), obfuscated.clone(), true);
189
190        assert_eq!(&obfuscated, scope.get_obfuscated_name(&real).unwrap());
191    }
192
193    #[test]
194    fn mapped_name_should_not_exist_after_pop() {
195        let mut scope = new_scope();
196        let real = "a".to_owned();
197        let obfuscated = "def".to_owned();
198
199        scope.push();
200        scope.add(real.clone(), obfuscated, true);
201        scope.pop();
202
203        assert_eq!(None, scope.get_obfuscated_name(&real));
204    }
205
206    #[test]
207    fn remapped_name_should_exist_after_pop() {
208        let mut scope = new_scope();
209        let real = "a".to_owned();
210        let obfuscated = "b".to_owned();
211        let other_obfuscated = "c".to_owned();
212
213        scope.add(real.clone(), obfuscated.clone(), true);
214
215        scope.push();
216        scope.add(real.clone(), other_obfuscated, true);
217        scope.pop();
218
219        assert_eq!(&obfuscated, scope.get_obfuscated_name(&real).unwrap());
220    }
221
222    #[test]
223    fn sort_char_digits() {
224        assert_eq!(sort_char('0', '1'), Ordering::Less);
225        assert_eq!(sort_char('1', '2'), Ordering::Less);
226        assert_eq!(sort_char('4', '2'), Ordering::Greater);
227        assert_eq!(sort_char('5', '5'), Ordering::Equal);
228    }
229
230    #[test]
231    fn sort_char_lowercase_letters() {
232        assert_eq!(sort_char('a', 'f'), Ordering::Less);
233        assert_eq!(sort_char('y', 'i'), Ordering::Greater);
234        assert_eq!(sort_char('t', 't'), Ordering::Equal);
235    }
236
237    #[test]
238    fn sort_char_uppercase_letters() {
239        assert_eq!(sort_char('A', 'F'), Ordering::Less);
240        assert_eq!(sort_char('Y', 'I'), Ordering::Greater);
241        assert_eq!(sort_char('T', 'T'), Ordering::Equal);
242    }
243
244    #[test]
245    fn sort_char_underscore_is_less_than_digit() {
246        for digit in "0123456789".chars() {
247            assert_eq!(sort_char('_', digit), Ordering::Less);
248        }
249    }
250
251    #[test]
252    fn sort_char_underscore_is_greather_than_letters() {
253        assert_eq!(sort_char('_', 'a'), Ordering::Greater);
254        assert_eq!(sort_char('_', 'A'), Ordering::Greater);
255        assert_eq!(sort_char('a', '_'), Ordering::Less);
256        assert_eq!(sort_char('A', '_'), Ordering::Less);
257    }
258
259    #[test]
260    fn sort_char_lowercase_is_less_than_uppercase() {
261        assert_eq!(sort_char('a', 'A'), Ordering::Less);
262        assert_eq!(sort_char('A', 'a'), Ordering::Greater);
263        assert_eq!(sort_char('z', 'A'), Ordering::Less);
264        assert_eq!(sort_char('A', 'z'), Ordering::Greater);
265    }
266
267    #[test]
268    fn sort_identifiers_compare_chars() {
269        assert_eq!(sort_identifiers("foo", "foo"), Ordering::Equal);
270        assert_eq!(sort_identifiers("aA", "ab"), Ordering::Greater);
271        assert_eq!(sort_identifiers("foo1", "foo9"), Ordering::Less);
272    }
273
274    #[test]
275    fn sort_identifiers_shorter_is_less() {
276        assert_eq!(sort_identifiers("foo", "fooo"), Ordering::Less);
277    }
278
279    #[test]
280    fn sort_identifiers_longer_is_greather() {
281        assert_eq!(sort_identifiers("foo", "fo"), Ordering::Greater);
282    }
283}
darklua_core/rules/rename_variables/rename_processor.rs

darklua_core/rules/rename_variables/
rename_processor.rs
