#!/bin/bash
finish() {
  sync
  umount -q /mnt
  [ -n "$REQUESTID" ] && [ -n "$TOKEN" ] && [ -n "$APIURL" ] && \
  nohup curl -q -X POST -H "Content-Type: application/json" -H "x-request-id: $REQUESTID" -H "x-trace-id: rescue-reboot" -d "{\"token\":\"$TOKEN\"}" "$APIURL" &
  shutdown -r 1
  exit
}

while getopts "p:r:s:t:u:y" option; do
  case $option in
    p) PASSWORD="$OPTARG"   ;;
    r) REQUESTID="$OPTARG"  ;;
    s) SSHKEY="$OPTARG"     ;;
    t) TOKEN="$OPTARG"      ;;
    u) APIURL="$OPTARG"     ;;
    y) CONTINUE=1           ;;
    *) ;;
  esac
done

if [ -z "$CONTINUE" ]; then
  echo "WARNING! This script looks for a Linux installation on a local disk, resets"
  echo "the root password and automatically reboots the instance."
  echo
  read -rsp $'Press any key to continue or CTRL+C to abort... ' -n1 key
  echo
fi

PARTARR=($(blkid -o device -t TYPE="ext4"))
PARTARR+=($(blkid -o device -t TYPE="ext3"))
PARTARR+=($(blkid -o device -t TYPE="xfs"))
PARTARR+=($(blkid -o device -t TYPE="btrfs"))

trap finish ERR

DATA="/var/lib/cloud/instance/user-data.txt"
[ -f "$DATA" ] && [ -z "$PASSWORD" ] && PASSWORD="$(yq -r '.password' $DATA)"
[ -n "$PASSWORD" ] || [ -n "$SSHKEY" ]

for PART in "${PARTARR[@]}"; do
  mount "$PART" /mnt || continue
  TARGET=""
  [ -f /mnt/etc/shadow ] && TARGET="/mnt"
  [ -f /mnt/root/etc/shadow ] && TARGET="/mnt/root"
  if [ -n "$TARGET" ]; then
    [ -n "$PASSWORD" ] && printf "root:%s\n" "$PASSWORD" | chpasswd -c SHA512 -R "$TARGET"
    [ -n "$SSHKEY" ] && mkdir -p "$TARGET"/root/.ssh && echo "$SSHKEY" > "$TARGET"/root/.ssh/authorized_keys && chmod 600 "$TARGET"/root/.ssh/authorized_keys
    finish
  fi
  umount /mnt
done

finish