From d5122757717dc3cffad7cfe8fcc7829915cc6471 Mon Sep 17 00:00:00 2001
From: Luis Cavalcante <luis@turbot.com>
Date: Wed, 14 Feb 2024 16:31:16 -0300
Subject: [PATCH] Add support for optional flags in
 generate_config_for_cross_account_roles.sh

Signed-off-by: Luis Cavalcante <luis@turbot.com>
---
 ...generate_config_for_cross_account_roles.sh | 40 ++++++++++++++++++-
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/all/aws-organizations-scripts/generate_config_for_cross_account_roles.sh b/all/aws-organizations-scripts/generate_config_for_cross_account_roles.sh
index 4abb131ec..6d9f304ad 100755
--- a/all/aws-organizations-scripts/generate_config_for_cross_account_roles.sh
+++ b/all/aws-organizations-scripts/generate_config_for_cross_account_roles.sh
@@ -15,13 +15,35 @@
 
 set -e
 
+# Initialize variables
+DURATION_SECONDS=""
+OU=""
+EXTERNAL_ID=""
+
+# Process flags
+while getopts d:o:e: flag
+do
+    case "${flag}" in
+        d) DURATION_SECONDS=${OPTARG};;
+        o) OU=${OPTARG};;
+        e) EXTERNAL_ID=${OPTARG};;
+    esac
+done
+
+# Shift positional parameters
+shift $((OPTIND -1))
+
+# Now $1, $2, etc. are positional parameters
 COMMAND=$1
 AUDIT_ROLE=$2
 AWS_CONFIG_FILE=$3
 SOURCE_PROFILE=$4
 
 usage () {
-  echo "Usage: $0 [IMDS | ECS | LOCAL ] <AUDIT_ROLE> <AWS_CONFIG_FILE> <SOURCE_PROFILE>"
+  echo "Usage: $0 [IMDS | ECS | LOCAL ] <AUDIT_ROLE> <AWS_CONFIG_FILE> <SOURCE_PROFILE> [-d DURATION_SECONDS] [-o OU] [-e EXTERNAL_ID]"
+  echo "  -d DURATION_SECONDS: Optional duration in seconds for the role session"
+  echo "  -o OU: Optional Organizational Unit ID to list accounts for"
+  echo "  -e EXTERNAL_ID: Optional External ID to use when assuming the role"
   exit 1
 }
 
@@ -49,6 +71,13 @@ if [ -z $SOURCE_PROFILE ] ; then
   fi
 fi
 
+# Check if OU is provided
+if [ ! -z $OU ] ; then
+  LIST_ACCOUNTS_COMMAND="aws organizations list-accounts-for-parent --parent-id $OU --query \"Accounts[?Status!='SUSPENDED'].[Name,Id,Status]\" --output text --profile $SOURCE_PROFILE | sort -f"
+else
+  LIST_ACCOUNTS_COMMAND="aws organizations list-accounts --query \"Accounts[?Status!='SUSPENDED'].[Name,Id,Status]\" --output text --profile $SOURCE_PROFILE | sort -f"
+fi
+
 # STEAMPIPE_INSTALL_DIR overrides the default steampipe directory of ~/.steampipe
 if [ -z $STEAMPIPE_INSTALL_DIR ] ; then
   echo "STEAMPIPE_INSTALL_DIR not defined, using the default location"
@@ -141,6 +170,13 @@ role_session_name = steampipe
 EOF
 fi
 
+if [ ! -z $DURATION_SECONDS ] ; then
+  echo "duration_seconds = $DURATION_SECONDS" >> $AWS_CONFIG_FILE
+fi
+
+if [ ! -z $EXTERNAL_ID ] ; then
+  echo "external_id = $EXTERNAL_ID" >> $AWS_CONFIG_FILE
+fi
 
 # And append an entry to the Steampipe config file
 cat <<EOF>>$SP_CONFIG_FILE
@@ -152,7 +188,7 @@ connection "aws_${SP_NAME}" {
 
 EOF
 
-done < <(aws organizations list-accounts --query "Accounts[?Status!='SUSPENDED'].[Name,Id,Status]" --output text --profile $SOURCE_PROFILE | sort -f) 
+done < <(eval $LIST_ACCOUNTS_COMMAND)
 
 if [ $COMMAND == "LOCAL" ] ; then
   echo "Append $AWS_CONFIG_FILE to your active AWS config file where profile $SOURCE_PROFILE is defined"