fix(aws): Only display AWS if there are credentials configured (#3504)

* only display aws on credential_process defined

* add check for both credential_process and valid credentials

* fix tests

* update aws module documentation

* add better explanation of requirements to documentation

* add support for AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN, and update docs

* remove credential_process env var

docs/config/README.md

@@ -270,11 +270,18 @@ format = "$all$directory$character"
 
 ## AWS
 
-The `aws` module shows the current AWS region and profile. This is based on
+The `aws` module shows the current AWS region and profile when
+credentials or a `credential_process` have been setup. This is based on
 `AWS_REGION`, `AWS_DEFAULT_REGION`, and `AWS_PROFILE` env var with
 `~/.aws/config` file. This module also shows an expiration timer when using temporary
 credentials.
 
+The module will display a profile only if its credentials are present in
+`~/.aws/credentials` or a `credential_process` is defined in
+`~/.aws/config`. Alternatively, having any of the `AWS_ACCESS_KEY_ID`,
+`AWS_SECRET_ACCESS_KEY`, or `AWS_SESSION_TOKEN` env vars defined will
+also suffice.
+
 When using [aws-vault](https://github.com/99designs/aws-vault) the profile
 is read from the `AWS_VAULT` env var and the credentials expiration date
 is read from the `AWS_SESSION_EXPIRATION` env var.

src/modules/aws.rs

@@ -118,6 +118,59 @@ fn alias_region(region: String, aliases: &HashMap<String, &str>) -> String {
     }
 }
 
+fn get_credential_process(context: &Context, aws_profile: Option<&Profile>) -> Option<String> {
+    let contents = read_file(get_config_file_path(context)?).ok()?;
+
+    let profile_line = if let Some(aws_profile) = aws_profile {
+        format!("[profile {}]", aws_profile)
+    } else {
+        "[default]".to_string()
+    };
+
+    let cred_proc_line = contents
+        .lines()
+        .skip_while(|line| line != &profile_line)
+        .skip(1)
+        .take_while(|line| !line.starts_with('['))
+        .find(|line| line.starts_with("credential_process"))?;
+
+    let cred_proc = cred_proc_line.split('=').nth(1)?.trim();
+    Some(cred_proc.to_string())
+}
+
+fn get_defined_credentials(context: &Context, aws_profile: Option<&Profile>) -> Option<String> {
+    let valid_env_vars = vec![
+        "AWS_ACCESS_KEY_ID",
+        "AWS_SECRET_ACCESS_KEY",
+        "AWS_SESSION_TOKEN",
+    ];
+
+    // accept if set through environment variable
+    if let Some(aws_identity_cred) = valid_env_vars
+        .iter()
+        .find_map(|env_var| context.get_env(env_var))
+    {
+        return Some(aws_identity_cred);
+    }
+
+    let contents = read_file(get_credentials_file_path(context)?).ok()?;
+
+    let profile_line = if let Some(aws_profile) = aws_profile {
+        format!("[{}]", aws_profile)
+    } else {
+        "[default]".to_string()
+    };
+
+    let aws_key_id_line = contents
+        .lines()
+        .skip_while(|line| line != &profile_line)
+        .skip(1)
+        .take_while(|line| !line.starts_with('['))
+        .find(|line| line.starts_with("aws_access_key_id"))?;
+    let aws_key_id = aws_key_id_line.split('=').nth(1)?.trim();
+    Some(aws_key_id.to_string())
+}
+
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let mut module = context.new_module("aws");
     let config: AwsConfig = AwsConfig::try_load(module.config);
@@ -127,6 +180,13 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         return None;
     }
 
+    // only display if credential_process is defined or has valid credentials
+    if get_credential_process(context, aws_profile.as_ref()).is_none()
+        && get_defined_credentials(context, aws_profile.as_ref()).is_none()
+    {
+        return None;
+    }
+
     let mapped_region = if let Some(aws_region) = aws_region {
         Some(alias_region(aws_region, &config.region_aliases))
     } else {
@@ -193,6 +253,7 @@ mod tests {
     fn region_set() {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_REGION", "ap-northeast-2")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -206,6 +267,7 @@ mod tests {
     fn region_set_with_alias() {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_REGION", "ap-southeast-2")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .config(toml::toml! {
                 [aws.region_aliases]
                 ap-southeast-2 = "au"
@@ -221,6 +283,7 @@ mod tests {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_REGION", "ap-northeast-2")
             .env("AWS_DEFAULT_REGION", "ap-northeast-1")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -234,6 +297,7 @@ mod tests {
     fn profile_set() {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_PROFILE", "astronauts")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -248,6 +312,7 @@ mod tests {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_VAULT", "astronauts-vault")
             .env("AWS_PROFILE", "astronauts-profile")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -262,6 +327,7 @@ mod tests {
         let actual = ModuleRenderer::new("aws")
             .env("AWSU_PROFILE", "astronauts-awsu")
             .env("AWS_PROFILE", "astronauts-profile")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -276,6 +342,7 @@ mod tests {
         let actual = ModuleRenderer::new("aws")
             .env("AWSUME_PROFILE", "astronauts-awsume")
             .env("AWS_PROFILE", "astronauts-profile")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -290,6 +357,7 @@ mod tests {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_PROFILE", "astronauts")
             .env("AWS_REGION", "ap-northeast-2")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -341,6 +409,7 @@ mod tests {
         file.write_all(
             "[default]
 region = us-east-1
+credential_process = /opt/bin/awscreds-retriever
 
 [profile astronauts]
 region = us-east-2
@@ -372,6 +441,7 @@ region = us-east-1
 
 [profile astronauts]
 region = us-east-2
+credential_process = /opt/bin/awscreds-retriever
 "
             .as_bytes(),
         )?;
@@ -397,6 +467,7 @@ region = us-east-2
         let actual = ModuleRenderer::new("aws")
             .env("AWS_PROFILE", "astronauts")
             .env("AWS_REGION", "ap-northeast-1")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -412,6 +483,7 @@ region = us-east-2
     fn profile_set_with_display_all() {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_PROFILE", "astronauts")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -425,6 +497,7 @@ region = us-east-2
     fn region_set_with_display_all() {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_REGION", "ap-northeast-1")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -439,6 +512,7 @@ region = us-east-2
         let actual = ModuleRenderer::new("aws")
             .env("AWS_PROFILE", "astronauts")
             .env("AWS_DEFAULT_REGION", "ap-northeast-1")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .config(toml::toml! {
                 [aws]
                 format = "on [$symbol$region]($style) "
@@ -457,6 +531,7 @@ region = us-east-2
         let actual = ModuleRenderer::new("aws")
             .env("AWS_PROFILE", "astronauts")
             .env("AWS_REGION", "ap-northeast-1")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .config(toml::toml! {
                 [aws]
                 format = "on [$symbol$profile]($style) "
@@ -474,6 +549,7 @@ region = us-east-2
     fn region_set_with_display_profile() {
         let actual = ModuleRenderer::new("aws")
             .env("AWS_REGION", "ap-northeast-1")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .config(toml::toml! {
                 [aws]
                 format = "on [$symbol$profile]($style) "
@@ -500,6 +576,7 @@ region = us-east-2
             })
             .env("AWS_PROFILE", "astronauts")
             .env("AWS_REGION", "ap-northeast-2")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .env(
                 "AWS_SESSION_EXPIRATION",
                 now_plus_half_hour.to_rfc3339_opts(SecondsFormat::Secs, true),
@@ -580,6 +657,7 @@ expiration={}
             })
             .env("AWS_PROFILE", "astronauts")
             .env("AWS_REGION", "ap-northeast-2")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .collect();
         let expected = Some(format!(
             "on {}",
@@ -610,6 +688,7 @@ expiration={}
             })
             .env("AWS_PROFILE", "astronauts")
             .env("AWS_REGION", "ap-northeast-2")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
             .env(
                 "AWS_SESSION_EXPIRATION",
                 now.to_rfc3339_opts(SecondsFormat::Secs, true),
@@ -638,4 +717,133 @@ expiration={}
 
         assert_eq!(expected, actual);
     }
+
+    #[test]
+    fn missing_any_credentials() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let config_path = dir.path().join("config");
+        let mut file = File::create(&config_path)?;
+
+        file.write_all(
+            "[default]
+region = us-east-1
+output = json
+
+[profile astronauts]
+region = us-east-2
+"
+            .as_bytes(),
+        )?;
+
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref())
+            .collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
+    #[test]
+    fn access_key_credential_set() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let credentials_path = dir.path().join("credentials");
+        let mut file = File::create(&credentials_path)?;
+
+        file.write_all(
+            "[astronauts]
+aws_access_key_id=dummy
+aws_secret_access_key=dummy
+"
+            .as_bytes(),
+        )?;
+
+        let actual = ModuleRenderer::new("aws")
+            .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) ")
+        ));
+
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn credential_process_set() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let config_path = dir.path().join("config");
+        let mut file = File::create(&config_path)?;
+
+        file.write_all(
+            "[default]
+region = ap-northeast-2
+credential_process = /opt/bin/awscreds-retriever
+"
+            .as_bytes(),
+        )?;
+
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref())
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  (ap-northeast-2) ")
+        ));
+
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
+    #[test]
+    fn access_key_env_var_set() {
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_PROFILE", "astronauts")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  astronauts ")
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn secret_access_key_env_var_set() {
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_PROFILE", "astronauts")
+            .env("AWS_SECRET_ACCESS_KEY", "dummy")
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  astronauts ")
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn session_token_env_var_set() {
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_PROFILE", "astronauts")
+            .env("AWS_SESSION_TOKEN", "dummy")
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  astronauts ")
+        ));
+
+        assert_eq!(expected, actual);
+    }
 }