#!/bin/bash -e
set -o pipefail

# Writes a raw image (possibly zipped)
# to a device (e.g., USB drive or SD card)
# using dd, with status monitoring via pv

ARGS=$(getopt -o fh -l "force,sparse,debug,removable,fixed,help" -n "$0" -- "$@")
eval set -- "$ARGS"

usage() {
    cat <<EOF
Usage:
   $0 [options] IMAGE DEVICE [BLOCK_SIZE]
Arguments:
   IMAGE        raw or zipped image file (.img, .img.gz or .img.xz)
                or ISO9660 image file (.iso)
   DEVICE       device name (e.g., '/dev/sdb' or '/dev/mmcblk0')
                or '-' to write uncompressed image to stdout
   BLOCK_SIZE   block size in bytes (default = '1M')
Options:
   -f,--force   don't ask to proceed with writing
   --sparse     use dd sparse conversion option
   --debug      turn on debugging messages
   --removable  device is removable (remove fallback from ESP);
                not supported if DEVICE is '-'
   --fixed      device is not removable (don't remove fallback from ESP)
   -h,--help    show this message
Note:
   If neither --removable nor --fixed is specified, 'lsblk -o HOTPLUG' is used
   to guess; the guess will be shown before you're asked to proceed with
   writing.
EOF
}

FORCE=
REMOVABLE=
DD_IN_OPTS=(iflag=nonblock)
DD_OUT_OPTS=(conv=fsync)
while true; do
    case "$1" in
        -f|--force)
            FORCE=true
            shift
            ;;
        --sparse)
            DD_OUT_OPTS+=(conv=sparse)
            shift
            ;;
        --debug)
            set -x
            shift
            ;;
        --removable)
            if [ "$REMOVABLE" = "0" ]; then
                echo "--removable and --fixed cannot both be used" >&2
                exit 1
            fi
            REMOVABLE=1
            shift
            ;;
        --fixed)
            if [ "$REMOVABLE" = "1" ]; then
                echo "--removable and --fixed cannot both be used" >&2
                exit 1
            fi
            REMOVABLE=0
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        --)
            shift
            break
            ;;
    esac
done

declare -A dependencies
dependencies=(
    [dd]='coreutils'
    [xzcat]='xz-utils'
    [zcat]='gzip'
)
for command in "${!dependencies[@]}"; do
    if ! which "$command" >/dev/null 2>&1; then
        echo "$command is not installed... aborting!" >&2
        echo "Try 'sudo apt-get install ${dependencies[$command]}'" >&2
        exit 1
    fi
done

if [ $# -lt 2 ] || [ $# -gt 3 ] ; then
    if [ $# -eq 0 ]; then
        echo "Error: missing IMAGE and DEVICE arguments" >&2
    elif [ $# -eq 1 ]; then
        echo "Error: missing DEVICE argument" >&2
    else
        echo "Error: extra arguments after IMAGE DEVICE BLOCK_SIZE" >&2
    fi
    echo >&2
    usage >&2
    exit 1
fi

USERID=$(id -u)
IMAGE="$1"
DEVICE="$2"

BLOCK_SIZE=1M
if [ $# -ge 3 ] ; then
    BLOCK_SIZE="$3"
fi
DD_IN_OPTS+=(bs="$BLOCK_SIZE")
DD_OUT_OPTS+=(bs="$BLOCK_SIZE")

if [ "$DEVICE" != "-" ] && [ "$USERID" != "0" ]; then
    echo "Program requires superuser privileges" >&2
    exit 1
fi

if [ ! -f "$IMAGE" ]; then
    echo "$IMAGE not found" >&2
    exit 1
fi

# fall back if pv isn't available
if ! which pv >/dev/null 2>&1; then
  pv() {
    echo "Proceeding with no pv, there will be no progress message!" >&2
    echo "Try 'sudo apt-get install pv'" >&2
    echo "Writing..." >&2
    cat
    echo "Done!" >&2
  }
fi

cat_image() {
    MIME_TYPE="$(file --brief --mime-type "$IMAGE")"
    if [[ "$MIME_TYPE" =~ xz$ ]]; then
        # Image is xzipped, xz has its own rate meter
        xzcat -v "$IMAGE"
    elif [[ "$MIME_TYPE" =~ gzip$ ]]; then
        # Image is gzipped
        # The following would calculate the original size
        # of the unzipped file if less than 4G:
        # IMAGE_SIZE=$(zcat -l "$IMAGE" | awk 'NR==2 { print $2 }')
        # But this doesn't help for large images,
        # so we just show relative progress unless .size is present
        if [ -f "$IMAGE.size" ]; then
            IMAGE_SIZE=$(cat $IMAGE.size)
            PV_ARGS="-s $IMAGE_SIZE"
        fi
        zcat "$IMAGE" | pv $PV_ARGS
    else
        # Image is uncompressed
        # We can show progress as percentage of total image size
        IMAGE_SIZE=$(ls -l "$IMAGE" | awk '{ print $5 }')
        dd if="$IMAGE" "${DD_IN_OPTS[@]}" | pv -s ${IMAGE_SIZE}

        # ISO images do not need to be mangled
        if [[ "$MIME_TYPE" =~ x-iso9660-image$ ]]; then
            REMOVABLE=0
        fi
    fi
}

if [ "$DEVICE" == "-" ]; then
    if [ "$REMOVABLE" = "1" ]; then
        echo "--removable not supported when writing to stdout" >&2
        exit 1
    fi
    cat_image
else
    if [ ! -e "$DEVICE" ]; then
        echo "$DEVICE does not exist... aborting!" >&2
        exit 1
    fi

    if [ ! -b "$DEVICE" ]; then
        echo "$DEVICE is not a block device... aborting!" >&2
        exit 1
    fi

    if grep -qs "$DEVICE" /proc/mounts; then
        # Protect against overwriting the device currently in use
        echo "$DEVICE is currently in use -- please unmount and try again" >&2
        exit 1
    fi

    if [ ! "$REMOVABLE" ]; then
        if ! REMOVABLE=$(lsblk --raw --nodeps --noheading --output HOTPLUG "$DEVICE"); then
            # This can happen if the disk has been "ejected" rather than
            # unmounting. Writing to it will fail in this case.
            echo "Unable to detect whether $DEVICE is removable;" \
                 "writing to it will probably fail." >&2
            REMOVABLE=1
        elif [ "$REMOVABLE" = 1 ]; then
            echo "Guessed that $DEVICE is a removable disk;" \
                 "pass --fixed if this is wrong." >&2
        else
            echo "Guessed that $DEVICE is a fixed disk;" \
                 "pass --removable if this is wrong." >&2
        fi
    fi

    if [ ! "$FORCE" ]; then
        read -p "Are you sure you want to overwrite all data on $DEVICE? [y/N] "
        response="${REPLY,,}" # to lower
        if [[ ! "$response" =~ ^(yes|y)$ ]]; then
            exit 1
        fi
    fi

    # Try to format / discard all the sectors on the device
    blkdiscard "$DEVICE" &>/dev/null || true
    cat_image | dd of="$DEVICE" "${DD_OUT_OPTS[@]}"

    # Probe new partitions
    partprobe "$DEVICE"
    udevadm settle

    # If removable device, delete fallback from the ESP
    if [ "$REMOVABLE" = "1" ]; then
        # Assume ESP is the first partition
        case "$DEVICE" in
            /dev/mmcblk?)
                ESP=${DEVICE}p1
                ;;
            /dev/sd?)
                ESP=${DEVICE}1
                ;;
        esac

        # Try to mount ESP
        udisksctl mount -b "$ESP" --no-user-interaction || exit 1

        # Grab the mountpoint
        ESP_MP=$(grep "$ESP" /proc/mounts | cut -f2 -d' ')

        if test -n "$ESP_MP" -a -d "$ESP_MP"
        then
            ESP_BOOT="${ESP_MP}/EFI/BOOT"
            ESP_ENDLESS="${ESP_MP}/EFI/endless"
            FALLBACK_OLD="${ESP_BOOT}/fallback.efi"
            FALLBACK_NEW="${ESP_BOOT}/fbx64.efi"
            MOKMANAGER_OLD="${ESP_ENDLESS}/MokManager.efi"
            MOKMANAGER_NEW="${ESP_ENDLESS}/mmx64.efi"
            GRUB="${ESP_ENDLESS}/grubx64.efi"

            echo "Checking ESP layout" >&2
            if [ -f "$FALLBACK_OLD" ]; then
                echo "Found $FALLBACK_OLD" >&2
                FALLBACK=$FALLBACK_OLD
                MOKMANAGER=$MOKMANAGER_OLD
            elif [ -f "$FALLBACK_NEW" ]; then
                echo "Found $FALLBACK_NEW" >&2
                FALLBACK=$FALLBACK_NEW
                MOKMANAGER=$MOKMANAGER_NEW
            fi

            if [ -z "$FALLBACK" ]; then
                echo "Neither $FALLBACK_OLD nor $FALLBACK_NEW were found" >&2
            else
                rm -v "$FALLBACK"
                mv -v "$MOKMANAGER" "$ESP_BOOT"
                mv -v "$GRUB" "$ESP_BOOT"
                rm -frv "$ESP_ENDLESS"
            fi
        fi

        # Unmount ESP
        udisksctl unmount -b "$ESP" --no-user-interaction
    fi
fi