From 6a5e4c3801b181a65d1ed71030cf101571efe31c Mon Sep 17 00:00:00 2001 From: mnv Date: Wed, 23 Oct 2024 14:33:19 +0530 Subject: [PATCH] fix(aws): improve shell functions safety and reliability 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 --- plugins/aws/aws.plugin.zsh | 285 +++++++++++++++++++------------------ 1 file changed, 149 insertions(+), 136 deletions(-) diff --git a/plugins/aws/aws.plugin.zsh b/plugins/aws/aws.plugin.zsh index 0c43031df..4ee2b91e5 100644 --- a/plugins/aws/aws.plugin.zsh +++ b/plugins/aws/aws.plugin.zsh @@ -1,23 +1,29 @@ function agp() { - echo $AWS_PROFILE + echo "${AWS_PROFILE:-}" } function agr() { - echo $AWS_REGION + echo "${AWS_REGION:-}" } # Update state file if enabled function _aws_update_state() { if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then - test -d $(dirname ${AWS_STATE_FILE}) || exit 1 - echo "${AWS_PROFILE} ${AWS_REGION}" > "${AWS_STATE_FILE}" + local state_dir="$(dirname "${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 } function _aws_clear_state() { if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then - test -d $(dirname ${AWS_STATE_FILE}) || exit 1 - echo -n > "${AWS_STATE_FILE}" + local state_dir="$(dirname "${AWS_STATE_FILE}")" + if [[ ! -d "$state_dir" ]]; then + mkdir -p "$state_dir" 2>/dev/null || return 1 + fi + : > "${AWS_STATE_FILE}" fi } @@ -26,35 +32,39 @@ function asp() { if [[ -z "$1" ]]; then unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE AWS_PROFILE_REGION _aws_clear_state - echo AWS profile cleared. - return + echo "AWS profile cleared." + return 0 fi local -a available_profiles available_profiles=($(aws_profiles)) if [[ -z "${available_profiles[(r)$1]}" ]]; then - echo "${fg[red]}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 "Profile '$1' not found in '${AWS_CONFIG_FILE:-$HOME/.aws/config}'" >&2 + echo "Available profiles: ${(j:, :)available_profiles:-no profiles found}" >&2 return 1 fi - export AWS_DEFAULT_PROFILE=$1 - export AWS_PROFILE=$1 - export AWS_EB_PROFILE=$1 + export AWS_DEFAULT_PROFILE="$1" + export AWS_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 - if [[ "$2" == "login" ]]; then - if [[ -n "$3" ]]; then - aws sso login --sso-session $3 - else - aws sso login - fi - elif [[ "$2" == "logout" ]]; then - aws sso logout - fi + case "$2" in + "login") + if [[ -n "$3" ]]; then + aws sso login --sso-session "$3" + else + aws sso login + fi + ;; + "logout") + aws sso logout + ;; + esac } # AWS region selection @@ -62,51 +72,57 @@ function asr() { if [[ -z "$1" ]]; then unset AWS_DEFAULT_REGION AWS_REGION _aws_update_state - echo AWS region cleared. - return + echo "AWS region cleared." + return 0 fi local -a available_regions available_regions=($(aws_regions)) 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 fi - export AWS_REGION=$1 - export AWS_DEFAULT_REGION=$1 + export AWS_DEFAULT_REGION="$1" + export AWS_REGION="$1" _aws_update_state } # AWS profile switch function acp() { if [[ -z "$1" ]]; then - unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE - unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN - echo AWS profile cleared. - return + unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN + echo "AWS profile cleared." + return 0 fi local -a available_profiles available_profiles=($(aws_profiles)) if [[ -z "${available_profiles[(r)$1]}" ]]; then - echo "${fg[red]}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 "Profile '$1' not found in '${AWS_CONFIG_FILE:-$HOME/.aws/config}'" >&2 + echo "Available profiles: ${(j:, :)available_profiles:-no profiles found}" >&2 return 1 fi local profile="$1" local mfa_token="$2" - # Get fallback credentials for if the aws command fails or no command is run - local aws_access_key_id="$(aws configure get aws_access_key_id --profile $profile)" - local aws_secret_access_key="$(aws configure get aws_secret_access_key --profile $profile)" - local aws_session_token="$(aws configure get aws_session_token --profile $profile)" + # Get fallback credentials + local aws_access_key_id + local aws_secret_access_key + 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")" - - # First, if the profile has MFA configured, lets get the token and session duration - local mfa_serial="$(aws configure get mfa_serial --profile $profile)" - local sess_duration="$(aws configure get duration_seconds --profile $profile)" + # Check for MFA configuration + local mfa_serial + local sess_duration + + mfa_serial="$(aws configure get mfa_serial --profile "$profile")" + sess_duration="$(aws configure get duration_seconds --profile "$profile")" if [[ -n "$mfa_serial" ]]; then local -a mfa_opt @@ -115,55 +131,52 @@ function acp() { read -r mfa_token fi 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 fi mfa_opt=(--serial-number "$mfa_serial" --token-code "$mfa_token" --duration-seconds "${sess_duration:-3600}") fi - # Now see whether we need to just MFA for the current role, or assume a different one - local role_arn="$(aws configure get role_arn --profile $profile)" - local sess_name="$(aws configure get role_session_name --profile $profile)" + # Check for role assumption + local role_arn + 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 - # Means we need to assume a specified role + local -a aws_command 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="$(aws configure get external_id --profile $profile)" + local external_id + external_id="$(aws configure get external_id --profile "$profile")" if [[ -n "$external_id" ]]; then aws_command+=(--external-id "$external_id") fi - # Get source profile to use to assume role - local source_profile="$(aws configure get source_profile --profile $profile)" + local source_profile + source_profile="$(aws configure get source_profile --profile "$profile")" if [[ -z "$sess_name" ]]; then - sess_name="${source_profile:-profile}" + sess_name="${source_profile:-$profile}" 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 - # Means we only need to do MFA + local -a aws_command aws_command=(aws sts get-session-token --profile="$profile" "${mfa_opt[@]}") echo "Obtaining session token for profile $profile" fi - # Format output of aws command for easier processing aws_command+=(--query '[Credentials.AccessKeyId,Credentials.SecretAccessKey,Credentials.SessionToken]' --output text) - # Run the aws command to obtain credentials - local -a credentials - credentials=(${(ps:\t:)"$(${aws_command[@]})"}) - + local credentials + credentials="$(${aws_command[@]})" || return 1 if [[ -n "$credentials" ]]; then - aws_access_key_id="${credentials[1]}" - aws_secret_access_key="${credentials[2]}" - aws_session_token="${credentials[3]}" + read -r aws_access_key_id aws_secret_access_key aws_session_token <<< "$credentials" 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_PROFILE="$profile" export AWS_EB_PROFILE="$profile" @@ -177,73 +190,79 @@ function acp() { fi echo "Switched to AWS Profile: $profile" + else + echo "Failed to obtain valid credentials" >&2 + return 1 fi } function aws_change_access_key() { if [[ -z "$1" ]]; then - echo "usage: $0 " + echo "usage: ${0} " >&2 return 1 fi local profile="$1" - # Get current access key - local original_aws_access_key_id="$(aws configure get aws_access_key_id --profile $profile)" + local original_aws_access_key_id + original_aws_access_key_id="$(aws configure get aws_access_key_id --profile "$profile")" - asp "$profile" || return 1 - 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\`" + if ! asp "$profile"; then return 1 fi - read -q "yn?Would you like to disable your previous access key (${original_aws_access_key_id}) now? " - case $yn in + echo "Generating a new access key pair..." + 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]*) - echo -n "\nDisabling 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 - echo "done." + 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 + echo "Access key disabled successfully." else - echo "\nFailed to disable ${original_aws_access_key_id} key." + echo "Failed to disable ${original_aws_access_key_id} key." >&2 fi ;; - *) - echo "" - ;; 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:" aws --no-cli-pager iam list-access-keys } function aws_regions() { - local region - if [[ $AWS_DEFAULT_REGION ]];then - region="$AWS_DEFAULT_REGION" - elif [[ $AWS_REGION ]];then - region="$AWS_REGION" - else - region="us-west-1" - fi + local region="${AWS_DEFAULT_REGION:-${AWS_REGION:-us-west-1}}" - if [[ $AWS_DEFAULT_PROFILE || $AWS_PROFILE ]];then - aws ec2 describe-regions --region $region |grep RegionName | awk -F ':' '{gsub(/"/, "", $2);gsub(/,/, "", $2);gsub(/ /, "", $2); print $2}' + if [[ -n "$AWS_DEFAULT_PROFILE" || -n "$AWS_PROFILE" ]]; then + aws ec2 describe-regions --region "$region" --query 'Regions[].RegionName' --output text | tr '\t' '\n' else - echo "You must specify a AWS profile." + echo "You must specify an AWS profile." >&2 + return 1 fi } function aws_profiles() { - aws --no-cli-pager configure list-profiles 2> /dev/null && return - [[ -r "${AWS_CONFIG_FILE:-$HOME/.aws/config}" ]] || return 1 - grep --color=never -Eo '\[.*\]' "${AWS_CONFIG_FILE:-$HOME/.aws/config}" | sed -E 's/^[[:space:]]*\[(profile)?[[:space:]]*([^[:space:]]+)\][[:space:]]*$/\2/g' + if aws --no-cli-pager configure list-profiles 2>/dev/null; then + return 0 + 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() { reply=($(aws_regions)) } @@ -256,83 +275,77 @@ compctl -K _aws_profiles asp acp aws_change_access_key # AWS prompt function aws_prompt_info() { - local _aws_to_show + local _aws_to_show="" local region="${AWS_REGION:-${AWS_DEFAULT_REGION:-$AWS_PROFILE_REGION}}" - if [[ -n "$AWS_PROFILE" ]];then - _aws_to_show+="${ZSH_THEME_AWS_PROFILE_PREFIX=""}" + if [[ -n "$AWS_PROFILE" ]]; then + _aws_to_show="${ZSH_THEME_AWS_PROFILE_PREFIX:-}" fi if [[ -n "$region" ]]; then - [[ -n "$_aws_to_show" ]] && _aws_to_show+="${ZSH_THEME_AWS_DIVIDER=" "}" - _aws_to_show+="${ZSH_THEME_AWS_REGION_PREFIX=""}" + [[ -n "$_aws_to_show" ]] && _aws_to_show+="${ZSH_THEME_AWS_DIVIDER:- }" + _aws_to_show+="${ZSH_THEME_AWS_REGION_PREFIX:-}" fi echo "$_aws_to_show" } +# Add AWS prompt to RPROMPT if enabled if [[ "$SHOW_AWS_PROMPT" != false && "$RPROMPT" != *'$(aws_prompt_info)'* ]]; then RPROMPT='$(aws_prompt_info)'"$RPROMPT" fi +# Initialize state from file if enabled if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then 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]}" - export AWS_PROFILE="$AWS_DEFAULT_PROFILE" - export AWS_EB_PROFILE="$AWS_DEFAULT_PROFILE" + if [[ -z "${aws_state[2]}" ]]; then + AWS_REGION="$(aws configure get region)" + else + AWS_REGION="${aws_state[2]}" + fi - test -z "${aws_state[2]}" && AWS_REGION=$(aws configure get region) - - export AWS_REGION=${AWS_REGION:-$aws_state[2]} - export AWS_DEFAULT_REGION="$AWS_REGION" + export AWS_REGION + export AWS_DEFAULT_REGION="$AWS_REGION" + fi fi -# Load awscli completions - -# AWS CLI v2 comes with its own autocompletion. Check if that is there, otherwise fall back -if command -v aws_completer &> /dev/null; then +# Load AWS CLI completions +if command -v aws_completer &>/dev/null; then autoload -Uz bashcompinit && bashcompinit complete -C aws_completer aws else function _awscli-homebrew-installed() { - # check if Homebrew is installed (( $+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 else - # ok, it is not in the default prefix - # this call to brew is expensive (about 400 ms), so at least let's make it only once - _brew_prefix=$(brew --prefix awscli) + _brew_prefix="$(brew --prefix awscli)" fi } - # get aws_zsh_completer.sh location from $PATH - _aws_zsh_completer_path="$commands[aws_zsh_completer.sh]" + local _aws_zsh_completer_path="$commands[aws_zsh_completer.sh]" - # otherwise check common locations - if [[ -z $_aws_zsh_completer_path ]]; then - # Homebrew + if [[ -z "$_aws_zsh_completer_path" ]]; then if _awscli-homebrew-installed; then - _aws_zsh_completer_path=$_brew_prefix/libexec/bin/aws_zsh_completer.sh - # Ubuntu + _aws_zsh_completer_path="$_brew_prefix/libexec/bin/aws_zsh_completer.sh" elif [[ -e /usr/share/zsh/vendor-completions/_awscli ]]; then _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 _aws_zsh_completer_path="${commands[aws]:P:h:h}/share/zsh/site-functions/aws_zsh_completer.sh" - # RPM else _aws_zsh_completer_path=/usr/share/zsh/site-functions/aws_zsh_completer.sh 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 fi -