feat(aws): add temporary credentials countdown (#2464)

5 files changed, 277 insertions, 72 deletions

diff --git a/docs/config/README.md b/docs/config/README.md
index 0aee51188..a3cb82545 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -250,35 +250,40 @@ $character"""
 
 The `aws` module shows the current AWS region and profile. This is based on
 `AWS_REGION`, `AWS_DEFAULT_REGION`, and `AWS_PROFILE` env var with
-`~/.aws/config` file.
+`~/.aws/config` file. This module also shows an expiration timer when using temporary
+credentials.
 
 When using [aws-vault](https://github.com/99designs/aws-vault) the profile
-is read from the `AWS_VAULT` env var.
+is read from the `AWS_VAULT` env var and the credentials expiration date
+is read from the `AWS_SESSION_EXPIRATION` env var.
 
 When using [awsu](https://github.com/kreuzwerker/awsu) the profile
 is read from the `AWSU_PROFILE` env var.
 
 When using [AWSume](https://awsu.me) the profile
-is read from the `AWSUME_PROFILE` env var.
+is read from the `AWSUME_PROFILE` env var and the credentials expiration
+date is read from the `AWSUME_EXPIRATION` env var.
 
 ### Options
 
-| Option           | Default                                           | Description                                                     |
-| ---------------- | ------------------------------------------------- | --------------------------------------------------------------- |
-| `format`         | `'on [$symbol($profile )(\($region\) )]($style)'` | The format for the module.                                      |
-| `symbol`         | `"☁️ "`                                            | The symbol used before displaying the current AWS profile.      |
-| `region_aliases` |                                                   | Table of region aliases to display in addition to the AWS name. |
-| `style`          | `"bold yellow"`                                   | The style for the module.                                       |
-| `disabled`       | `false`                                           | Disables the `aws` module.                                      |
+| Option              | Default                                                          | Description                                                       |
+| ------------------- | ---------------------------------------------------------------- | ----------------------------------------------------------------- |
+| `format`            | `'on [$symbol($profile )(\($region\) )(\[$duration\])]($style)'` | The format for the module.                                        |
+| `symbol`            | `"☁️ "`                                                           | The symbol used before displaying the current AWS profile.        |
+| `region_aliases`    |                                                                  | Table of region aliases to display in addition to the AWS name.   |
+| `style`             | `"bold yellow"`                                                  | The style for the module.                                         |
+| `expiration_symbol` | `X`                                                              | The symbol displayed when the temporary credentials have expired. |
+| `disabled`          | `false`                                                          | Disables the `AWS` module.                                        |
 
 ### Variables
 
-| Variable | Example          | Description                          |
-| -------- | ---------------- | ------------------------------------ |
-| region   | `ap-northeast-1` | The current AWS region               |
-| profile  | `astronauts`     | The current AWS profile              |
-| symbol   |                  | Mirrors the value of option `symbol` |
-| style\*  |                  | Mirrors the value of option `style`  |
+| Variable | Example          | Description                                 |
+| -------- | ---------------- | ------------------------------------------- |
+| region   | `ap-northeast-1` | The current AWS region                      |
+| profile  | `astronauts`     | The current AWS profile                     |
+| duration | `2h27m20s`       | The temporary credentials validity duration |
+| symbol   |                  | Mirrors the value of option `symbol`        |
+| style\*  |                  | Mirrors the value of option `style`         |
 
 \*: This variable can only be used as a part of a style string
 
diff --git a/src/configs/aws.rs b/src/configs/aws.rs
index 8cef66362..8cb2b2dd5 100644
--- a/src/configs/aws.rs
+++ b/src/configs/aws.rs
@@ -10,16 +10,18 @@ pub struct AwsConfig<'a> {
     pub style: &'a str,
     pub disabled: bool,
     pub region_aliases: HashMap<String, &'a str>,
+    pub expiration_symbol: &'a str,
 }
 
 impl<'a> Default for AwsConfig<'a> {
     fn default() -> Self {
         AwsConfig {
-            format: "on [$symbol($profile )(\\($region\\) )]($style)",
+            format: "on [$symbol($profile )(\\($region\\) )(\\[$duration\\])]($style)",
             symbol: "☁️  ",
             style: "bold yellow",
             disabled: false,
             region_aliases: HashMap::new(),
+            expiration_symbol: "X",
         }
     }
 }
diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index 637ecc7fc..6f35b512e 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -4,23 +4,41 @@ use std::io::{BufRead, BufReader};
 use std::path::PathBuf;
 use std::str::FromStr;
 
+use chrono::DateTime;
+
 use super::{Context, Module, RootModuleConfig};
 
 use crate::configs::aws::AwsConfig;
 use crate::formatter::StringFormatter;
+use crate::utils::render_time;
 
 type Profile = String;
 type Region = String;
 
-fn get_aws_region_from_config(context: &Context, aws_profile: Option<&str>) -> Option<Region> {
-    let config_location = context
+fn get_credentials_file_path(context: &Context) -> Option<PathBuf> {
+    context
+        .get_env("AWS_CREDENTIALS_FILE")
+        .and_then(|path| PathBuf::from_str(&path).ok())
+        .or_else(|| {
+            let mut home = context.get_home()?;
+            home.push(".aws/credentials");
+            Some(home)
+        })
+}
+
+fn get_config_file_path(context: &Context) -> Option<PathBuf> {
+    context
         .get_env("AWS_CONFIG_FILE")
         .and_then(|path| PathBuf::from_str(&path).ok())
         .or_else(|| {
             let mut home = context.get_home()?;
             home.push(".aws/config");
             Some(home)
-        })?;
+        })
+}
+
+fn get_aws_region_from_config(context: &Context, aws_profile: Option<&str>) -> Option<Region> {
+    let config_location = get_config_file_path(context)?;
 
     let file = File::open(&config_location).ok()?;
     let reader = BufReader::new(file);
@@ -65,6 +83,39 @@ fn get_aws_profile_and_region(context: &Context) -> (Option<Profile>, Option<Reg
     }
 }
 
+fn get_credentials_duration(context: &Context, aws_profile: Option<&Profile>) -> Option<i64> {
+    let expiration_env_vars = vec!["AWS_SESSION_EXPIRATION", "AWSUME_EXPIRATION"];
+    let expiration_date = if let Some(expiration_date) = expiration_env_vars
+        .iter()
+        .find_map(|env_var| context.get_env(env_var))
+    {
+        chrono::DateTime::parse_from_rfc3339(&expiration_date).ok()
+    } else {
+        let credentials_location = get_credentials_file_path(context)?;
+
+        let file = File::open(&credentials_location).ok()?;
+        let reader = BufReader::new(file);
+        let lines = reader.lines().filter_map(Result::ok);
+
+        let profile_line = if let Some(aws_profile) = aws_profile {
+            format!("[{}]", aws_profile)
+        } else {
+            "[default]".to_string()
+        };
+
+        let expiration_date_line = lines
+            .skip_while(|line| line != &profile_line)
+            .skip(1)
+            .take_while(|line| !line.starts_with('['))
+            .find(|line| line.starts_with("expiration"))?;
+
+        let expiration_date = expiration_date_line.split('=').nth(1)?.trim();
+        DateTime::parse_from_rfc3339(expiration_date).ok()
+    }?;
+
+    Some(expiration_date.timestamp() - chrono::Local::now().timestamp())
+}
+
 fn alias_region(region: String, aliases: &HashMap<String, &str>) -> String {
     match aliases.get(&region) {
         None => region,
@@ -87,6 +138,16 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         None
     };
 
+    let duration = {
+        get_credentials_duration(context, aws_profile.as_ref()).map(|duration| {
+            if duration > 0 {
+                render_time((duration * 1000) as u128, false)
+            } else {
+                config.expiration_symbol.to_string()
+            }
+        })
+    };
+
     let parsed = StringFormatter::new(config.format).and_then(|formatter| {
         formatter
             .map_meta(|variable, _| match variable {
@@ -100,6 +161,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
             .map(|variable| match variable {
                 "profile" => aws_profile.as_ref().map(Ok),
                 "region" => mapped_region.as_ref().map(Ok),
+                "duration" => duration.as_ref().map(Ok),
                 _ => None,
             })
             .parse(None)
@@ -397,6 +459,142 @@ region = us-east-2
     }
 
     #[test]
+    fn expiration_date_set() {
+        use chrono::{DateTime, NaiveDateTime, SecondsFormat, Utc};
+
+        let now_plus_half_hour: DateTime<Utc> = chrono::DateTime::from_utc(
+            NaiveDateTime::from_timestamp(chrono::Local::now().timestamp() + 1800, 0),
+            Utc,
+        );
+
+        let actual = ModuleRenderer::new("aws")
+            .config(toml::toml! {
+                [aws]
+                show_duration = true
+            })
+            .env("AWS_PROFILE", "astronauts")
+            .env("AWS_REGION", "ap-northeast-2")
+            .env(
+                "AWS_SESSION_EXPIRATION",
+                now_plus_half_hour.to_rfc3339_opts(SecondsFormat::Secs, true),
+            )
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow
+                .bold()
+                .paint("☁️  astronauts (ap-northeast-2) [30m]")
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn expiration_date_set_from_file() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let credentials_path = dir.path().join("credentials");
+        let mut file = File::create(&credentials_path)?;
+
+        use chrono::{DateTime, NaiveDateTime, Utc};
+
+        let now_plus_half_hour: DateTime<Utc> = chrono::DateTime::from_utc(
+            NaiveDateTime::from_timestamp(chrono::Local::now().timestamp() + 1800, 0),
+            Utc,
+        );
+
+        let expiration_date = now_plus_half_hour.to_rfc3339_opts(chrono::SecondsFormat::Secs, true);
+
+        file.write_all(
+            format!(
+                "[astronauts]
+aws_access_key_id=dummy
+aws_secret_access_key=dummy
+expiration={}
+",
+                expiration_date
+            )
+            .as_bytes(),
+        )?;
+
+        let actual = ModuleRenderer::new("aws")
+            .config(toml::toml! {
+                [aws]
+                show_duration = true
+            })
+            .env("AWS_PROFILE", "astronauts")
+            .env("AWS_REGION", "ap-northeast-2")
+            .env(
+                "AWS_CREDENTIALS_FILE",
+                credentials_path.to_string_lossy().as_ref(),
+            )
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow
+                .bold()
+                .paint("☁️  astronauts (ap-northeast-2) [30m]")
+        ));
+
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn profile_and_region_set_show_duration() {
+        let actual = ModuleRenderer::new("aws")
+            .config(toml::toml! {
+                [aws]
+                show_duration = true
+            })
+            .env("AWS_PROFILE", "astronauts")
+            .env("AWS_REGION", "ap-northeast-2")
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow
+                .bold()
+                .paint("☁️  astronauts (ap-northeast-2) ")
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn expiration_date_set_expired() {
+        use chrono::{DateTime, NaiveDateTime, SecondsFormat, Utc};
+
+        let now: DateTime<Utc> = chrono::DateTime::from_utc(
+            NaiveDateTime::from_timestamp(chrono::Local::now().timestamp() - 1800, 0),
+            Utc,
+        );
+
+        let symbol = "!!!";
+
+        let actual = ModuleRenderer::new("aws")
+            .config(toml::toml! {
+                [aws]
+                show_duration = true
+                expiration_symbol = symbol
+            })
+            .env("AWS_PROFILE", "astronauts")
+            .env("AWS_REGION", "ap-northeast-2")
+            .env(
+                "AWS_SESSION_EXPIRATION",
+                now.to_rfc3339_opts(SecondsFormat::Secs, true),
+            )
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow
+                .bold()
+                .paint(format!("☁️  astronauts (ap-northeast-2) [{}]", symbol))
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
     #[ignore]
     fn region_not_set_with_display_region() {
         let actual = ModuleRenderer::new("aws")
diff --git a/src/modules/cmd_duration.rs b/src/modules/cmd_duration.rs
index 1d732d299..a6556ea62 100644
--- a/src/modules/cmd_duration.rs
+++ b/src/modules/cmd_duration.rs
@@ -2,6 +2,7 @@ use super::{Context, Module, RootModuleConfig};
 
 use crate::configs::cmd_duration::CmdDurationConfig;
 use crate::formatter::StringFormatter;
+use crate::utils::render_time;
 
 /// Outputs the time it took the last command to execute
 ///
@@ -50,36 +51,6 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     Some(undistract_me(module, &config, elapsed))
 }
 
-// Render the time into a nice human-readable string
-fn render_time(raw_millis: u128, show_millis: bool) -> String {
-    // Calculate a simple breakdown into days/hours/minutes/seconds/milliseconds
-    let (millis, raw_seconds) = (raw_millis % 1000, raw_millis / 1000);
-    let (seconds, raw_minutes) = (raw_seconds % 60, raw_seconds / 60);
-    let (minutes, raw_hours) = (raw_minutes % 60, raw_minutes / 60);
-    let (hours, days) = (raw_hours % 24, raw_hours / 24);
-
-    let components = [days, hours, minutes, seconds];
-    let suffixes = ["d", "h", "m", "s"];
-
-    let mut rendered_components: Vec<String> = components
-        .iter()
-        .zip(&suffixes)
-        .map(render_time_component)
-        .collect();
-    if show_millis || raw_millis < 1000 {
-        rendered_components.push(render_time_component((&millis, &"ms")));
-    }
-    rendered_components.join("")
-}
-
-/// Render a single component of the time string, giving an empty string if component is zero
-fn render_time_component((component, suffix): (&u128, &&str)) -> String {
-    match component {
-        0 => String::new(),
-        n => format!("{}{}", n, suffix),
-    }
-}
-
 #[cfg(not(feature = "notify-rust"))]
 fn undistract_me<'a, 'b>(
     module: Module<'a>,
@@ -125,32 +96,10 @@ fn undistract_me<'a, 'b>(
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use crate::test::ModuleRenderer;
     use ansi_term::Color;
 
     #[test]
-    fn test_500ms() {
-        assert_eq!(render_time(500_u128, true), "500ms")
-    }
-    #[test]
-    fn test_10s() {
-        assert_eq!(render_time(10_000_u128, true), "10s")
-    }
-    #[test]
-    fn test_90s() {
-        assert_eq!(render_time(90_000_u128, true), "1m30s")
-    }
-    #[test]
-    fn test_10110s() {
-        assert_eq!(render_time(10_110_000_u128, true), "2h48m30s")
-    }
-    #[test]
-    fn test_1d() {
-        assert_eq!(render_time(86_400_000_u128, true), "1d")
-    }
-
-    #[test]
     fn config_blank_duration_1s() {
         let actual = ModuleRenderer::new("cmd_duration")
             .cmd_duration(1000)
diff --git a/src/utils.rs b/src/utils.rs
index 7ef829a31..a70192e31 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -363,11 +363,62 @@ fn internal_exec_cmd(cmd: &str, args: &[&str], time_limit: Duration) -> Option<C
     }
 }
 
+// Render the time into a nice human-readable string
+pub fn render_time(raw_millis: u128, show_millis: bool) -> String {
+    // Calculate a simple breakdown into days/hours/minutes/seconds/milliseconds
+    let (millis, raw_seconds) = (raw_millis % 1000, raw_millis / 1000);
+    let (seconds, raw_minutes) = (raw_seconds % 60, raw_seconds / 60);
+    let (minutes, raw_hours) = (raw_minutes % 60, raw_minutes / 60);
+    let (hours, days) = (raw_hours % 24, raw_hours / 24);
+
+    let components = [days, hours, minutes, seconds];
+    let suffixes = ["d", "h", "m", "s"];
+
+    let mut rendered_components: Vec<String> = components
+        .iter()
+        .zip(&suffixes)
+        .map(render_time_component)
+        .collect();
+    if show_millis || raw_millis < 1000 {
+        rendered_components.push(render_time_component((&millis, &"ms")));
+    }
+    rendered_components.join("")
+}
+
+/// Render a single component of the time string, giving an empty string if component is zero
+fn render_time_component((component, suffix): (&u128, &&str)) -> String {
+    match component {
+        0 => String::new(),
+        n => format!("{}{}", n, suffix),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
+    fn test_500ms() {
+        assert_eq!(render_time(500_u128, true), "500ms")
+    }
+    #[test]
+    fn test_10s() {
+        assert_eq!(render_time(10_000_u128, true), "10s")
+    }
+    #[test]
+    fn test_90s() {
+        assert_eq!(render_time(90_000_u128, true), "1m30s")
+    }
+    #[test]
+    fn test_10110s() {
+        assert_eq!(render_time(10_110_000_u128, true), "2h48m30s")
+    }
+    #[test]
+    fn test_1d() {
+        assert_eq!(render_time(86_400_000_u128, true), "1d")
+    }
+
+    #[test]
     fn exec_mocked_command() {
         let result = exec_cmd("dummy_command", &[], Duration::from_millis(500));
         let expected = Some(CommandOutput {