custom.rs

use std::env;
use std::io::Write;
use std::path::Path;
use std::process::{Command, Stdio};
use std::time::Duration;
use std::time::Instant;

use process_control::{ChildExt, Control, Output};

use super::{Context, Module, ModuleConfig};

use crate::{
    config::Either, configs::custom::CustomConfig, formatter::StringFormatter,
    utils::create_command,
};

/// Creates a custom module with some configuration
///
/// The relevant TOML config will set the files, extensions, and directories needed
/// for the module to be displayed. If none of them match, and optional "when"
/// command can be run -- if its result is 0, the module will be shown.
///
/// Finally, the content of the module itself is also set by a command.
pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
    let start: Instant = Instant::now();
    let toml_config = context.config.get_custom_module_config(name).expect(
        "modules::custom::module should only be called after ensuring that the module exists",
    );
    let config = CustomConfig::load(toml_config);

    if let Some(os) = config.os {
        if os != env::consts::OS && !(os == "unix" && cfg!(unix)) {
            return None;
        }
    }

    let mut is_match = context
        .try_begin_scan()?
        .set_extensions(&config.detect_extensions)
        .set_files(&config.detect_files)
        .set_folders(&config.detect_folders)
        .is_match();

    if !is_match {
        is_match = match config.when {
            Either::First(b) => b,
            Either::Second(s) => exec_when(s, &config, context),
        };

        if !is_match {
            return None;
        }
    }

    let mut module = Module::new(name, config.description, Some(toml_config));

    let parsed = StringFormatter::new(config.format).and_then(|formatter| {
        formatter
            .map_meta(|var, _| match var {
                "symbol" => Some(config.symbol),
                _ => None,
            })
            .map_style(|variable| match variable {
                "style" => Some(Ok(config.style)),
                _ => None,
            })
            .map_no_escaping(|variable| match variable {
                "output" => {
                    let output = exec_command(config.command, context, &config)?;
                    let trimmed = output.trim();

                    if trimmed.is_empty() {
                        None
                    } else {
                        Some(Ok(trimmed.to_string()))
                    }
                }
                _ => None,
            })
            .parse(None, Some(context))
    });

    match parsed {
        Ok(segments) => module.set_segments(segments),
        Err(error) => {
            log::warn!("Error in module `custom.{}`:\n{}", name, error);
        }
    };
    let elapsed = start.elapsed();
    log::trace!("Took {:?} to compute custom module {:?}", elapsed, name);
    module.duration = elapsed;
    Some(module)
}

/// Return the invoking shell, using `shell` and fallbacking in order to `STARSHIP_SHELL` and "sh"/"cmd"
fn get_shell<'a, 'b>(
    shell_args: &'b [&'a str],
    context: &Context,
) -> (std::borrow::Cow<'a, str>, &'b [&'a str]) {
    if !shell_args.is_empty() {
        (shell_args[0].into(), &shell_args[1..])
    } else if let Some(env_shell) = context.get_env("STARSHIP_SHELL") {
        (env_shell.into(), &[] as &[&str])
    } else if cfg!(windows) {
        // `/C` is added by `handle_shell`
        ("cmd".into(), &[] as &[&str])
    } else {
        ("sh".into(), &[] as &[&str])
    }
}

/// Attempt to run the given command in a shell by passing it as either `stdin` or an argument to `get_shell()`,
/// depending on the configuration or by invoking a platform-specific falback shell if `shell` is empty.
fn shell_command(cmd: &str, config: &CustomConfig, context: &Context) -> Option<Output> {
    let (shell, shell_args) = get_shell(config.shell.0.as_ref(), context);
    let mut use_stdin = config.use_stdin;

    let mut command = match create_command(shell.as_ref()) {
        Ok(command) => command,
        // Don't attempt to use fallback shell if the user specified a shell
        Err(error) if !shell_args.is_empty() => {
            log::debug!(
                "Error creating command with STARSHIP_SHELL, falling back to fallback shell: {}",
                error
            );

            // Skip `handle_shell` and just set the shell and command
            use_stdin = Some(!cfg!(windows));

            if cfg!(windows) {
                let mut c = create_command("cmd").ok()?;
                c.arg("/C");
                c
            } else {
                let mut c = create_command("/usr/bin/env").ok()?;
                c.arg("sh");
                c
            }
        }
        _ => return None,
    };

    command
        .current_dir(&context.current_dir)
        .args(shell_args)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());

    let use_stdin = use_stdin.unwrap_or_else(|| handle_shell(&mut command, &shell, shell_args));

    if !use_stdin {
        command.arg(cmd);
    }

    let mut child = match command.spawn() {
        Ok(child) => child,
        Err(error) => {
            log::debug!(
                "Failed to run command with given shell or STARSHIP_SHELL env variable:: {}",
                error
            );
            return None;
        }
    };

    if use_stdin {
        child.stdin.as_mut()?.write_all(cmd.as_bytes()).ok()?;
    }

    let mut output = child.controlled_with_output();

    if !config.ignore_timeout {
        output = output
            .time_limit(Duration::from_millis(context.root_config.command_timeout))
            .terminate_for_timeout()
    }

    match output.wait().ok()? {
        None => {
            log::warn!("Executing custom command {cmd:?} timed out.");
            log::warn!("You can set command_timeout in your config to a higher value or set ignore_timeout to true for this module to allow longer-running commands to keep executing.");
            None
        }
        Some(status) => Some(status),
    }
}

/// Execute the given command capturing all output, and return whether it return 0
fn exec_when(cmd: &str, config: &CustomConfig, context: &Context) -> bool {
    log::trace!("Running '{}'", cmd);

    if let Some(output) = shell_command(cmd, config, context) {
        if !output.status.success() {
            log::trace!("non-zero exit code '{:?}'", output.status.code());
            log::trace!(
                "stdout: {}",
                std::str::from_utf8(&output.stdout).unwrap_or("<invalid utf8>")
            );
            log::trace!(
                "stderr: {}",
                std::str::from_utf8(&output.stderr).unwrap_or("<invalid utf8>")
            );
        }

        output.status.success()
    } else {
        log::debug!("Cannot start command");

        false
    }
}

/// Execute the given command, returning its output on success
fn exec_command(cmd: &str, context: &Context, config: &CustomConfig) -> Option<String> {
    log::trace!("Running '{cmd}'");

    if let Some(output) = shell_command(cmd, config, context) {
        if !output.status.success() {
            log::trace!("Non-zero exit code '{:?}'", output.status.code());
            log::trace!(
                "stdout: {}",
                std::str::from_utf8(&output.stdout).unwrap_or("<invalid utf8>")
            );
            log::trace!(
                "stderr: {}",
                std::str::from_utf8(&output.stderr).unwrap_or("<invalid utf8>")
            );
            return None;
        }

        Some(String::from_utf8_lossy(&output.stdout).into())
    } else {
        None
    }
}

/// If the specified shell refers to `PowerShell`, adds the arguments "-Command -" to the
/// given command.
/// Retruns `false` if the shell shell expects scripts as arguments, `true` if as `stdin`.
fn handle_shell(command: &mut Command, shell: &str, shell_args: &[&str]) -> bool {
    let shell_exe = Path::new(shell).file_stem();
    let no_args = shell_args.is_empty();

    match shell_exe.and_then(std::ffi::OsStr::to_str) {
        Some("pwsh" | "powershell") => {
            if no_args {
                command.arg("-NoProfile").arg("-Command").arg("-");
            }
            true
        }
        Some("cmd") => {
            if no_args {
                command.arg("/C");
            }
            false
        }
        Some("nu") => {
            if no_args {
                command.arg("-c");
            }
            false
        }
        _ => true,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::test::ModuleRenderer;
    use ansi_term::Color;
    use std::fs::File;
    use std::io;

    #[cfg(not(windows))]
    const SHELL: &[&str] = &["/bin/sh"];
    #[cfg(windows)]
    const SHELL: &[&str] = &["cmd"];

    #[cfg(not(windows))]
    const FAILING_COMMAND: &str = "false";
    #[cfg(windows)]
    const FAILING_COMMAND: &str = "color 00";

    const UNKNOWN_COMMAND: &str = "ydelsyiedsieudleylse dyesdesl";

    fn render_cmd(cmd: &str) -> io::Result<Option<String>> {
        let dir = tempfile::tempdir()?;
        let cmd = cmd.to_owned();
        let shell = SHELL
            .iter()
            .map(std::borrow::ToOwned::to_owned)
            .collect::<Vec<_>>();
        let out = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "$output"
                command = cmd
                shell = shell
                when = true
                ignore_timeout = true
            })
            .collect();
        dir.close()?;
        Ok(out)
    }

    fn render_when(cmd: &str) -> io::Result<bool> {
        let dir = tempfile::tempdir()?;
        let cmd = cmd.to_owned();
        let shell = SHELL
            .iter()
            .map(std::borrow::ToOwned::to_owned)
            .collect::<Vec<_>>();
        let out = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                when = cmd
                shell = shell
                ignore_timeout = true
            })
            .collect()
            .is_some();
        dir.close()?;
        Ok(out)
    }

    #[test]
    fn when_returns_right_value() -> io::Result<()> {
        assert!(render_cmd("echo hello")?.is_some());
        assert!(render_cmd(FAILING_COMMAND)?.is_none());
        Ok(())
    }

    #[test]
    fn when_returns_false_if_invalid_command() -> io::Result<()> {
        assert!(!render_when(UNKNOWN_COMMAND)?);
        Ok(())
    }

    #[test]
    #[cfg(not(windows))]
    fn command_returns_right_string() -> io::Result<()> {
        assert_eq!(render_cmd("echo hello")?, Some("hello".into()));
        assert_eq!(render_cmd("echo 강남스타일")?, Some("강남스타일".into()));
        Ok(())
    }

    #[test]
    #[cfg(windows)]
    fn command_returns_right_string() -> io::Result<()> {
        assert_eq!(render_cmd("echo hello")?, Some("hello".into()));
        assert_eq!(render_cmd("echo 강남스타일")?, Some("강남스타일".into()));
        Ok(())
    }

    #[test]
    #[cfg(not(windows))]
    fn command_ignores_stderr() -> io::Result<()> {
        assert_eq!(render_cmd("echo foo 1>&2; echo bar")?, Some("bar".into()));
        assert_eq!(render_cmd("echo foo; echo bar 1>&2")?, Some("foo".into()));
        Ok(())
    }

    #[test]
    #[cfg(windows)]
    fn command_ignores_stderr() -> io::Result<()> {
        assert_eq!(render_cmd("echo foo 1>&2 & echo bar")?, Some("bar".into()));
        assert_eq!(render_cmd("echo foo& echo bar 1>&2")?, Some("foo".into()));
        Ok(())
    }

    #[test]
    fn command_can_fail() -> io::Result<()> {
        assert_eq!(render_cmd(FAILING_COMMAND)?, None);
        assert_eq!(render_cmd(UNKNOWN_COMMAND)?, None);
        Ok(())
    }

    #[test]
    fn cwd_command() -> io::Result<()> {
        let dir = tempfile::tempdir()?;

        let mut f = File::create(dir.path().join("a.txt"))?;
        write!(f, "hello")?;
        f.sync_all()?;

        let cat = if cfg!(windows) { "type" } else { "cat" };
        let cmd = format!("{cat} a.txt");

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                command = cmd
                when = true
                ignore_timeout = true
            })
            .collect();
        let expected = Some(format!("{}", Color::Green.bold().paint("hello ")));

        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    fn cwd_when() -> io::Result<()> {
        let dir = tempfile::tempdir()?;

        File::create(dir.path().join("a.txt"))?.sync_all()?;

        let cat = if cfg!(windows) { "type" } else { "cat" };
        let cmd = format!("{cat} a.txt");

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                when = cmd
                ignore_timeout = true
            })
            .collect();
        let expected = Some("test".to_owned());

        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    fn use_stdin_false() -> io::Result<()> {
        let dir = tempfile::tempdir()?;

        let shell = if cfg!(windows) {
            vec![
                "powershell".to_owned(),
                "-NoProfile".to_owned(),
                "-Command".to_owned(),
            ]
        } else {
            vec!["sh".to_owned(), "-c".to_owned()]
        };

        // `use_stdin = false` doesn't like Korean on Windows
        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                command = "echo test"
                when = true
                use_stdin = false
                shell = shell
                ignore_timeout = true
            })
            .collect();
        let expected = Some(format!("{}", Color::Green.bold().paint("test ")));

        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    fn use_stdin_true() -> io::Result<()> {
        let dir = tempfile::tempdir()?;

        let shell = if cfg!(windows) {
            vec![
                "powershell".to_owned(),
                "-NoProfile".to_owned(),
                "-Command".to_owned(),
                "-".to_owned(),
            ]
        } else {
            vec!["sh".to_owned()]
        };

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                command = "echo 강남스타일"
                when = true
                use_stdin = true
                ignore_timeout = true
                shell = shell
            })
            .collect();
        let expected = Some(format!("{}", Color::Green.bold().paint("강남스타일 ")));

        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    #[cfg(not(windows))]
    fn when_true_with_string() -> std::io::Result<()> {
        let dir = tempfile::tempdir()?;

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                shell = ["sh"]
                when = "true"
                ignore_timeout = true
            })
            .collect();
        let expected = Some("test".to_string());
        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    #[cfg(not(windows))]
    fn when_false_with_string() -> std::io::Result<()> {
        let dir = tempfile::tempdir()?;

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                shell = ["sh"]
                when = "false"
                ignore_timeout = true
            })
            .collect();
        let expected = None;
        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    fn when_true_with_bool() -> std::io::Result<()> {
        let dir = tempfile::tempdir()?;

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                when = true
            })
            .collect();
        let expected = Some("test".to_string());
        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    #[cfg(not(windows))]
    fn when_false_with_bool() -> std::io::Result<()> {
        let dir = tempfile::tempdir()?;

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                when = false
            })
            .collect();
        let expected = None;
        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    fn timeout_short_cmd() -> std::io::Result<()> {
        let dir = tempfile::tempdir()?;

        let shell = if cfg!(windows) {
            "powershell".to_owned()
        } else {
            "sh".to_owned()
        };

        let when = if cfg!(windows) {
            "$true".to_owned()
        } else {
            "true".to_owned()
        };

        // Use a long timeout to ensure that the test doesn't fail
        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                command_timeout = 100_000
                [custom.test]
                format = "test"
                when = when
                shell = shell
                ignore_timeout = false
            })
            .collect();
        let expected = Some("test".to_owned());
        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    fn timeout_cmd() -> std::io::Result<()> {
        let dir = tempfile::tempdir()?;

        let shell = if cfg!(windows) {
            "powershell".to_owned()
        } else {
            "sh".to_owned()
        };

        // Use a long timeout to ensure that the test doesn't fail
        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                when = "sleep 3"
                shell = shell
                ignore_timeout = false
            })
            .collect();
        let expected = None;
        assert_eq!(expected, actual);

        dir.close()
    }

    #[test]
    fn config_aliases_work() -> std::io::Result<()> {
        let dir = tempfile::tempdir()?;

        File::create(dir.path().join("a.txt"))?;
        std::fs::create_dir(dir.path().join("dir"))?;

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                files = ["a.txt"]
            })
            .collect();
        let expected = Some("test".to_string());
        assert_eq!(expected, actual);

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                extensions = ["txt"]
            })
            .collect();
        let expected = Some("test".to_string());
        assert_eq!(expected, actual);

        let actual = ModuleRenderer::new("custom.test")
            .path(dir.path())
            .config(toml::toml! {
                [custom.test]
                format = "test"
                directories = ["dir"]
            })
            .collect();
        let expected = Some("test".to_string());
        assert_eq!(expected, actual);

        dir.close()
    }
}