use crate::config::paths::Paths; use crate::conversation::message::ToolRequest; use anyhow::Result; use blake3::Hasher; use chrono::Utc; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::time::Duration; use std::{fs::File, path::PathBuf}; #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ToolPermissionRecord { tool_name: String, allowed: bool, context_hash: String, // Hash of the tool's arguments/context to differentiate similar calls #[serde(skip_serializing_if = "Option::is_none")] // Don't serialize if None readable_context: Option, timestamp: i64, expiry: Option, // Optional expiry timestamp } #[derive(Debug, Serialize, Deserialize)] pub struct ToolPermissionStore { permissions: HashMap>, version: u32, // For future schema migrations #[serde(skip)] // Don't serialize this field permissions_dir: PathBuf, } impl Default for ToolPermissionStore { fn default() -> Self { Self::new() } } impl ToolPermissionStore { pub fn new() -> Self { Self { permissions: HashMap::new(), version: 1, permissions_dir: Paths::config_dir().join("permissions"), } } pub fn load() -> Result { let store = Self::new(); let file_path = store.permissions_dir.join("tool_permissions.json"); if !file_path.exists() { return Ok(store); } let file = File::open(file_path)?; let mut permissions: ToolPermissionStore = serde_json::from_reader(file)?; permissions.permissions_dir = store.permissions_dir; // Clean up expired entries on load permissions.cleanup_expired()?; Ok(permissions) } pub fn save(&self) -> anyhow::Result<()> { std::fs::create_dir_all(&self.permissions_dir)?; let path = self.permissions_dir.join("tool_permissions.json"); let temp_path = path.with_extension("tmp"); // Write complete content to temporary file let content = serde_json::to_string_pretty(self)?; std::fs::write(&temp_path, &content)?; // Atomically rename temp file to target file std::fs::rename(temp_path, path)?; Ok(()) } pub fn check_permission(&self, tool_request: &ToolRequest) -> Option { let context_hash = self.hash_tool_context(tool_request); let tool_call = tool_request.tool_call.as_ref().unwrap(); let key = format!("{}:{}", tool_call.name, context_hash); self.permissions.get(&key).and_then(|records| { records .iter() .rfind(|record| record.expiry.is_none_or(|exp| exp > Utc::now().timestamp())) .map(|record| record.allowed) }) } pub fn record_permission( &mut self, tool_request: &ToolRequest, allowed: bool, expiry_duration: Option, ) -> anyhow::Result<()> { let context_hash = self.hash_tool_context(tool_request); let tool_call = tool_request.tool_call.as_ref().unwrap(); let key = format!("{}:{}", tool_call.name, context_hash); let record = ToolPermissionRecord { tool_name: tool_call.name.to_string().clone(), allowed, context_hash, readable_context: Some(tool_request.to_readable_string()), timestamp: Utc::now().timestamp(), expiry: expiry_duration.map(|d| Utc::now().timestamp() + d.as_secs() as i64), }; self.permissions.entry(key).or_default().push(record); self.save()?; Ok(()) } fn hash_tool_context(&self, tool_request: &ToolRequest) -> String { // Create a hash of the tool's arguments to differentiate similar calls // This helps identify when the same tool is being used in a different context let mut hasher = Hasher::new(); hasher.update( serde_json::to_string(&tool_request.tool_call.as_ref().unwrap().arguments) .unwrap_or_default() .as_bytes(), ); hasher.finalize().to_hex().to_string() } pub fn cleanup_expired(&mut self) -> anyhow::Result<()> { let now = Utc::now().timestamp(); let mut changed = false; self.permissions.retain(|_, records| { records.retain(|record| record.expiry.is_none_or(|exp| exp > now)); changed = changed || records.is_empty(); !records.is_empty() }); if changed { self.save()?; } Ok(()) } }