Skip to main content

thalir_core/obfuscation/
deobfuscator.rs

1use super::ObfuscationMapping;
2use std::collections::HashMap;
3
4pub struct VulnerabilityMapper {
5    mapping: HashMap<String, String>,
6}
7
8impl VulnerabilityMapper {
9    pub fn from_mapping(mapping: ObfuscationMapping) -> Self {
10        Self {
11            mapping: mapping.mapping,
12        }
13    }
14
15    pub fn deobfuscate_identifier(&self, obfuscated: &str) -> Option<String> {
16        self.mapping.get(obfuscated).cloned()
17    }
18
19    pub fn deobfuscate_report(&self, report: &str) -> String {
20        let mut result = report.to_string();
21
22        let mut mappings: Vec<_> = self.mapping.iter().collect();
23        mappings.sort_by(|a, b| b.0.len().cmp(&a.0.len()));
24
25        for (obfuscated, original) in mappings {
26            result = result.replace(obfuscated, original);
27        }
28
29        result
30    }
31
32    pub fn deobfuscate_reports(&self, reports: &[String]) -> Vec<String> {
33        reports.iter().map(|r| self.deobfuscate_report(r)).collect()
34    }
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40    use crate::obfuscation::MappingMetadata;
41
42    fn create_test_mapping() -> ObfuscationMapping {
43        let mut mapping = HashMap::new();
44        mapping.insert("contract_0".to_string(), "NovelBondingCurve".to_string());
45        mapping.insert("fn_0".to_string(), "calculateBondingCurve".to_string());
46        mapping.insert("var_0".to_string(), "liquidityPoolReserves".to_string());
47        mapping.insert("fn_1".to_string(), "transfer".to_string());
48
49        ObfuscationMapping {
50            mapping,
51            metadata: MappingMetadata {
52                created_at: "2024-01-01T00:00:00Z".to_string(),
53                obfuscation_level: "minimal".to_string(),
54                hash_salt: None,
55            },
56        }
57    }
58
59    #[test]
60    fn test_deobfuscate_identifier() {
61        let mapping = create_test_mapping();
62        let mapper = VulnerabilityMapper::from_mapping(mapping);
63
64        assert_eq!(
65            mapper.deobfuscate_identifier("contract_0"),
66            Some("NovelBondingCurve".to_string())
67        );
68        assert_eq!(
69            mapper.deobfuscate_identifier("fn_0"),
70            Some("calculateBondingCurve".to_string())
71        );
72        assert_eq!(
73            mapper.deobfuscate_identifier("var_0"),
74            Some("liquidityPoolReserves".to_string())
75        );
76        assert_eq!(mapper.deobfuscate_identifier("unknown"), None);
77    }
78
79    #[test]
80    fn test_deobfuscate_simple_report() {
81        let mapping = create_test_mapping();
82        let mapper = VulnerabilityMapper::from_mapping(mapping);
83
84        let obfuscated = "Reentrancy in contract_0::fn_0 at position [7]";
85        let expected = "Reentrancy in NovelBondingCurve::calculateBondingCurve at position [7]";
86
87        assert_eq!(mapper.deobfuscate_report(obfuscated), expected);
88    }
89
90    #[test]
91    fn test_deobfuscate_complex_report() {
92        let mapping = create_test_mapping();
93        let mapper = VulnerabilityMapper::from_mapping(mapping);
94
95        let obfuscated = "Function fn_1 in contract_0 accesses var_0 unsafely";
96        let expected =
97            "Function transfer in NovelBondingCurve accesses liquidityPoolReserves unsafely";
98
99        assert_eq!(mapper.deobfuscate_report(obfuscated), expected);
100    }
101
102    #[test]
103    fn test_deobfuscate_multiple_reports() {
104        let mapping = create_test_mapping();
105        let mapper = VulnerabilityMapper::from_mapping(mapping);
106
107        let reports = vec![
108            "Issue in fn_0".to_string(),
109            "Problem with var_0".to_string(),
110            "contract_0 has vulnerability".to_string(),
111        ];
112
113        let deobfuscated = mapper.deobfuscate_reports(&reports);
114
115        assert_eq!(deobfuscated[0], "Issue in calculateBondingCurve");
116        assert_eq!(deobfuscated[1], "Problem with liquidityPoolReserves");
117        assert_eq!(deobfuscated[2], "NovelBondingCurve has vulnerability");
118    }
119
120    #[test]
121    fn test_deobfuscate_preserves_unknown_identifiers() {
122        let mapping = create_test_mapping();
123        let mapper = VulnerabilityMapper::from_mapping(mapping);
124
125        let obfuscated = "Unknown identifier unknown_fn in contract_0";
126        let result = mapper.deobfuscate_report(obfuscated);
127
128        assert!(result.contains("NovelBondingCurve"));
129        assert!(result.contains("unknown_fn"));
130    }
131
132    #[test]
133    fn test_deobfuscate_handles_position_markers() {
134        let mapping = create_test_mapping();
135        let mapper = VulnerabilityMapper::from_mapping(mapping);
136
137        let obfuscated = "At [5] in fn_0, call to contract_0::fn_1 detected";
138        let result = mapper.deobfuscate_report(obfuscated);
139
140        assert_eq!(
141            result,
142            "At [5] in calculateBondingCurve, call to NovelBondingCurve::transfer detected"
143        );
144    }
145}