kindly_tools/config/
wrap.rs

1// Copyright 2025 Kindly Software Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Configuration support for the wrap command
16
17use anyhow::{Context, Result};
18use serde::{Deserialize, Serialize};
19use std::collections::HashSet;
20use std::path::{Path, PathBuf};
21
22/// Default configuration file name
23const CONFIG_FILE_NAME: &str = "wrap.toml";
24
25/// Default commands to auto-wrap
26const DEFAULT_AUTO_WRAP_COMMANDS: &[&str] = &[
27    "claude",
28    "openai",
29    "gemini",
30    "anthropic",
31    "gpt",
32    "ai",
33    "llm",
34    "copilot",
35    "codewhisperer",
36    "tabnine",
37    "codium",
38    "bard",
39    "perplexity",
40];
41
42/// Wrap command configuration
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct WrapConfig {
45    /// Whether auto-wrapping is enabled
46    #[serde(default = "default_enabled")]
47    pub enabled: bool,
48
49    /// Protection mode: "warning" or "blocking"
50    #[serde(default = "default_mode")]
51    pub mode: WrapMode,
52
53    /// List of commands to auto-wrap
54    #[serde(default = "default_commands")]
55    pub commands: HashSet<String>,
56
57    /// Additional user-defined commands to wrap
58    #[serde(default)]
59    pub custom_commands: HashSet<String>,
60
61    /// Server URL for threat detection
62    #[serde(default = "default_server")]
63    pub server: String,
64
65    /// Whether to show detailed threat information
66    #[serde(default = "default_verbose")]
67    pub verbose: bool,
68
69    /// Whether to log wrapped sessions
70    #[serde(default)]
71    pub log_sessions: bool,
72
73    /// Path to session log directory
74    #[serde(default)]
75    pub log_directory: Option<PathBuf>,
76}
77
78/// Protection mode for wrapped commands
79#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
80#[serde(rename_all = "lowercase")]
81pub enum WrapMode {
82    /// Show warnings but allow threats through
83    Warning,
84    /// Block detected threats
85    Blocking,
86}
87
88impl Default for WrapConfig {
89    fn default() -> Self {
90        Self {
91            enabled: default_enabled(),
92            mode: default_mode(),
93            commands: default_commands(),
94            custom_commands: HashSet::new(),
95            server: default_server(),
96            verbose: default_verbose(),
97            log_sessions: false,
98            log_directory: None,
99        }
100    }
101}
102
103impl WrapConfig {
104    /// Load configuration from the default location
105    pub fn load() -> Result<Self> {
106        let config_path = Self::default_config_path()?;
107
108        if config_path.exists() {
109            Self::load_from_path(&config_path)
110        } else {
111            // Return default config if file doesn't exist
112            Ok(Self::default())
113        }
114    }
115
116    /// Load configuration from a specific path
117    pub fn load_from_path(path: &Path) -> Result<Self> {
118        let content = std::fs::read_to_string(path)
119            .with_context(|| format!("Failed to read configuration from {}", path.display()))?;
120
121        let mut config: Self = toml::from_str(&content)
122            .with_context(|| format!("Failed to parse configuration from {}", path.display()))?;
123
124        // Merge custom commands into the main commands set
125        config.merge_custom_commands();
126
127        Ok(config)
128    }
129
130    /// Save configuration to the default location
131    pub fn save(&self) -> Result<()> {
132        let config_path = Self::default_config_path()?;
133        self.save_to_path(&config_path)
134    }
135
136    /// Save configuration to a specific path
137    pub fn save_to_path(&self, path: &Path) -> Result<()> {
138        // Ensure parent directory exists
139        if let Some(parent) = path.parent() {
140            std::fs::create_dir_all(parent).with_context(|| {
141                format!("Failed to create config directory: {}", parent.display())
142            })?;
143        }
144
145        let content = toml::to_string_pretty(self).context("Failed to serialize configuration")?;
146
147        std::fs::write(path, content)
148            .with_context(|| format!("Failed to write configuration to {}", path.display()))?;
149
150        Ok(())
151    }
152
153    /// Get the default configuration file path
154    pub fn default_config_path() -> Result<PathBuf> {
155        let config_dir = crate::config_dir()?;
156        Ok(config_dir.join(CONFIG_FILE_NAME))
157    }
158
159    /// Check if a command should be wrapped
160    pub fn should_wrap(&self, command: &str) -> bool {
161        if !self.enabled {
162            return false;
163        }
164
165        // Check if the command or its basename is in our wrap list
166        let basename = Path::new(command)
167            .file_name()
168            .and_then(|n| n.to_str())
169            .unwrap_or(command);
170
171        self.commands.contains(basename) || self.commands.contains(command)
172    }
173
174    /// Add a custom command to wrap
175    pub fn add_command(&mut self, command: String) {
176        self.custom_commands.insert(command.clone());
177        self.commands.insert(command);
178    }
179
180    /// Remove a command from the wrap list
181    pub fn remove_command(&mut self, command: &str) -> bool {
182        self.custom_commands.remove(command);
183        self.commands.remove(command)
184    }
185
186    /// Merge custom commands into the main commands set
187    fn merge_custom_commands(&mut self) {
188        for cmd in &self.custom_commands {
189            self.commands.insert(cmd.clone());
190        }
191    }
192
193    /// Create a default configuration file with example settings
194    pub fn create_default_config() -> Result<()> {
195        let config_path = Self::default_config_path()?;
196
197        if config_path.exists() {
198            anyhow::bail!(
199                "Configuration file already exists at {}",
200                config_path.display()
201            );
202        }
203
204        let default_config = Self::default();
205        default_config.save_to_path(&config_path)?;
206
207        println!(
208            "Created default configuration at: {}",
209            config_path.display()
210        );
211        Ok(())
212    }
213
214    /// Get effective mode as a string
215    pub fn mode_string(&self) -> &'static str {
216        match self.mode {
217            WrapMode::Warning => "warning",
218            WrapMode::Blocking => "blocking",
219        }
220    }
221}
222
223// Default value functions for serde
224fn default_enabled() -> bool {
225    true
226}
227
228fn default_mode() -> WrapMode {
229    WrapMode::Warning
230}
231
232fn default_commands() -> HashSet<String> {
233    DEFAULT_AUTO_WRAP_COMMANDS
234        .iter()
235        .map(|&s| s.to_string())
236        .collect()
237}
238
239fn default_server() -> String {
240    "http://localhost:8080".to_string()
241}
242
243fn default_verbose() -> bool {
244    false
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use tempfile::TempDir;
251
252    #[test]
253    fn test_default_config() {
254        let config = WrapConfig::default();
255        assert!(config.enabled);
256        assert_eq!(config.mode, WrapMode::Warning);
257        assert!(config.commands.contains("claude"));
258        assert!(config.commands.contains("openai"));
259    }
260
261    #[test]
262    fn test_should_wrap() {
263        let mut config = WrapConfig::default();
264
265        // Should wrap known commands
266        assert!(config.should_wrap("claude"));
267        assert!(config.should_wrap("openai"));
268
269        // Should handle full paths
270        assert!(config.should_wrap("/usr/bin/claude"));
271        assert!(config.should_wrap("./openai"));
272
273        // Should not wrap unknown commands
274        assert!(!config.should_wrap("ls"));
275        assert!(!config.should_wrap("cat"));
276
277        // Should respect enabled flag
278        config.enabled = false;
279        assert!(!config.should_wrap("claude"));
280    }
281
282    #[test]
283    fn test_add_remove_commands() {
284        let mut config = WrapConfig::default();
285
286        // Add custom command
287        config.add_command("mycli".to_string());
288        assert!(config.should_wrap("mycli"));
289        assert!(config.custom_commands.contains("mycli"));
290
291        // Remove command
292        assert!(config.remove_command("mycli"));
293        assert!(!config.should_wrap("mycli"));
294        assert!(!config.custom_commands.contains("mycli"));
295    }
296
297    #[test]
298    fn test_save_load_config() -> Result<()> {
299        let temp_dir = TempDir::new()?;
300        let config_path = temp_dir.path().join("wrap.toml");
301
302        // Create and save config
303        let mut config = WrapConfig::default();
304        config.mode = WrapMode::Blocking;
305        config.add_command("custom-ai".to_string());
306        config.verbose = true;
307        config.save_to_path(&config_path)?;
308
309        // Load and verify
310        let loaded = WrapConfig::load_from_path(&config_path)?;
311        assert_eq!(loaded.mode, WrapMode::Blocking);
312        assert!(loaded.should_wrap("custom-ai"));
313        assert!(loaded.verbose);
314
315        Ok(())
316    }
317
318    #[test]
319    fn test_mode_serialization() -> Result<()> {
320        let warning_toml = r#"mode = "warning""#;
321        let config: WrapConfig = toml::from_str(warning_toml)?;
322        assert_eq!(config.mode, WrapMode::Warning);
323
324        let blocking_toml = r#"mode = "blocking""#;
325        let config: WrapConfig = toml::from_str(blocking_toml)?;
326        assert_eq!(config.mode, WrapMode::Blocking);
327
328        Ok(())
329    }
330}
kindly_tools/config/wrap.rs

kindly_tools/config/
wrap.rs
