From 1dc6c271dd6e73e50f58aa73fa3b0374dcfb75b8 Mon Sep 17 00:00:00 2001
From: Cosimo Matteini <dev.matteini@gmail.com>
Date: Thu, 26 Jan 2023 23:46:30 +0100
Subject: [PATCH] feat(aws): add support for source_profile (#3834)

---
 src/modules/aws.rs | 115 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 114 insertions(+), 1 deletion(-)

diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index ec21408e3307b..4c806871ec077 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -200,6 +200,29 @@ fn has_defined_credentials(
     Some(section.contains_key("aws_access_key_id"))
 }
 
+// https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html#cli-configure-files-settings
+fn has_source_profile(
+    context: &Context,
+    aws_profile: Option<&Profile>,
+    aws_config: &AwsConfigFile,
+    aws_creds: &AwsCredsFile,
+) -> Option<bool> {
+    let config = get_config(context, aws_config)?;
+
+    let config_section = get_profile_config(config, aws_profile)?;
+    let source_profile = config_section
+        .get("source_profile")
+        .map(std::borrow::ToOwned::to_owned);
+
+    let has_credential_process =
+        has_credential_process_or_sso(context, source_profile.as_ref(), aws_config, aws_creds)
+            .unwrap_or(false);
+    let has_credentials =
+        has_defined_credentials(context, source_profile.as_ref(), aws_creds).unwrap_or(false);
+
+    Some(has_credential_process || has_credentials)
+}
+
 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);
@@ -212,10 +235,12 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         return None;
     }
 
-    // only display if credential_process is defined or has valid credentials
+    // only display in the presence of credential_process, source_profile or valid credentials
     if !config.force_display
         && !has_credential_process_or_sso(context, aws_profile.as_ref(), &aws_config, &aws_creds)
             .unwrap_or(false)
+        && !has_source_profile(context, aws_profile.as_ref(), &aws_config, &aws_creds)
+            .unwrap_or(false)
         && !has_defined_credentials(context, aws_profile.as_ref(), &aws_creds).unwrap_or(false)
     {
         return None;
@@ -1034,4 +1059,92 @@ sso_role_name = <AWS-ROLE-NAME>
 
         assert_eq!(expected, actual);
     }
+
+    #[test]
+    fn source_profile_set() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let config_path = dir.path().join("config");
+        let credential_path = dir.path().join("credentials");
+        let mut config = File::create(&config_path)?;
+        config.write_all(
+            "[profile astronauts]
+source_profile = starship
+"
+            .as_bytes(),
+        )?;
+        let mut credentials = File::create(&credential_path)?;
+        credentials.write_all(
+            "[starship]
+aws_access_key_id=dummy
+aws_secret_access_key=dummy
+"
+            .as_bytes(),
+        )?;
+
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref())
+            .env(
+                "AWS_CREDENTIALS_FILE",
+                credential_path.to_string_lossy().as_ref(),
+            )
+            .env("AWS_PROFILE", "astronauts")
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  astronauts ")
+        ));
+
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
+    #[test]
+    fn source_profile_not_exists() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let config_path = dir.path().join("config");
+        let mut config = File::create(&config_path)?;
+        config.write_all(
+            "[profile astronauts]
+source_profile = starship
+"
+            .as_bytes(),
+        )?;
+
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref())
+            .env("AWS_PROFILE", "astronauts")
+            .collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
+    #[test]
+    fn source_profile_uses_credential_process() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let config_path = dir.path().join("config");
+        let mut config = File::create(&config_path)?;
+        config.write_all(
+            "[profile starship]
+credential_process = /opt/bin/awscreds-retriever --username starship
+
+[profile astronauts]
+source_profile = starship
+"
+            .as_bytes(),
+        )?;
+
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref())
+            .env("AWS_PROFILE", "astronauts")
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  astronauts ")
+        ));
+
+        assert_eq!(expected, actual);
+        dir.close()
+    }
 }