vtcode_core/tools/registry/
shell_policy.rs

1use anyhow::{Result, anyhow};
2use regex::Regex;
3use std::collections::hash_map::DefaultHasher;
4use std::hash::{Hash, Hasher};
5use tracing::warn;
6
7#[derive(Clone, Debug)]
8pub struct ShellPolicyCacheEntry {
9    pub signature: u64,
10    pub deny_regexes: Vec<(String, Regex)>,
11    pub deny_globs: Vec<(String, Regex)>,
12}
13
14pub struct ShellPolicyChecker {
15    cache: Option<ShellPolicyCacheEntry>,
16    commands_config: Option<crate::config::CommandsConfig>,
17}
18
19impl ShellPolicyChecker {
20    pub fn new() -> Self {
21        Self {
22            cache: None,
23            commands_config: None,
24        }
25    }
26}
27
28impl Default for ShellPolicyChecker {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl ShellPolicyChecker {
35    pub fn set_commands_config(&mut self, commands_config: &crate::config::CommandsConfig) {
36        self.commands_config = Some(commands_config.clone());
37        self.reset_cache();
38    }
39
40    pub fn commands_config(&self) -> Option<&crate::config::CommandsConfig> {
41        self.commands_config.as_ref()
42    }
43
44    pub fn check_command(
45        &mut self,
46        command: &str,
47        agent_type: &str,
48        deny_regex_patterns: &[String],
49        deny_glob_patterns: &[String],
50    ) -> Result<()> {
51        let mut hasher = DefaultHasher::new();
52        deny_regex_patterns.hash(&mut hasher);
53        deny_glob_patterns.hash(&mut hasher);
54        let signature = hasher.finish();
55
56        let entry = if let Some(ref entry) = self.cache
57            && entry.signature == signature
58        {
59            entry
60        } else {
61            let compiled_regexes = deny_regex_patterns
62                .iter()
63                .filter_map(|pattern| {
64                    if pattern.is_empty() { return None; }
65                    match Regex::new(pattern) {
66                        Ok(re) => Some((pattern.clone(), re)),
67                        Err(err) => {
68                            warn!(agent = agent_type, pattern, error = %err, "Invalid deny regex pattern skipped");
69                            None
70                        }
71                    }
72                })
73                .collect::<Vec<_>>();
74
75            let compiled_globs = deny_glob_patterns
76                .iter()
77                .filter_map(|pattern| {
78                    if pattern.is_empty() { return None; }
79                    let re_pattern = format!("^{}$", regex::escape(pattern).replace(r"\\*", ".*"));
80                    match Regex::new(&re_pattern) {
81                        Ok(re) => Some((pattern.clone(), re)),
82                        Err(err) => {
83                            warn!(agent = agent_type, pattern, error = %err, "Invalid deny glob pattern skipped");
84                            None
85                        }
86                    }
87                })
88                .collect::<Vec<_>>();
89
90            let new_entry = ShellPolicyCacheEntry {
91                signature,
92                deny_regexes: compiled_regexes,
93                deny_globs: compiled_globs,
94            };
95            self.cache = Some(new_entry);
96            self.cache
97                .as_ref()
98                .ok_or_else(|| anyhow!("Failed to initialize shell policy cache entry"))?
99        };
100
101        for (pattern, compiled) in &entry.deny_regexes {
102            if compiled.is_match(command) {
103                return Err(anyhow!(
104                    "Shell command denied by agent regex policy: {}",
105                    pattern
106                ));
107            }
108        }
109
110        for (pattern, compiled) in &entry.deny_globs {
111            if compiled.is_match(command) {
112                return Err(anyhow!(
113                    "Shell command denied by agent glob policy: {}",
114                    pattern
115                ));
116            }
117        }
118
119        Ok(())
120    }
121
122    pub fn reset_cache(&mut self) {
123        self.cache = None;
124    }
125}
vtcode_core/tools/registry/shell_policy.rs

vtcode_core/tools/registry/
shell_policy.rs
