From 3af8c488b92c79c7ee715234b9fc65df9e013254 Mon Sep 17 00:00:00 2001
From: mgeeky <mb@binary-offensive.com>
Date: Fri, 29 Apr 2022 11:29:41 +0200
Subject: [PATCH] Updated assume-role-helper.sh

---
 clouds/aws/assume-role-helper.sh | 52 ++++++++++++++++++--------------
 1 file changed, 30 insertions(+), 22 deletions(-)

diff --git a/clouds/aws/assume-role-helper.sh b/clouds/aws/assume-role-helper.sh
index 64a72f3..5b81325 100644
--- a/clouds/aws/assume-role-helper.sh
+++ b/clouds/aws/assume-role-helper.sh
@@ -3,7 +3,7 @@
 # This script simply calls `aws sts assume-role` using hardcoded parameters, in order
 # to retrieve set of session credentials and reformat it into ~/.aws/credentials file format.
 #
-# Mariusz Banach, mgeeky '19-20
+# Mariusz B., mgeeky '19-20
 #
 
 
@@ -13,9 +13,12 @@
 #
 
 # Below two values are REQUIRED
-PROFILE_NAME=
+PROFILE_NAME=default
 ROLE_NAME=
 
+# Printed output role name
+OUTPUT_ROLE_NAME=
+
 # If left empty, will be deduced from `aws sts get-caller-identity` output.
 ACCOUNT_NUMBER=
 
@@ -38,8 +41,8 @@ DURATION=3600
 # regular commands sent first.
 out=$(aws --profile $PROFILE_NAME sts get-caller-identity)
 if [ $? -ne 0 ]; then
-    echo "[!] Could not get caller's identity: "
-    echo "$out"
+    >&2 echo "[!] Could not get caller's identity: "
+    >&2 echo "$out"
     exit 1
 fi
 
@@ -57,38 +60,43 @@ fi
 
 ROLE_ARN=arn:aws:iam::$ACCOUNT_NUMBER:role/$ROLE_NAME
 
-echo "[.] Using Role ARN: $ROLE_ARN"
+>&2 echo "[.] Using Role ARN: $ROLE_ARN"
 
-read -p "Type your AWS MFA Code (leave empty if not needed): " code
-echo
+code=""
 
 if [[ "$code" = "" ]] || [[ "$SERIAL_MFA" == "" ]]; then
-    echo "[.] MFA not provided, will attempt to assume role without it."
+    >&2 echo "[.] MFA not provided, will attempt to assume role without it."
     out=$(aws --profile $PROFILE_NAME sts assume-role --role-arn $ROLE_ARN --role-session-name $SESSION_NAME --duration-seconds $DURATION 2>&1)
 else
-    echo "[.] Will attempt to assume role with MFA provided."
+    >&2 echo "[.] Will attempt to assume role with MFA provided."
     out=$(aws --profile $PROFILE_NAME sts assume-role --serial-number $SERIAL_MFA --role-arn $ROLE_ARN --role-session-name $ROLE_NAME --duration-seconds $DURATION --token-code $code 2>&1)
 fi
 
+rolename=$PROFILE_NAME-$SESSION_NAME
+
+if [[ "$OUTPUT_ROLE_NAME" != "" ]]; then
+    rolename=$OUTPUT_ROLE_NAME
+fi
+
 if [ $? -eq 0 ]; then
     valid=$(printf '%dh:%dm:%ds\n' $(($DURATION/3600)) $(($DURATION%3600/60)) $(($DURATION%60)))
-    echo "[+] Collected session credentials. They will be valid for: $valid. "
-    echo -e "\tPaste below lines to your '~/.aws/credentials' file:"
+    >&2 echo "[+] Collected session credentials. They will be valid for: $valid. "
+    >&2 echo -e "\tPaste below lines to your '~/.aws/credentials' file:"
     echo
-    echo "[$PROFILE_NAME-$SESSION_NAME]"
+    echo "[$rolename]"
     echo "$out" | python3 -c 'import sys,json; foo=json.loads(sys.stdin.read()); print("aws_access_key_id={}\naws_secret_access_key={}\naws_session_token={}".format(foo["Credentials"]["AccessKeyId"],foo["Credentials"]["SecretAccessKey"],foo["Credentials"]["SessionToken"]))'
-    echo
+    >&2 echo
 else
-    echo "[!] Could not obtain assume-role session credentials:"
-    echo "$out"
-    echo
+    >&2 echo "[!] Could not obtain assume-role session credentials:"
+    >&2 echo "$out"
+    >&2 echo
     out2=$(env | grep -E 'AWS_[^=]+')
     if [[ "$out2" != "" ]]; then
-        echo "[!] Your command could fail because of pre-set AWS-related environment variables."
-        echo -e "\tPlease review them, correct any problems and re-launch that script."
-        echo
-        echo "$out2"
-        echo
+        >&2 echo "[!] Your command could fail because of pre-set AWS-related environment variables."
+        >&2 echo -e "\tPlease review them, correct any problems and re-launch that script."
+        >&2 echo
+        >&2 echo "$out2"
+        >&2 echo
     fi
     exit 1
-fi
\ No newline at end of file
+fi