#!/bin/bash

# Syntax:
# sm-local-start-ssh <instance_id> <extra_ssh_args>

INSTANCE_ID=$1
shift

INSTANCE_ID=$(echo "$INSTANCE_ID" | tr '\n' ' ' | grep -o 'mi-.*$' | awk -F' ' '{print $1}')

if [ -z "${INSTANCE_ID}" ]; then
  echo "INSTANCE_ID is not provided or incorrect (should be in the form mi-1234567890abcdef0)"
  exit 1
fi

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

read -r -d '' program << EOF
import logging
import sagemaker
logging.getLogger('sagemaker.config').setLevel(logging.WARNING)
print(sagemaker.Session().default_bucket())
EOF

if [ -z "${SSH_AUTHORIZED_KEYS_PATH}" ]; then
  bucket=$(python -c "$program")
  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