#!/bin/bash -efu

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See LICENSE for more details.
#
# Copyright: Red Hat Inc. 2018
# Author: Andrei Stepanov <astepano@redhat.com>


PROG="${PROG:-${0##*/}}"
DEFAULT_SIZE="20G"
DEBUG=

debug() {
    if [ -n "$DEBUG" ]; then
        echo "$*" >&2
    fi
}

msg_usage() {
    cat << EOF
Usage:
$PROG <options>

You must run this script with root credentials.

Options:
-i, --image=IMAGE                            one of: install, update, downgrade, remove
-s, --size=SIZE                              grow image for desired size. Use suffixes: K, M, G, Default: $DEFAULT_SIZE
-h, --help                                   display this help and exit
-v, --verbose                                print debug messages
EOF
}

box_out() {
    local s=("$@")
    local b=
    local w=
    for l in "${s[@]}"; do
        ((w<${#l})) && { b="$l"; w="${#l}"; }
    done
    echo -e " -${b//?/-}-\n| ${b//?/ } |"
    for l in "${s[@]}"; do
        printf '| %*s |\n' "-$w" "$l"
    done
    echo -e "| ${b//?/ } |\n-${b//?/-}-"
}

# http://wiki.bash-hackers.org/howto/getopts_tutorial
opt_str="$@"
opt=$(getopt -n "$0" --options "hvi:s:" --longoptions "help,verbose,image:,size:" -- "$@")
eval set -- "$opt"
while [[ $# -gt 0 ]]; do
    case "$1" in
        -i|--image)
            IMAGE="$2"
            shift 2
            ;;
        -s|--size)
            SIZE="$2"
            shift 2
            ;;
        -v|--verbose)
            DEBUG="yes"
            shift
            ;;
        -h|--help)
            msg_usage
            exit 0
            ;;
        --)
            shift
            ;;
        *)
            msg_usage
            exit 1
    esac
done

# Entry

IMAGE="${IMAGE:-}"
SIZE="${SIZE:-$DEFAULT_SIZE}"
ACTIVE_UID="$(id -u)"

debug "IMAGE: $IMAGE"
debug "SIZE: $SIZE"
debug "ACTIVE_UID: $ACTIVE_UID"

if [ -z "$IMAGE" ] || [ "$ACTIVE_UID" -ne 0 ]; then
    echo "Use: $PROG -h for help."
    exit
fi

kernel_config="/boot/config-$(uname -r)"
if [ -e "$kernel_config"  ]; then
    ret=
    grep -q -s "^CONFIG_BLK_DEV_NBD" "$kernel_config" || ret="not_supported"
    if [ -n "$ret" ]; then
        echo "Kernel was compiled without CONFIG_BLK_DEV_NBD support. Use other distro."
        exit 1
    fi
fi

req_pkgs=()
req_pkgs+=('qemu-img')
req_pkgs+=('xfsprogs')
req_pkgs+=('e2fsprogs')
req_pkgs+=('util-linux')

dnf=dnf
[ -f "/usr/bin/yum" ] && dnf=yum

for pkg in ${req_pkgs[@]}; do
    present=
    info_msg=
    rpm -q --qf "" "$pkg" >/dev/null 2>&1 && present="yes" || :
    if [ -z "$present" ]; then
        [ -z "$info_msg" ] && box_out "Install missing packages"; info_msg="run_once"
        "$dnf" install -y "$pkg"
    fi
done

nbd_present=
lsmod | grep -q '^nbd[[:space:]]' && nbd_present="yes"
if [ -z "$nbd_present" ]; then
    box_out \
        "Loading module: nbd"
    modprobe nbd max_part=8
fi

nbd_dev_index=0
last_nbd=$(cat "/proc/partitions" | tr -s ' ' | sed -n -e 's/^[[:space:]]*//; /[[:digit:]]/p' | cut -f '4' -d ' ' | sed -n -e '/^nbd[[:digit:]]\+$/p' | sort -V | sed -n -e '$ s/[^[:digit:]]\+// p')
[ -n "$last_nbd}" ] && nbd_dev_index=$(($last_nbd + 1))
nbd_dev="/dev/nbd${nbd_dev_index}"
mntpoint="$(mktemp -d -p /mnt qcow2-grow-XXXX)"

do_cleanup() {
    rc=$?;
    trap - SIGINT SIGTERM SIGABRT EXIT # clear the trap
    if mountpoint -q $mntpoint; then
        echo "Cleanup: umount $mntpoint"
        umount "$mntpoint" || :
    fi
    echo "Cleanup: remove $mntpoint"
    rmdir "$mntpoint" || :
    echo "Cleanup: disconnect $nbd_dev"
    qemu-nbd -d "$nbd_dev" || :
    exit $rc
}

trap do_cleanup SIGINT SIGTERM SIGABRT EXIT

box_out \
    "Before extending" \
    "----------------" \
    "" \
    "Image: $IMAGE" \
    "Device: $nbd_dev" \
    "Add: $SIZE"

(set -x; qemu-img "info" "$IMAGE")
qemu-img "resize" "$IMAGE" "+${SIZE}"
qemu-nbd --connect="$nbd_dev" "$IMAGE"
sync
sleep 2
partition="$(cat "/proc/partitions" | \
    tr -s ' ' | \
    sed -n -e 's/^[[:space:]]*//; /[[:digit:]]/p' | \
    cut -f '4' -d ' ' | \
    sed -n -e "/^nbd${nbd_dev_index}p[[:digit:]]\+$/p" | \
    sort -V | \
    sed -n -e '$ p')"
(set -x; fdisk -l "$nbd_dev")
[ -z "partition" ] && echo "Bad qcow2 image." && exit 1 || :
partition="/dev/$partition"
box_out "Grow the latest partition: $partition"
step=0
while true; do
    # Sometimes fdisk can fail with: Re-reading the partition table failed.: Device or resource busy.
    # Retry 5 times, untill success.
    failed=
    echo -e 'd\n\nn\n\n\n\n\nw' | fdisk "$nbd_dev" || failed="yes"
    [ -z "$failed" ] && break
    step+=1
    sleep 1
    if [ $step -eq 5 ]; then
        echo "fdisk commands failed."
        exit 1
    fi
done
(set -x; fdisk -l "$nbd_dev")

mount -o rw "$partition" "$mntpoint"
sync
sleep 1
fs_type="$(cat /proc/mounts | tr -s ' '  | grep "$partition " | cut -f '3' -d ' ')"

box_out "Grow $partition with fs: $fs_type"

if [ "$fs_type" = "xfs" ]; then
    xfs_growfs -d "$mntpoint"
elif [ "${fs_type%?}" = "ext" ]; then
    resize2fs "$partition"
fi

(set -x; df -h "$mntpoint")

exit $ret
