1use std::borrow::Cow;
7use std::collections::HashMap;
8
9const SKIP_VARS: &[&str] = &[
12 "BASH",
13 "BASHOPTS",
14 "BASHPID",
15 "BASH_ALIASES",
16 "BASH_ARGC",
17 "BASH_ARGV",
18 "BASH_CMDS",
19 "BASH_COMMAND",
20 "BASH_EXECUTION_STRING",
21 "BASH_LINENO",
22 "BASH_LOADABLES_PATH",
23 "BASH_REMATCH",
24 "BASH_SOURCE",
25 "BASH_SUBSHELL",
26 "BASH_VERSINFO",
27 "BASH_VERSION",
28 "COLUMNS",
29 "COMP_WORDBREAKS",
30 "DIRSTACK",
31 "EUID",
32 "FUNCNAME",
33 "GROUPS",
34 "HISTCMD",
35 "HISTFILE",
36 "HOSTNAME",
37 "HOSTTYPE",
38 "IFS",
39 "LINES",
40 "MACHTYPE",
41 "MAILCHECK",
42 "OLDPWD",
43 "OPTERR",
44 "OPTIND",
45 "OSTYPE",
46 "PIPESTATUS",
47 "PPID",
48 "PS1",
49 "PS2",
50 "PS4",
51 "PWD",
52 "RANDOM",
53 "SECONDS",
54 "SHELL",
55 "SHELLOPTS",
56 "SHLVL",
57 "UID",
58 "_",
59];
60
61#[derive(Debug, Clone)]
63pub struct EnvSnapshot {
64 vars: HashMap<String, String>,
65 cwd: String,
66}
67
68impl EnvSnapshot {
69 #[must_use]
71 pub fn new(vars: HashMap<String, String>, cwd: String) -> Self {
72 EnvSnapshot { vars, cwd }
73 }
74
75 #[must_use]
77 pub fn capture_current() -> Self {
78 let vars: HashMap<String, String> = std::env::vars()
79 .filter(|(k, _)| !should_skip_var(k))
80 .collect();
81 let cwd = std::env::current_dir()
82 .map(|p| p.to_string_lossy().into_owned())
83 .unwrap_or_default();
84 EnvSnapshot { vars, cwd }
85 }
86
87 #[must_use]
97 #[allow(dead_code)] pub fn vars(&self) -> &HashMap<String, String> {
99 &self.vars
100 }
101
102 #[must_use]
112 #[allow(dead_code)] pub fn cwd(&self) -> &str {
114 &self.cwd
115 }
116
117 pub fn diff_into(&self, after: &EnvSnapshot, out: &mut String) {
123 for (key, new_val) in &after.vars {
125 if should_skip_var(key) {
126 continue;
127 }
128
129 let changed = match self.vars.get(key) {
130 Some(old_val) => old_val != new_val,
131 None => true,
132 };
133
134 if changed {
135 out.push_str("set -gx ");
136 out.push_str(key);
137 out.push(' ');
138 if key.ends_with("PATH") && new_val.contains(':') {
140 for (i, part) in new_val.split(':').enumerate() {
141 if i > 0 {
142 out.push(' ');
143 }
144 out.push_str(part);
145 }
146 } else {
147 out.push_str(&shell_escape(new_val));
148 }
149 out.push('\n');
150 }
151 }
152
153 for key in self.vars.keys() {
155 if should_skip_var(key) {
156 continue;
157 }
158 if !after.vars.contains_key(key) {
159 out.push_str("set -e ");
160 out.push_str(key);
161 out.push('\n');
162 }
163 }
164
165 if !after.cwd.is_empty() && self.cwd != after.cwd {
167 out.push_str("cd ");
168 out.push_str(&shell_escape(&after.cwd));
169 out.push('\n');
170 }
171 }
172
173 #[must_use]
178 #[allow(dead_code)]
179 pub fn diff(&self, after: &EnvSnapshot) -> String {
180 let mut out = String::new();
181 self.diff_into(after, &mut out);
182 out
183 }
184}
185
186#[must_use]
188pub fn parse_null_separated_env(data: &str) -> HashMap<String, String> {
189 let mut vars = HashMap::new();
190
191 for entry in data.split('\0') {
193 let entry = entry.trim_start_matches('\n');
194 if entry.is_empty() {
195 continue;
196 }
197 if let Some(eq_pos) = entry.find('=') {
198 let key = &entry[..eq_pos];
199 let value = &entry[eq_pos + 1..];
200 if !key.is_empty() && key.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_') {
202 vars.insert(key.to_string(), value.to_string());
203 }
204 }
205 }
206
207 vars
208}
209
210#[must_use]
212pub(crate) fn should_skip_var(name: &str) -> bool {
213 SKIP_VARS.binary_search(&name).is_ok()
214}
215
216fn shell_escape(s: &str) -> Cow<'_, str> {
219 if s.bytes().all(|b| {
221 b.is_ascii_alphanumeric()
222 || matches!(b, b'/' | b'.' | b'-' | b'_' | b':' | b'~' | b'+' | b',')
223 }) {
224 return Cow::Borrowed(s);
225 }
226 let mut result = String::with_capacity(s.len() + 2);
228 result.push('\'');
229 for &b in s.as_bytes() {
230 if b == b'\'' {
231 result.push_str("'\\''");
232 } else {
233 result.push(b as char);
234 }
235 }
236 result.push('\'');
237 Cow::Owned(result)
238}
239
240#[cfg(test)]
241mod tests {
242 use super::*;
243
244 #[test]
245 fn skip_vars_sorted() {
246 for pair in SKIP_VARS.windows(2) {
247 assert!(
248 pair[0] < pair[1],
249 "SKIP_VARS not sorted: {:?} >= {:?}",
250 pair[0],
251 pair[1]
252 );
253 }
254 }
255
256 #[test]
257 fn parse_null_env() {
258 let data = "FOO=bar\0BAZ=qux\0MULTI=hello world\0";
259 let vars = parse_null_separated_env(data);
260 assert_eq!(vars.get("FOO").unwrap(), "bar");
261 assert_eq!(vars.get("BAZ").unwrap(), "qux");
262 assert_eq!(vars.get("MULTI").unwrap(), "hello world");
263 }
264
265 #[test]
266 fn diff_new_var() {
267 let before = EnvSnapshot::new(HashMap::new(), "/home".to_string());
268 let mut after_vars = HashMap::new();
269 after_vars.insert("NEW_VAR".to_string(), "hello".to_string());
270 let after = EnvSnapshot::new(after_vars, "/home".to_string());
271
272 let out = before.diff(&after);
273 assert!(out.contains("set -gx NEW_VAR"));
274 }
275
276 #[test]
277 fn diff_removed_var() {
278 let mut before_vars = HashMap::new();
279 before_vars.insert("OLD_VAR".to_string(), "gone".to_string());
280 let before = EnvSnapshot::new(before_vars, "/home".to_string());
281 let after = EnvSnapshot::new(HashMap::new(), "/home".to_string());
282
283 let out = before.diff(&after);
284 assert!(out.lines().any(|l| l == "set -e OLD_VAR"));
285 }
286
287 #[test]
288 fn diff_changed_cwd() {
289 let before = EnvSnapshot::new(HashMap::new(), "/home".to_string());
290 let after = EnvSnapshot::new(HashMap::new(), "/tmp".to_string());
291
292 let out = before.diff(&after);
293 assert!(out.contains("cd /tmp"));
294 }
295
296 #[test]
297 fn diff_path_split() {
298 let before = EnvSnapshot::new(HashMap::new(), "/home".to_string());
299 let mut after_vars = HashMap::new();
300 after_vars.insert("PATH".to_string(), "/usr/bin:/usr/local/bin".to_string());
301 let after = EnvSnapshot::new(after_vars, "/home".to_string());
302
303 let out = before.diff(&after);
304 let path_line = out.lines().find(|l| l.contains("PATH")).unwrap();
305 assert!(path_line.contains("/usr/bin /usr/local/bin"));
306 }
307
308 #[test]
309 fn skip_bash_internal_vars() {
310 let before = EnvSnapshot::new(HashMap::new(), "/home".to_string());
311 let mut after_vars = HashMap::new();
312 after_vars.insert("BASH_VERSION".to_string(), "5.2.0".to_string());
313 after_vars.insert("REAL_VAR".to_string(), "keep".to_string());
314 let after = EnvSnapshot::new(after_vars, "/home".to_string());
315
316 let out = before.diff(&after);
317 assert!(!out.contains("BASH_VERSION"));
318 assert!(out.contains("REAL_VAR"));
319 }
320
321 #[test]
322 fn shell_escape_simple() {
323 assert_eq!(shell_escape("/usr/bin"), "/usr/bin");
324 assert_eq!(shell_escape("hello"), "hello");
325 }
326
327 #[test]
328 fn shell_escape_spaces() {
329 assert_eq!(shell_escape("hello world"), "'hello world'");
330 }
331
332 #[test]
333 fn shell_escape_quotes() {
334 assert_eq!(shell_escape("it's"), "'it'\\''s'");
335 }
336
337 #[test]
338 fn capture_current_env() {
339 let snap = EnvSnapshot::capture_current();
340 assert!(!snap.vars().is_empty());
341 assert!(!snap.cwd().is_empty());
342 assert!(snap.vars().contains_key("HOME"));
343 }
344}