aws-ssm-ec2-proxycommand.sh

#!/usr/bin/env bash
# Example usage in ~/.ssh/config:
#
#	host aws:*
#		IdentityFile ~/.ssh/id_ed25519
#		User ec2-user
#		ProxyCommand ~/bin/aws-ssm-ec2-proxycommand.sh %h %r %p ~/.ssh/id_ed25519.pub
#
# Example SSH command:
#
#	ssh aws.my-profile.i-123456789

set -euo pipefail

IFS='.'; read -ra split_host <<< "$1"
aws_profile="${split_host[1]}"
instance_id="${split_host[2]}"
ssh_user="$2"
ssh_port="$3"
ssh_public_key_path="$4"
ssh_public_key="$(cat "${ssh_public_key_path}")"
ssh_public_key_timeout=60

if [[ $instance_id != i-* ]]; then
	instance_name="$instance_id"
	echo >&2 "looking up instance ID for name ${instance_name}"
	instance_id="$(
		aws-vault exec "$aws_profile" -- \
		aws ec2 describe-instances \
			--filters \
				"Name=tag:Name,Values=${instance_name}" \
				"Name=instance-state-name,Values=running" \
			--output text \
			--query 'Reservations[0].Instances[0].InstanceId'
	)"
	if [[ -z $instance_id || "$instance_id" == "None" ]]; then
		echo >&2 "failed to find instance ID for name ${instance_name}"
		exit 1
	fi
fi
# echo >&2 "aws_profile=${aws_profile}"
# echo >&2 "instance_id=${instance_id}"
# echo >&2 "ssh_user=${ssh_user}"
# echo >&2 "ssh_port=${ssh_port}"
# echo >&2 "ssh_public_key_path=${ssh_public_key_path}"
# set -x

echo >&2 "registering public key in /home/${ssh_user}/.ssh/authorized_keys on ${instance_id}"
command_id="$(
	aws-vault exec "$aws_profile" -- \
	aws ssm send-command \
	--instance-ids "${instance_id}" \
	--document-name 'AWS-RunShellScript' \
	--comment "Add an SSH public key to authorized_keys for ${ssh_public_key_timeout} seconds" \
	--parameters "commands=\"
cd /home/${ssh_user} || { echo >&2 '/home/${ssh_user} does not exist'; exit 1; }
mkdir -p .ssh && cd .ssh || exit 1
authorized_key=\\\"expiry-time=\\\\\\\"\$(date -d '+${ssh_public_key_timeout} seconds' -u +%Y%m%d%H%M%S)\\\\\\\" ${ssh_public_key} ssm-session\\\"
echo \\\"\${authorized_key}\\\" >> authorized_keys
sleep ${ssh_public_key_timeout}
grep -v -F \\\"\${authorized_key}\\\" authorized_keys > authorized_keys.tmp
mv authorized_keys.tmp authorized_keys
\"" \
	--query 'Command.CommandId' \
	--output text
)"

sleep 0.5
command_status="$(
	aws-vault exec "$aws_profile" -- \
	aws ssm get-command-invocation \
	--command "$command_id" \
	--instance-id "$instance_id" \
	--query 'Status' \
	--output text
)"
if [[ "$command_status" != "Success" && "$command_status" != "InProgress" ]]; then
	echo >&2 "command ${command_id} failed; stderr:"
	stderr="$(
		aws-vault exec "$aws_profile" -- \
		aws ssm get-command-invocation \
		--command "$command_id" \
		--instance-id "$instance_id" \
		--query 'StandardErrorContent' \
		--output text
	)"
	echo >&2 "$stderr"
	exit 1
fi

echo >&2 "starting ssm session on ${instance_id}"
aws-vault exec "$aws_profile" -- \
	aws ssm start-session \
	--target "${instance_id}" \
	--document-name 'AWS-StartSSHSession' \
	--parameters "portNumber=${ssh_port}"