#!/bin/bash

[[ $EUID == 0 ]] || exec sudo "$BASH_SOURCE" "$@"

x="$(readlink -f "$BASH_SOURCE")"; source "${x%/*}/bash-trace"

usage() {
    cat <<EOF
Usage: ${0##*/} [OPTIONS] -- +/-SIZE[g] swap|boot

Assuming Ian Kelling's partition scheme and we are currently into one of
it's encrypted oses (we it's btrfs filesystem to be mounted), resize
swap or boot, expanding or shrinking the root fs and partition to
compensate. If it changes the partition tables incorrectly, the
originals are stored in /root/backup_partition_table_<device_names>.

TODO: only tested on stretch. deactivation of swap on reboot
probably needs to be fixed on other oses. Even on stretch,
we get 1.5 minutes of waiting for the crypt_dev and normal
boot .device units.

Warning!!! Backup your data. This script could have bugs in it.

-n         Dry run. Note, this likely won't be the exact commands,
           for example, if you are running outside a vm, there will
           probably be a reboot required in the middle so the kernel
           can know about partition changes.
-r         Reboot right away if it's needed.
-f         Force running on a distro that has not been tested.
-h|--help  Print help and exit.

SIZE is MiB, or if g is specified, GiB.

If using multiple devices, SIZE is applied to each device, so total change is
SIZE * devices.

Note: Uses GNU getopt options parsing style
EOF
    exit $1
}

reboot_not=false
dry_run=false
force=false

temp=$(getopt -l help rnfh "$@") || usage 1
eval set -- "$temp"
while true; do
    case $1 in
        -r) reboot_now=true; shift ;;
        -n) dry_run=true; shift ;;
        -f) force=true; shift ;;
        -h|--help) usage ;;
        --) shift; break ;;
        *) echo "$0: Internal error!" ; exit 1 ;;
    esac
done

#### begin arg error checking ####

if [[ $# != 2 ]]; then
    echo "$0: error: expected 2 arguments"
    usage 1
fi

if [[ $1 != [+-][0-9]* ]]; then
    echo "$0: error: bad 1st arg: $1"
    usage 1
fi

if ! which parted &>/dev/null; then
    echo "$0: error: install parted"
    exit 1
fi

case $2 in swap|boot) : ;; *) echo "$0: error: bad 2nd arg"; usage 1 ;; esac

if ! $force && ! grep -q 'VERSION=.*stretch' /etc/os-release; then
    echo "$0: error: This distro is untested. Only tested version atm is Stretch."
    exit 1
fi

#### end arg error checking ####


boot=true
[[ $2 == boot ]] || boot=false

op_size=$1 # operator plus size
if [[ $op_size == *g ]]; then
    op_size=${op_size%g}
    size=${op_size#?}
    op=${op_size%$size}
    size=$(( $size * 1024 ))
    op_size=$op$size
else
    size=${op_size#?}
fi

if [[ $op_size == +* ]]; then
    op_size_rev=-$size # rev = reverse
    grow=true
else
    op_size_rev=+$size
    grow=false
fi

##### end command line parsing ########

rootn=1
swapn=2
bootn=3
needs_reboot=false
reboot_script_initialized=false

pmk() { # partition make
    part=$1
    start_op=$2
    end_op=$3
    fs_type="$4"

    # This fails outside a vm, but actually succeeds. also prints this
    # message:
    # Error: Partition(s) 2 on /dev/sda have been written, but
    # we have been unable to inform the kernel of the change, probably
    # because it/they are in use.  As a result, the old partition(s)
    # will remain in use.  You should reboot now before making further
    # changes.

    if ! p mkpart primary "$fs_type" \
         $((${ptable[start$part]} $start_op)) $((${ptable[end$part]} $end_op)); then
        echo "$0: warning: ignoring failure return of mkpart"
    fi
}

def-e() {
    if $dry_run; then
        e() { echo "+ $@"; }
    else
        e() { echo "+ $@"; "$@"; }
    fi
}

def-e


while read devid dev; do
    case $dev in
        /dev/dm-[0-9])
            # older oses, it points to /dev/dm-x
            dev=$(dmsetup info $dev | sed -rn 's/^\s*Name:\s*(\S*)/\1/p')
            ;;
        /dev/mapper/*)
            dev=${dev#/dev/mapper/}
            ;;
        *)
            echo "$0: error: could not find devicemapper root dev,
make sure you are running from a encrypted root this script is resizing"
            exit 1
            ;;
    esac
    if [[ $dev != crypt_dev_*-part$rootn ]]; then
        echo "$0: error: unexpected root device name,
make sure you are running from a encrypted root this script is resizing"
        exit 1
    fi
    dev=${dev#crypt_dev_}
    dev=${dev%-part$rootn}
    devpath=/dev/disk/by-id/$dev
    echo skip=$size
    def-e
    declare -A ptable
    while IFS=: read id start end psize _; do
        [[ $id == [0-9] ]] || continue
        ptable[start$id]=start=${start%%[^0-9]*}
        ptable[end$id]=${end%%[^0-9]*}
        ptable[size$id]=${psize%%[^0-9]*}
    done < <(parted -m $devpath unit MiB print)
    parted $devpath unit MiB print | tee /root/backup_partition_table_$dev
    p() { e parted -a optimal -s -- $devpath unit MiB "$@"; }
    unit=systemd-cryptsetup@crypt_dev_$dev-part$swapn
    # note systemctl show can test if a unit exists.
    if ! e systemctl stop $unit; then
        e swapoff -a
    fi
    # there is a bug in jessie. this and the .swap unit are
    # generated from /etc/fstab, and it escapes - to x2d, then doesn't escape it
    # when looking for the file to use as swap. so, no swap is working on jessie
    # right now.
    sleep 1 # dunno if this is needed,
    # but systemd likes to do these kind of things in the background.

    # These partition comments seems a little verbose now, but I bet they
    # will be helpfull if I read this in more than a week from now.
    # <> = deleted partition, () = partition
    p rm $swapn # ( root )< swap >( boot )

    root_resize_cmd="e btrfs fi resize $devid:${op_size_rev}M /"
    if $grow; then $root_resize_cmd; fi
    # if $grow; then
    # < root >< swap >( boot )
    # ( root ) >< swap >( boot )
    # else
    # < root >< swap >( boot )
    # ( root >< ) swap >( boot )

    out=$(p rm $rootn 2>&1)
    echo "$out"

    pmk $rootn "" $op_size_rev

    if echo "$out" | \
            grep "but we have been unable to inform the kernel" &>/dev/null; then
        needs_reboot=true
    fi
    if $needs_reboot; then
        # note: even if these units don't exist, this will succeed.
        e systemctl mask dev-mapper-crypt_swap_$dev$swapn.swap
        e systemctl mask systemd-cryptsetup@crypt_swap_$dev$swapn.service
        e() { echo "$@" >> /root/finish-resize; }
        if ! $reboot_script_initialized; then
            reboot_script_initialized=true
            rm -rf /root/finish-resize
            cat >/root/finish-resize <<'EOF'
#!/bin/bash -x
set -eE -o pipefail
trap 'echo "$0:$LINENO:error: \"$BASH_COMMAND\" returned $?" >&2' ERR
EOF
            chmod +x /root/finish-resize
        fi
        e swapoff -a
        e systemctl unmask systemd-cryptsetup@crypt_swap_$dev$swapn.service
        e systemctl unmask dev-mapper-crypt_swap_$dev$swapn.swap
    fi
    if ! $grow; then
        $root_resize_cmd
    fi
    if $boot; then
        # non by-id path, to match what btrfs fi show will tell us
        boot_dev_path=$(readlink -f $devpath-part$bootn)
        boot_devid=$(btrfs fi show /boot | \
                         sed -rn "s#^\s*devid\s+(\S+)\s.*$boot_dev_path#\1#p")

        if ! $grow; then
            # shrink boot, move it to a temp file
            e btrfs fi resize $boot_devid:${op_size}M /boot
            e umount /boot
            temp_boot=/root/temp_boot_$dev
            e dd bs=1M if=$boot_dev_path of=$temp_boot \
              count=$((${ptable[size$bootn]} $op_size))
        else
            e umount /boot
        fi
        # if $grow; then
        # ( root ) >< swap >< boot >
        # ( root )( >< swap ) >< boot >
        # ( root )( >< swap )( >< boot )
        # else
        # ( root >< ) swap >< boot >
        # ( root >< )( swap >< ) boot >
        # ( root >< )( swap >< )( boot )
        p rm $bootn
        pmk $swapn $op_size_rev $op_size_rev "linux-swap"
        pmk $bootn $op_size_rev ""

        if $grow; then
            e dd bs=1M if=$boot_dev_path of=$boot_dev_path skip=$size
            e mount /boot
            e btrfs fi resize $boot_devid:${op_size}M /boot
        else
            e dd bs=1M if=$temp_boot of=$boot_dev_path
            e mount /boot
        fi
    else
        # if $grow; then  ( root )( >< swap )( boot )
        # else            ( root >< )( swap )( boot )
        pmk $swapn $op_size_rev "" "linux-swap"
        e systemctl start systemd-cryptsetup@crypt_swap_$dev$swapn
    fi
done < <(btrfs fi show / | sed -nr 's#^\s*devid\s*(\S+).* path (.*)$#\1 \2#p')


if $boot; then
    e rm -rf "/root/temp_boot_*"
    e rm -f /root/finish-resize
fi

if $needs_reboot; then
    echo "$0: Reboot, run /root/finish-resize. It's contents:"
    cat /root/finish-resize
    if $reboot_now; then
        echo "$0: rebooting now"
        reboot now
        exit
    fi
fi