cortex-agent 0.5.2

Self-learning AI agent with persistent memory, skills, project awareness, and a beautiful terminal UI
use std::collections::HashSet;
use std::path::Path;
use std::sync::Mutex;

/// Simple in-memory permission store with optional file persistence.
pub struct PermissionStore {
    /// Tools allowed for the current session (resets on restart)
    session_allowed: Mutex<HashSet<String>>,
    /// Tools permanently allowed (persisted to disk)
    permanently_allowed: Mutex<HashSet<String>>,
    /// Path to persistent permissions file
    db_path: String,
}

impl PermissionStore {
    pub fn new(db_dir: &str) -> Self {
        let db_path = format!("{}/permissions.json", db_dir);
        let existing = std::fs::read_to_string(&db_path)
            .ok()
            .and_then(|s| serde_json::from_str::<Vec<String>>(&s).ok())
            .unwrap_or_default();
        Self {
            session_allowed: Mutex::new(HashSet::new()),
            permanently_allowed: Mutex::new(existing.into_iter().collect()),
            db_path,
        }
    }

    /// Check if a tool is allowed to run.
    pub fn is_allowed(&self, tool_name: &str) -> bool {
        let session = self.session_allowed.lock().unwrap();
        if session.contains(tool_name) {
            return true;
        }
        let permanent = self.permanently_allowed.lock().unwrap();
        permanent.contains(tool_name)
    }

    /// Allow a tool for the current session only.
    pub fn allow_session(&self, tool_name: &str) {
        self.session_allowed.lock().unwrap().insert(tool_name.to_string());
    }

    /// Permanently allow a tool (persisted).
    pub fn allow_permanently(&self, tool_name: &str) -> Result<(), String> {
        {
            let mut perm = self.permanently_allowed.lock().unwrap();
            perm.insert(tool_name.to_string());
        }
        self.save()
    }

    /// Revoke permission for a tool.
    pub fn revoke(&self, tool_name: &str) -> Result<(), String> {
        {
            self.session_allowed.lock().unwrap().remove(tool_name);
            self.permanently_allowed.lock().unwrap().remove(tool_name);
        }
        self.save()
    }

    /// List all permanently allowed tools.
    pub fn list_allowed(&self) -> Vec<String> {
        let mut all: Vec<String> = self.session_allowed.lock().unwrap().iter().cloned().collect();
        for tool in self.permanently_allowed.lock().unwrap().iter() {
            if !all.contains(tool) {
                all.push(tool.clone());
            }
        }
        all.sort();
        all
    }

    /// Persist permanent permissions to disk.
    fn save(&self) -> Result<(), String> {
        let perm = self.permanently_allowed.lock().unwrap();
        let data = serde_json::to_string_pretty(&perm.iter().cloned().collect::<Vec<_>>())
            .map_err(|e| format!("Serialize: {}", e))?;
        if let Some(parent) = Path::new(&self.db_path).parent() {
            let _ = std::fs::create_dir_all(parent);
        }
        std::fs::write(&self.db_path, &data).map_err(|e| format!("Write: {}", e))
    }

    /// Check if a tool is considered "dangerous" and needs permission.
    pub fn is_dangerous_tool(name: &str) -> bool {
        let dangerous = [
            "execute_bash", "execute_command", "delete_file", "remove_file",
            "write_file", "create_file", "patch_file", "kill_process",
            "restart_service", "install_package", "uninstall_package",
            "format_disk", "mkfs", "dd", "shutdown", "reboot",
        ];
        dangerous.contains(&name)
    }
}