Skip to content

Instantly share code, notes, and snippets.

@Szpadel
Last active November 13, 2025 13:39
Show Gist options
  • Select an option

  • Save Szpadel/1c169eca69d865111a866762ce3b892a to your computer and use it in GitHub Desktop.

Select an option

Save Szpadel/1c169eca69d865111a866762ce3b892a to your computer and use it in GitHub Desktop.
aws-sts-activate.fish - virtualenv like command for switching to different aws account using assume role. AI converted verions (untested) for bash and zsh
# Source this file in bash: source /path/to/aws-sts.bash
# Provides:
# - aws-sts-activate (alias for aws_sts_activate)
# - aws-sts-deactivate (alias for aws_sts_deactivate)
# Behavior mirrors your Fish version:
# * Loads/writes a cache file with credentials
# * Assumes role when cache is missing/near expiry
# * Shows the active profile in your prompt (left side, bash has no RPROMPT)
# * Auto-refreshes creds when <=30m to expiry (pre-exec and pre-prompt)
_aws_sts_usage() {
cat <<'USAGE'
Usage: aws-sts-activate [-n <name>|--name <name>] [-r <region>|--region <region>] [-h|--help] <account_id> <role_name>
Activate shell with AWS STS credentials
Options:
-n, --name <name> Set the profile/session name
-r, --region <region> Set AWS_DEFAULT_REGION
-h, --help Show this help
USAGE
}
__aws_sts_load_cache() {
local cf="$1"
[[ -f "$cf" ]] || return 1
export AWS_ACCESS_KEY_ID
export AWS_SECRET_ACCESS_KEY
export AWS_SESSION_TOKEN
export STS_EXPIRATION
AWS_ACCESS_KEY_ID="$(grep '^export AWS_ACCESS_KEY_ID=' "$cf" | cut -d'"' -f2)"
AWS_SECRET_ACCESS_KEY="$(grep '^export AWS_SECRET_ACCESS_KEY=' "$cf" | cut -d'"' -f2)"
AWS_SESSION_TOKEN="$(grep '^export AWS_SESSION_TOKEN=' "$cf" | cut -d'"' -f2)"
STS_EXPIRATION="$(grep '^EXPIRATION=' "$cf" | cut -d'=' -f2)"
[[ -n "$STS_EXPIRATION" ]]
}
__aws_sts_assume_role() {
local account_id="$1" role_name="$2" session_name="$3" out
# Ensure prerequisites exist
command -v aws >/dev/null 2>&1 || { echo "aws CLI not found in PATH" >&2; return 1; }
command -v jq >/dev/null 2>&1 || { echo "jq not found in PATH" >&2; return 1; }
unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN STS_EXPIRATION
if ! out="$(aws sts assume-role \
--role-arn "arn:aws:iam::$account_id:role/$role_name" \
--role-session-name "$session_name")"; then
return 1
fi
export AWS_ACCESS_KEY_ID
export AWS_SECRET_ACCESS_KEY
export AWS_SESSION_TOKEN
export STS_EXPIRATION
AWS_ACCESS_KEY_ID="$(jq -r .Credentials.AccessKeyId <<<"$out")"
AWS_SECRET_ACCESS_KEY="$(jq -r .Credentials.SecretAccessKey <<<"$out")"
AWS_SESSION_TOKEN="$(jq -r .Credentials.SessionToken <<<"$out")"
STS_EXPIRATION="$(( $(date +%s) + 3600 ))" # match fish: now + 1h
}
__aws_sts_write_cache() {
local cf="$1"
mkdir -p "$(dirname "$cf")"
printf 'export AWS_ACCESS_KEY_ID="%s"\nexport AWS_SECRET_ACCESS_KEY="%s"\nexport AWS_SESSION_TOKEN="%s"\nEXPIRATION=%s\n' \
"$AWS_ACCESS_KEY_ID" "$AWS_SECRET_ACCESS_KEY" "$AWS_SESSION_TOKEN" "$STS_EXPIRATION" >"$cf"
}
_aws_sts_check_expiration() {
[[ -n "$STS_EXPIRATION" ]] || return 0
local now; now="$(date +%s)"
if (( STS_EXPIRATION <= now + 1800 )); then # 30 minutes
printf '\e[38;5;208mRefreshing AWS STS credentials for %s …\e[0m\n' "${STS_PROFILE_NAME:-?}"
__aws_sts_assume_role "$STS_ACCOUNT_ID" "$STS_ROLE_NAME" "$STS_PROFILE_NAME" \
&& __aws_sts_write_cache "$STS_CACHE_FILE"
fi
}
aws_sts_activate() {
local profile_name region account_id role_name cache_file
# Option parsing (supports -n/--name, -r/--region, -h/--help)
while (($#)); do
case "$1" in
-n|--name) profile_name="$2"; shift 2 ;;
-r|--region) region="$2"; shift 2 ;;
-h|--help) _aws_sts_usage; return 0 ;;
--) shift; break ;;
-*) echo "Unknown option: $1" >&2; _aws_sts_usage; return 1 ;;
*) break ;;
esac
done
if (($# != 2)); then _aws_sts_usage; return 1; fi
account_id="$1"; role_name="$2"
[[ -n "$profile_name" ]] || profile_name="${account_id}-${role_name}"
cache_file="$HOME/.cache/aws-sts/${account_id}_${role_name}.env"
export STS_CACHE_FILE="$cache_file" STS_ACCOUNT_ID="$account_id" STS_ROLE_NAME="$role_name" STS_PROFILE_NAME="$profile_name"
if __aws_sts_load_cache "$cache_file"; then
local now; now="$(date +%s)"
if (( STS_EXPIRATION <= now + 900 )); then # <= 15 minutes -> refresh
__aws_sts_assume_role "$account_id" "$role_name" "$profile_name" && __aws_sts_write_cache "$cache_file"
fi
else
__aws_sts_assume_role "$account_id" "$role_name" "$profile_name" && __aws_sts_write_cache "$cache_file"
fi
if [[ -n "$region" ]]; then export AWS_DEFAULT_REGION="$region"; fi
# Prompt customization and refresh hooks (interactive only)
if [[ -z "$AWS_STS_ACTIVE" ]]; then
export AWS_STS_OLD_PS1="$PS1"
export AWS_STS_OLD_PROMPT_COMMAND="$PROMPT_COMMAND"
# Make profile visible on the left of the prompt (bash has no RPROMPT)
PS1="\[\e[48;5;208;97m\] ${STS_PROFILE_NAME} \[\e[0m\] ${PS1}"
if [[ $- == *i* ]]; then
# Save & set DEBUG trap (pre-exec)
export AWS_STS_OLD_DEBUG_TRAP
AWS_STS_OLD_DEBUG_TRAP="$(trap -p DEBUG)"
trap '_aws_sts_check_expiration' DEBUG
# Also check just before the prompt displays
PROMPT_COMMAND="_aws_sts_check_expiration; ${PROMPT_COMMAND:-:}"
fi
export AWS_STS_ACTIVE=1
fi
aws_sts_deactivate() {
# Restore traps/hooks/prompts
if [[ -n "$AWS_STS_OLD_DEBUG_TRAP" ]]; then
eval "$AWS_STS_OLD_DEBUG_TRAP"
else
trap - DEBUG 2>/dev/null
fi
if [[ -n "$AWS_STS_OLD_PROMPT_COMMAND" ]]; then
PROMPT_COMMAND="$AWS_STS_OLD_PROMPT_COMMAND"
else
PROMPT_COMMAND=
fi
[[ -n "$AWS_STS_OLD_PS1" ]] && PS1="$AWS_STS_OLD_PS1"
# Unset env and clean up
unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN AWS_DEFAULT_REGION
unset STS_PROFILE_NAME STS_CACHE_FILE STS_ACCOUNT_ID STS_ROLE_NAME STS_EXPIRATION
unset AWS_STS_ACTIVE AWS_STS_OLD_PS1 AWS_STS_OLD_PROMPT_COMMAND AWS_STS_OLD_DEBUG_TRAP
unset -f _aws_sts_check_expiration aws_sts_deactivate
}
alias aws-sts-deactivate='aws_sts_deactivate'
printf '\e[38;5;208mActivated AWS STS profile %s, type '\''aws-sts-deactivate'\'' to deactivate\e[0m\n' "$profile_name"
}
alias aws-sts-activate='aws_sts_activate'
# Save in ~/.config/fish/functions/aws-sts-activate.fish
# $argv[1] - cache file path
function __aws_sts_load_cache
set -l cf $argv[1]
if not test -f $cf
return 1
end
set -gx AWS_ACCESS_KEY_ID (grep '^export AWS_ACCESS_KEY_ID=' $cf | cut -d\" -f2)
set -gx AWS_SECRET_ACCESS_KEY (grep '^export AWS_SECRET_ACCESS_KEY=' $cf | cut -d\" -f2)
set -gx AWS_SESSION_TOKEN (grep '^export AWS_SESSION_TOKEN=' $cf | cut -d\" -f2)
set -gx STS_EXPIRATION (grep '^EXPIRATION=' $cf | cut -d= -f2)
# basic sanity
test -n "$STS_EXPIRATION"
end
# $argv[1] - account id, $argv[2] - role, $argv[3] - session name
function __aws_sts_assume_role
# Unset AWS session variables before assuming role again
set -e AWS_ACCESS_KEY_ID
set -e AWS_SECRET_ACCESS_KEY
set -e AWS_SESSION_TOKEN
set -e STS_EXPIRATION
set -l out (aws sts assume-role --role-arn arn:aws:iam::$argv[1]:role/$argv[2] --role-session-name $argv[3])
or return 1
set -gx AWS_ACCESS_KEY_ID (echo $out | jq -r .Credentials.AccessKeyId)
set -gx AWS_SECRET_ACCESS_KEY (echo $out | jq -r .Credentials.SecretAccessKey)
set -gx AWS_SESSION_TOKEN (echo $out | jq -r .Credentials.SessionToken)
set -gx STS_EXPIRATION (math (date +%s) + 3600)
end
# $argv[1] - cache file path
function __aws_sts_write_cache
mkdir -p (dirname $argv[1])
printf 'export AWS_ACCESS_KEY_ID="%s"\nexport AWS_SECRET_ACCESS_KEY="%s"\nexport AWS_SESSION_TOKEN="%s"\nEXPIRATION=%s\n' \
$AWS_ACCESS_KEY_ID $AWS_SECRET_ACCESS_KEY $AWS_SESSION_TOKEN $STS_EXPIRATION >$argv[1]
end
function aws-sts-activate --description "Activate shell with AWS STS credentials"
argparse n/name= h/help r/region= -- $argv
or return 1
function _print_help
echo "Usage: aws-sts-activate [-n <name>|--name <name>] [-r <region>| --region <region>] "
echo " [-h|--help] <account_id> <role_name>"
echo "Activate shell with AWS STS credentials"
echo ""
echo "Options:"
echo " -n, --name <name> Set tjhe profile name"
echo " -r, --region <region> Set the region"
echo " -h, --help Show this help"
end
if not test (count $argv) = 2
_print_help
return 1
end
if set -q _flag_help
_print_help
return 0
end
set -l account_id $argv[1]
set -l role_name $argv[2]
set -l profile_name
if set -q _flag_name
set profile_name $_flag_name
else
set profile_name (printf "%s-%s" $account_id $role_name)
end
set -l cache_file (printf "%s/.cache/aws-sts/%s_%s.env" $HOME $account_id $role_name)
set -gx STS_CACHE_FILE $cache_file
set -gx STS_ACCOUNT_ID $account_id
set -gx STS_ROLE_NAME $role_name
# try cache first
if __aws_sts_load_cache $cache_file
set -l now (date +%s)
if test $STS_EXPIRATION -le (math "$now + 900")
# almost expired - force refresh
__aws_sts_assume_role $account_id $role_name $profile_name
and __aws_sts_write_cache $cache_file
end
else
# no cache, or unreadable - get fresh credentials
__aws_sts_assume_role $account_id $role_name $profile_name
and __aws_sts_write_cache $cache_file
end
if set -q _flag_region
set -gx AWS_DEFAULT_REGION $_flag_region
end
set -gx STS_PROFILE_NAME $profile_name
if functions -q _sts_old_fish_right_prompt
functions -e _sts_old_fish_right_prompt
end
if functions -q fish_right_prompt
functions -c fish_right_prompt _sts_old_fish_right_prompt
set -g _STS_PROMPT_HAD_ORIGINAL 1
else
function _sts_old_fish_right_prompt
end
set -g _STS_PROMPT_HAD_ORIGINAL 0
end
function new_fish_right_prompt
set -l prompt (_sts_old_fish_right_prompt)
printf '%s %s %s ' (set_color -b F90 -o brwhite) $STS_PROFILE_NAME (set_color normal)
string join -- \n $prompt # handle multi-line prompts
end
function _aws-sts-deactivate --description "Deactivate shell with AWS STS credentials"
set -e AWS_ACCESS_KEY_ID
set -e AWS_SECRET_ACCESS_KEY
set -e AWS_SESSION_TOKEN
set -e AWS_DEFAULT_REGION
set -e STS_PROFILE_NAME
set -e STS_CACHE_FILE
set -e STS_ACCOUNT_ID
set -e STS_ROLE_NAME
set -e STS_EXPIRATION
functions -e fish_right_prompt
if set -q _STS_PROMPT_HAD_ORIGINAL
if test $_STS_PROMPT_HAD_ORIGINAL -eq 1
functions -c _sts_old_fish_right_prompt fish_right_prompt
end
set -e _STS_PROMPT_HAD_ORIGINAL
end
functions -e _sts_old_fish_right_prompt
functions -e __aws_sts_check_expiration
functions -e aws-sts-deactivate
end
if functions -q aws-sts-deactivate
functions -e aws-sts-deactivate
end
functions -c _aws-sts-deactivate aws-sts-deactivate
functions -e fish_right_prompt
functions -c new_fish_right_prompt fish_right_prompt
function __aws_sts_check_expiration --on-event fish_preexec
if not set -q STS_EXPIRATION
return
end
set -l now (date +%s)
if test $STS_EXPIRATION -le (math "$now + 1800")
echo (set_color F90)"Refreshing AWS STS credentials for $STS_PROFILE_NAME …"(set_color normal)
__aws_sts_assume_role $STS_ACCOUNT_ID $STS_ROLE_NAME $STS_PROFILE_NAME
and __aws_sts_write_cache $STS_CACHE_FILE
end
end
echo (set_color F90) "Activated AWS STS profile $profile_name, type 'aws-sts-deactivate' to deactivate" (set_color normal)
end
# Source this file in zsh: source /path/to/aws-sts.zsh
# Provides:
# - aws-sts-activate (alias for aws_sts_activate)
# - aws-sts-deactivate (alias for aws_sts_deactivate)
# Behavior mirrors your Fish version:
# * Loads/writes a cache file with credentials
# * Assumes role when cache is missing/near expiry
# * Shows the active profile in the right prompt (RPROMPT)
# * Auto-refreshes creds when <=30m to expiry (preexec + precmd hooks)
_aws_sts_usage() {
cat <<'USAGE'
Usage: aws-sts-activate [-n <name>|--name <name>] [-r <region>|--region <region>] [-h|--help] <account_id> <role_name>
Activate shell with AWS STS credentials
Options:
-n, --name <name> Set the profile/session name
-r, --region <region> Set AWS_DEFAULT_REGION
-h, --help Show this help
USAGE
}
__aws_sts_load_cache() {
local cf="$1"
[[ -f "$cf" ]] || return 1
export AWS_ACCESS_KEY_ID
export AWS_SECRET_ACCESS_KEY
export AWS_SESSION_TOKEN
export STS_EXPIRATION
AWS_ACCESS_KEY_ID="$(grep '^export AWS_ACCESS_KEY_ID=' "$cf" | cut -d'"' -f2)"
AWS_SECRET_ACCESS_KEY="$(grep '^export AWS_SECRET_ACCESS_KEY=' "$cf" | cut -d'"' -f2)"
AWS_SESSION_TOKEN="$(grep '^export AWS_SESSION_TOKEN=' "$cf" | cut -d'"' -f2)"
STS_EXPIRATION="$(grep '^EXPIRATION=' "$cf" | cut -d'=' -f2)"
[[ -n "$STS_EXPIRATION" ]]
}
__aws_sts_assume_role() {
local account_id="$1" role_name="$2" session_name="$3" out
# Ensure prerequisites exist
command -v aws >/dev/null 2>&1 || { print -r -- "aws CLI not found in PATH" >&2; return 1; }
command -v jq >/dev/null 2>&1 || { print -r -- "jq not found in PATH" >&2; return 1; }
unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN STS_EXPIRATION
if ! out="$(aws sts assume-role \
--role-arn "arn:aws:iam::$account_id:role/$role_name" \
--role-session-name "$session_name")"; then
return 1
fi
export AWS_ACCESS_KEY_ID
export AWS_SECRET_ACCESS_KEY
export AWS_SESSION_TOKEN
export STS_EXPIRATION
AWS_ACCESS_KEY_ID="$(jq -r .Credentials.AccessKeyId <<<"$out")"
AWS_SECRET_ACCESS_KEY="$(jq -r .Credentials.SecretAccessKey <<<"$out")"
AWS_SESSION_TOKEN="$(jq -r .Credentials.SessionToken <<<"$out")"
STS_EXPIRATION="$(( $(date +%s) + 3600 ))" # match fish: now + 1h
}
__aws_sts_write_cache() {
local cf="$1"
mkdir -p "$(dirname "$cf")"
printf 'export AWS_ACCESS_KEY_ID="%s"\nexport AWS_SECRET_ACCESS_KEY="%s"\nexport AWS_SESSION_TOKEN="%s"\nEXPIRATION=%s\n' \
"$AWS_ACCESS_KEY_ID" "$AWS_SECRET_ACCESS_KEY" "$AWS_SESSION_TOKEN" "$STS_EXPIRATION" >"$cf"
}
_aws_sts_check_expiration() {
[[ -n "$STS_EXPIRATION" ]] || return 0
local now; now="$(date +%s)"
if (( STS_EXPIRATION <= now + 1800 )); then # 30 minutes
print -P "%F{208}Refreshing AWS STS credentials for $STS_PROFILE_NAME …%f"
__aws_sts_assume_role "$STS_ACCOUNT_ID" "$STS_ROLE_NAME" "$STS_PROFILE_NAME" \
&& __aws_sts_write_cache "$STS_CACHE_FILE"
fi
}
aws_sts_activate() {
local profile_name region account_id role_name cache_file
# Option parsing (supports -n/--name, -r/--region, -h/--help)
while (( $# )); do
case "$1" in
-n|--name) profile_name="$2"; shift 2 ;;
-r|--region) region="$2"; shift 2 ;;
-h|--help) _aws_sts_usage; return 0 ;;
--) shift; break ;;
-*) print -r -- "Unknown option: $1" >&2; _aws_sts_usage; return 1 ;;
*) break ;;
esac
done
if (( $# != 2 )); then _aws_sts_usage; return 1; fi
account_id="$1"; role_name="$2"
[[ -n "$profile_name" ]] || profile_name="${account_id}-${role_name}"
cache_file="$HOME/.cache/aws-sts/${account_id}_${role_name}.env"
export STS_CACHE_FILE="$cache_file" STS_ACCOUNT_ID="$account_id" STS_ROLE_NAME="$role_name" STS_PROFILE_NAME="$profile_name"
if __aws_sts_load_cache "$cache_file"; then
local now; now="$(date +%s)"
if (( STS_EXPIRATION <= now + 900 )); then # <= 15 minutes -> refresh
__aws_sts_assume_role "$account_id" "$role_name" "$profile_name" && __aws_sts_write_cache "$cache_file"
fi
else
__aws_sts_assume_role "$account_id" "$role_name" "$profile_name" && __aws_sts_write_cache "$cache_file"
fi
if [[ -n "$region" ]]; then export AWS_DEFAULT_REGION="$region"; fi
# Prompt customization & hooks
if [[ -z "$AWS_STS_ACTIVE" ]]; then
export AWS_STS_OLD_RPROMPT="$RPROMPT"
# Right prompt: orange background (208), bold white text, then existing RPROMPT
RPROMPT="%K{208}%B%F{white} $STS_PROFILE_NAME %f%b%k ${RPROMPT}"
# Hook before each command (preexec) and before each prompt (precmd)
autoload -Uz add-zsh-hook 2>/dev/null
if whence -w add-zsh-hook >/dev/null 2>&1; then
add-zsh-hook preexec _aws_sts_check_expiration
add-zsh-hook precmd _aws_sts_check_expiration
else
# Fallback for very old zsh
typeset -ga preexec_functions precmd_functions
preexec_functions+=(_aws_sts_check_expiration)
precmd_functions+=(_aws_sts_check_expiration)
fi
export AWS_STS_ACTIVE=1
fi
aws_sts_deactivate() {
# Remove hooks and restore RPROMPT
if whence -w add-zsh-hook >/dev/null 2>&1; then
add-zsh-hook -d preexec _aws_sts_check_expiration 2>/dev/null
add-zsh-hook -d precmd _aws_sts_check_expiration 2>/dev/null
else
preexec_functions=(${preexec_functions:#_aws_sts_check_expiration})
precmd_functions=(${precmd_functions:#_aws_sts_check_expiration})
fi
RPROMPT="$AWS_STS_OLD_RPROMPT"
unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN AWS_DEFAULT_REGION
unset STS_PROFILE_NAME STS_CACHE_FILE STS_ACCOUNT_ID STS_ROLE_NAME STS_EXPIRATION
unset AWS_STS_ACTIVE AWS_STS_OLD_RPROMPT
unset -f _aws_sts_check_expiration aws_sts_deactivate
}
alias aws-sts-deactivate='aws_sts_deactivate'
print -P "%F{208}Activated AWS STS profile $profile_name, type 'aws-sts-deactivate' to deactivate%f"
}
alias aws-sts-activate='aws_sts_activate'
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment