#!/usr/bin/env sh
set -euo pipefail

######## Usage #################################################################
#
# #1 Install the AWS CLI
#   https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-install.html
#
# #2 Install the Session Manager Plugin for the AWS CLI
#   https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html
#
# #3 Install ProxyCommand
#   - Move this script to ~/.ssh/aws-ssm-ec2-proxy-command.sh
#   - Make it executable (chmod +x ~/.ssh/aws-ssm-ec2-proxy-command.sh)
#
# #4 Setup SSH Config
#   - Add foolowing entry to your ~/.ssh/config
#
#   host i-* mi-*
#     ProxyCommand ~/.ssh/aws-ssm-ec2-proxy-command.sh %h %r %p
#
# #5 Ensure SSM Permissions of Target Instance Profile
#
#   https://docs.aws.amazon.com/systems-manager/latest/userguide/setup-instance-profile.html
#
# #6 Ensure latest SSM Agent on Target Instance
#
#   Is preinstalled on all amazon linux AMIs, however may needs to be updated
#   yum install -y https://s3.amazonaws.com/ec2-downloads-windows/SSMAgent/latest/linux_amd64/amazon-ssm-agent.rpm & service amazon-ssm-agent restart
#
# #7 Finally connect to ec2 instance
#
#   AWS_PROFILE='default' ssh ec2-user@i-xxxxxxxxxxxxxxxx
#
################################################################################

ec2_instance_id="$1"
ssh_user="${2}"
ssh_port="${3}"
ssh_public_key_path="${HOME}/.ssh/id_rsa.pub"

# Temporary add your public SSH key to authorized_keys on target instance
ssh_public_key_timeout=10

# Try to get an public ssh key from 'ssh agent'
ssh_public_key="$(keys="$(ssh-add -L)" && echo $keys | head -1)"
if [[ -n "$ssh_public_key" ]]; then
  ssh_public_key_source='ssh agent'
else
  # Try read public ssh key from '${ssh_public_key_path}'
  ssh_public_key="$([[ -e "${ssh_public_key_path}" ]] && cat "${ssh_public_key_path}")"
  if [[ -n "$ssh_public_key" ]]; then
    ssh_public_key_source="${ssh_public_key_path}"
  else
    echo "No ssh key present in ssh agent nor at ${ssh_public_key_path}"
    exit 1
  fi
fi

echo "Temporary add your public ssh key from '$ssh_public_key_source' to authorized_keys on target instance ${ec2_instance_id}"
aws ssm send-command \
  --instance-ids "${ec2_instance_id}" \
  --document-name 'AWS-RunShellScript' \
  --parameters commands="\"
    cd ~${ssh_user}/.ssh || exit 1
    grep -F '${ssh_public_key}' authorized_keys || echo '${ssh_public_key} ssm-session' >> authorized_keys
    sleep ${ssh_public_key_timeout}
    grep -v -F '${ssh_public_key}' authorized_keys > .tmp.authorized_keys
    mv .tmp.authorized_keys authorized_keys
  \"" \
  --comment "grant ssh access for ${ssh_public_key_timeout} seconds"

# Start SSM SSH session
aws ssm start-session \
  --target "${ec2_instance_id}" \
  --document-name 'AWS-StartSSHSession' \
  --parameters "portNumber=${ssh_port}"