linux/ohmyzsh
1
0
Fork 0
mirror of https://github.com/ohmyzsh/ohmyzsh.git synced 2024-11-19 21:41:07 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(aws): improve shell functions safety and reliability

Browse source 
Modified AWS profile state file handling to create directories
if they don't exist. Users with custom state file locations should ensure
parent directories exist or have proper permissions.

- Added proper variable quoting and parameter expansion
- Improved error handling and messaging
- Enhanced credentials validation
- Made state file handling more robust
- Improved MFA token and session handling
- Fixed profile switching reliability issues
- Added consistent error output to stderr
- Improved AWS regions parsing using native query
- Added fallbacks for theme variables
This commit is contained in:
mnv 2024-10-23 14:33:19 +05:30 committed by GitHub
parent 7bbebcd520
commit 6a5e4c3801
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 149 additions and 136 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

285
plugins/aws/aws.plugin.zsh
View file

@ -1,23 +1,29 @@
function agp() { function agp() {
echo $AWS_PROFILE echo "${AWS_PROFILE:-}"
} }
function agr() { function agr() {
echo $AWS_REGION echo "${AWS_REGION:-}"
} }
# Update state file if enabled # Update state file if enabled
function _aws_update_state() { function _aws_update_state() {
if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then
test -d $(dirname ${AWS_STATE_FILE}) || exit 1 local state_dir="$(dirname "${AWS_STATE_FILE}")"
echo "${AWS_PROFILE} ${AWS_REGION}" > "${AWS_STATE_FILE}" if [[ ! -d "$state_dir" ]]; then
mkdir -p "$state_dir" 2>/dev/null || return 1
fi
echo "${AWS_PROFILE:-} ${AWS_REGION:-}" > "${AWS_STATE_FILE}"
fi fi
} }
function _aws_clear_state() { function _aws_clear_state() {
if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then
test -d $(dirname ${AWS_STATE_FILE}) || exit 1 local state_dir="$(dirname "${AWS_STATE_FILE}")"
echo -n > "${AWS_STATE_FILE}" if [[ ! -d "$state_dir" ]]; then
mkdir -p "$state_dir" 2>/dev/null || return 1
fi
: > "${AWS_STATE_FILE}"
fi fi
} }
@ -26,35 +32,39 @@ function asp() {
if [[ -z "$1" ]]; then if [[ -z "$1" ]]; then
unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE AWS_PROFILE_REGION unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE AWS_PROFILE_REGION
_aws_clear_state _aws_clear_state
echo AWS profile cleared. echo "AWS profile cleared."
return return 0
fi fi
local -a available_profiles local -a available_profiles
available_profiles=($(aws_profiles)) available_profiles=($(aws_profiles))
if [[ -z "${available_profiles[(r)$1]}" ]]; then if [[ -z "${available_profiles[(r)$1]}" ]]; then
echo "${fg[red]}Profile '$1' not found in '${AWS_CONFIG_FILE:-$HOME/.aws/config}'" >&2 echo "Profile '$1' not found in '${AWS_CONFIG_FILE:-$HOME/.aws/config}'" >&2
echo "Available profiles: ${(j:, :)available_profiles:-no profiles found}${reset_color}" >&2 echo "Available profiles: ${(j:, :)available_profiles:-no profiles found}" >&2
return 1 return 1
fi fi
export AWS_DEFAULT_PROFILE=$1 export AWS_DEFAULT_PROFILE="$1"
export AWS_PROFILE=$1 export AWS_PROFILE="$1"
export AWS_EB_PROFILE=$1 export AWS_EB_PROFILE="$1"
export AWS_PROFILE_REGION=$(aws configure get region) AWS_PROFILE_REGION="$(aws configure get region)" || AWS_PROFILE_REGION=""
export AWS_PROFILE_REGION
_aws_update_state _aws_update_state
if [[ "$2" == "login" ]]; then case "$2" in
if [[ -n "$3" ]]; then "login")
aws sso login --sso-session $3 if [[ -n "$3" ]]; then
else aws sso login --sso-session "$3"
aws sso login else
fi aws sso login
elif [[ "$2" == "logout" ]]; then fi
aws sso logout ;;
fi "logout")
aws sso logout
;;
esac
} }
# AWS region selection # AWS region selection
@ -62,51 +72,57 @@ function asr() {
if [[ -z "$1" ]]; then if [[ -z "$1" ]]; then
unset AWS_DEFAULT_REGION AWS_REGION unset AWS_DEFAULT_REGION AWS_REGION
_aws_update_state _aws_update_state
echo AWS region cleared. echo "AWS region cleared."
return return 0
fi fi
local -a available_regions local -a available_regions
available_regions=($(aws_regions)) available_regions=($(aws_regions))
if [[ -z "${available_regions[(r)$1]}" ]]; then if [[ -z "${available_regions[(r)$1]}" ]]; then
echo "${fg[red]}Available regions: \n$(aws_regions)" echo "Invalid region. Available regions:" >&2
aws_regions >&2
return 1 return 1
fi fi
export AWS_REGION=$1 export AWS_DEFAULT_REGION="$1"
export AWS_DEFAULT_REGION=$1 export AWS_REGION="$1"
_aws_update_state _aws_update_state
} }
# AWS profile switch # AWS profile switch
function acp() { function acp() {
if [[ -z "$1" ]]; then if [[ -z "$1" ]]; then
unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN echo "AWS profile cleared."
echo AWS profile cleared. return 0
return
fi fi
local -a available_profiles local -a available_profiles
available_profiles=($(aws_profiles)) available_profiles=($(aws_profiles))
if [[ -z "${available_profiles[(r)$1]}" ]]; then if [[ -z "${available_profiles[(r)$1]}" ]]; then
echo "${fg[red]}Profile '$1' not found in '${AWS_CONFIG_FILE:-$HOME/.aws/config}'" >&2 echo "Profile '$1' not found in '${AWS_CONFIG_FILE:-$HOME/.aws/config}'" >&2
echo "Available profiles: ${(j:, :)available_profiles:-no profiles found}${reset_color}" >&2 echo "Available profiles: ${(j:, :)available_profiles:-no profiles found}" >&2
return 1 return 1
fi fi
local profile="$1" local profile="$1"
local mfa_token="$2" local mfa_token="$2"
# Get fallback credentials for if the aws command fails or no command is run # Get fallback credentials
local aws_access_key_id="$(aws configure get aws_access_key_id --profile $profile)" local aws_access_key_id
local aws_secret_access_key="$(aws configure get aws_secret_access_key --profile $profile)" local aws_secret_access_key
local aws_session_token="$(aws configure get aws_session_token --profile $profile)" local aws_session_token
aws_access_key_id="$(aws configure get aws_access_key_id --profile "$profile")"
aws_secret_access_key="$(aws configure get aws_secret_access_key --profile "$profile")"
aws_session_token="$(aws configure get aws_session_token --profile "$profile")"
# Check for MFA configuration
# First, if the profile has MFA configured, lets get the token and session duration local mfa_serial
local mfa_serial="$(aws configure get mfa_serial --profile $profile)" local sess_duration
local sess_duration="$(aws configure get duration_seconds --profile $profile)"
mfa_serial="$(aws configure get mfa_serial --profile "$profile")"
sess_duration="$(aws configure get duration_seconds --profile "$profile")"
if [[ -n "$mfa_serial" ]]; then if [[ -n "$mfa_serial" ]]; then
local -a mfa_opt local -a mfa_opt
@ -115,55 +131,52 @@ function acp() {
read -r mfa_token read -r mfa_token
fi fi
if [[ -z "$sess_duration" ]]; then if [[ -z "$sess_duration" ]]; then
echo -n "Please enter the session duration in seconds (900-43200; default: 3600, which is the default maximum for a role): " echo -n "Please enter the session duration in seconds (900-43200; default: 3600): "
read -r sess_duration read -r sess_duration
fi fi
mfa_opt=(--serial-number "$mfa_serial" --token-code "$mfa_token" --duration-seconds "${sess_duration:-3600}") mfa_opt=(--serial-number "$mfa_serial" --token-code "$mfa_token" --duration-seconds "${sess_duration:-3600}")
fi fi
# Now see whether we need to just MFA for the current role, or assume a different one # Check for role assumption
local role_arn="$(aws configure get role_arn --profile $profile)" local role_arn
local sess_name="$(aws configure get role_session_name --profile $profile)" local sess_name
role_arn="$(aws configure get role_arn --profile "$profile")"
sess_name="$(aws configure get role_session_name --profile "$profile")"
if [[ -n "$role_arn" ]]; then if [[ -n "$role_arn" ]]; then
# Means we need to assume a specified role local -a aws_command
aws_command=(aws sts assume-role --role-arn "$role_arn" "${mfa_opt[@]}") aws_command=(aws sts assume-role --role-arn "$role_arn" "${mfa_opt[@]}")
# Check whether external_id is configured to use while assuming the role local external_id
local external_id="$(aws configure get external_id --profile $profile)" external_id="$(aws configure get external_id --profile "$profile")"
if [[ -n "$external_id" ]]; then if [[ -n "$external_id" ]]; then
aws_command+=(--external-id "$external_id") aws_command+=(--external-id "$external_id")
fi fi
# Get source profile to use to assume role local source_profile
local source_profile="$(aws configure get source_profile --profile $profile)" source_profile="$(aws configure get source_profile --profile "$profile")"
if [[ -z "$sess_name" ]]; then if [[ -z "$sess_name" ]]; then
sess_name="${source_profile:-profile}" sess_name="${source_profile:-$profile}"
fi fi
aws_command+=(--profile="${source_profile:-profile}" --role-session-name "${sess_name}") aws_command+=(--profile="${source_profile:-$profile}" --role-session-name "$sess_name")
echo "Assuming role $role_arn using profile ${source_profile:-profile}" echo "Assuming role $role_arn using profile ${source_profile:-$profile}"
else else
# Means we only need to do MFA local -a aws_command
aws_command=(aws sts get-session-token --profile="$profile" "${mfa_opt[@]}") aws_command=(aws sts get-session-token --profile="$profile" "${mfa_opt[@]}")
echo "Obtaining session token for profile $profile" echo "Obtaining session token for profile $profile"
fi fi
# Format output of aws command for easier processing
aws_command+=(--query '[Credentials.AccessKeyId,Credentials.SecretAccessKey,Credentials.SessionToken]' --output text) aws_command+=(--query '[Credentials.AccessKeyId,Credentials.SecretAccessKey,Credentials.SessionToken]' --output text)
# Run the aws command to obtain credentials local credentials
local -a credentials credentials="$(${aws_command[@]})" || return 1
credentials=(${(ps:\t:)"$(${aws_command[@]})"})
if [[ -n "$credentials" ]]; then if [[ -n "$credentials" ]]; then
aws_access_key_id="${credentials[1]}" read -r aws_access_key_id aws_secret_access_key aws_session_token <<< "$credentials"
aws_secret_access_key="${credentials[2]}"
aws_session_token="${credentials[3]}"
fi fi
# Switch to AWS profile if [[ -n "$aws_access_key_id" && -n "$aws_secret_access_key" ]]; then
if [[ -n "${aws_access_key_id}" && -n "$aws_secret_access_key" ]]; then
export AWS_DEFAULT_PROFILE="$profile" export AWS_DEFAULT_PROFILE="$profile"
export AWS_PROFILE="$profile" export AWS_PROFILE="$profile"
export AWS_EB_PROFILE="$profile" export AWS_EB_PROFILE="$profile"
@ -177,73 +190,79 @@ function acp() {
fi fi
echo "Switched to AWS Profile: $profile" echo "Switched to AWS Profile: $profile"
else
echo "Failed to obtain valid credentials" >&2
return 1
fi fi
} }
function aws_change_access_key() { function aws_change_access_key() {
if [[ -z "$1" ]]; then if [[ -z "$1" ]]; then
echo "usage: $0 <profile>" echo "usage: ${0} <profile>" >&2
return 1 return 1
fi fi
local profile="$1" local profile="$1"
# Get current access key local original_aws_access_key_id
local original_aws_access_key_id="$(aws configure get aws_access_key_id --profile $profile)" original_aws_access_key_id="$(aws configure get aws_access_key_id --profile "$profile")"
asp "$profile" || return 1 if ! asp "$profile"; then
echo "Generating a new access key pair for you now."
if aws --no-cli-pager iam create-access-key; then
echo "Insert the newly generated credentials when asked."
aws --no-cli-pager configure --profile $profile
else
echo "Current access keys:"
aws --no-cli-pager iam list-access-keys
echo "Profile \"${profile}\" is currently using the $original_aws_access_key_id key. You can delete an old access key by running \`aws --profile $profile iam delete-access-key --access-key-id AccessKeyId\`"
return 1 return 1
fi fi
read -q "yn?Would you like to disable your previous access key (${original_aws_access_key_id}) now? " echo "Generating a new access key pair..."
case $yn in if aws --no-cli-pager iam create-access-key; then
echo "Insert the newly generated credentials when asked."
aws --no-cli-pager configure --profile "$profile"
else
echo "Current access keys:"
aws --no-cli-pager iam list-access-keys
echo "Profile \"${profile}\" is currently using the $original_aws_access_key_id key."
echo "You can delete an old access key by running: aws --profile $profile iam delete-access-key --access-key-id AccessKeyId"
return 1
fi
echo -n "Would you like to disable your previous access key (${original_aws_access_key_id}) now? [y/N] "
read -r yn
case "$yn" in
[Yy]*) [Yy]*)
echo -n "\nDisabling access key ${original_aws_access_key_id}..." echo "Disabling access key ${original_aws_access_key_id}..."
if aws --no-cli-pager iam update-access-key --access-key-id ${original_aws_access_key_id} --status Inactive; then if aws --no-cli-pager iam update-access-key --access-key-id "${original_aws_access_key_id}" --status Inactive; then
echo "done." echo "Access key disabled successfully."
else else
echo "\nFailed to disable ${original_aws_access_key_id} key." echo "Failed to disable ${original_aws_access_key_id} key." >&2
fi fi
;; ;;
*)
echo ""
;;
esac esac
echo "You can now safely delete the old access key by running \`aws --profile $profile iam delete-access-key --access-key-id ${original_aws_access_key_id}\`"
echo "You can now safely delete the old access key by running:"
echo "aws --profile $profile iam delete-access-key --access-key-id ${original_aws_access_key_id}"
echo "Your current keys are:" echo "Your current keys are:"
aws --no-cli-pager iam list-access-keys aws --no-cli-pager iam list-access-keys
} }
function aws_regions() { function aws_regions() {
local region local region="${AWS_DEFAULT_REGION:-${AWS_REGION:-us-west-1}}"
if [[ $AWS_DEFAULT_REGION ]];then
region="$AWS_DEFAULT_REGION"
elif [[ $AWS_REGION ]];then
region="$AWS_REGION"
else
region="us-west-1"
fi
if [[ $AWS_DEFAULT_PROFILE || $AWS_PROFILE ]];then if [[ -n "$AWS_DEFAULT_PROFILE" || -n "$AWS_PROFILE" ]]; then
aws ec2 describe-regions --region $region |grep RegionName | awk -F ':' '{gsub(/"/, "", $2);gsub(/,/, "", $2);gsub(/ /, "", $2); print $2}' aws ec2 describe-regions --region "$region" --query 'Regions[].RegionName' --output text | tr '\t' '\n'
else else
echo "You must specify a AWS profile." echo "You must specify an AWS profile." >&2
return 1
fi fi
} }
function aws_profiles() { function aws_profiles() {
aws --no-cli-pager configure list-profiles 2> /dev/null && return if aws --no-cli-pager configure list-profiles 2>/dev/null; then
[[ -r "${AWS_CONFIG_FILE:-$HOME/.aws/config}" ]] || return 1 return 0
grep --color=never -Eo '\[.*\]' "${AWS_CONFIG_FILE:-$HOME/.aws/config}" | sed -E 's/^[[:space:]]*\[(profile)?[[:space:]]*([^[:space:]]+)\][[:space:]]*$/\2/g' fi
local config_file="${AWS_CONFIG_FILE:-$HOME/.aws/config}"
[[ -r "$config_file" ]] || return 1
grep -Eo '\[.*\]' "$config_file" | sed -E 's/^[[:space:]]*\[(profile)?[[:space:]]*([^[:space:]]+)\][[:space:]]*$/\2/g'
} }
# Completion functions
function _aws_regions() { function _aws_regions() {
reply=($(aws_regions)) reply=($(aws_regions))
} }
@ -256,83 +275,77 @@ compctl -K _aws_profiles asp acp aws_change_access_key
# AWS prompt # AWS prompt
function aws_prompt_info() { function aws_prompt_info() {
local _aws_to_show local _aws_to_show=""
local region="${AWS_REGION:-${AWS_DEFAULT_REGION:-$AWS_PROFILE_REGION}}" local region="${AWS_REGION:-${AWS_DEFAULT_REGION:-$AWS_PROFILE_REGION}}"
if [[ -n "$AWS_PROFILE" ]];then if [[ -n "$AWS_PROFILE" ]]; then
_aws_to_show+="${ZSH_THEME_AWS_PROFILE_PREFIX="<aws:"}${AWS_PROFILE}${ZSH_THEME_AWS_PROFILE_SUFFIX=">"}" _aws_to_show="${ZSH_THEME_AWS_PROFILE_PREFIX:-<aws:}${AWS_PROFILE}${ZSH_THEME_AWS_PROFILE_SUFFIX:->}"
fi fi
if [[ -n "$region" ]]; then if [[ -n "$region" ]]; then
[[ -n "$_aws_to_show" ]] && _aws_to_show+="${ZSH_THEME_AWS_DIVIDER=" "}" [[ -n "$_aws_to_show" ]] && _aws_to_show+="${ZSH_THEME_AWS_DIVIDER:- }"
_aws_to_show+="${ZSH_THEME_AWS_REGION_PREFIX="<region:"}${region}${ZSH_THEME_AWS_REGION_SUFFIX=">"}" _aws_to_show+="${ZSH_THEME_AWS_REGION_PREFIX:-<region:}${region}${ZSH_THEME_AWS_REGION_SUFFIX:->}"
fi fi
echo "$_aws_to_show" echo "$_aws_to_show"
} }
# Add AWS prompt to RPROMPT if enabled
if [[ "$SHOW_AWS_PROMPT" != false && "$RPROMPT" != *'$(aws_prompt_info)'* ]]; then if [[ "$SHOW_AWS_PROMPT" != false && "$RPROMPT" != *'$(aws_prompt_info)'* ]]; then
RPROMPT='$(aws_prompt_info)'"$RPROMPT" RPROMPT='$(aws_prompt_info)'"$RPROMPT"
fi fi
# Initialize state from file if enabled
if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then
AWS_STATE_FILE="${AWS_STATE_FILE:-/tmp/.aws_current_profile}" AWS_STATE_FILE="${AWS_STATE_FILE:-/tmp/.aws_current_profile}"
test -s "${AWS_STATE_FILE}" || return if [[ -s "$AWS_STATE_FILE" ]]; then
local -a aws_state
aws_state=($(< "$AWS_STATE_FILE"))
aws_state=($(cat $AWS_STATE_FILE)) export AWS_DEFAULT_PROFILE="${aws_state[1]}"
export AWS_PROFILE="$AWS_DEFAULT_PROFILE"
export AWS_EB_PROFILE="$AWS_DEFAULT_PROFILE"
export AWS_DEFAULT_PROFILE="${aws_state[1]}" if [[ -z "${aws_state[2]}" ]]; then
export AWS_PROFILE="$AWS_DEFAULT_PROFILE" AWS_REGION="$(aws configure get region)"
export AWS_EB_PROFILE="$AWS_DEFAULT_PROFILE" else
AWS_REGION="${aws_state[2]}"
fi
test -z "${aws_state[2]}" && AWS_REGION=$(aws configure get region) export AWS_REGION
export AWS_DEFAULT_REGION="$AWS_REGION"
export AWS_REGION=${AWS_REGION:-$aws_state[2]} fi
export AWS_DEFAULT_REGION="$AWS_REGION"
fi fi
# Load awscli completions # Load AWS CLI completions
if command -v aws_completer &>/dev/null; then
# AWS CLI v2 comes with its own autocompletion. Check if that is there, otherwise fall back
if command -v aws_completer &> /dev/null; then
autoload -Uz bashcompinit && bashcompinit autoload -Uz bashcompinit && bashcompinit
complete -C aws_completer aws complete -C aws_completer aws
else else
function _awscli-homebrew-installed() { function _awscli-homebrew-installed() {
# check if Homebrew is installed
(( $+commands[brew] )) || return 1 (( $+commands[brew] )) || return 1
# speculatively check default brew prefix if [[ -h /usr/local/opt/awscli ]]; then
if [ -h /usr/local/opt/awscli ]; then
_brew_prefix=/usr/local/opt/awscli _brew_prefix=/usr/local/opt/awscli
else else
# ok, it is not in the default prefix _brew_prefix="$(brew --prefix awscli)"
# this call to brew is expensive (about 400 ms), so at least let's make it only once
_brew_prefix=$(brew --prefix awscli)
fi fi
} }
# get aws_zsh_completer.sh location from $PATH local _aws_zsh_completer_path="$commands[aws_zsh_completer.sh]"
_aws_zsh_completer_path="$commands[aws_zsh_completer.sh]"
# otherwise check common locations if [[ -z "$_aws_zsh_completer_path" ]]; then
if [[ -z $_aws_zsh_completer_path ]]; then
# Homebrew
if _awscli-homebrew-installed; then if _awscli-homebrew-installed; then
_aws_zsh_completer_path=$_brew_prefix/libexec/bin/aws_zsh_completer.sh _aws_zsh_completer_path="$_brew_prefix/libexec/bin/aws_zsh_completer.sh"
# Ubuntu
elif [[ -e /usr/share/zsh/vendor-completions/_awscli ]]; then elif [[ -e /usr/share/zsh/vendor-completions/_awscli ]]; then
_aws_zsh_completer_path=/usr/share/zsh/vendor-completions/_awscli _aws_zsh_completer_path=/usr/share/zsh/vendor-completions/_awscli
# NixOS
elif [[ -e "${commands[aws]:P:h:h}/share/zsh/site-functions/aws_zsh_completer.sh" ]]; then elif [[ -e "${commands[aws]:P:h:h}/share/zsh/site-functions/aws_zsh_completer.sh" ]]; then
_aws_zsh_completer_path="${commands[aws]:P:h:h}/share/zsh/site-functions/aws_zsh_completer.sh" _aws_zsh_completer_path="${commands[aws]:P:h:h}/share/zsh/site-functions/aws_zsh_completer.sh"
# RPM
else else
_aws_zsh_completer_path=/usr/share/zsh/site-functions/aws_zsh_completer.sh _aws_zsh_completer_path=/usr/share/zsh/site-functions/aws_zsh_completer.sh
fi fi
fi fi
[[ -r $_aws_zsh_completer_path ]] && source $_aws_zsh_completer_path [[ -r "$_aws_zsh_completer_path" ]] && source "$_aws_zsh_completer_path"
unset _aws_zsh_completer_path _brew_prefix unset _aws_zsh_completer_path _brew_prefix
fi fi