#!/bin/bash

# Creates a new hybrid activation in SSM and reports back the managed instance ID
# If successful, the log line with the instance ID will look like this:
#   Successfully registered the instance with AWS SSM using Managed instance-id: mi-01234567890abcdef

# Requires environment variable SSH_SSM_ROLE to be passed as an argument
# The role for SSM is not a full IAM ARN, but only the last part of it such as 'service-role/SageMakerRole'

set -e

dir=$(dirname "$0")
source "$dir"/sm-helper-functions

CURRENT_REGION=$(aws configure get region || echo "$AWS_REGION")

SSH_CREATOR=$(aws sts get-caller-identity | jq --raw-output '.UserId')
SSH_TIMESTAMP=$(date +%s)

if [ -f /opt/ml/metadata/resource-metadata.json ]; then
  # SageMaker Studio and notebook instances
  RESOURCE_NAME=$(jq --raw-output '.ResourceName' < /opt/ml/metadata/resource-metadata.json)
  RESOURCE_ARN=$(jq --raw-output '.ResourceArn' < /opt/ml/metadata/resource-metadata.json)
elif [ -f /opt/ml/config/processingjobconfig.json ]; then
  # Processing job
  RESOURCE_NAME=$(jq --raw-output '.ProcessingJobName' < /opt/ml/config/processingjobconfig.json)
  RESOURCE_ARN=$(jq --raw-output '.ProcessingJobArn' < /opt/ml/config/processingjobconfig.json)
elif [[ "$TRAINING_JOB_NAME" != "" ]]; then
  # Training job
  RESOURCE_NAME=$TRAINING_JOB_NAME
  RESOURCE_ARN=$TRAINING_JOB_ARN  # empty for local mode
elif [[ "$TRANSFORM_JOB_ARN" != "" ]]; then
  # Transform job
  RESOURCE_NAME=$(echo $TRANSFORM_JOB_ARN | awk -F/ '{print $2}')
  RESOURCE_ARN=$TRANSFORM_JOB_ARN
else
  # Probably, endpoint
  RESOURCE_NAME=""
  RESOURCE_ARN=""
fi

echo "sm-init-ssm: Detected SageMaker resource: $RESOURCE_NAME [$RESOURCE_ARN]"

SSH_SSM_TAGS="[{\"Key\": \"SSHOwner\", \"Value\": \"$SSH_OWNER_TAG\"}, {\"Key\": \"SSHCreator\", \"Value\": \"$SSH_CREATOR\"}, {\"Key\": \"SSHTimestamp\", \"Value\": \"$SSH_TIMESTAMP\"}, {\"Key\": \"SSHResourceName\", \"Value\": \"$RESOURCE_NAME\"}, {\"Key\": \"SSHResourceArn\", \"Value\": \"$RESOURCE_ARN\"}]"

response=$(aws ssm create-activation \
  --description "Activation for Amazon SageMaker integration with SSH and IDEs" \
  --iam-role "$SSH_SSM_ROLE" \
  --registration-limit 1 \
  --region "$CURRENT_REGION" \
  --tags "$SSH_SSM_TAGS")

acode=$(echo $response | jq --raw-output '.ActivationCode')
aid=$(echo $response | jq --raw-output '.ActivationId')

echo Yes | amazon-ssm-agent -register -id "$aid" -code "$acode" -region "$CURRENT_REGION"