Skip to main content

victauri_plugin/
privacy.rs

1use std::collections::HashSet;
2
3use crate::redaction::Redactor;
4
5/// Privacy controls for the MCP server. Blocklist takes precedence over allowlist.
6#[derive(Default)]
7pub struct PrivacyConfig {
8    /// If set, only these commands can be invoked (positive allowlist).
9    pub command_allowlist: Option<HashSet<String>>,
10    /// Commands that are always blocked, even if on the allowlist.
11    pub command_blocklist: HashSet<String>,
12    /// MCP tool names that are disabled (e.g., `"eval_js"`, `"screenshot"`).
13    pub disabled_tools: HashSet<String>,
14    /// Output redactor with regex and JSON-key matching.
15    pub redactor: Redactor,
16    /// Whether output redaction is active.
17    pub redaction_enabled: bool,
18}
19
20impl PrivacyConfig {
21    /// Returns `true` if the command passes both the allowlist and blocklist checks.
22    pub fn is_command_allowed(&self, command: &str) -> bool {
23        if self.command_blocklist.contains(command) {
24            return false;
25        }
26        match &self.command_allowlist {
27            Some(allow) => allow.contains(command),
28            None => true,
29        }
30    }
31
32    /// Returns `true` unless the tool is in [`disabled_tools`](Self::disabled_tools).
33    pub fn is_tool_enabled(&self, tool_name: &str) -> bool {
34        !self.disabled_tools.contains(tool_name)
35    }
36
37    /// Apply redaction rules to the output string if redaction is enabled, otherwise pass through.
38    pub fn redact_output(&self, output: &str) -> String {
39        if self.redaction_enabled {
40            self.redactor.redact(output)
41        } else {
42            output.to_string()
43        }
44    }
45}
46
47const STRICT_DISABLED_TOOLS: &[&str] = &[
48    "eval_js",
49    "screenshot",
50    "inject_css",
51    "set_storage",
52    "delete_storage",
53    "navigate",
54    "set_dialog_response",
55    "fill",
56    "type_text",
57];
58
59/// Create a [`PrivacyConfig`] that disables dangerous tools (eval, screenshot, mutations) and enables redaction.
60pub fn strict_privacy_config() -> PrivacyConfig {
61    PrivacyConfig {
62        command_allowlist: None,
63        command_blocklist: HashSet::new(),
64        disabled_tools: STRICT_DISABLED_TOOLS
65            .iter()
66            .map(|s| s.to_string())
67            .collect(),
68        redactor: Redactor::default(),
69        redaction_enabled: true,
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn default_allows_all_commands() {
79        let config = PrivacyConfig::default();
80        assert!(config.is_command_allowed("get_settings"));
81        assert!(config.is_command_allowed("anything"));
82    }
83
84    #[test]
85    fn blocklist_blocks() {
86        let mut config = PrivacyConfig::default();
87        config.command_blocklist.insert("save_api_key".to_string());
88        assert!(!config.is_command_allowed("save_api_key"));
89        assert!(config.is_command_allowed("get_settings"));
90    }
91
92    #[test]
93    fn allowlist_restricts() {
94        let mut allow = HashSet::new();
95        allow.insert("get_settings".to_string());
96        allow.insert("get_monitoring_status".to_string());
97        let config = PrivacyConfig {
98            command_allowlist: Some(allow),
99            ..Default::default()
100        };
101        assert!(config.is_command_allowed("get_settings"));
102        assert!(!config.is_command_allowed("save_api_key"));
103    }
104
105    #[test]
106    fn blocklist_wins_over_allowlist() {
107        let mut allow = HashSet::new();
108        allow.insert("save_api_key".to_string());
109        let mut block = HashSet::new();
110        block.insert("save_api_key".to_string());
111        let config = PrivacyConfig {
112            command_allowlist: Some(allow),
113            command_blocklist: block,
114            ..Default::default()
115        };
116        assert!(!config.is_command_allowed("save_api_key"));
117    }
118
119    #[test]
120    fn tool_disabling() {
121        let mut disabled = HashSet::new();
122        disabled.insert("eval_js".to_string());
123        let config = PrivacyConfig {
124            disabled_tools: disabled,
125            ..Default::default()
126        };
127        assert!(!config.is_tool_enabled("eval_js"));
128        assert!(config.is_tool_enabled("dom_snapshot"));
129    }
130
131    #[test]
132    fn strict_mode_disables_dangerous_tools() {
133        let config = strict_privacy_config();
134        assert!(!config.is_tool_enabled("eval_js"));
135        assert!(!config.is_tool_enabled("screenshot"));
136        assert!(!config.is_tool_enabled("inject_css"));
137        assert!(!config.is_tool_enabled("navigate"));
138        assert!(config.is_tool_enabled("dom_snapshot"));
139        assert!(config.is_tool_enabled("get_window_state"));
140        assert!(config.redaction_enabled);
141    }
142
143    #[test]
144    fn redaction_when_enabled() {
145        let config = PrivacyConfig {
146            redaction_enabled: true,
147            ..Default::default()
148        };
149        let output = config.redact_output("key is sk-abc123def456ghi789jkl012mno");
150        assert!(output.contains("[REDACTED]"));
151    }
152
153    #[test]
154    fn no_redaction_when_disabled() {
155        let config = PrivacyConfig::default();
156        let input = "key is sk-abc123def456ghi789jkl012mno";
157        assert_eq!(config.redact_output(input), input);
158    }
159}