From 852a44094a3bb4df39f8f778bc7ada2ddda09727 Mon Sep 17 00:00:00 2001
From: Maksym <maksyms@users.noreply.github.com>
Date: Thu, 29 Oct 2020 21:13:36 +0000
Subject: [PATCH] aws: split setting profile from changing profile (#9402)

the change to assume a role when it is specified in configuration broke some workflows. This fix addresses that

Fixes #9394
---
 plugins/aws/README.md      |  7 +++++++
 plugins/aws/aws.plugin.zsh | 24 +++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/plugins/aws/README.md b/plugins/aws/README.md
index 4ceb71425..851f586dd 100644
--- a/plugins/aws/README.md
+++ b/plugins/aws/README.md
@@ -15,6 +15,13 @@ plugins=(... aws)
   It also sets `$AWS_EB_PROFILE` to `<profile>` for the Elastic Beanstalk CLI.
   Run `asp` without arguments to clear the profile.
 
+* `acp [<profile>]`: in addition to `asp` functionality, it actually changes the profile by
+   assuming the role specified in the `<profile>` configuration. It supports MFA and sets
+   `$AWS_ACCESS_KEY_ID`, `$AWS_SECRET_ACCESS_KEY` and `$AWS_SESSION_TOKEN`, if obtained. It
+   requires the roles to be configured as per the
+   [official guide](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-role.html).
+   Run `acp` without arguments to clear the profile.
+
 * `agp`: gets the current value of `$AWS_PROFILE`.
 
 * `aws_change_access_key`: changes the AWS access key of a profile.
diff --git a/plugins/aws/aws.plugin.zsh b/plugins/aws/aws.plugin.zsh
index fe1f098e8..8149ba121 100644
--- a/plugins/aws/aws.plugin.zsh
+++ b/plugins/aws/aws.plugin.zsh
@@ -4,6 +4,27 @@ function agp() {
 
 # AWS profile selection
 function asp() {
+  if [[ -z "$1" ]]; then
+    unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE
+    echo AWS profile cleared.
+    return
+  fi
+
+  local -a available_profiles
+  available_profiles=($(aws_profiles))
+  if [[ -z "${available_profiles[(r)$1]}" ]]; then
+    echo "${fg[red]}Profile '$1' not found in '${AWS_CONFIG_FILE:-$HOME/.aws/config}'" >&2
+    echo "Available profiles: ${(j:, :)available_profiles:-no profiles found}${reset_color}" >&2
+    return 1
+  fi
+
+  export AWS_DEFAULT_PROFILE=$1
+  export AWS_PROFILE=$1
+  export AWS_EB_PROFILE=$1
+}
+
+# AWS profile switch
+function acp() {
   if [[ -z "$1" ]]; then
     unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
     echo AWS profile cleared.
@@ -34,7 +55,7 @@ function asp() {
         echo "Please enter the session duration in seconds (900-43200; default: 3600, which is the default maximum for a role):"
         read sess_duration
         if [[ -z $sess_duration ]]; then
-          sess_duration = 3600
+          sess_duration="3600"
         fi
         mfa_opt="--serial-number $mfa_serial --token-code $mfa_token --duration-seconds $sess_duration"
       fi
@@ -100,6 +121,7 @@ function _aws_profiles() {
   reply=($(aws_profiles))
 }
 compctl -K _aws_profiles asp aws_change_access_key
+compctl -K _aws_profiles acp aws_change_access_key
 
 # AWS prompt
 function aws_prompt_info() {