#!/usr/bin/env sh
# Description
# Bootstrap SSH Session to an SSM-managed instance
# by temporarily adding a public SSH key available on the local machine (ssh-agent or in ~/.ssh)
#
# 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-* ssm-jumphost
#     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:-ec2-user}"
ssh_port="${3-22}"

ssh_public_key_paths="${HOME}/.ssh/id_rsa.pub ${HOME}/.ssh/id_ed25519.pub"

# Cleanup any stale running session managers - these may cause us to hang later when trying to connect
pkill session-manager-plugin

# A magic hostname when we don't know/care about the intermediate jump host
# just want to retrieve him via ec2 instance tag
jumphost_hostname="ssm-jumphost"

# An EC2 instance Tag that should have vale = True to denote this as an SSM Jumpbox
jumphost_tagname_true="SSM-Jumphost"

if [[ "$ec2_instance_id" == $jumphost_hostname ]] ; then
  ec2_instance_id=`aws --output text ec2 describe-instances --filters Name=tag:$jumphost_tagname_true,Values=True --query "Reservations[0].Instances[0].InstanceId"`
fi

# 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_paths}'
  for ssh_public_key_path in $ssh_public_key_paths; do
    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}"
    fi
  done
fi

# Try getting ANY ssh key in ~/.ssh
# If found - start ssh-agent and add it
if [[ -z "${ssh_public_key}" ]]; then
  for ssh_public_key_path in $(ls ~/.ssh/*.pub 2>/dev/null); do
    # Check we have (what looks like) the corresponding private key
    ssh_private_key_path="$(dirname $ssh_public_key_path)/$(basename $ssh_public_key_path .pub)"
    if [[ -e "$ssh_private_key_path" ]] ; then
      ssh_public_key="$(cat ${ssh_public_key_path})"
      ssh-agent
      echo "Adding ${ssh_private_key_path} to ssh-agent ..."
      ssh-add ${ssh_private_key_path}
      break
    fi
  done
fi

if [[ -z "$ssh_public_key" ]]; then
  echo "No ssh key present in ssh agent nor ~/.ssh/"
  exit 1
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}"