keeper_secrets_manager_core/
storage.rs

1// -*- coding: utf-8 -*-
2//  _  __
3// | |/ /___ ___ _ __  ___ _ _ (R)
4// | ' </ -_) -_) '_ \/ -_) '_|
5// |_|\_\___\___| .__/\___|_|
6//              |_|
7//
8// Keeper Secrets Manager
9// Copyright 2024 Keeper Security Inc.
10// Contact: sm@keepersecurity.com
11//
12
13use crate::config_keys::ConfigKeys;
14use crate::custom_error::KSMRError;
15use crate::enums::KvStoreType;
16use base64::{engine::general_purpose::{STANDARD, STANDARD_NO_PAD}, Engine as _};
17use serde_json::{self};
18use std::collections::HashMap;
19use std::fs::{File, OpenOptions};
20use std::io::{BufReader, Read, Write};
21use std::path::Path;
22use std::{env, fs};
23
24pub trait KeyValueStorage {
25    fn read_storage(&self) -> Result<HashMap<ConfigKeys, String>, KSMRError>;
26    fn save_storage(
27        &mut self,
28        updated_config: HashMap<ConfigKeys, String>,
29    ) -> Result<bool, KSMRError>;
30    fn get(&self, key: ConfigKeys) -> Result<Option<String>, KSMRError>;
31    fn set(
32        &mut self,
33        key: ConfigKeys,
34        value: String,
35    ) -> Result<HashMap<ConfigKeys, String>, KSMRError>;
36    fn delete(&mut self, key: ConfigKeys) -> Result<HashMap<ConfigKeys, String>, KSMRError>;
37    fn delete_all(&mut self) -> Result<HashMap<ConfigKeys, String>, KSMRError>;
38    fn contains(&self, key: ConfigKeys) -> Result<bool, KSMRError>;
39    fn create_config_file_if_missing(&self) -> Result<(), KSMRError>;
40    fn is_empty(&self) -> Result<bool, KSMRError>;
41}
42
43#[derive(Clone)]
44pub struct FileKeyValueStorage {
45    config_file_location: String,
46}
47
48impl FileKeyValueStorage {
49    const DEFAULT_CONFIG_FILE_LOCATION: &str = "client-config.json";
50    pub fn new(config_file_location: Option<String>) -> Result<Self, KSMRError> {
51        let location = config_file_location
52            .or_else(|| env::var("KSM_CONFIG_FILE").ok())
53            .unwrap_or_else(|| Self::DEFAULT_CONFIG_FILE_LOCATION.to_string());
54
55        Ok(FileKeyValueStorage {
56            config_file_location: location,
57        })
58    }
59
60    pub fn new_config_storage(file_name: String) -> Result<KvStoreType, KSMRError> {
61        let file_storage = FileKeyValueStorage::new(Some(file_name.to_string()))?;
62        Ok(KvStoreType::File(file_storage))
63    }
64}
65
66impl KeyValueStorage for FileKeyValueStorage {
67    fn read_storage(&self) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
68        // Check if the config file exists, create it if necessary
69        self.create_config_file_if_missing().map_err(|err| {
70            KSMRError::StorageError(format!("Failed to ensure config file exists: {}", err))
71        })?;
72
73        // Check if the file can be opened
74        let file = File::open(&self.config_file_location).map_err(|err| {
75            KSMRError::StorageError(format!(
76                "Unable to open config file {}: {}",
77                self.config_file_location, err
78            ))
79        })?;
80
81        // Read file contents into buffer
82        let mut reader = BufReader::new(file);
83        let mut contents = String::new();
84        reader
85            .read_to_string(&mut contents)
86            .map_err(|err| KSMRError::StorageError(format!("Failed to read file: {}", err)))?;
87
88        // Deserialize the string to JSON
89        let config_result: Result<HashMap<ConfigKeys, String>, KSMRError> =
90            serde_json::from_str(&contents)
91                .map_err(|err| KSMRError::StorageError(format!("Failed to parse JSON: {}", err)));
92
93        match config_result {
94            Ok(config) => Ok(config),
95            Err(err) => {
96                // Print the error details in case JSON parsing fails
97                eprintln!("Failed to parse JSON: {}", err);
98                Err(KSMRError::StorageError(format!(
99                    "Failed to parse JSON: {}",
100                    err
101                )))
102            }
103        }
104    }
105
106    fn save_storage(
107        &mut self,
108        updated_config: HashMap<ConfigKeys, String>,
109    ) -> Result<bool, KSMRError> {
110        // Ensure the config file exists, create it if missing
111        self.create_config_file_if_missing().map_err(|err| {
112            KSMRError::StorageError(format!("Failed to ensure config file exists: {}", err))
113        })?;
114
115        // Open the file in write mode and truncate it
116        let mut file = OpenOptions::new()
117            .write(true)
118            .truncate(true) // Clear the file before writing
119            .open(&self.config_file_location)
120            .map_err(|err| {
121                KSMRError::StorageError(format!("Failed to open config file for writing: {}", err))
122            })?;
123
124        // Serialize the updated config to JSON
125        let json_data = serde_json::to_string_pretty(&updated_config).map_err(|err| {
126            KSMRError::StorageError(format!("Failed to serialize config to JSON: {}", err))
127        })?;
128
129        // Write the JSON data to the file
130        file.write_all(json_data.as_bytes()).map_err(|err| {
131            KSMRError::StorageError(format!("Failed to write JSON to config file: {}", err))
132        })?;
133
134        Ok(true)
135    }
136
137    fn get(&self, key: ConfigKeys) -> Result<Option<String>, KSMRError> {
138        let config: HashMap<ConfigKeys, String> = self
139            .read_storage()
140            .map_err(|err| KSMRError::StorageError(format!("Failed to Read storage: {}", err)))?;
141
142        // Return the value corresponding to the key, cloning the String to give ownership
143        Ok(config.get(&key).cloned())
144    }
145
146    fn set(
147        &mut self,
148        key: ConfigKeys,
149        value: String,
150    ) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
151        // Check if the key is valid
152        if ConfigKeys::get_enum(key.value()).is_none() {
153            return Err(KSMRError::StorageError(format!("Invalid key: {:?}", key)));
154        }
155
156        // Read the current configuration
157        let mut config = self
158            .read_storage()
159            .map_err(|err| KSMRError::StorageError(format!("Failed to read storage: {}", err)))?;
160
161        // Update the value for the given key
162        config.insert(key, value);
163
164        // Save the updated configuration
165        self.save_storage(config.clone()).map_err(|err| {
166            KSMRError::StorageError(format!("Failed to save updated config: {}", err))
167        })?;
168
169        Ok(config) // Return the updated config
170    }
171
172    fn delete(&mut self, key: ConfigKeys) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
173        // Read the current configuration
174        let mut config = self
175            .read_storage()
176            .map_err(|err| KSMRError::StorageError(format!("Failed to read storage: {}", err)))?;
177
178        // Check if the key exists in the config and remove it
179        if config.remove(&key).is_some() {
180            log::debug!("Removed key {}", key);
181        } else {
182            log::debug!("No key {} was found in config", key);
183        }
184
185        // Save the updated configuration
186        self.save_storage(config.clone()).map_err(|err| {
187            KSMRError::StorageError(format!("Failed to save updated config: {}", err))
188        })?;
189
190        Ok(config) // Return the updated config
191    }
192
193    fn delete_all(&mut self) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
194        // Check if we are able to read storage and read from storage
195        let mut config = self
196            .read_storage()
197            .map_err(|e| KSMRError::StorageError(format!("Failed to read storage: {}", e)))?;
198
199        // Clear the configuration
200        config.clear();
201
202        // Save the cleared configuration
203        self.save_storage(config.clone()).map_err(|e| {
204            KSMRError::StorageError(format!("Failed to save cleared config: {}", e))
205        })?;
206
207        Ok(config) // Return the cleared config
208    }
209
210    fn contains(&self, key: ConfigKeys) -> Result<bool, KSMRError> {
211        // Read the current configuration
212        let config = self
213            .read_storage()
214            .map_err(|e| KSMRError::StorageError(format!("Failed to read storage: {}", e)))?;
215
216        // Check if the key exists in the config
217        Ok(config.contains_key(&key))
218    }
219
220    fn create_config_file_if_missing(&self) -> Result<(), KSMRError> {
221        // Check if parent directory exists, if not, create it.
222        if let Some(parent) = Path::new(&self.config_file_location).parent() {
223            fs::create_dir_all(parent)
224                .map_err(|e| KSMRError::DirectoryCreationError(parent.display().to_string(), e))?;
225        }
226
227        // Check if the configuration file exists, if not, create it.
228        let config_path = Path::new(&self.config_file_location);
229        if !config_path.exists() {
230            let mut file = File::create(config_path)
231                .map_err(|e| KSMRError::FileCreationError(config_path.display().to_string(), e))?;
232
233            // Attempt to write an empty JSON object to the file
234            let empty_json_string = b"{}";
235            file.write_all(empty_json_string)
236                .map_err(|e| KSMRError::FileWriteError(config_path.display().to_string(), e))?;
237        }
238
239        Ok(())
240    }
241
242    fn is_empty(&self) -> Result<bool, KSMRError> {
243        // Attempt to read the storage and handle errors using custom KSMRError
244        let config = self
245            .read_storage()
246            .map_err(|e| KSMRError::StorageError(format!("Failed to read storage: {}", e)))?;
247
248        // Check if the config is empty and return the result
249        Ok(config.is_empty())
250    }
251}
252
253#[derive(Clone)]
254pub struct InMemoryKeyValueStorage {
255    config: HashMap<ConfigKeys, String>,
256}
257
258impl InMemoryKeyValueStorage {
259    pub fn new(config: Option<String>) -> Result<Self, KSMRError> {
260        let mut config_map: HashMap<ConfigKeys, String> = HashMap::new();
261
262        if let Some(cfg) = config {
263            if Self::is_base64(&cfg) {
264                // Try decoding as padded, then un-padded
265                let decoded_bytes = STANDARD
266                    .decode(&cfg)
267                    .or_else(|_| STANDARD_NO_PAD.decode(&cfg))
268                    .map_err(|e| {
269                        KSMRError::DecodeError(format!("Failed to decode Base64 string: {}", e))
270                    })?;
271
272                let decoded_string = String::from_utf8(decoded_bytes).map_err(|e| {
273                    KSMRError::StringConversionError(format!(
274                        "Failed to convert decoded bytes to string: {}",
275                        e
276                    ))
277                })?;
278
279                config_map = Self::json_to_dict(&decoded_string)?;
280            } else {
281                // Directly parse the JSON string
282                config_map = Self::json_to_dict(&cfg)?;
283            }
284        }
285        Ok(InMemoryKeyValueStorage { config: config_map })
286    }
287
288    pub fn new_config_storage(config: Option<String>) -> Result<KvStoreType, KSMRError> {
289        let in_memory = InMemoryKeyValueStorage::new(config)?;
290        Ok(KvStoreType::InMemory(in_memory))
291    }
292
293    fn is_base64(s: &str) -> bool {
294        // Accept either padded or un-padded Base64
295        STANDARD.decode(s).is_ok() || STANDARD_NO_PAD.decode(s).is_ok()
296    }
297
298    pub fn json_to_dict(json_str: &str) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
299        // Handle empty string as an empty JSON object
300        let json_str = if json_str.is_empty() { "{}" } else { json_str };
301
302        // Deserialize the JSON string
303        let value: serde_json::Value = serde_json::from_str(json_str)
304            .map_err(|e| KSMRError::SerializationError(format!("Failed to parse JSON: {}", e)))?;
305
306        let mut result = HashMap::new();
307
308        // Ensure we are dealing with a JSON object
309        if let serde_json::Value::Object(obj) = value {
310            for (k, v) in obj {
311                if let serde_json::Value::String(s) = v {
312                    // Attempt to convert the key to a ConfigKeys enum
313                    if let Some(key) = ConfigKeys::get_enum(&k) {
314                        result.insert(key, s);
315                    } else {
316                        return Err(KSMRError::SerializationError(format!(
317                            "Invalid key in JSON: {}",
318                            k
319                        )));
320                    }
321                } else {
322                    return Err(KSMRError::SerializationError(format!(
323                        "Expected string value for key: {}",
324                        k
325                    )));
326                }
327            }
328        } else {
329            return Err(KSMRError::SerializationError(
330                "Expected JSON object".to_string(),
331            ));
332        }
333
334        Ok(result) // Return the populated HashMap
335    }
336}
337
338impl KeyValueStorage for InMemoryKeyValueStorage {
339    fn read_storage(&self) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
340        Ok(self.config.clone()) // Return a clone of the in-memory storage
341    }
342
343    fn save_storage(
344        &mut self,
345        _updated_config: HashMap<ConfigKeys, String>,
346    ) -> Result<bool, KSMRError> {
347        // Since this is in-memory, we just replace the storage
348        self.config = _updated_config;
349        Ok(true)
350    }
351
352    fn get(&self, key: ConfigKeys) -> Result<Option<String>, KSMRError> {
353        Ok(self.config.get(&key).cloned()) // Get the value for the given key
354    }
355
356    fn set(
357        &mut self,
358        key: ConfigKeys,
359        value: String,
360    ) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
361        self.config.insert(key, value.clone()); // Insert or update the key
362        Ok(self.config.clone()) // Return the updated storage
363    }
364
365    fn delete(&mut self, key: ConfigKeys) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
366        self.config.remove(&key); // Remove the key if it exists
367        Ok(self.config.clone()) // Return the updated storage
368    }
369
370    fn delete_all(&mut self) -> Result<HashMap<ConfigKeys, String>, KSMRError> {
371        self.config.clear(); // Clear all entries
372        Ok(self.config.clone()) // Return the cleared storage
373    }
374
375    fn contains(&self, key: ConfigKeys) -> Result<bool, KSMRError> {
376        Ok(self.config.contains_key(&key)) // Check if the key exists
377    }
378
379    fn create_config_file_if_missing(&self) -> Result<(), KSMRError> {
380        // No file to create for in-memory storage
381        Ok(())
382    }
383
384    fn is_empty(&self) -> Result<bool, KSMRError> {
385        Ok(self.config.is_empty()) // Check if storage is empty
386    }
387}