From 414f7d13fa4e1efcee9bddb4c74920075483ccdd Mon Sep 17 00:00:00 2001
From: Fabio Kreusch <fabiokr@gmail.com>
Date: Tue, 2 Aug 2022 13:55:24 -0300
Subject: [PATCH] Allows MFA through sts get-session-token

---
 awsp_functions.sh | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/awsp_functions.sh b/awsp_functions.sh
index 7add4c8..62f2ea3 100755
--- a/awsp_functions.sh
+++ b/awsp_functions.sh
@@ -26,10 +26,12 @@ function _awsListProfile() {
 function _awsSwitchProfile() {
    if [ -z $1 ]; then  echo "Usage: awsp profilename"; return; fi
    exists="$(aws configure get aws_access_key_id --profile $1)"
-   role_arn="$(aws configure get role_arn --profile $1)"
-   if [[ -n $exists || -n $role_arn ]]; then
-       if [[ -n $role_arn ]]; then
-           mfa_serial="$(aws configure get mfa_serial --profile $1)"
+
+   if [[ -n $exists ]]; then
+       role_arn="$(aws configure get role_arn --profile $1)"
+       mfa_serial="$(aws configure get mfa_serial --profile $1)"
+
+       if [[ -n $role_arn || -n $mfa_serial ]]; then
            if [[ -n $mfa_serial ]]; then
                echo "Please enter your MFA token for $mfa_serial:"
                read mfa_token
@@ -42,11 +44,15 @@ function _awsSwitchProfile() {
                profile=$1
            fi
 
-           echo "Assuming role $role_arn using profile $profile"
-           if [[ -n $mfa_serial ]]; then
-               JSON="$(aws sts assume-role --profile=$profile --role-arn $role_arn --role-session-name "$profile" --serial-number $mfa_serial --token-code $mfa_token)"
+           if [[ -n $role_arn ]]; then
+              echo "Assuming role $role_arn using profile $profile"
+              if [[ -n $mfa_serial ]]; then
+                  JSON="$(aws sts assume-role --profile=$profile --role-arn $role_arn --role-session-name "$profile" --serial-number $mfa_serial --token-code $mfa_token)"
+              else
+                  JSON="$(aws sts assume-role --profile=$profile --role-arn $role_arn --role-session-name "$profile")"
+              fi
            else
-               JSON="$(aws sts assume-role --profile=$profile --role-arn $role_arn --role-session-name "$profile")"
+              JSON="$(aws sts get-session-token --profile=$profile --serial-number $mfa_serial --token-code $mfa_token)"
            fi
 
            aws_access_key_id="$(echo $JSON | jq -r '.Credentials.AccessKeyId')"