From 69a754573d1321032c82083f86097550ea7920a4 Mon Sep 17 00:00:00 2001 From: Carl-Louis Van Brandt Date: Thu, 13 May 2021 02:43:46 +0200 Subject: feat(aws): add temporary credentials countdown (#2464) --- docs/config/README.md | 37 ++++---- src/configs/aws.rs | 4 +- src/modules/aws.rs | 204 +++++++++++++++++++++++++++++++++++++++++++- src/modules/cmd_duration.rs | 53 +----------- src/utils.rs | 51 +++++++++++ 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, + 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 { - let config_location = context +fn get_credentials_file_path(context: &Context) -> Option { + 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 { + 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 { + 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, Option) -> Option { + 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 { match aliases.get(®ion) { None => region, @@ -87,6 +138,16 @@ pub fn module<'a>(context: &'a Context) -> Option> { 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> { .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) @@ -396,6 +458,142 @@ region = us-east-2 assert_eq!(expected, actual); } + #[test] + fn expiration_date_set() { + use chrono::{DateTime, NaiveDateTime, SecondsFormat, Utc}; + + let now_plus_half_hour: DateTime = 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 = 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 = 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() { 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> { 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 = 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,31 +96,9 @@ 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") diff --git a/src/utils.rs b/src/utils.rs index 7ef829a31..a70192e31 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -363,10 +363,61 @@ fn internal_exec_cmd(cmd: &str, args: &[&str], time_limit: Duration) -> Option 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 = 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)); -- cgit v1.2.3