thalir_core/obfuscation/
deobfuscator.rs1use 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}