thalir_core/obfuscation/
mapping_store.rs

1use super::NameObfuscator;
2use anyhow::Result;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::path::Path;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ObfuscationMapping {
9    pub mapping: HashMap<String, String>,
10    pub metadata: MappingMetadata,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct MappingMetadata {
15    pub created_at: String,
16    pub obfuscation_level: String,
17    pub hash_salt: Option<String>,
18}
19
20impl ObfuscationMapping {
21    pub fn from_obfuscator(obfuscator: &NameObfuscator) -> Self {
22        let mapping = obfuscator.export_mapping();
23
24        Self {
25            mapping,
26            metadata: MappingMetadata {
27                created_at: chrono::Utc::now().to_rfc3339(),
28                obfuscation_level: "standard".to_string(),
29                hash_salt: None,
30            },
31        }
32    }
33
34    pub fn save_to_file(&self, path: &Path) -> Result<()> {
35        let json = serde_json::to_string_pretty(self)?;
36        std::fs::write(path, json)?;
37        Ok(())
38    }
39
40    pub fn load_from_file(path: &Path) -> Result<Self> {
41        let json = std::fs::read_to_string(path)?;
42        let mapping: ObfuscationMapping = serde_json::from_str(&json)?;
43        Ok(mapping)
44    }
45
46    pub fn deobfuscate(&self, obfuscated: &str) -> Option<&str> {
47        self.mapping.get(obfuscated).map(|s| s.as_str())
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54    use crate::obfuscation::{ObfuscationConfig, ObfuscationLevel};
55    use std::fs;
56    use tempfile::NamedTempFile;
57
58    #[test]
59    fn test_mapping_serialization() {
60        let mut mapping = HashMap::new();
61        mapping.insert("contract_0".to_string(), "MyContract".to_string());
62        mapping.insert("fn_0".to_string(), "transfer".to_string());
63
64        let obf_mapping = ObfuscationMapping {
65            mapping,
66            metadata: MappingMetadata {
67                created_at: "2024-01-01T00:00:00Z".to_string(),
68                obfuscation_level: "minimal".to_string(),
69                hash_salt: None,
70            },
71        };
72
73        let json = serde_json::to_string_pretty(&obf_mapping).unwrap();
74        assert!(json.contains("MyContract"));
75        assert!(json.contains("transfer"));
76
77        let deserialized: ObfuscationMapping = serde_json::from_str(&json).unwrap();
78        assert_eq!(deserialized.deobfuscate("contract_0"), Some("MyContract"));
79        assert_eq!(deserialized.deobfuscate("fn_0"), Some("transfer"));
80    }
81
82    #[test]
83    fn test_file_save_and_load() {
84        let temp_file = NamedTempFile::new().unwrap();
85        let temp_path = temp_file.path();
86
87        let mut mapping = HashMap::new();
88        mapping.insert("c_abc123".to_string(), "TestContract".to_string());
89        mapping.insert("f_def456".to_string(), "testFunction".to_string());
90
91        let obf_mapping = ObfuscationMapping {
92            mapping,
93            metadata: MappingMetadata {
94                created_at: "2024-01-01T00:00:00Z".to_string(),
95                obfuscation_level: "standard".to_string(),
96                hash_salt: Some("test-salt".to_string()),
97            },
98        };
99
100        obf_mapping.save_to_file(temp_path).unwrap();
101
102        let loaded = ObfuscationMapping::load_from_file(temp_path).unwrap();
103
104        assert_eq!(loaded.deobfuscate("c_abc123"), Some("TestContract"));
105        assert_eq!(loaded.deobfuscate("f_def456"), Some("testFunction"));
106        assert_eq!(loaded.metadata.hash_salt, Some("test-salt".to_string()));
107    }
108
109    #[test]
110    fn test_from_obfuscator() {
111        let config = ObfuscationConfig {
112            level: ObfuscationLevel::Minimal,
113            retain_mapping: true,
114            ..Default::default()
115        };
116
117        let mut obfuscator = NameObfuscator::new(config);
118        obfuscator.obfuscate_contract_name("Contract1");
119        obfuscator.obfuscate_function_name("function1");
120        obfuscator.obfuscate_storage_name("storage1");
121
122        let mapping = ObfuscationMapping::from_obfuscator(&obfuscator);
123
124        assert_eq!(mapping.deobfuscate("contract_0"), Some("Contract1"));
125        assert_eq!(mapping.deobfuscate("fn_0"), Some("function1"));
126        assert_eq!(mapping.deobfuscate("var_0"), Some("storage1"));
127    }
128}
thalir_core/obfuscation/mapping_store.rs

thalir_core/obfuscation/
mapping_store.rs
