Lots more updates
Also begin adding rust building capabilities to be able to write rust binaries for some commands.
This commit is contained in:
parent
624508dd14
commit
dd1cfa79e7
GPG key ID:
C7D890804F01E9F0
52 changed files with 2509 additions and 150 deletions
355
rust/lib/common/src/proc.rs
Normal file
355
rust/lib/common/src/proc.rs
Normal file
|
|
@ -0,0 +1,355 @@
|
|||
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<u8>),
|
||||
Inherit,
|
||||
}
|
||||
|
||||
impl Stdin {
|
||||
fn into_data(self) -> Option<Vec<u8>> {
|
||||
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(_) => "<REDACTED>",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Command {
|
||||
program: String,
|
||||
args: Vec<String>,
|
||||
env: BTreeMap<String, Option<EnvData>>,
|
||||
is_sudo: bool,
|
||||
show_command: bool,
|
||||
stdin: Stdin,
|
||||
stderr: Stderr,
|
||||
}
|
||||
|
||||
impl Command {
|
||||
pub fn new(program: impl Into<String>) -> 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<Item = &str> {
|
||||
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<Item = (&str, Option<&EnvData>)> {
|
||||
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<S: Serialize>(&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<S: Serialize>(&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<String>) -> &mut Self {
|
||||
self.stdin = Stdin::Pipe(Vec::from(stdin.into()));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn stdin_bytes(&mut self, stdin: impl Into<Vec<u8>>) -> &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<String>) -> &mut Self {
|
||||
self.args.push(arg.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn args<A>(&mut self, args: impl IntoIterator<Item = A>) -> &mut Self
|
||||
where
|
||||
A: Into<String>,
|
||||
{
|
||||
self.args.extend(args.into_iter().map(Into::into));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn env(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
|
||||
let _ = self
|
||||
.env
|
||||
.insert(key.into(), Some(EnvData::Insensitive(value.into())));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn env_sensitive(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
|
||||
let _ = self
|
||||
.env
|
||||
.insert(key.into(), Some(EnvData::Sensitive(value.into())));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn env_remove(&mut self, key: impl Into<String>) -> &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<Vec<u8>> {
|
||||
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<String> {
|
||||
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<D: for<'de> Deserialize<'de>>(&mut self) -> anyhow::Result<D> {
|
||||
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<ExitStatus> {
|
||||
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<Vec<u8>> {
|
||||
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")
|
||||
);
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue