#!/usr/bin/env bash

show_help () {
  cat << EOF
checkvolumesize 1.0.0

Automatically expand Digital Ocean block storage volumes.

Usage: checkvolumesize [options]

Options:

  --help | -h | help    [Optional] Show this help message
  --token               [Required] Digital Ocean API token
  --device              [Required] The volume's block device or mount point
  --volume-name         [Required] The volume's name
  --volume-region       [Required] The volume's region
  --buffer              [Optional] The amount of GB to keep available    Default: 10
  --log                 [Optional] Prepend timestamps to stdout

Example:

  checkvolumesize /
    --token bc99be9f73b037da64074472fe58f643e619328b85c5467615964a59abd12029 /
    --device /mnt/block-storage /
    --volume-name volume-sgp1-01 /
    --volume-region sgp1 /
    --buffer 10

GitHub: https://github.com/lukechilds/autoscale-do-block-storage
EOF
}

# Parse options
POSITIONAL=()
while [[ $# -gt 0 ]]; do
  key="$1"

  BUFFER=10
  LOG=false

  case $key in
      --help|-h|help)
        show_help
        exit
        ;;

      --token)
        TOKEN="$2"
        shift
        shift
        ;;

      --device)
        DEVICE="$2"
        shift
        shift
        ;;

      --volume-name)
        VOLUME_NAME="$2"
        shift
        shift
        ;;

      --volume-region)
        VOLUME_REGION="$2"
        shift
        shift
        ;;

      --buffer)
        BUFFER="$2"
        shift
        shift
        ;;

      --log)
        LOG=true
        shift
        shift
        ;;

      *)
      POSITIONAL+=("$1")
        shift
        ;;
  esac
done
set -- "${POSITIONAL[@]}"

log () {
  prefix=""
  if [[ "$LOG" = true ]]; then
    prefix="$(date "+%Y-%m-%d %H:%M:%S") "
  fi
  echo $prefix$@
}

check_root () {
  if [[ $UID != 0 ]]; then
    log "This script must be run as root"
    exit 1
  fi
}

check_dependencies () {
  for cmd in "$@"; do
    if ! command -v $cmd >/dev/null 2>&1; then
      log "This script requires \"${cmd}\" to be installed"
      exit 1
    fi
  done
}

check_options () {
  if [[ -z $TOKEN || -z $DEVICE || -z $VOLUME_NAME || -z $VOLUME_REGION ]]; then
    log 'You are missing required options'
    echo
    show_help
    exit 1
  fi
}

digital_ocean_api () {
  endpoint=$1
  post_data=$2

  method="GET"
  if [[ $post_data ]]; then
    method="POST"
  fi

  args=()
  args+=("--silent")
  args+=("-X" $method)
  args+=("-H" "Content-Type: application/json")
  args+=("-H" "Authorization: Bearer ${TOKEN}")
  if [[ $post_data ]]; then
    args+=("-d" $post_data)
  fi
  args+=("https://api.digitalocean.com/v2/${endpoint}")

  curl "${args[@]}"
}

bytes_to_gigabytes () {
  bytes=$1

  expr $bytes / 1024 / 1024
}

check_device_free_space () {
  device=$1

  bytes_free=$(df --output=avail $device | tail -n 1)
  bytes_to_gigabytes $bytes_free
}

mountpoint_to_device () {
  mountpoint=$1

  df --output=source $mountpoint | tail -n 1
}

main () {
  check_root
  check_dependencies curl jq expr resize2fs
  check_options

  # Check available volume space
  log "Checking available space on device \"${DEVICE}\" is above ${BUFFER}GB requirement..."
  volume_free_space=$(check_device_free_space $DEVICE)
  if [[ $volume_free_space -gt $BUFFER ]]; then
    log "${volume_free_space}GB available, all good!"
    return
  fi
  log "Only ${volume_free_space}GB available, volume resize required"

  # Save start time
  start_time=$SECONDS

  # Get volume data
  log "Getting data for volume \"${VOLUME_NAME}\" in region \"${VOLUME_REGION}\"..."
  volume_json=$(digital_ocean_api "volumes?region=${VOLUME_REGION}" | jq -r --arg VOLUME_NAME "${VOLUME_NAME}" '.volumes | .[] | select(.name==$VOLUME_NAME)')
  volume_id=$(echo $volume_json | jq -r .id)
  volume_size=$(echo $volume_json | jq -r .size_gigabytes)
  log "Volume ID is \"${volume_id}\" and is currently ${volume_size}GB"

  # Increase volume size
  new_volume_size=$(expr $volume_size + $BUFFER)
  log "Increasing volume size to ${new_volume_size}GB..."
  action_json=$(digital_ocean_api "volumes/${volume_id}/actions" '{"type":"resize","size_gigabytes":'$new_volume_size'}')
  action_id=$(echo $action_json | jq -r '.action.id')
  log "Created action \"${action_id}\""

  # Wait for volume resize
  log "Waiting for action to complete..."
  while true; do
    resize_status=$(digital_ocean_api "actions/$action_id" | jq -r '.action.status')
    if [[ "${resize_status}" == "completed" ]]; then
      break
    fi
    sleep 1
  done
  log "Volume resize complete!"

  # Resize filesystem
  log "Resizing filesystem..."
  sudo resize2fs $(mountpoint_to_device $DEVICE) 2>&1 | while read line; do
    if [[ "${line}" != "" ]]; then
      log $line
    fi
  done
  log "Filesystem resize complete!"

  # Show new available space on completion
  new_free_space=$(check_device_free_space $DEVICE)
  log "Device \"${DEVICE}\" now has ${new_free_space}GB available"

  # Show duration
  duration=$(expr $SECONDS - $start_time)
  log "Completed in ${duration} seconds"
}

main