use crate::commands; use commands::autoenv; use indexmap::{IndexMap, IndexSet}; use nu_errors::ShellError; use serde::Deserialize; use std::env::*; use std::process::Command; use std::{ ffi::OsString, fmt::Debug, path::{Path, PathBuf}, }; type EnvKey = String; type EnvVal = OsString; #[derive(Debug, Default)] pub struct DirectorySpecificEnvironment { pub last_seen_directory: PathBuf, //If an environment var has been added from a .nu in a directory, we track it here so we can remove it when the user leaves the directory. //If setting the var overwrote some value, we save the old value in an option so we can restore it later. added_vars: IndexMap>>, //We track directories that we have read .nu-env from. This is different from the keys in added_vars since sometimes a file only wants to run scripts. visited_dirs: IndexSet, exitscripts: IndexMap>, } #[derive(Deserialize, Debug, Default)] pub struct NuEnvDoc { pub env: Option>, pub scriptvars: Option>, pub scripts: Option>>, pub entryscripts: Option>, pub exitscripts: Option>, } impl DirectorySpecificEnvironment { pub fn new() -> DirectorySpecificEnvironment { let root_dir = if cfg!(target_os = "windows") { PathBuf::from("c:\\") } else { PathBuf::from("/") }; DirectorySpecificEnvironment { last_seen_directory: root_dir, added_vars: IndexMap::new(), visited_dirs: IndexSet::new(), exitscripts: IndexMap::new(), } } fn toml_if_trusted(&mut self, nu_env_file: &PathBuf) -> Result { let content = std::fs::read(&nu_env_file)?; if autoenv::file_is_trusted(&nu_env_file, &content)? { let mut doc: NuEnvDoc = toml::de::from_slice(&content) .map_err(|e| ShellError::untagged_runtime_error(format!("{:?}", e)))?; if let Some(scripts) = doc.scripts.as_ref() { for (k, v) in scripts { if k == "entryscripts" { doc.entryscripts = Some(v.clone()); } else if k == "exitscripts" { doc.exitscripts = Some(v.clone()); } } } return Ok(doc); } Err(ShellError::untagged_runtime_error( format!("{:?} is untrusted. Run 'autoenv trust {:?}' to trust it.\nThis needs to be done after each change to the file.", nu_env_file, nu_env_file.parent().unwrap_or_else(|| &Path::new(""))))) } pub fn maintain_autoenv(&mut self) -> Result<(), ShellError> { let mut dir = current_dir()?; if self.last_seen_directory == dir { return Ok(()); } //We track which keys we set as we go up the directory hierarchy, so that we don't overwrite a value we set in a subdir. let mut added_keys = IndexSet::new(); //We note which directories we pass so we can clear unvisited dirs later. let mut seen_directories = IndexSet::new(); //Add all .nu-envs until we reach a dir which we have already added, or we reached the root. let mut popped = true; while popped && !self.visited_dirs.contains(&dir) { let nu_env_file = dir.join(".nu-env"); if nu_env_file.exists() { let nu_env_doc = self.toml_if_trusted(&nu_env_file)?; //add regular variables from the [env section] if let Some(env) = nu_env_doc.env { for (env_key, env_val) in env { self.maybe_add_key(&mut added_keys, &dir, &env_key, &env_val); } } self.visited_dirs.insert(dir.clone()); //Add variables that need to evaluate scripts to run, from [scriptvars] section if let Some(sv) = nu_env_doc.scriptvars { for (key, script) in sv { self.maybe_add_key( &mut added_keys, &dir, &key, value_from_script(&script)?.as_str(), ); } } if let Some(es) = nu_env_doc.entryscripts { for s in es { run(s.as_str())?; } } if let Some(es) = nu_env_doc.exitscripts { self.exitscripts.insert(dir.clone(), es); } } seen_directories.insert(dir.clone()); popped = dir.pop(); } //Time to clear out vars set by directories that we have left. let mut new_vars = IndexMap::new(); for (dir, dirmap) in self.added_vars.drain(..) { if seen_directories.contains(&dir) { new_vars.insert(dir, dirmap); } else { for (k, v) in dirmap { if let Some(v) = v { std::env::set_var(k, v); } else { std::env::remove_var(k); } } } } let mut new_visited = IndexSet::new(); for dir in self.visited_dirs.drain(..) { if seen_directories.contains(&dir) { new_visited.insert(dir); } } //Run exitscripts, can not be done in same loop as new vars as some files can contain only exitscripts let mut new_exitscripts = IndexMap::new(); for (dir, scripts) in self.exitscripts.drain(..) { if seen_directories.contains(&dir) { new_exitscripts.insert(dir, scripts); } else { for s in scripts { run(s.as_str())?; } } } self.visited_dirs = new_visited; self.exitscripts = new_exitscripts; self.added_vars = new_vars; self.last_seen_directory = current_dir()?; Ok(()) } pub fn maybe_add_key( &mut self, seen_vars: &mut IndexSet, dir: &PathBuf, key: &str, val: &str, ) { //This condition is to make sure variables in parent directories don't overwrite variables set by subdirectories. if !seen_vars.contains(key) { seen_vars.insert(key.to_string()); self.added_vars .entry(dir.clone()) .or_insert(IndexMap::new()) .insert(key.to_string(), var_os(key)); std::env::set_var(key, val); } } // If the user recently ran autoenv untrust on a file, we clear the environment variables it set and make sure to not run any possible exitscripts. pub fn clear_recently_untrusted_file(&mut self) -> Result<(), ShellError> { // Figure out which file was untrusted // Remove all vars set by it let current_trusted_files: IndexSet = autoenv::read_trusted()? .files .iter() .map(|(k, _)| PathBuf::from(k)) .collect(); // We figure out which file(s) the user untrusted by taking the set difference of current trusted files in .config/nu/nu-env.toml and the files tracked by self.added_env_vars // If a file is in self.added_env_vars but not in nu-env.toml, it was just untrusted. let untrusted_files: IndexSet = self .added_vars .iter() .filter_map(|(path, _)| { if !current_trusted_files.contains(path) { return Some(path.clone()); } None }) .collect(); for path in untrusted_files { if let Some(added_keys) = self.added_vars.get(&path) { for (key, _) in added_keys { remove_var(key); } } self.exitscripts.remove(&path); self.added_vars.remove(&path); } Ok(()) } } fn run(cmd: &str) -> Result<(), ShellError> { if cfg!(target_os = "windows") { Command::new("cmd").args(&["/C", cmd]).output()? } else { Command::new("sh").arg("-c").arg(&cmd).output()? }; Ok(()) } fn value_from_script(cmd: &str) -> Result { let command = if cfg!(target_os = "windows") { Command::new("cmd").args(&["/C", cmd]).output()? } else { Command::new("sh").arg("-c").arg(&cmd).output()? }; if command.stdout.is_empty() { return Err(ShellError::untagged_runtime_error(format!( "{:?} did not return any output", cmd ))); } let response = std::str::from_utf8(&command.stdout[..command.stdout.len()]).map_err(|e| { ShellError::untagged_runtime_error(format!( "Couldn't parse stdout from command {:?}: {:?}", command, e )) })?; Ok(response.trim().to_string()) }