ssh-into

#!/usr/bin/env bash

display_usage() {
    echo -e 'ssh-into connects to an instance over SSM either directly or by using ssh over SSM.
Usage: ssh-into [user@]<destination> [options...]
Destination: Either a hostname (server.env) or an instance-id

Optional Arguments:
 -i, --identity_file      Identity file (~/.ssh/id_rsa)
 -p, --profile            AWS Profile (default)
 -h, --help               Displays this help text'
}

AWS_PROFILE='default'

params=( $@ )
INPUT="$1"

# parse args
i=0
for arg in "${params[@]}"; do
    if [[ "$arg" == "-h" || "$arg" == "--help" ]]; then
        display_usage
        exit 0
    fi
    if [[ "$arg" == "-i" || "$arg" == "--identity_file" ]]; then
        SSH_KEY="${params[$i+1]}"
    fi
    if [[ "$arg" == "-p" || "$arg" == "--profile" ]]; then
        AWS_PROFILE="${params[$i+1]}"
    fi
    ((i++))
done

# check required args were provided
if [[ -z "$INPUT" ]]; then
    echo 'ERROR: No server specified, please specify a server to connect to!' 1>&2
    display_usage
    exit 1
fi

# extract username, if provided
if [[ $INPUT =~ "@" ]]; then
    USERNAME="${INPUT%%@*}"
fi
REMAINING_INPUT="${INPUT##*@}"

instance_id_pattern='^m?i-[[:xdigit:]]+$'
if [[ $REMAINING_INPUT =~ $instance_id_pattern ]]; then
    # we already have the instance id
    INSTANCE_ID="$REMAINING_INPUT"
else
    # we need to look up the instance id
    SERVER="${REMAINING_INPUT%%.*}"
    DOMAIN="${REMAINING_INPUT#*.}"

    HOSTED_ZONE_ID="$(aws route53 list-hosted-zones --query 'HostedZones[?Name ==`'$DOMAIN'.`].Id' --profile $AWS_PROFILE --output json 2> /dev/null)"
    if (( $? != 0 )); then
        echo -e "You don't appear to have the correct permissions. \nPlease ensure the AWS profile that you're using has the correct Route 53 permissions."
        exit 1
    fi

    HOSTED_ZONE_ID="$(python -c 'import sys, json; print json.load(sys.stdin)[0].split("/")[2]' <<< "$HOSTED_ZONE_ID" 2> /dev/null)"
    if [[ "$HOSTED_ZONE_ID" == "null" || -z "$HOSTED_ZONE_ID" ]]; then
        echo "$DOMAIN doesn't appear to be Route 53 Private Hosted Zone in your account. Please double check the domain."
        exit 1
    fi

    PRIVATE_IP="$(aws route53 list-resource-record-sets --hosted-zone-id $HOSTED_ZONE_ID --query 'ResourceRecordSets[?contains(Name,`'$SERVER'`)].ResourceRecords[].Value' --profile $AWS_PROFILE --output json 2> /dev/null | python -c 'import sys, json; print json.load(sys.stdin)[0]' 2> /dev/null)"
    INSTANCE_ID="$(aws ec2 describe-instances --profile $AWS_PROFILE --filters 'Name=private-ip-address,Values='$PRIVATE_IP'' --query 'Reservations[].Instances[0].InstanceId' --output text 2> /dev/null)"

    if [[ -z "$INSTANCE_ID" ]]; then
        echo -e "The domain name provided doesn't exist. Please double check the name and try again.\nOR Ensure the the route 53 entry contains a PRIVATE ip address."
        exit 1
    fi
    if [[ ! "$INSTANCE_ID" =~ $instance_id_pattern ]]; then
        echo "Retrieved instance id $INSTANCE_ID does not seem valid!"
        exit 1
    fi

    echo "[INFO] Found Private IP for $INSTANCE_ID: $PRIVATE_IP"
fi

if [[ -z "$USERNAME" ]]; then
    # connect with ssm directly
    echo "[INFO] Username not provided, using aws ssm start-session command."
    # echo "aws ssm start-session --target $INSTANCE_ID --profile $AWS_PROFILE"
    aws ssm start-session --target $INSTANCE_ID --profile $AWS_PROFILE
else
    # connect with ssh over ssm
    unset SSH_COMMAND
    if [[ -n "$SSH_KEY" ]]; then
        SSH_COMMAND=" -i $SSH_KEY"
    fi
    if [[ -n "$SERVER" && -n "$DOMAIN" ]]; then
        SSH_COMMAND="${SSH_COMMAND} -o "HostKeyAlias=${SERVER}.${DOMAIN}""
    fi
    SSH_COMMAND="ssh${SSH_COMMAND} $USERNAME@$INSTANCE_ID"
    echo "$SSH_COMMAND"
    $SSH_COMMAND
fi