#!/bin/sh

function log { logger -t "user-data.${_FUNC}" -- $@; }

function die { log "$@"; exit_trap 1 1 / "$@"; }

# msg used for sns event, last one wins
function msg { MSG="$@"; log "$@"; }

# Generic retry command wrapper, incl. timeout of 30s
# $1 = number of tries; 0 = forever
# $2 = number of seconds to sleep between tries
# $@ actual command
retry() {
  local tries=$1
  local waitfor=$2
  shift 2
  while true; do
    # Only use timeout of $1 is an executable, call directly if function
    type -tf $1 >/dev/null && { timeout 30 $@ && return; } || { $@ && return; }
    ((tries=tries-1))
    [ $tries -eq 0 ] && return 1
    sleep $waitfor
  done
}

function add_swap() {
  [ -f /.swapfile ] || { dd if=/dev/zero of=/.swapfile bs=1M count=$1 && chmod 600 /.swapfile && mkswap /.swapfile && swapon /.swapfile; }
  grep -q "/.swapfile" /etc/fstab || echo "/.swapfile none swap sw 0 0" >> /etc/fstab
  sysctl -w vm.swappiness=10
}

# Get SSM secure string base64 decoded
# $0 SSM_PATH, value to stdout
function get_secret() {
  aws ssm get-parameter --name ${1,,} --with-decryption --query 'Parameter.Value' | base64 -d
}

# Store values as base64 on SSM
# $0 SSM_PATH VALUE
function put_secret() {
  aws ssm put-parameter --name ${1,,} --type SecureString --value "$(echo "$2" | base64 -w0)" --overwrite
}

# Gets existing passphrase or creates new passphrase and stores it
function init_passphrase() {
  local _URL=$1
  local _PPFILE=$2

  # If secret already exists noop
  [ -f $_PPFILE ] && return 0

  get_secret $_URL > $_PPFILE && chmod 600 $_PPFILE || \
    { xxd -l16 -p /dev/random > $_PPFILE; chmod 600 $_PPFILE; put_secret $_URL "$(cat $_PPFILE)"; }
}

function asg_heartbeat {
  [ -n "$LAUNCH_HOOK" ] && aws autoscaling record-lifecycle-action-heartbeat --instance-id $INSTANCE_ID --lifecycle-hook-name $LAUNCH_HOOK --auto-scaling-group-name $AWS_AUTOSCALING_GROUPNAME || true
}

function setup_sns_alarms() {
  # store SNS message json template
  cat <<EOF > /etc/cloudbender/sns_alarm.json
{
  "Source": "CloudBender",
  "AWSAccountId": "$AWS_ACCOUNT_ID",
  "Region": "$REGION",
  "Artifact": "$ARTIFACT",
  "Asg": "$AWS_AUTOSCALING_GROUPNAME",
  "Instance": "$INSTANCE_ID",
  "ip": "$IP_ADDRESS"
}
EOF
  mkdir -p /var/lib/cloudbender
  cat <<EOF > /var/lib/cloudbender/sns_alarm.sh
#!/bin/bash

SUBJECT=\$1
MSG=\$2
LEVEL=\${3:-Info}
ATTACHMENT=\${4:-""}
EMOJI=\${5:-""}

jq -M --arg subject "\$SUBJECT" --arg level "\$LEVEL" --arg msg "\$MSG" --arg attachment "\$ATTACHMENT" --arg emoji "\$EMOJI" --arg hostname "\$HOSTNAME" '.Subject = \$subject | .Level = \$level | .Message = \$msg | .Attachment = \$attachment | .Emoji = \$emoji | .Hostname = \$hostname' < /etc/cloudbender/sns_alarm.json | sed -e 's/\\\\\\\\/\\\\/g' > /tmp/sns.json
aws sns publish --region ${REGION} --target-arn $ALARMSNSARN --message file:///tmp/sns.json
EOF
  chmod +x /var/lib/cloudbender/sns_alarm.sh
}

function exit_trap {
  set +e
  trap - ERR EXIT
  local ERR_CODE=$1
  local ERR_LINE="$2"
  local ERR_FUNC="$3"
  local ERR_CMD="$4"

  if [ $ERR_CODE -ne 0 ]; then
    CFN_STATUS="FAILURE"
    RESULT="ABANDON"
  else
    CFN_STATUS="SUCCESS"
    RESULT="CONTINUE"
  fi

  # Add SNS events on demand
  if [ "x${ALARMSNSARN}" != 'x' ]; then
    if [ $ERR_CODE -ne 0 ]; then
      LEVEL="Error"
      SUBJECT="Error during cloud-init."
      if [ $ERR_LINE -ne 1 ]; then
        MSG="$ERR_CMD failed in $ERR_FUNC at $ERR_LINE. Return: $ERR_CODE"
        ATTACHMENT="$(pr -tn $0 | tail -n+$((ERR_LINE - 3)) | head -n7)"
      else
        MSG="$ERR_CMD"
      fi

      if [ -n "$DEBUG" ]; then
        SUBJECT="$SUBJECT Instance kept running for debug."
      else
        SUBJECT="$SUBJECT Instance terminated by ASG lifecycle hook."
      fi
    else
      LEVEL="Info"
      SUBJECT="ZDT Alpine Instance launched."
    fi

    if [ -z "${DISABLECLOUDBENDERSNSSCALINGEVENTS}" ] || [ "$LEVEL" != "Info" ]; then
      /var/lib/cloudbender/sns_alarm.sh "$SUBJECT" "$MSG" "$LEVEL" "$ATTACHMENT"
    fi

    # Disable scaling events during shutdown
    [ -n "${DISABLECLOUDBENDERSNSSCALINGEVENTS}" ] && echo "DISABLE_SCALING_EVENTS=1" >> /etc/cloudbender/rc.conf
  fi

  [ -n "$LAUNCH_HOOK" ] && aws autoscaling complete-lifecycle-action --lifecycle-action-result $RESULT --instance-id $INSTANCE_ID --lifecycle-hook-name $LAUNCH_HOOK --auto-scaling-group-name ${AWS_AUTOSCALING_GROUPNAME} || true

  if [ -n "${AWS_CLOUDFORMATION_LOGICAL_ID}" ]; then
    aws cloudformation signal-resource --stack-name ${AWS_CLOUDFORMATION_STACK_NAME} --logical-resource-id ${AWS_CLOUDFORMATION_LOGICAL_ID} --unique-id ${INSTANCE_ID} --status ${CFN_STATUS}
  fi

  # timestamp being done
  end_uptime=$(awk '{print $1}' < /proc/uptime)
  log "Exiting user-data. Duration: $(echo "$end_uptime-$start_uptime" | bc) seconds"

  # Shutdown / poweroff if we ran into error and not DEBUG
  [ $ERR_CODE -ne 0 -a -z "$DEBUG" ] && poweroff

  exit 0
}