Llewellyn/starship
1
0
Fork 0
mirror of https://github.com/Llewellynvdm/starship.git synced 2024-11-11 07:40:57 +00:00
Code Issues Releases Activity

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

Browse Source
This commit is contained in:
Carl-Louis Van Brandt 2021-05-13 02:43:46 +02:00 committed by GitHub
parent baabc7743d
commit 69a754573d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 277 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
docs/config/README.md
View File

@ -250,33 +250,38 @@ $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. |
| ------------------- | ---------------------------------------------------------------- | ----------------------------------------------------------------- |
| `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. |
| `disabled` | `false` | Disables the `aws` 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 |
| duration | `2h27m20s` | The temporary credentials validity duration |
| symbol | | Mirrors the value of option `symbol` |
| style\* | | Mirrors the value of option `style` |

4
src/configs/aws.rs
View File

@ -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",
}
}
}

204
src/modules/aws.rs
View File

@ -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)
@ -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<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() {

53
src/modules/cmd_duration.rs
View File

@ -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,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")

51
src/utils.rs
View File

@ -363,10 +363,61 @@ 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));