From b3275d8ddf4ab00a120a1f6edeb954987b549cc9 Mon Sep 17 00:00:00 2001 From: Thomas Lee Date: Tue, 15 Oct 2019 00:05:03 +0900 Subject: [PATCH] feat: Show AWS region in `aws` module (#482) --- docs/config/README.md | 5 ++- src/configs/aws.rs | 2 + src/modules/aws.rs | 74 ++++++++++++++++++++++++++++++++++- tests/testsuite/aws.rs | 87 +++++++++++++++++++++++++++++++++++++++++- 4 files changed, 162 insertions(+), 6 deletions(-) diff --git a/docs/config/README.md b/docs/config/README.md index 490b042d9..4e19a31a4 100644 --- a/docs/config/README.md +++ b/docs/config/README.md @@ -114,8 +114,9 @@ prompt_order = [ ## AWS -The `aws` module shows the current AWS profile. This is based on the -`AWS_PROFILE` env var. +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. ### Options diff --git a/src/configs/aws.rs b/src/configs/aws.rs index 2c9589508..a4dc4ca71 100644 --- a/src/configs/aws.rs +++ b/src/configs/aws.rs @@ -7,6 +7,7 @@ use starship_module_config_derive::ModuleConfig; pub struct AwsConfig<'a> { pub symbol: SegmentConfig<'a>, pub profile: SegmentConfig<'a>, + pub region: SegmentConfig<'a>, pub style: Style, pub disabled: bool, } @@ -16,6 +17,7 @@ impl<'a> RootModuleConfig<'a> for AwsConfig<'a> { AwsConfig { symbol: SegmentConfig::new("☁️ "), profile: SegmentConfig::default(), + region: SegmentConfig::default(), style: Color::Yellow.bold(), disabled: false, } diff --git a/src/modules/aws.rs b/src/modules/aws.rs index 8e6fd50f9..9225c0fe7 100644 --- a/src/modules/aws.rs +++ b/src/modules/aws.rs @@ -1,16 +1,85 @@ use std::env; +use std::fs::File; +use std::io::{BufRead, BufReader}; +use std::path::PathBuf; +use std::str::FromStr; + +use dirs::home_dir; use super::{Context, Module, RootModuleConfig}; use crate::configs::aws::AwsConfig; +fn get_aws_region_from_config(aws_profile: &Option) -> Option { + let config_location = env::var("AWS_CONFIG_FILE") + .ok() + .and_then(|path| PathBuf::from_str(&path).ok()) + .or_else(|| { + let mut home = home_dir()?; + home.push(".aws/config"); + Some(home) + })?; + + let file = File::open(&config_location).ok()?; + let reader = BufReader::new(file); + let lines = reader.lines().filter_map(Result::ok); + + let region_line = if let Some(ref aws_profile) = aws_profile { + lines + .skip_while(|line| line != &format!("[profile {}]", aws_profile)) + .skip(1) + .take_while(|line| !line.starts_with('[')) + .find(|line| line.starts_with("region")) + } else { + lines + .skip_while(|line| line != "[default]") + .skip(1) + .take_while(|line| !line.starts_with('[')) + .find(|line| line.starts_with("region")) + }?; + + let region = region_line.split('=').nth(1)?; + let region = region.trim(); + + Some(region.to_string()) +} + +fn get_aws_region() -> Option<(String, String)> { + env::var("AWS_DEFAULT_REGION") + .ok() + .map(|region| (String::new(), region)) + .or_else(|| { + let aws_profile = env::var("AWS_PROFILE").ok(); + let aws_region = get_aws_region_from_config(&aws_profile); + + if aws_profile.is_none() && aws_region.is_none() { + None + } else { + Some(( + aws_profile.unwrap_or_default(), + aws_region.unwrap_or_default(), + )) + } + }) + .or_else(|| { + env::var("AWS_REGION") + .ok() + .map(|region| (String::new(), region)) + }) +} + pub fn module<'a>(context: &'a Context) -> Option> { const AWS_PREFIX: &str = "on "; - let aws_profile = env::var("AWS_PROFILE").ok()?; - if aws_profile.is_empty() { + let (aws_profile, aws_region) = get_aws_region()?; + if aws_profile.is_empty() && aws_region.is_empty() { return None; } + let aws_region = if aws_profile.is_empty() || aws_region.is_empty() { + aws_region + } else { + format!("({})", aws_region) + }; let mut module = context.new_module("aws"); let config: AwsConfig = AwsConfig::try_load(module.config); @@ -21,6 +90,7 @@ pub fn module<'a>(context: &'a Context) -> Option> { module.create_segment("symbol", &config.symbol); module.create_segment("profile", &config.profile.with_value(&aws_profile)); + module.create_segment("region", &config.profile.with_value(&aws_region)); Some(module) } diff --git a/tests/testsuite/aws.rs b/tests/testsuite/aws.rs index 65a7bce71..d057a2432 100644 --- a/tests/testsuite/aws.rs +++ b/tests/testsuite/aws.rs @@ -1,10 +1,12 @@ +use std::fs::File; +use std::io::{self, Write}; + use ansi_term::Color; -use std::io; use crate::common; #[test] -fn no_profile_set() -> io::Result<()> { +fn no_region_set() -> io::Result<()> { let output = common::render_module("aws").env_clear().output()?; let expected = ""; let actual = String::from_utf8(output.stdout).unwrap(); @@ -12,6 +14,31 @@ fn no_profile_set() -> io::Result<()> { Ok(()) } +#[test] +fn region_set() -> io::Result<()> { + let output = common::render_module("aws") + .env_clear() + .env("AWS_REGION", "ap-northeast-2") + .output()?; + let expected = format!("on {} ", Color::Yellow.bold().paint("☁️ ap-northeast-2")); + let actual = String::from_utf8(output.stdout).unwrap(); + assert_eq!(expected, actual); + Ok(()) +} + +#[test] +fn default_region_set() -> io::Result<()> { + let output = common::render_module("aws") + .env_clear() + .env("AWS_REGION", "ap-northeast-2") + .env("AWS_DEFAULT_REGION", "ap-northeast-1") + .output()?; + let expected = format!("on {} ", Color::Yellow.bold().paint("☁️ ap-northeast-1")); + let actual = String::from_utf8(output.stdout).unwrap(); + assert_eq!(expected, actual); + Ok(()) +} + #[test] fn profile_set() -> io::Result<()> { let output = common::render_module("aws") @@ -23,3 +50,59 @@ fn profile_set() -> io::Result<()> { assert_eq!(expected, actual); Ok(()) } + +#[test] +fn default_profile_set() -> io::Result<()> { + let dir = common::new_tempdir()?; + let config_path = dir.path().join("config"); + let mut file = File::create(&config_path)?; + + file.write_all( + "[default] +region = us-east-1 + +[profile astronauts] +region = us-east-2 +" + .as_bytes(), + )?; + + let output = common::render_module("aws") + .env_clear() + .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref()) + .output()?; + let expected = format!("on {} ", Color::Yellow.bold().paint("☁️ us-east-1")); + let actual = String::from_utf8(output.stdout).unwrap(); + assert_eq!(expected, actual); + Ok(()) +} + +#[test] +fn profile_and_config_set() -> io::Result<()> { + let dir = common::new_tempdir()?; + let config_path = dir.path().join("config"); + let mut file = File::create(&config_path)?; + + file.write_all( + "[default] +region = us-east-1 + +[profile astronauts] +region = us-east-2 +" + .as_bytes(), + )?; + + let output = common::render_module("aws") + .env_clear() + .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref()) + .env("AWS_PROFILE", "astronauts") + .output()?; + let expected = format!( + "on {} ", + Color::Yellow.bold().paint("☁️ astronauts(us-east-2)") + ); + let actual = String::from_utf8(output.stdout).unwrap(); + assert_eq!(expected, actual); + Ok(()) +}