#!/usr/bin/env bash
set -euo pipefail
usage() {
cat << EOF
"""
EXAMPLE:
# ssh to bastion
$ IAM=zhengjiajin hack/bastion.sh ssh

# port-forward bastion, local_port is 9999, remote_port is 80
$ IAM=zhengjiajin hack/bastion.sh port-forward 9999 80

Available Bastion machine command:
  ssh               ssh bastion machine
  port-forward      port forward local port to bastion machine port (the default local_port is 9999, remote_port is 80)

Available Environment:
  ENV      dev or prod
  IAM      special iam user name
"""
EOF
}

ENV=${ENV:-""}
IAM=${IAM:-""}
REGION=${AWS_DEFAULT_REGION:-"us-west-2"}
[ -z "$IAM" ] && echo "must specify IAM environment variable, eg: IAM=zhengjiajin" && exit 1
[ -z "$ENV" ] && echo "must specify ENV environment variable, eg: ENV=dev" && exit 1
CMD=${1:-"ssh"}

ROOT=$(cd $(dirname ${BASH_SOURCE[0]})/.. && pwd)

if [ "${CMD}" != "ssh" ] && [ "${CMD}" != "port-forward" ];then
    usage
    exit
fi

printenv ENV
printenv IAM

INSTANCE_ID=$(aws-vault exec ${ENV} --region ${REGION} -- aws ec2 describe-instances \
                                                          --filters "Name=tag:Name,Values=bastion-${IAM}" \
                                                          --query "Reservations[].Instances[?State.Name == 'running'].InstanceId[]" \
                                                          --output text)

if [[ -z $INSTANCE_ID ]] ; then
  echo "Bastion instance not found" >&2
  exit 1
fi

case "${CMD}" in
  ssh)
    aws-vault exec ${ENV} --region ${REGION} -- aws ssm start-session --target $INSTANCE_ID
    ;;
  port-forward)
    local_port=${2:-"9999"}
    remote_port=${3:-"80"}
    aws-vault exec ${ENV} --region ${REGION} -- aws ssm start-session \
                                                --target $INSTANCE_ID \
                                                --document-name AWS-StartPortForwardingSession \
                                                --parameters '{"portNumber":["'${remote_port}'"],"localPortNumber":["'${local_port}'"]}' \
    ;;
esac