machines/rust/lib/common/src/proc.rs

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")
        );
    }
}