path: root/subplotlib/src/steplibrary/runcmd.rs

//! Step library for running subprocesses as part of scenarios

use regex::RegexBuilder;

pub use super::datadir::Datadir;
pub use crate::prelude::*;

use std::collections::HashMap;
use std::env;
use std::ffi::{OsStr, OsString};
use std::io::Write;
use std::path::PathBuf;
use std::process::{Command, Stdio};

#[derive(Default)]
pub struct Runcmd {
    env: HashMap<OsString, OsString>,
    // push to "prepend", order reversed when added to env
    paths: Vec<OsString>,
    // The following are the result of any executed command
    exitcode: Option<i32>,
    stdout: Vec<u8>,
    stderr: Vec<u8>,
}

// Note, this prefix requires that the injection env vars must have
// names which are valid unicode (and ideally ASCII)
const ENV_INJECTION_PREFIX: &str = "SUBPLOT_ENV_";

#[cfg(not(windows))]
static DEFAULT_PATHS: &[&str] = &["/usr/bin", "/bin"];

// Note, this comes from https://www.computerhope.com/issues/ch000549.htm#defaultpath
#[cfg(windows)]
static DEFAULT_PATHS: &[&str] = &[
    r"%SystemRoot%\system32",
    r"%SystemRoot%",
    r"%SystemRoot%\System32\Wbem",
];

// This us used internally to force CWD for running commands
const USE_CWD: &str = "\0USE_CWD";

impl ContextElement for Runcmd {
    fn scenario_starts(&mut self) -> StepResult {
        self.env.drain();
        self.paths.drain(..);
        self.env.insert("SHELL".into(), "/bin/sh".into());
        self.env.insert(
            "PATH".into(),
            env::var_os("PATH")
                .map(Ok)
                .unwrap_or_else(|| env::join_paths(DEFAULT_PATHS.iter()))?,
        );

        // Having assembled the 'default' environment, override it with injected
        // content from the calling environment.
        for (k, v) in env::vars_os() {
            if let Some(k) = k.to_str() {
                if let Some(k) = k.strip_prefix(ENV_INJECTION_PREFIX) {
                    self.env.insert(k.into(), v);
                }
            }
        }
        Ok(())
    }
}

impl Runcmd {
    /// Prepend the given location to the run path
    pub fn prepend_to_path<S: Into<OsString>>(&mut self, element: S) {
        self.paths.push(element.into());
    }

    /// Retrieve the last run command's stdout as a string.
    ///
    /// This does a lossy conversion from utf8 so should always succeed.
    pub fn stdout_as_string(&self) -> String {
        String::from_utf8_lossy(&self.stdout).into_owned()
    }

    /// Retrieve the last run command's stderr as a string.
    ///
    /// This does a lossy conversion from utf8 so should always succeed.
    pub fn stderr_as_string(&self) -> String {
        String::from_utf8_lossy(&self.stderr).into_owned()
    }
}

#[step]
pub fn helper_script(context: &Datadir, script: SubplotDataFile) {
    context
        .open_write(script.name())?
        .write_all(script.data())?;
}

#[step]
pub fn helper_srcdir_path(context: &mut Runcmd) {
    context.prepend_to_path(env!("CARGO_MANIFEST_DIR"));
}

#[step]
#[context(Datadir)]
#[context(Runcmd)]
pub fn run(context: &ScenarioContext, argv0: &str, args: &str) {
    try_to_run::call(context, argv0, args)?;
    exit_code_is::call(context, 0)?;
}

#[step]
#[context(Datadir)]
#[context(Runcmd)]
pub fn run_in(context: &ScenarioContext, dirname: &str, argv0: &str, args: &str) {
    try_to_run_in::call(context, dirname, argv0, args)?;
    exit_code_is::call(context, 0)?;
}

#[step]
#[context(Datadir)]
#[context(Runcmd)]
pub fn try_to_run(context: &ScenarioContext, argv0: &str, args: &str) {
    try_to_run_in::call(context, USE_CWD, argv0, args)?;
}

#[step]
#[context(Datadir)]
#[context(Runcmd)]
pub fn try_to_run_in(context: &ScenarioContext, dirname: &str, argv0: &str, args: &str) {
    // This is the core of runcmd and is how we handle things
    let argv0: PathBuf = if argv0.starts_with('.') {
        context.with(
            |datadir: &Datadir| datadir.canonicalise_filename(argv0),
            false,
        )?
    } else {
        argv0.into()
    };
    let mut datadir = context.with(
        |datadir: &Datadir| Ok(datadir.base_path().to_path_buf()),
        false,
    )?;
    if dirname != USE_CWD {
        datadir = datadir.join(dirname);
    }
    let mut proc = Command::new(argv0);
    proc.args(&shell_words::split(args)?);
    proc.current_dir(&datadir);
    proc.env("HOME", &datadir);
    proc.env("TMPDIR", &datadir);

    let mut curpath = context.with(
        |runcmd: &Runcmd| {
            let mut curpath = None;
            for (k, v) in runcmd.env.iter() {
                proc.env(k, v);
                if k == "PATH" {
                    curpath = Some(v.to_owned());
                }
            }
            Ok(curpath)
        },
        false,
    )?;

    let path = context.with(
        |runcmd: &Runcmd| {
            Ok(env::join_paths(
                runcmd
                    .paths
                    .iter()
                    .rev()
                    .map(PathBuf::from)
                    .chain(env::split_paths(
                        curpath.as_deref().unwrap_or_else(|| OsStr::new("")),
                    )),
            )?)
        },
        false,
    )?;
    proc.env("PATH", path);
    proc.stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    let mut output = proc.output()?;
    context.with_mut(
        |runcmd: &mut Runcmd| {
            std::mem::swap(&mut runcmd.stdout, &mut output.stdout);
            std::mem::swap(&mut runcmd.stderr, &mut output.stderr);
            runcmd.exitcode = output.status.code();
            Ok(())
        },
        false,
    )?;
}

#[step]
pub fn exit_code_is(context: &Runcmd, exit: i32) {
    if context.exitcode != Some(exit) {
        throw!(format!(
            "expected exit code {}, but had {:?}",
            exit, context.exitcode
        ));
    }
}

#[step]
pub fn exit_code_is_not(context: &Runcmd, exit: i32) {
    if context.exitcode.is_none() || context.exitcode == Some(exit) {
        throw!(format!("Expected exit code to not equal {}", exit));
    }
}

#[step]
#[context(Runcmd)]
pub fn exit_code_is_zero(context: &ScenarioContext) {
    exit_code_is::call(context, 0)
}

#[step]
#[context(Runcmd)]
pub fn exit_code_is_nonzero(context: &ScenarioContext) {
    exit_code_is_not::call(context, 0)
}

enum Stream {
    Stdout,
    Stderr,
}
enum MatchKind {
    Exact,
    Contains,
    Regex,
}

#[throws(StepError)]
fn check_matches(runcmd: &Runcmd, which: Stream, how: MatchKind, against: &str) -> bool {
    let stream = match which {
        Stream::Stdout => &runcmd.stdout,
        Stream::Stderr => &runcmd.stderr,
    };
    let against = if matches!(how, MatchKind::Regex) {
        against.to_string()
    } else {
        unescape::unescape(against).ok_or("unable to unescape input")?
    };
    match how {
        MatchKind::Exact => stream.as_slice() == against.as_bytes(),
        MatchKind::Contains => stream
            .windows(against.len())
            .any(|window| window == against.as_bytes()),
        MatchKind::Regex => {
            let stream = String::from_utf8_lossy(stream);
            let regex = RegexBuilder::new(&against).multi_line(true).build()?;
            regex.is_match(&stream)
        }
    }
}

#[step]
pub fn stdout_is(runcmd: &Runcmd, text: &str) {
    if !check_matches(runcmd, Stream::Stdout, MatchKind::Exact, text)? {
        throw!(format!("stdout is not {:?}", text));
    }
}

#[step]
pub fn stdout_isnt(runcmd: &Runcmd, text: &str) {
    if check_matches(runcmd, Stream::Stdout, MatchKind::Exact, text)? {
        throw!(format!("stdout is exactly {:?}", text));
    }
}

#[step]
pub fn stderr_is(runcmd: &Runcmd, text: &str) {
    if !check_matches(runcmd, Stream::Stderr, MatchKind::Exact, text)? {
        throw!(format!("stderr is not {:?}", text));
    }
}

#[step]
pub fn stderr_isnt(runcmd: &Runcmd, text: &str) {
    if check_matches(runcmd, Stream::Stderr, MatchKind::Exact, text)? {
        throw!(format!("stderr is exactly {:?}", text));
    }
}

#[step]
pub fn stdout_contains(runcmd: &Runcmd, text: &str) {
    if !check_matches(runcmd, Stream::Stdout, MatchKind::Contains, text)? {
        throw!(format!("stdout does not contain {:?}", text));
    }
}

#[step]
pub fn stdout_doesnt_contain(runcmd: &Runcmd, text: &str) {
    if check_matches(runcmd, Stream::Stdout, MatchKind::Contains, text)? {
        throw!(format!("stdout contains {:?}", text));
    }
}

#[step]
pub fn stderr_contains(runcmd: &Runcmd, text: &str) {
    if !check_matches(runcmd, Stream::Stderr, MatchKind::Contains, text)? {
        throw!(format!("stderr does not contain {:?}", text));
    }
}

#[step]
pub fn stderr_doesnt_contain(runcmd: &Runcmd, text: &str) {
    if check_matches(runcmd, Stream::Stderr, MatchKind::Contains, text)? {
        throw!(format!("stderr contains {:?}", text));
    }
}

#[step]
pub fn stdout_matches_regex(runcmd: &Runcmd, regex: &str) {
    if !check_matches(runcmd, Stream::Stdout, MatchKind::Regex, regex)? {
        throw!(format!("stdout does not match {:?}", regex));
    }
}

#[step]
pub fn stdout_doesnt_match_regex(runcmd: &Runcmd, regex: &str) {
    if check_matches(runcmd, Stream::Stdout, MatchKind::Regex, regex)? {
        throw!(format!("stdout matches {:?}", regex));
    }
}

#[step]
pub fn stderr_matches_regex(runcmd: &Runcmd, regex: &str) {
    if !check_matches(runcmd, Stream::Stderr, MatchKind::Regex, regex)? {
        throw!(format!("stderr does not match {:?}", regex));
    }
}

#[step]
pub fn stderr_doesnt_match_regex(runcmd: &Runcmd, regex: &str) {
    if check_matches(runcmd, Stream::Stderr, MatchKind::Regex, regex)? {
        throw!(format!("stderr matches {:?}", regex));
    }
}