AWS SSO: Add ALL Profiles to ~/.aws/config, and login to all Roles to ~/.aws/credentials

aws_sso_helper.sh

#!/bin/bash

# Exit on error
set -e

# Function to check if jq is installed
check_jq() {
    if ! command -v jq &> /dev/null; then
        echo "Error: jq is required but not installed. Please install jq first."
        exit 1
    fi
}

# Function to get SSO configuration
get_sso_config() {
    local config_file="$HOME/.aws/config"
    if [ ! -f "$config_file" ]; then
        echo "Error: AWS config file not found. Please run 'aws configure sso' first."
        exit 1
    fi

    # Get the first SSO configuration found
    local sso_start_url=$(awk -F' = ' '/sso_start_url/{print $2; exit}' "$config_file")
    local sso_region=$(awk -F' = ' '/sso_region/{print $2; exit}' "$config_file")

    if [ -z "$sso_start_url" ] || [ -z "$sso_region" ]; then
        echo "Error: No SSO configuration found. Please run 'aws configure sso' first."
        exit 1
    fi

    echo "$sso_start_url|$sso_region"
}

# Function to get cached token
get_cached_token() {
    local sso_start_url="$1"
    local sso_region="$2"
    local cache_dir="$HOME/.aws/sso/cache"

    if [ ! -d "$cache_dir" ]; then
        echo "Error: AWS SSO cache directory not found. Please run 'aws sso login' first."
        exit 1
    fi

    local latest_token=""
    local latest_time=0

    for token_file in "$cache_dir"/*.json; do
        if [ -f "$token_file" ]; then
            local file_url=$(jq -r '.startUrl' "$token_file")
            local file_region=$(jq -r '.region' "$token_file")
            
            if [ "$file_url" = "$sso_start_url" ] && [ "$file_region" = "$sso_region" ]; then
                local mtime=$(stat -f "%m" "$token_file")
                if [ "$mtime" -gt "$latest_time" ]; then
                    latest_time=$mtime
                    latest_token=$token_file
                fi
            fi
        fi
    done

    if [ -z "$latest_token" ]; then
        echo "Error: No valid SSO token found for $sso_start_url. Please run 'aws sso login'."
        exit 1
    fi

    # Check token expiration
    local token_expiration=$(jq -r '.expiresAt' "$latest_token")
    local current_time=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
    
    if [[ "$token_expiration" < "$current_time" ]]; then
        echo "Error: SSO token has expired. Please run 'aws sso login' to refresh it."
        exit 1
    fi

    jq -r '.accessToken' "$latest_token"
}

# Function to update AWS credentials file
update_credentials() {
    local account_id="$1"
    local role_name="$2"
    local access_key="$3"
    local secret_key="$4"
    local session_token="$5"
    local creds_file="$HOME/.aws/credentials"

    # Create or update the profile in credentials file
    if [ ! -f "$creds_file" ]; then
        mkdir -p "$(dirname "$creds_file")"
        touch "$creds_file"
    fi

    # Remove existing profile if it exists
    sed -i '' "/\[$account_id\_$role_name\]/,/^$/d" "$creds_file"

    # Add new profile
    {
        echo "[$account_id\_$role_name]"
        echo "aws_access_key_id=$access_key"
        echo "aws_secret_access_key=$secret_key"
        echo "aws_session_token=$session_token"
        echo ""
    } >> "$creds_file"
}

# Function to update AWS config file
update_config() {
    local account_id="$1"
    local role_name="$2"
    local sso_start_url="$3"
    local sso_region="$4"
    local config_file="$HOME/.aws/config"

    # Create or update the profile in config file
    if [ ! -f "$config_file" ]; then
        mkdir -p "$(dirname "$config_file")"
        touch "$config_file"
    fi

    # Remove existing profile if it exists
    sed -i '' "/\[profile $account_id\_$role_name\]/,/^$/d" "$config_file"

    # Add new profile
    {
        echo "[profile ${account_id}_${role_name}]"
        echo "sso_start_url = $sso_start_url"
        echo "sso_region = $sso_region"
        echo "sso_account_id = $account_id"
        echo "sso_role_name = $role_name"
        echo ""
    } >> "$config_file"
}

main() {
    check_jq

    # Get SSO configuration
    local sso_config=$(get_sso_config)
    local sso_start_url=$(echo "$sso_config" | cut -d'|' -f1)
    local sso_region=$(echo "$sso_config" | cut -d'|' -f2)

    # Get access token
    local access_token=$(get_cached_token "$sso_start_url" "$sso_region")

    echo "Fetching AWS accounts..."

    # Get list of accounts
    local accounts=$(aws sso list-accounts \
        --access-token "$access_token" \
        --region "$sso_region" \
        --output json)

    local total_accounts=$(echo "$accounts" | jq '.accountList | length')
    echo "Found $total_accounts accounts"

    local count=0
    echo "$accounts" | jq -c '.accountList[]' | while read -r account; do
        ((count++))
        local account_id=$(echo "$account" | jq -r '.accountId')
        local account_name=$(echo "$account" | jq -r '.accountName // "Unknown Account"')

        echo -ne "\rProgress: [$count/$total_accounts] $(( count * 100 / total_accounts ))% - Processing $account_name"

        # Get roles for the account
        local roles=$(aws sso list-account-roles \
            --access-token "$access_token" \
            --account-id "$account_id" \
            --region "$sso_region" \
            --output json)

        echo "$roles" | jq -c '.roleList[]' | while read -r role; do
            local role_name=$(echo "$role" | jq -r '.roleName')

            # Get credentials for the role
            local credentials=$(aws sso get-role-credentials \
                --access-token "$access_token" \
                --account-id "$account_id" \
                --role-name "$role_name" \
                --region "$sso_region" \
                --output json)

            # Extract credential values
            local access_key=$(echo "$credentials" | jq -r '.roleCredentials.accessKeyId')
            local secret_key=$(echo "$credentials" | jq -r '.roleCredentials.secretAccessKey')
            local session_token=$(echo "$credentials" | jq -r '.roleCredentials.sessionToken')

            # Update credentials and config files
            update_credentials "$account_id" "$role_name" "$access_key" "$secret_key" "$session_token"
            update_config "$account_id" "$role_name" "$sso_start_url" "$sso_region"
        done
    done

    echo -e "\nCredential fetching completed!"
    echo "Credentials have been updated in $HOME/.aws/credentials"
    echo "Profiles have been updated in $HOME/.aws/config"
}

main