okbot_agent/
env.rs

1//! Keeping okbot's secrets out of agent processes.
2
3use std::ffi::OsString;
4
5/// Decides which environment variables are removed before an agent CLI starts.
6///
7/// Agent CLIs inherit okbot's environment, which holds platform tokens. A prompt-injected agent
8/// must not be able to read them, so every variable named here, and every variable whose name
9/// starts with one of the prefixes, is removed from the child.
10#[derive(Debug, Clone)]
11pub struct EnvScrub {
12    names: Vec<String>,
13    prefixes: Vec<String>,
14}
15
16impl Default for EnvScrub {
17    fn default() -> Self {
18        Self {
19            names: vec![
20                "GOOGLE_APPLICATION_CREDENTIALS".into(),
21                "DISCORD_TOKEN".into(),
22            ],
23            prefixes: vec![
24                "OKBOT_".into(),
25                "ZALO_".into(),
26                "SLACK_".into(),
27                "DISCORD_".into(),
28                "GCHAT_".into(),
29                "TELEGRAM_".into(),
30            ],
31        }
32    }
33}
34
35impl EnvScrub {
36    /// Also removes `name` (for example a configured `token_env`).
37    pub fn with_name(mut self, name: impl Into<String>) -> Self {
38        let name = name.into();
39        if !self.names.contains(&name) {
40            self.names.push(name);
41        }
42        self
43    }
44
45    /// Also removes every variable starting with `prefix`.
46    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
47        self.prefixes.push(prefix.into());
48        self
49    }
50
51    /// The variables to remove, given the current environment's variable names.
52    pub fn names_to_remove(&self, current: impl IntoIterator<Item = OsString>) -> Vec<OsString> {
53        let mut out: Vec<OsString> = self.names.iter().map(OsString::from).collect();
54        for name in current {
55            let Some(text) = name.to_str() else { continue };
56            if self.prefixes.iter().any(|p| text.starts_with(p.as_str())) && !out.contains(&name) {
57                out.push(name);
58            }
59        }
60        out
61    }
62
63    /// The variables to remove from the current process environment.
64    pub fn current(&self) -> Vec<OsString> {
65        self.names_to_remove(std::env::vars_os().map(|(k, _)| k))
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn removes_names_and_prefixes() {
75        let scrub = EnvScrub::default().with_name("MY_TOKEN");
76        let removed = scrub.names_to_remove(
77            ["PATH", "ZALO_BOT_TOKEN", "OKBOT_X", "HOME", "MY_TOKEN"]
78                .into_iter()
79                .map(OsString::from),
80        );
81        assert!(removed.contains(&OsString::from("ZALO_BOT_TOKEN")));
82        assert!(removed.contains(&OsString::from("OKBOT_X")));
83        assert!(removed.contains(&OsString::from("MY_TOKEN")));
84        assert!(!removed.contains(&OsString::from("PATH")));
85        assert_eq!(
86            removed.iter().filter(|n| *n == "MY_TOKEN").count(),
87            1,
88            "no duplicates"
89        );
90    }
91}
okbot_agent/env.rs

okbot_agent/
env.rs
