jump

#!/usr/bin/env bash

set -o errexit   # Exit on most errors (see the manual)
set -o errtrace  # Make sure any error trap is inherited
set -o nounset   # Disallow expansion of unset variables
set -o pipefail  # Use last non-zero exit code in a pipeline

readonly __blue="\033[36m"
readonly __red="\033[31m"
readonly __reset="\033[0m"
readonly __yellow="\033[33m"

highlight() { echo -e "${__yellow}$*${__reset}"; }
fail() { echo -e "${__red}✖ $*${__reset}" >&2; exit 1; }
info() { echo -e "${__blue}‣ $*${__reset}"; }

INSTANCE_ID=""
INSTANCE_NAME=""
PORT=22
PROFILE="${AWS_PROFILE:-}"
REGION="${AWS_DEFAULT_REGION:-eu-west-2}"
SESSION=""
SSH_KEY=""
USER="ec2-user"

processArguments() {
  while [[ $# -gt 0 ]]
  do
    case $1 in
      -i | --instance-id)  INSTANCE_ID="$2"; shift; ;;
      -k | --key)          SSH_KEY=$2; shift ;;
      -n | --name)         INSTANCE_NAME="$2"; shift; ;;
      -p | --profile)      PROFILE="$2"; shift ;;
      -P | --port)         PORT="$2"; shift ;;
      -r | --region)       REGION="$2"; shift ;;
      -s | --session)      SESSION="true" ;;
      -u | --user)         USER=$2; shift ;;
      --* | -?)        fail "Unknown option \"$1\"" ;;
    esac
    shift
  done

  validateArguments
}

validateArguments() {
  [[ -n "$PROFILE" ]] || fail "AWS profile has not been specified"
  [[ -n "$REGION" ]] || fail "AWS region has not been specified"
  [[ -n "$INSTANCE_NAME" ]] || [[ -n "$INSTANCE_ID" ]] || fail "Instance name or ID has not been specified"
}

ssoLogin() {
  info "Logging into AWS account"
  echo "  - Profile: $(highlight "$PROFILE")"
  echo "  - Region: $(highlight "$REGION")"

  if ! aws sts --profile "$PROFILE" get-caller-identity > /dev/null 2>&1; then
    aws sso --profile "$PROFILE" login
  fi
}

getInstanceId() {
  INSTANCE_ID=$(
    aws ec2 describe-instances \
      --profile "$PROFILE" \
      --region "$REGION" \
      --query "Reservations[*].Instances[*].[InstanceId]" \
      --filters "Name=tag:Name,Values=$INSTANCE_NAME" "Name=instance-state-name,Values=running" \
      --output text | head -n 1
  )

  [[ -n "$INSTANCE_ID" ]] || fail "Unknown instance name '$INSTANCE_NAME'"

  echo "  - Instance ID: $(highlight "$INSTANCE_ID")"
}

createTemporarySshKey() {
  info "Creating temporary SSH key"
  SSH_KEY="/tmp/jump_ssh_key"

  if [[ ! -e "$SSH_KEY" ]]; then
    ssh-keygen -t rsa -b 4096 -f "$SSH_KEY" -q -N ""
  fi
}

sendSshPublicKey() {
  info "Uploading SSH key"
  echo "  - SSH key: $(highlight "$SSH_KEY")"
  echo "  - User: $(highlight "$USER")"

  aws ec2-instance-connect send-ssh-public-key \
    --profile "$PROFILE" \
    --region "$REGION" \
    --instance-id "$INSTANCE_ID" \
    --instance-os-user "$USER" \
    --ssh-public-key "file://${SSH_KEY}.pub" > /dev/null
}

startSsmSessionCommand() {
  cat << EOF  
    aws ssm start-session \
      --profile "$PROFILE" \
      --region "$REGION" \
      --target "$INSTANCE_ID" \
      --document AWS-StartSSHSession \
      --parameters "portNumber=$PORT"
EOF
}

startSsmSession() {
  info "Starting SSM session\n"
  eval "$(startSsmSessionCommand)"
}

main() {
  processArguments "$@"
  ssoLogin
  [[ -n "$INSTANCE_ID" ]] || getInstanceId
  [[ -e "$SSH_KEY" ]] || createTemporarySshKey
  sendSshPublicKey
  if [[ "$SESSION" == "true" ]]; then
    startSsmSession
  else
    ssh -i "$SSH_KEY" -p "$PORT" -o ProxyCommand="$(startSsmSessionCommand)" "$USER@$INSTANCE_ID"
  fi 
}

main "$@"