feat: Add pulumi module (#3055)

7 files changed, 357 insertions, 1 deletions

diff --git a/src/configs/mod.rs b/src/configs/mod.rs
index 733c8ae35..80d47cacf 100644
--- a/src/configs/mod.rs
+++ b/src/configs/mod.rs
@@ -48,6 +48,7 @@ pub mod openstack;
 pub mod package;
 pub mod perl;
 pub mod php;
+pub mod pulumi;
 pub mod purescript;
 pub mod python;
 pub mod red;
@@ -125,6 +126,7 @@ pub struct FullConfig<'a> {
     package: package::PackageConfig<'a>,
     perl: perl::PerlConfig<'a>,
     php: php::PhpConfig<'a>,
+    pulumi: pulumi::PulumiConfig<'a>,
     purescript: purescript::PureScriptConfig<'a>,
     python: python::PythonConfig<'a>,
     red: red::RedConfig<'a>,
@@ -200,6 +202,7 @@ impl<'a> Default for FullConfig<'a> {
             package: Default::default(),
             perl: Default::default(),
             php: Default::default(),
+            pulumi: Default::default(),
             purescript: Default::default(),
             python: Default::default(),
             red: Default::default(),
diff --git a/src/configs/pulumi.rs b/src/configs/pulumi.rs
new file mode 100644
index 000000000..5a796dd79
--- /dev/null
+++ b/src/configs/pulumi.rs
@@ -0,0 +1,25 @@
+use crate::config::ModuleConfig;
+
+use serde::Serialize;
+use starship_module_config_derive::ModuleConfig;
+
+#[derive(Clone, ModuleConfig, Serialize)]
+pub struct PulumiConfig<'a> {
+    pub format: &'a str,
+    pub version_format: &'a str,
+    pub symbol: &'a str,
+    pub style: &'a str,
+    pub disabled: bool,
+}
+
+impl<'a> Default for PulumiConfig<'a> {
+    fn default() -> Self {
+        PulumiConfig {
+            format: "via [$symbol$stack]($style) ",
+            version_format: "v${raw}",
+            symbol: " ",
+            style: "bold 5",
+            disabled: false,
+        }
+    }
+}
diff --git a/src/configs/starship_root.rs b/src/configs/starship_root.rs
index dfbf87224..04378f651 100644
--- a/src/configs/starship_root.rs
+++ b/src/configs/starship_root.rs
@@ -53,6 +53,7 @@ pub const PROMPT_ORDER: &[&str] = &[
     "ocaml",
     "perl",
     "php",
+    "pulumi",
     "purescript",
     "python",
     "rlang",
diff --git a/src/module.rs b/src/module.rs
index 2c88b5533..c886b6a30 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -53,6 +53,7 @@ pub const ALL_MODULES: &[&str] = &[
     "package",
     "perl",
     "php",
+    "pulumi",
     "purescript",
     "python",
     "red",
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 100b05a2a..61640f78b 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -43,6 +43,7 @@ mod openstack;
 mod package;
 mod perl;
 mod php;
+mod pulumi;
 mod purescript;
 mod python;
 mod red;
@@ -125,6 +126,7 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
             "package" => package::module(context),
             "perl" => perl::module(context),
             "php" => php::module(context),
+            "pulumi" => pulumi::module(context),
             "purescript" => purescript::module(context),
             "python" => python::module(context),
             "rlang" => rlang::module(context),
@@ -212,6 +214,7 @@ pub fn description(module: &str) -> &'static str {
         "package" => "The package version of the current directory's project",
         "perl" => "The currently installed version of Perl",
         "php" => "The currently installed version of PHP",
+        "pulumi" => "The current stack and installed version of Pulumi",
         "purescript" => "The currently installed version of PureScript",
         "python" => "The currently installed version of Python",
         "red" => "The currently installed version of Red",
diff --git a/src/modules/pulumi.rs b/src/modules/pulumi.rs
new file mode 100644
index 000000000..78b7567f9
--- /dev/null
+++ b/src/modules/pulumi.rs
@@ -0,0 +1,296 @@
+#![warn(missing_docs)]
+use sha1::{Digest, Sha1};
+use std::ffi::OsStr;
+use std::fs::File;
+use std::io::Read;
+use std::path::{Path, PathBuf};
+use std::str::FromStr;
+use yaml_rust::{Yaml, YamlLoader};
+
+use super::{Context, Module, RootModuleConfig};
+use crate::configs::pulumi::PulumiConfig;
+use crate::formatter::{StringFormatter, VersionFormatter};
+
+static PULUMI_HOME: &str = "PULUMI_HOME";
+
+/// Creates a module with the current Pulumi version and stack name.
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let mut module = context.new_module("pulumi");
+    let config = PulumiConfig::try_load(module.config);
+
+    let project_file = find_package_file(&context.logical_dir)?;
+
+    let parsed = StringFormatter::new(config.format).and_then(|formatter| {
+        formatter
+            .map_meta(|variable, _| match variable {
+                "symbol" => Some(config.symbol),
+                _ => None,
+            })
+            .map_style(|variable| match variable {
+                "style" => Some(Ok(config.style)),
+                _ => None,
+            })
+            .map(|variable| match variable {
+                "version" => {
+                    let stdout = context.exec_cmd("pulumi", &["version"])?.stdout;
+                    VersionFormatter::format_module_version(
+                        module.get_name(),
+                        parse_version(&stdout),
+                        config.version_format,
+                    )
+                }
+                .map(Ok),
+                "stack" => stack_name(&project_file, context).map(Ok),
+                _ => None,
+            })
+            .parse(None)
+    });
+
+    match parsed {
+        Ok(x) => {
+            module.set_segments(x);
+            Some(module)
+        }
+        Err(e) => {
+            log::warn!("Error in module `pulumi`:\n{}", e);
+            None
+        }
+    }
+}
+
+/// Parse the output of `pulumi version` into just the version string.
+///
+/// Normally, this just means returning it. When Pulumi is being developed, it
+/// can return results like `3.12.0-alpha.1630554544+f89e9a29.dirty`, which we
+/// don't want to see. Instead we display that as `3.12.0-alpha`.
+fn parse_version(version: &str) -> &str {
+    let mut periods = 0;
+    for (i, c) in version.as_bytes().iter().enumerate() {
+        if *c == b'.' {
+            if periods == 2 {
+                return &version[0..i];
+            } else {
+                periods += 1;
+            }
+        }
+    }
+    // We didn't hit 3 periods, so we just return the whole string.
+    version
+}
+
+/// Find a file describing a Pulumi package in the current directory (or any parrent directory).
+fn find_package_file(path: &Path) -> Option<PathBuf> {
+    for path in path.ancestors() {
+        log::trace!("Looking for package file in {:?}", path);
+        let dir = std::fs::read_dir(path).ok()?;
+        let goal = dir.filter_map(Result::ok).find(|path| {
+            path.file_name() == OsStr::new("Pulumi.yaml")
+                || path.file_name() == OsStr::new("Pulumi.yml")
+        });
+        if let Some(goal) = goal {
+            return Some(goal.path());
+        }
+    }
+    log::trace!("Did not find a Pulumi package file");
+    None
+}
+
+/// We get the name of the current stack.
+///
+/// Pulumi has no CLI option that is fast enough to get this for us, but finding
+/// the location is simple. We get it ourselves.
+fn stack_name(project_file: &Path, context: &Context) -> Option<String> {
+    let mut file = File::open(&project_file).ok()?;
+
+    let mut contents = String::new();
+    file.read_to_string(&mut contents).ok()?;
+    let name = YamlLoader::load_from_str(&contents).ok().and_then(
+        |yaml| -> Option<Option<String>> {
+            log::trace!("Parsed {:?} into yaml", project_file);
+            let yaml = yaml.into_iter().next()?;
+            yaml.into_hash().map(|mut hash| -> Option<String> {
+                hash.remove(&Yaml::String("name".to_string()))?
+                    .into_string()
+            })
+        },
+    )??;
+    log::trace!("Found project name: {:?}", name);
+
+    let workspace_file = get_pulumi_workspace(context, &name, project_file)
+        .map(File::open)?
+        .ok()?;
+    log::trace!("Trying to read workspace_file: {:?}", workspace_file);
+    let workspace: serde_json::Value = match serde_json::from_reader(workspace_file) {
+        Ok(k) => k,
+        Err(e) => {
+            log::debug!("Failed to parse workspace file: {}", e);
+            return None;
+        }
+    };
+    log::trace!("Read workspace_file: {:?}", workspace);
+    workspace
+        .as_object()?
+        .get("stack")?
+        .as_str()
+        .map(ToString::to_string)
+}
+
+/// Calculates the path of the workspace settings file for a given pulumi stack.
+fn get_pulumi_workspace(context: &Context, name: &str, project_file: &Path) -> Option<PathBuf> {
+    let project_file = if cfg!(test) {
+        // Because this depends on the absolute path of the file, it changes in
+        // each test run. We thus mock it.
+        "test".to_string()
+    } else {
+        let mut hasher = Sha1::new();
+        hasher.update(project_file.to_str()?.as_bytes());
+        crate::utils::encode_to_hex(&hasher.finalize().to_vec())
+    };
+    let unique_file_name = format!("{}-{}-workspace.json", name, project_file);
+    let mut path = pulumi_home_dir(context)?;
+    path.push("workspaces");
+    path.push(unique_file_name);
+    Some(path)
+}
+
+/// Get the Pulumi home directory. We first check `PULUMI_HOME`. If that isn't
+/// set, we return `$HOME/.pulumi`.
+fn pulumi_home_dir(context: &Context) -> Option<PathBuf> {
+    if let Some(k) = context.get_env(PULUMI_HOME) {
+        std::path::PathBuf::from_str(&k).ok()
+    } else {
+        context.get_home().map(|p| p.join(".pulumi"))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io;
+
+    use super::*;
+    use crate::test::ModuleRenderer;
+    use ansi_term::Color;
+    use clap::ArgMatches;
+
+    #[test]
+    fn pulumi_version_release() {
+        let input = "3.12.0";
+        assert_eq!(parse_version(input), input);
+    }
+
+    #[test]
+    fn pulumi_version_prerelease() {
+        let input = "3.12.0-alpha";
+        assert_eq!(parse_version(input), input);
+    }
+
+    #[test]
+    fn pulumi_version_dirty() {
+        let input = "3.12.0-alpha.1630554544+f89e9a29.dirty";
+        assert_eq!(parse_version(input), "3.12.0-alpha");
+    }
+
+    #[test]
+    fn get_home_dir() {
+        let mut context = Context::new(ArgMatches::default());
+        context.env.insert("HOME", "/home/sweet/home".to_string());
+        assert_eq!(
+            pulumi_home_dir(&context),
+            Some(PathBuf::from("/home/sweet/home/.pulumi"))
+        );
+        context.env.insert("PULUMI_HOME", "/a/dir".to_string());
+        assert_eq!(pulumi_home_dir(&context), Some(PathBuf::from("/a/dir")))
+    }
+
+    #[test]
+    fn test_get_pulumi_workspace() {
+        let mut context = Context::new(ArgMatches::default());
+        context.env.insert("HOME", "/home/sweet/home".to_string());
+        let name = "foobar";
+        let project_file = PathBuf::from("/hello/Pulumi.yaml");
+        assert_eq!(
+            get_pulumi_workspace(&context, name, &project_file),
+            Some("/home/sweet/home/.pulumi/workspaces/foobar-test-workspace.json")
+                .map(PathBuf::from)
+        );
+    }
+
+    #[test]
+    fn version_render() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let pulumi_file = File::create(dir.path().join("Pulumi.yaml"))?;
+        pulumi_file.sync_all()?;
+        let rendered = ModuleRenderer::new("pulumi")
+            .path(dir.path())
+            .config(toml::toml! {
+                [pulumi]
+                format = "with [$version]($style) "
+            })
+            .collect();
+        dir.close()?;
+        let expected = format!("with {} ", Color::Fixed(5).bold().paint("v1.2.3-ver"));
+
+        assert_eq!(expected, rendered.expect("a result"));
+        Ok(())
+    }
+
+    #[test]
+    /// This test confirms a full render. This means finding a Pulumi.yml file,
+    /// tracing back to the backing workspace settings file, and printing the
+    /// stack name.
+    fn render_valid_paths() -> io::Result<()> {
+        use io::Write;
+        let dir = tempfile::tempdir()?;
+        let root = std::fs::canonicalize(dir.path())?;
+        let mut yaml = File::create(root.join("Pulumi.yml"))?;
+        yaml.write_all("name: starship\nruntime: nodejs\ndescription: A thing\n".as_bytes())?;
+        yaml.sync_all()?;
+
+        let workspace_path = root.join(".pulumi").join("workspaces");
+        let _ = std::fs::create_dir_all(&workspace_path)?;
+        let workspace_path = &workspace_path.join("starship-test-workspace.json");
+        let mut workspace = File::create(&workspace_path)?;
+        serde_json::to_writer_pretty(
+            &mut workspace,
+            &serde_json::json!(
+                {
+                    "stack": "launch"
+                }
+            ),
+        )?;
+        workspace.sync_all()?;
+        let rendered = ModuleRenderer::new("pulumi")
+            .path(root.clone())
+            .logical_path(root.clone())
+            .config(toml::toml! {
+                [pulumi]
+                format = "in [$symbol($stack)]($style) "
+            })
+            .env("HOME", root.to_str().unwrap())
+            .collect();
+        let expected = format!("in {} ", Color::Fixed(5).bold().paint(" launch"));
+        assert_eq!(expected, rendered.expect("a result"));
+        dir.close()?;
+        Ok(())
+    }
+
+    #[test]
+    fn empty_config_file() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let yaml = File::create(dir.path().join("Pulumi.yaml"))?;
+        yaml.sync_all()?;
+
+        let rendered = ModuleRenderer::new("pulumi")
+            .path(dir.path())
+            .logical_path(dir.path())
+            .config(toml::toml! {
+                [pulumi]
+                format = "in [$symbol($stack)]($style) "
+            })
+            .collect();
+        let expected = format!("in {} ", Color::Fixed(5).bold().paint(" "));
+        assert_eq!(expected, rendered.expect("a result"));
+        dir.close()?;
+        Ok(())
+    }
+}
diff --git a/src/utils.rs b/src/utils.rs
index 910bf9be1..4a41c7dcf 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -224,7 +224,11 @@ active boot switches: -d:release\n",
                 stdout: String::from("7.3.8"),
                 stderr: String::default(),
             })
-        }
+        },
+        "pulumi version" => Some(CommandOutput{
+            stdout: String::from("1.2.3-ver.1631311768+e696fb6c"),
+            stderr: String::default(),
+        }),
         "purs --version" => Some(CommandOutput {
             stdout: String::from("0.13.5\n"),
             stderr: String::default(),
@@ -499,6 +503,21 @@ pub fn home_dir() -> Option<PathBuf> {
     directories_next::BaseDirs::new().map(|base_dirs| base_dirs.home_dir().to_owned())
 }
 
+const HEXTABLE: &[char] = &[
+    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f',
+];
+
+/// Encode a u8 slice into a hexadecimal string.
+pub fn encode_to_hex(slice: &[u8]) -> String {
+    // let mut j = 0;
+    let mut dst = Vec::with_capacity(slice.len() * 2);
+    for &v in slice {
+        dst.push(HEXTABLE[(v >> 4) as usize] as u8);
+        dst.push(HEXTABLE[(v & 0x0f) as usize] as u8);
+    }
+    String::from_utf8(dst).unwrap()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -709,4 +728,12 @@ mod tests {
         };
         assert_eq!(get_command_string_output(case2), "stderr");
     }
+
+    #[test]
+    fn sha1_hex() {
+        assert_eq!(
+            encode_to_hex(&[8, 13, 9, 189, 129, 94]),
+            "080d09bd815e".to_string()
+        );
+    }
 }