essh

#!/usr/bin/env bash

verbose=

error() {
  printf '%s\n' "$*" >&2
  exit 1
}

usage_error() {
  [ "$#" -gt 0 ] && printf '%s\n\n' "$*" >&2
  usage
  exit 1
}

verbose() {
  [ -n "$verbose" ] || return
  printf '%s\n' "$*" >&2
}

usage() {
  cat >&2 <<EOF
Usage: $(basename "$0") [OPTIONS] [SEARCH] [--] [SSH_OPT...]

SSH to an AWS EC2 instance by Name tag.

  essh web

OPTIONS
  -c number   Chose the given number instead of showing choices
  -d          Use private DNS name
  -D          Use public DNS name
  -h          Show this help message
  -l          Force showing the list even for only 1 option
  -n          Print the host name instead of sshing to the host
  -p profile  An AWS CLI configuration profile. Multiple allowed.
  -r region   An AWS region. Multiple allowed.
  -R          Search all AWS regions. Overrides -r.
  -u user     An ssh login user
  -v          Show verbose messages

SEARCH
  Searches for matches in the EC2 instance Name tag. Substrings are allowed.
  If search is omitted, all hosts will be selected.

SSH_OPTS
  If arguments are given after the search, they are passed to SSH
EOF
}

while getopts :c:dDhlnp:r:Ru:v opt; do
  case "$opt" in
    c) choose="$OPTARG";;
    d) dns_type=private;;
    D) dns_type=public;;
    h) usage && exit;;
    l) list=1;;
    n) host_only=1;;
    p) profiles+=("$OPTARG");;
    r) regions+=("$OPTARG");;
    R) all_regions=1;;
    u) user="$OPTARG";;
    v) verbose=1;;
    \?) usage_error "Invalid option -${OPTARG}";;
    :) usage_error "Option -${OPTARG} requires an argument" ;;
  esac
done

# Preserve -- option if it is given first
if [ "${*:$((OPTIND - 1)):1}" == '--' ]; then
  OPTIND=$((OPTIND - 1))
fi

# Shift off options
shift $((OPTIND - 1))

if [ "$1" == '--' ]; then
  shift
elif [ "$#" -gt 0 ]; then
  target_name="$1"
  shift
fi

if ! hash aws 2>/dev/null; then
  error 'The aws CLI is required'
fi

query="Reservations[*].Instances[*].["
query+="PrivateDnsName,"
query+="PublicDnsName,"
query+="State.Name,"
query+="Tags[?Key=='Name'] | [0].Value"
query+="]"

filter="Name=instance-state-name,Values=pending,running,shutting-down,stopping"

add_candidates() {
  local profile="$1"
  local region="$2"

  aws_cmd=(aws ec2 describe-instances \
    --output text \
    --query "$query" \
    --filter "$filter" \
  )

  [ -n "$profile" ] && aws_cmd+=(--profile "$profile")
  [ -n "$region" ] && aws_cmd+=(--region "$region")

  while read -r line; do
    private="$(echo "$line" | cut -d$'\t' -f1)"
    public="$(echo "$line" | cut -d$'\t' -f2)"
    state="$(echo "$line" | cut -d$'\t' -f3)"
    name="$(echo "$line" | cut -d$'\t' -f4-)"

    [[ "$name" == *"$target_name"* ]] || continue

    if [ "$dns_type" = 'private' ]; then
      host="$private"
    elif [ "$dns_type" = 'public' ]; then
      host="$public"
    else
      [ -n "$public" ] && host="$public" || host="$private"
    fi

    [ -n "$host" ] || continue

    candidates+=("${host} ${state} ${name}")
  done < <("${aws_cmd[@]}")
}

[ "${#profiles[@]}" = 0 ] && profiles+=('')
[ "${#regions[@]}" = 0 ] && regions+=('')

if [ -n "$all_regions" ]; then
  regions=()
  while read -r region; do
    regions+=("$region")
  done < <( \
    aws ec2 describe-regions --output text --query 'Regions[*].[RegionName]' \
  )
fi

for profile in "${profiles[@]}"; do
  for region in "${regions[@]}"; do
    add_candidates "$profile" "$region"
  done
done

# Sort candidates by name and host
while IFS='' read -r candidate; do
  sorted_candidates+=("$candidate")
done < <(printf '%s\n' "${candidates[@]}" | sort -s -k3 -k1)
candidates=("${sorted_candidates[@]}")

if [ "${#candidates[@]}" -eq 0 ]; then
  error "No hosts found"
elif [ "${#candidates[@]}" -eq 1 ] && [ -z "$list" ]; then
  IFS=' ' read -ra candidate <<< "${candidates[0]}"
  host="${candidate[0]}"
else
  if [ -z "$choose" ]; then
    number=1
    for candidate in "${candidates[@]}"; do
      IFS=' ' read -ra parts <<< "${candidate}"
      host="${parts[0]}"
      state="${parts[1]}"
      name="${parts[*]:2}"
      [ "${#name}" -gt 40 ] && name="${name:0:37}..."
      [ "${#host}" -gt 45 ] && host="${host:0:27}..."
      [ "$state" = 'running' ] && state=
      printf '%-3s %-40s %-45s %-13s\n' \
	"${number})" "$name" "$host" "$state" >&2
      number=$((number + 1))
    done

    echo -n '> ' >&2
    read -r choose
  fi

  if [ -z "$choose" ]; then
    exit
  fi

  if ! [[ "$choose" =~ ^[0-9]+$ ]]; then
    error "Invalid choice ${choose}"
  fi

  if [ "$choose" -lt 1 ] || [ "$choose" -gt "${#candidates[@]}" ]; then
    error "Choice ${choose} out of range"
  fi

  IFS=' ' read -ra chosen <<< "${candidates[(($choose - 1))]}"
  host="${chosen[0]}"
fi

if [ -n "$host_only" ]; then
  echo "$host"
  exit
fi

[ -n "$user" ] && host="${user}@${host}"

verbose "Connecting to ${host}"

ssh_cmd=(ssh)
[ -n "$verbose" ] && ssh_cmd+=(-v)
ssh_cmd+=("$host" "$@")

verbose "${ssh_cmd[@]}"
exec "${ssh_cmd[@]}"