From 8e36f307248e80580ffbc61707845bfefe1446ac 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 | 83 +++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 82 insertions(+), 1 deletion(-)

diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index ec21408e3307b..c249cc7bc8855 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -176,6 +176,25 @@ fn has_credential_process_or_sso(
     )
 }
 
+// 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 credentials = get_creds(context, aws_creds)?;
+
+    let config_section = get_profile_config(config, aws_profile)?;
+    let source_profile = config_section
+        .get("source_profile")
+        .map(std::borrow::ToOwned::to_owned);
+
+    let section = get_profile_creds(credentials, source_profile.as_ref())?;
+    Some(section.contains_key("aws_access_key_id"))
+}
+
 fn has_defined_credentials(
     context: &Context,
     aws_profile: Option<&Profile>,
@@ -212,10 +231,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 if credential_process, source_profile are defined or has 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 +1055,64 @@ 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()
+    }
 }