aws sso has never been this easy!

aws_sso.sh

#!/bin/bash

SSO_CACHE_DIR="$HOME/.aws/sso/cache"
ALIAS_FILE="$HOME/.aws/alias"

# !!! UPDATE THIS !!!
SSO_START_URL="https://COMPANY_SSO_ALIAS.awsapps.com/start/"
SSO_REGION="us-east-1"
SSO_REGISTRATION_SCOPES="sso:account:access"

# !!! UPDATE THIS !!!
export AWS_PROFILE_PREFIX="COMPANY_NAME"

# Detect current shell
CURRENT_SHELL=$(basename "$SHELL")
RC_FILE=""
if [ "$CURRENT_SHELL" = "zsh" ]; then
    RC_FILE="$HOME/.zshrc"
elif [ "$CURRENT_SHELL" = "bash" ]; then
    RC_FILE="$HOME/.bashrc"
else
    echo "Unsupported shell: $CURRENT_SHELL"
    exit 1
fi

cat > ~/.aws/config <<EOF
[sso-session $AWS_PROFILE_PREFIX]
sso_start_url = $SSO_START_URL
sso_region = $SSO_REGION
sso_registration_scopes = $SSO_REGISTRATION_SCOPES
EOF


aws sso login --sso-session "$AWS_PROFILE_PREFIX"

for file in "$SSO_CACHE_DIR"/*; do
    if [ -f "$file" ]; then
        TOKEN=$(cat "$file" | jq -r '.accessToken')
        if [ "$TOKEN" == "null" ]; then
            echo "No token found in $file"
            continue
        else
            break
        fi


    fi
done

function add_profile(){
    
    local account_id=$1
    local role_name=$2
    local profile_name=$3

    cat >> ~/.aws/config <<EOF
[profile $profile_name]
sso_session = $AWS_PROFILE_PREFIX
sso_account_id = $account_id
sso_role_name = $role_name
region = $SSO_REGION
EOF
    echo "[+] Profile $profile_name added!"
}


ACCOUNTS=$(aws sso list-accounts --access-token "$TOKEN" --region $SSO_REGION | jq -r '.accountList[] | "\(.accountId)-\(.accountName | gsub(" "; ""))"')

for account in $ACCOUNTS; do
    
    ACCOUNT_ID=$(echo "$account" | cut -d- -f 1)
    ACCOUNT_NAME=$(echo "$account" | cut -d- -f 2-)


    echo "----------------------------------------"
    echo "ACCOUNT: $ACCOUNT_NAME - $ACCOUNT_ID"
    echo "----------------------------------------"

    
    ROLES=$(aws sso list-account-roles --access-token "$TOKEN" --region $SSO_REGION --account-id $ACCOUNT_ID | jq -r ".roleList[].roleName")
    for role in $ROLES; do
        PROFILE_NAME=$(echo "$AWS_PROFILE_PREFIX-$ACCOUNT_NAME-$role" | tr '[:upper:]' '[:lower:]')
        add_profile $ACCOUNT_ID $role $PROFILE_NAME
    done
done

if ! grep -q "#DONOTEDIT:AWS_HELPERS" "$RC_FILE"; then
    cat >> "$RC_FILE" <<EOF

#DONOTEDIT:AWS_HELPERS
aws-selector() {
    export AWS_PROFILE=$(aws configure list-profiles | fzf)
}
#DONOTEDIT:AWS_HELPERS
EOF
    echo "[+] aws-selector function added to $RC_FILE"
    echo "[!] Please run 'source $RC_FILE' or restart your shell to use the aws-selector function"
else
    echo "[*] AWS helpers section already exists in $RC_FILE"
fi