[prompt] AWS profile, cluster name shortening (cloudposse#601)

rootfs/etc/profile.d/aws.sh

@@ -22,22 +22,39 @@ fi
 # Asks AWS what the currently active identity is and
 # sets environment variables accordingly
 function export_current_aws_role() {
-	local role_arn=$(aws sts get-caller-identity --output text --query 'Arn' | sed 's/:sts:/:iam:/g' | sed 's,:assumed-role/,:role/,' | cut -d/ -f1-2)
-	if [[ -z $role_arn ]]; then
+	local role_name
+	# Could be a primary or assumed role. If we have assumed a role, cut off the session name.
+	local current_role=$(aws sts get-caller-identity --output text --query 'Arn' | cut -d/ -f1-2 2>/dev/null)
+	if [[ -z $current_role ]]; then
 		unset ASSUME_ROLE
-	else
-		local role_name=$(crudini --get --format=lines "$AWS_CONFIG_FILE" | grep "$role_arn" | cut -d' ' -f 3)
-		if [[ -z $role_name ]]; then
-			if [[ "$role_arn" =~ "role/OrganizationAccountAccessRole" ]]; then
-				role_name="$(printf "%s" "$role_arn" | cut -d: -f 5):OrgAccess"
-				echo "* $(red Could not find profile name for ${role_arn}\; calling it \"${role_name}\")"
-			else
-				role_name="$(printf "%s" "$role_arn" | cut -d/ -f 2)"
-				echo "* $(green Could not find profile name for ${role_arn}\; calling it \"${role_name}\")"
-			fi
+		return 0
+	fi
+
+	# saml2aws will store the assumed role from sign-in as x_principal_arn in credentials file
+	# Default values from https://awscli.amazonaws.com/v2/documentation/api/latest/topic/config-vars.html
+	local creds_file="${AWS_SHARED_CREDENTIALS_FILE:-\~/.aws/credentials}"
+	if [[ -r $creds_file ]]; then
+		role_name=$(crudini --get --format=lines "${creds_file}" | grep "$current_role" | cut -d' ' -f 2)
+	fi
+
+	# Assumed roles are normally found in AWS config file, but using the role ARN,
+	# not the assumed role ARN. google2aws also puts login role in this file.
+	local config_file="${AWS_CONFIG_FILE:-\~/.aws/config}"
+	if [[ -z $role_name ]] && [[ -r $config_file ]]; then
+		local role_arn=$(printf "%s" "$current_role" | sed 's/:sts:/:iam:/g' | sed 's,:assumed-role/,:role/,')
+		role_name=$(crudini --get --format=lines "$config_file" | grep "$role_arn" | cut -d' ' -f 3)
+	fi
+
+	if [[ -z $role_name ]]; then
+		if [[ "$role_arn" =~ "role/OrganizationAccountAccessRole" ]]; then
+			role_name="$(printf "%s" "$role_arn" | cut -d: -f 5):OrgAccess"
+			echo "* $(red Could not find profile name for ${role_arn}\; calling it \"${role_name}\")" >&2
+		else
+			role_name="$(printf "%s" "$role_arn" | cut -d/ -f 2)"
+			echo "* $(green Could not find profile name for ${role_arn}\; calling it \"${role_name}\")" >&2
 		fi
-		export ASSUME_ROLE="$role_name"
 	fi
+	export ASSUME_ROLE="$role_name"
 }
 
 # Keep track of AWS credentials and updates to AWS role environment variables.

rootfs/etc/profile.d/geodesic.kube-ps1.sh

@@ -20,10 +20,22 @@ function kube_ps1_helper() {
 	fi
 }
 
-# This shortens the cluster name based on our EKS cluster naming pattern,
-# taking just the characters between the first and second dashes after "cluster/".
+# This shortens the cluster name of EKS clusters.
 # It should not affect other cluster names, so should be safe as default.
+# Users can override it if they want to.
 function short_cluster_name_from_eks() {
-	printf "%s" "$1" | sed -e 's%arn.*:cluster/[^-]\+-\([^-]\+\)-.*$%\1%'
+	# If it is not a cluster ARN, leave it alone
+	if ! [[ $1 =~ ^arn:.*:cluster/ ]]; then
+		printf "%s" "$1"
+		return 0
+	fi
+	local full_name=$(printf "%s" "$1" | cut -d/ -f2)
+	# remove namespace prefix if present
+	full_name=${full_name#${NAMESPACE}-}
+	# remove eks and everything after it, if present
+	full_name=${full_name%-eks-*}
+	printf "%s" "${full_name}"
+	# If NAMESPACE is unset, delete everything before and including the first dash
+	# printf "%s" "$1" | sed -e 's%arn.*:cluster/'"${NAMESPACE:-[^-]\+}"'-\([^-]\+\)-eks-.*$%\1%'
 }
-KUBE_PS1_CLUSTER_FUNCTION=short_cluster_name_from_eks
+[[ -z $KUBE_PS1_CLUSTER_FUNCTION ]] && KUBE_PS1_CLUSTER_FUNCTION=short_cluster_name_from_eks

rootfs/etc/profile.d/set-cluster.sh

@@ -14,7 +14,7 @@ function _update_cluster_config() {
 	local current_namespace
 	local set_namespace=1
 
-	current_namespace=$(KUBECONFIG="$new_config"kubens -c 2>/dev/null)
+	current_namespace=$(KUBECONFIG="$new_config" kubens -c 2>/dev/null)
 	set_namespace=$?
 	if ! KUBECONFIG="$new_config" kubectl auth can-i -Aq create selfsubjectaccessreviews.authorization.k8s.io >/dev/null 2>&1 </dev/null; then
 		eks-update-kubeconfig "$@"

rootfs/usr/local/bin/eks-update-kubeconfig

@@ -21,7 +21,7 @@ Usage:
 
   With "off", deletes the currently active kubecfg file.
 
-  NOTE: This tool assumes you are using Cloud Posses standard naming conventions:
+  NOTE: This tool assumes you are using Cloud Posse's standard naming conventions:
   * Cluster name "corp" expands to "${NAMESPACE}-corp-eks-cluster"
   * Role name "admin" expands to "${NAMESPACE}-corp-admin"