stakpak_shared/
secret_manager.rs

1use crate::local_store::LocalStore;
2use crate::secrets::{redact_password, redact_secrets, restore_secrets};
3use serde_json;
4use std::collections::HashMap;
5use tracing::{error, warn};
6
7/// Handles secret redaction and restoration across different tool types
8#[derive(Clone)]
9pub struct SecretManager {
10    redact_secrets: bool,
11    privacy_mode: bool,
12}
13
14impl SecretManager {
15    pub fn new(redact_secrets: bool, privacy_mode: bool) -> Self {
16        Self {
17            redact_secrets,
18            privacy_mode,
19        }
20    }
21
22    /// Load the redaction map from the session file
23    pub fn load_session_redaction_map(&self) -> HashMap<String, String> {
24        match LocalStore::read_session_data("secrets.json") {
25            Ok(content) => {
26                if content.trim().is_empty() {
27                    return HashMap::new();
28                }
29
30                match serde_json::from_str::<HashMap<String, String>>(&content) {
31                    Ok(map) => map,
32                    Err(e) => {
33                        error!("Failed to parse session redaction map JSON: {}", e);
34                        HashMap::new()
35                    }
36                }
37            }
38            Err(e) => {
39                warn!("Failed to read session redaction map file: {}", e);
40                HashMap::new()
41            }
42        }
43    }
44
45    /// Save the redaction map to the session file
46    pub fn save_session_redaction_map(&self, redaction_map: &HashMap<String, String>) {
47        match serde_json::to_string_pretty(redaction_map) {
48            Ok(json_content) => {
49                if let Err(e) = LocalStore::write_session_data("secrets.json", &json_content) {
50                    error!("Failed to save session redaction map: {}", e);
51                }
52            }
53            Err(e) => {
54                error!("Failed to serialize session redaction map to JSON: {}", e);
55            }
56        }
57    }
58
59    /// Add new redactions to the session map
60    pub fn add_to_session_redaction_map(&self, new_redactions: &HashMap<String, String>) {
61        if new_redactions.is_empty() {
62            return;
63        }
64
65        let mut existing_map = self.load_session_redaction_map();
66        existing_map.extend(new_redactions.clone());
67        self.save_session_redaction_map(&existing_map);
68    }
69
70    /// Restore secrets in a string using the session redaction map
71    pub fn restore_secrets_in_string(&self, input: &str) -> String {
72        let redaction_map = self.load_session_redaction_map();
73        if redaction_map.is_empty() {
74            return input.to_string();
75        }
76        restore_secrets(input, &redaction_map)
77    }
78
79    /// Redact secrets and add to session map
80    pub fn redact_and_store_secrets(&self, content: &str, path: Option<&str>) -> String {
81        if !self.redact_secrets {
82            return content.to_string();
83        }
84
85        // TODO: this is not thread safe, we need to use a mutex or an actor to protect the redaction map
86        let existing_redaction_map = self.load_session_redaction_map();
87        let redaction_result =
88            redact_secrets(content, path, &existing_redaction_map, self.privacy_mode);
89
90        // Add new redactions to session map
91        self.add_to_session_redaction_map(&redaction_result.redaction_map);
92
93        redaction_result.redacted_string
94    }
95
96    pub fn redact_and_store_password(&self, content: &str, password: &str) -> String {
97        if !self.redact_secrets {
98            return content.to_string();
99        }
100
101        // TODO: this is not thread safe, we need to use a mutex or an actor to protect the redaction map
102        let existing_redaction_map = self.load_session_redaction_map();
103        let redaction_result = redact_password(content, password, &existing_redaction_map);
104
105        // Add new redactions to session map
106        self.add_to_session_redaction_map(&redaction_result.redaction_map);
107
108        redaction_result.redacted_string
109    }
110}