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    #[must_use]
23    pub fn is_command_allowed(&self, command: &str) -> bool {
24        if self.command_blocklist.contains(command) {
25            return false;
26        }
27        match &self.command_allowlist {
28            Some(allow) => allow.contains(command),
29            None => true,
30        }
31    }
32
33    /// Returns `true` unless the tool is in [`disabled_tools`](Self::disabled_tools).
34    #[must_use]
35    pub fn is_tool_enabled(&self, tool_name: &str) -> bool {
36        !self.disabled_tools.contains(tool_name)
37    }
38
39    /// Apply redaction rules to the output string if redaction is enabled, otherwise pass through.
40    #[must_use]
41    pub fn redact_output(&self, output: &str) -> String {
42        if self.redaction_enabled {
43            self.redactor.redact(output)
44        } else {
45            output.to_string()
46        }
47    }
48}
49
50const STRICT_DISABLED_TOOLS: &[&str] = &[
51    "eval_js",
52    "screenshot",
53    "inject_css",
54    "set_storage",
55    "delete_storage",
56    "navigate",
57    "set_dialog_response",
58    "fill",
59    "type_text",
60];
61
62/// Create a [`PrivacyConfig`] that disables dangerous tools (eval, screenshot, mutations) and enables redaction.
63pub fn strict_privacy_config() -> PrivacyConfig {
64    PrivacyConfig {
65        command_allowlist: None,
66        command_blocklist: HashSet::new(),
67        disabled_tools: STRICT_DISABLED_TOOLS
68            .iter()
69            .map(std::string::ToString::to_string)
70            .collect(),
71        redactor: Redactor::default(),
72        redaction_enabled: true,
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn default_allows_all_commands() {
82        let config = PrivacyConfig::default();
83        assert!(config.is_command_allowed("get_settings"));
84        assert!(config.is_command_allowed("anything"));
85    }
86
87    #[test]
88    fn blocklist_blocks() {
89        let mut config = PrivacyConfig::default();
90        config.command_blocklist.insert("save_api_key".to_string());
91        assert!(!config.is_command_allowed("save_api_key"));
92        assert!(config.is_command_allowed("get_settings"));
93    }
94
95    #[test]
96    fn allowlist_restricts() {
97        let mut allow = HashSet::new();
98        allow.insert("get_settings".to_string());
99        allow.insert("get_monitoring_status".to_string());
100        let config = PrivacyConfig {
101            command_allowlist: Some(allow),
102            ..Default::default()
103        };
104        assert!(config.is_command_allowed("get_settings"));
105        assert!(!config.is_command_allowed("save_api_key"));
106    }
107
108    #[test]
109    fn blocklist_wins_over_allowlist() {
110        let mut allow = HashSet::new();
111        allow.insert("save_api_key".to_string());
112        let mut block = HashSet::new();
113        block.insert("save_api_key".to_string());
114        let config = PrivacyConfig {
115            command_allowlist: Some(allow),
116            command_blocklist: block,
117            ..Default::default()
118        };
119        assert!(!config.is_command_allowed("save_api_key"));
120    }
121
122    #[test]
123    fn tool_disabling() {
124        let mut disabled = HashSet::new();
125        disabled.insert("eval_js".to_string());
126        let config = PrivacyConfig {
127            disabled_tools: disabled,
128            ..Default::default()
129        };
130        assert!(!config.is_tool_enabled("eval_js"));
131        assert!(config.is_tool_enabled("dom_snapshot"));
132    }
133
134    #[test]
135    fn strict_mode_disables_dangerous_tools() {
136        let config = strict_privacy_config();
137        assert!(!config.is_tool_enabled("eval_js"));
138        assert!(!config.is_tool_enabled("screenshot"));
139        assert!(!config.is_tool_enabled("inject_css"));
140        assert!(!config.is_tool_enabled("navigate"));
141        assert!(config.is_tool_enabled("dom_snapshot"));
142        assert!(config.is_tool_enabled("get_window_state"));
143        assert!(config.redaction_enabled);
144    }
145
146    #[test]
147    fn redaction_when_enabled() {
148        let config = PrivacyConfig {
149            redaction_enabled: true,
150            ..Default::default()
151        };
152        let output = config.redact_output("key is sk-abc123def456ghi789jkl012mno");
153        assert!(output.contains("[REDACTED]"));
154    }
155
156    #[test]
157    fn no_redaction_when_disabled() {
158        let config = PrivacyConfig::default();
159        let input = "key is sk-abc123def456ghi789jkl012mno";
160        assert_eq!(config.redact_output(input), input);
161    }
162}