#!/bin/bash

INSTANCE_ID=$1
shift

if [ -z "${INSTANCE_ID}" ]; then
  echo "INSTANCE_ID not provided"
  exit 1
else
  echo "INSTANCE_ID -> ${INSTANCE_ID}"
fi

# Useful for port forwarding and debugging, eg. passing -vvv option
EXTRA_SSH_ARGS=$*

if [ -z "${SSH_AUTHORIZED_KEYS_PATH}" ]; then
  bucket=$(python -c 'import sagemaker; print(sagemaker.Session().default_bucket())')
  SSH_AUTHORIZED_KEYS_PATH="s3://$bucket/ssh-authorized-keys/"
fi

echo "SSH authorized keys path -> ${SSH_AUTHORIZED_KEYS_PATH}"

# shellcheck disable=SC2086  # extra args have to be unquoted to be parsed from the inner script
sm-connect-ssh-proxy "${INSTANCE_ID}" \
    "${SSH_AUTHORIZED_KEYS_PATH}" \
    $EXTRA_SSH_ARGS