use std::{ collections::BTreeMap, io::Write as _, ops::Deref, process::{Child, ExitStatus, Stdio}, }; use anyhow::Context; use log::Level; use serde::{Deserialize, Serialize}; use crate::{json, proc::util::command_to_string}; mod util; #[derive(Debug)] enum Stdin { Null, Pipe(Vec), Inherit, } impl Stdin { fn into_data(self) -> Option> { match self { Self::Pipe(d) => Some(d), _ => None, } } fn as_std_stdio(&self) -> Stdio { match self { Self::Inherit => Stdio::inherit(), Self::Pipe(_) => Stdio::piped(), Self::Null => Stdio::null(), } } } #[derive(Debug)] enum Stderr { Pipe, Inherit, } impl Stderr { fn as_std_stdio(&self) -> Stdio { match self { Self::Inherit => Stdio::inherit(), Self::Pipe => Stdio::piped(), } } } #[derive(Debug)] pub enum EnvData { Insensitive(String), Sensitive(String), } impl EnvData { const fn as_str(&self) -> &str { match self { Self::Insensitive(d) | Self::Sensitive(d) => d.as_str(), } } fn as_potentially_redacted_str(&self) -> &str { match self { Self::Insensitive(s) => s.as_str(), Self::Sensitive(_) => "", } } } #[derive(Debug)] pub struct Command { program: String, args: Vec, env: BTreeMap>, is_sudo: bool, show_command: bool, stdin: Stdin, stderr: Stderr, } impl Command { pub fn new(program: impl Into) -> Self { Self { program: program.into(), args: Vec::new(), env: BTreeMap::new(), is_sudo: false, show_command: false, stdin: Stdin::Null, stderr: Stderr::Pipe, } } pub fn get_program(&self) -> &str { if self.is_sudo { "sudo" } else { &self.program } } pub fn get_args(&self) -> impl Iterator { let prefix_vec = if self.is_sudo { Vec::from([self.program.as_str()]) } else { Vec::new() }; prefix_vec .into_iter() .chain(self.args.iter().map(Deref::deref)) } pub fn get_envs(&self) -> impl Iterator)> { self.env.iter().map(|(k, v)| (k.as_str(), v.as_ref())) } /// Sudo will automatically set stdin and stderr to inherit, to allow the user to enter the sudo password pub fn sudo(&mut self) -> &mut Self { self.is_sudo = true; self.stdin = Stdin::Inherit; self.stderr = Stderr::Inherit; self } pub fn announce(&mut self) -> &mut Self { self.show_command = true; self } pub fn stdin_json(&mut self, stdin: &S) -> anyhow::Result<&mut Self> { self.stdin = Stdin::Pipe(json::to_vec(stdin).context("Could not convert stdin to json")?); Ok(self) } pub fn stdin_json_base64(&mut self, stdin: &S) -> anyhow::Result<&mut Self> { let json = json::to_string(stdin).context("Could not convert stdin to json")?; let base64 = crate::base64::encode(json.as_bytes()); self.stdin = Stdin::Pipe(base64.into()); Ok(self) } pub fn stdin_string(&mut self, stdin: impl Into) -> &mut Self { self.stdin = Stdin::Pipe(Vec::from(stdin.into())); self } pub fn stdin_bytes(&mut self, stdin: impl Into>) -> &mut Self { self.stdin = Stdin::Pipe(stdin.into()); self } pub fn stdin_inherit(&mut self) -> &mut Self { self.stdin = Stdin::Inherit; self } pub fn stderr_inherit(&mut self) -> &mut Self { self.stderr = Stderr::Inherit; self } pub fn arg(&mut self, arg: impl Into) -> &mut Self { self.args.push(arg.into()); self } pub fn args(&mut self, args: impl IntoIterator) -> &mut Self where A: Into, { self.args.extend(args.into_iter().map(Into::into)); self } pub fn env(&mut self, key: impl Into, value: impl Into) -> &mut Self { let _ = self .env .insert(key.into(), Some(EnvData::Insensitive(value.into()))); self } pub fn env_sensitive(&mut self, key: impl Into, value: impl Into) -> &mut Self { let _ = self .env .insert(key.into(), Some(EnvData::Sensitive(value.into()))); self } pub fn env_remove(&mut self, key: impl Into) -> &mut Self { self.env.insert(key.into(), None); self } pub fn env_clear(&mut self) -> &mut Self { self.env.clear(); for k in std::env::args() { self.env.insert(k, None); } self } fn as_command(&self) -> std::process::Command { let is_sudo = self.is_sudo; let show_command = self.show_command; let mut cmd = if self.is_sudo { let mut cmd = std::process::Command::new("sudo"); cmd.arg(&self.program); cmd } else { std::process::Command::new(&self.program) }; cmd.args(self.args.iter()); for (k, v) in self.env.iter() { if let Some(v) = v { cmd.env(k, v.as_str()); } else { cmd.env_remove(k); } } if !self.env.contains_key("LOGEVEL") { if let Ok(loglevel) = std::env::var("LOGLEVEL") { match loglevel.to_ascii_lowercase().as_str() { "trace" | "debug" => cmd.env("LOGLEVEL", "debug"), "verbose" => cmd.env("LOGLEVEL", "verbose"), "info" => cmd.env("LOGLEVEL", "info"), "warning" | "warn" => cmd.env("LOGLEVEL", "warning"), "error" => cmd.env("LOGLEVEL", "error"), _ => cmd.env_remove("LOGLEVEL"), }; } } if is_sudo { log::info!("[SUDO] Running command {}", command_to_string(self)); } else { let level = if show_command { Level::Info } else { Level::Trace }; log::log!(level, "Running command {}", command_to_string(self)); } cmd } /// WARNING: Spawning a command consumes all of its stdin, as such spawning it twice may not yield the same result pub fn try_spawn_to_bytes(&mut self) -> anyhow::Result> { let mut cmd = self.as_command(); let mut child = cmd .stderr(self.stderr.as_std_stdio()) .stdin(self.stdin.as_std_stdio()) .stdout(Stdio::piped()) .spawn() .with_context(|| format!("Could not spawn command: {}", command_to_string(self)))?; let mut stdin = Stdin::Null; std::mem::swap(&mut self.stdin, &mut stdin); let join_handle = if let Some(data) = stdin.into_data() { let mut stdin_pipe = child.stdin.take().expect("Child has no stdin"); Some(std::thread::spawn(move || { stdin_pipe .write_all(data.as_slice()) .expect("Could not write to child"); })) } else { None }; let output = wait_with_output(child, || command_to_string(self)); if let Some(join_handle) = join_handle { join_handle .join() .map_err(|e| anyhow::format_err!("Thread sending stdin panicked: {e:?}"))?; } output } /// WARNING: Spawning a command consumes all of its stdin, as such spawning it twice may not yield the same result /// Using try_spawn_to_string will trim a single trailing newline, if you don't want this, use to_bytes and convert the string manually. pub fn try_spawn_to_string(&mut self) -> anyhow::Result { let output = self.try_spawn_to_bytes()?; let mut output: String = output.try_into().map_err(|_| { anyhow::format_err!( "Command {} didn't produce valid utf-8 output", command_to_string(self) ) })?; if output.ends_with("\n") { output.truncate(output.len() - 1); } Ok(output) } /// WARNING: Spawning a command consumes all of its stdin, as such spawning it twice may not yield the same result pub fn try_spawn_to_json Deserialize<'de>>(&mut self) -> anyhow::Result { let output = self.try_spawn_to_string()?; json::from_str(&output).with_context(|| { format!( "Could not parse output of {} as json", command_to_string(self) ) }) } /// WARNING: Spawning a command consumes all of its stdin, as such spawning it twice may not yield the same result pub fn try_spawn_stdout_inherit(&mut self) -> anyhow::Result { let mut cmd = self.as_command(); cmd.stdout(Stdio::inherit()); cmd.status() .with_context(|| format!("Could not spawn command: {}", command_to_string(self))) } } fn wait_with_output(child: Child, cmd_str: impl Fn() -> String) -> anyhow::Result> { let output = child .wait_with_output() .with_context(|| format!("Could not wait for output of commnad: {}", cmd_str()))?; if !output.status.success() { return Err(anyhow::format_err!( "Command {}, exited unexpectedly: {:?}. With stderr: {}", cmd_str(), output.status, String::from_utf8_lossy(&output.stderr), )); } Ok(output.stdout) } #[cfg(test)] mod tests { use super::Command; #[test] fn test_spawn() { let mut echo = Command::new("echo"); echo.args(["Hello", "World"]); assert_eq!( "Hello World", echo.try_spawn_to_string() .expect("Should be able to echo Hello World") ); } #[test] fn test_spawn_stdin() { let mut rev = Command::new("rev"); assert_eq!( "dlroW olleH", rev.stdin_string("Hello World".to_string()) .try_spawn_to_string() .expect("Should be able to rev Hello World") ); } }