From 9924af55db9e17c07aea2c64861c8d05632aa2a1 Mon Sep 17 00:00:00 2001
From: Kevin Pullin <kevin.pullin@gmail.com>
Date: Mon, 6 Apr 2020 07:59:56 -0700
Subject: [PATCH] feat(aws): Use AWS_VAULT as the profile if set (#984)

* Use AWS_VAULT as the profile if set

[aws-vault](https://github.com/99designs/aws-vault) sets
the `AWS_VAULT` env var instead of `AWS_PROFILE` when an
aws-vault session is active.

This PR adds support for reading the AWS profile value from
the `AWS_VAULT` env var, giving priority to `AWS_VAULT` if both
it and `AWS_PROFILE` are set.

* Update docs for AWS_VAULT
---
 docs/config/README.md  |  3 +++
 src/modules/aws.rs     |  4 +++-
 tests/testsuite/aws.rs | 12 ++++++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/docs/config/README.md b/docs/config/README.md
index 8c0131fb..7e95048b 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -136,6 +136,9 @@ 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.
 
+When using [aws-vault](https://github.com/99designs/aws-vault) the profile
+is read from the `AWS_VAULT` env var.
+
 ### Options
 
 | Variable          | Default         | Description                                                                 |
diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index 7962516d..ce409bdb 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -50,7 +50,9 @@ fn get_aws_region_from_config(aws_profile: Option<&str>) -> Option<Region> {
 
 fn get_aws_profile_and_region() -> (Option<Profile>, Option<Region>) {
     match (
-        env::var("AWS_PROFILE").ok(),
+        env::var("AWS_VAULT")
+            .or_else(|_| env::var("AWS_PROFILE"))
+            .ok(),
         env::var("AWS_REGION").ok(),
         env::var("AWS_DEFAULT_REGION").ok(),
     ) {
diff --git a/tests/testsuite/aws.rs b/tests/testsuite/aws.rs
index 01f518f8..988238b9 100644
--- a/tests/testsuite/aws.rs
+++ b/tests/testsuite/aws.rs
@@ -67,6 +67,18 @@ fn profile_set() -> io::Result<()> {
     Ok(())
 }
 
+#[test]
+fn profile_set_from_aws_vault() -> io::Result<()> {
+    let output = common::render_module("aws")
+        .env("AWS_VAULT", "astronauts-vault")
+        .env("AWS_PROFILE", "astronauts-profile")
+        .output()?;
+    let expected = format!("on {} ", Color::Yellow.bold().paint("☁️  astronauts-vault"));
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
 #[test]
 fn profile_and_region_set() -> io::Result<()> {
     let output = common::render_module("aws")