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

# Writes an Endless OS image to a device (e.g., USB drive or SD card) so that it
# can be used both as a live system, and to install itself. The image is
# extracted to an exFAT partition, and is booted using a GRUB and
# initramfs which know how to loopback-mount the image.
#
# NOTE: the partition table will not be in physical order so Windows can mount
# the image partition. Additional images copied to the image partition will be
# detected by the installer, but will not be bootable.

MOUNTS=/run/media/eos-write-live-image
ISO_MOUNTPOINT=$MOUNTS/iso
ESP_MOUNTPOINT=$MOUNTS/esp
EOSLIVE_MOUNTPOINT=$MOUNTS/eoslive


EOS_WRITE_IMAGE=$(dirname "$0")/eos-write-image
if [ ! -f "$EOS_WRITE_IMAGE" ]; then
    EOS_WRITE_IMAGE='eos-write-image'
fi
EOS_DOWNLOAD_IMAGE=$(dirname "$0")/eos-download-image
if [ ! -f "$EOS_DOWNLOAD_IMAGE" ]; then
    EOS_DOWNLOAD_IMAGE='eos-download-image'
fi

ARGS=$(getopt -o o:x:lp:r:nms:wL:fPS:e:h \
    -l os-image: \
    -l windows-tool: \
    -l latest \
    -l personality: \
    -l product: \
    -l ntfs \
    -l mbr \
    -l size: \
    -l fill \
    -l writable \
    -l no-bios \
    -l label: \
    -l force \
    -l debug \
    -l persistent \
    -l free-space: \
    -l extra-data: \
    -l help \
    -n "$0" -- "$@")
eval set -- "$ARGS"

NTFS=
MBR=
EXPAND=
SIZE=
WRITABLE=
BIOS=true
WINDOWS_TOOL_PROVIDED=
EXTRA_DATA_PATHS=()
FORCE=
PERSONALITY=base
PRODUCT=eos
LABEL=
PERSISTENT=
PERSISTENT_FREE_SPACE=

usage() {
    local SELF
    SELF=$(basename "$0")
    cat <<EOF
Usage:
    $SELF [options] OUTPUT

    # equivalent to $SELF --latest --os-image IMAGE OUTPUT
    $SELF IMAGE OUTPUT

Arguments:
    OUTPUT                   Device path (e.g. '/dev/sdb')

Options:
    -o, --os-image IMAGE     Path to Endless OS image
    -x, --windows-tool PATH  Path to Endless OS installer tool for Windows
    -l, --latest             Fetch latest OS image and/or Windows tool
    -p, --personality        Fetch a different personality (default: $PERSONALITY)
    -r, --product            Fetch a different product (default: $PRODUCT)
    -L, --label LABEL        Autorun.inf label (default: Endless OS)
    -s, --size SIZE          Expand the image to desired size in bytes
        --fill               Expand the image to fill the target disk
    -w, --writable           Allow image to be modified when run (which
                             prevents installing it later)
        --no-bios            Do not create a BIOS boot partition
    -f, --force              don't ask to proceed before writing
    -P, --persistent         Allocate space for persistent storage, leaving a
                             few megabytes for logs from the Endless OS
                             installer tool for Windows
    -S, --free-space SIZE    With --persistent, leave SIZE megabytes free on
                             the partition, rather than consuming almost all
                             free space
    -e, --extra-data PATH    A local path to copy to the live partition, it can
                             be used multiple times
    -h, --help               Show this message

Developer options (you probably don't want to use these):
    -n, --ntfs               Format the image partition as NTFS, not exFAT
    -m, --mbr                Format DEVICE as MBR, not GPT
        --debug              Enable debugging output
EOF
}

function check_exists() {
    if [ ! -f "$1" ]; then
        echo "$2 $1 does not exist or is not a file" >&2
        exit 1
    fi
}

function cleanup_mountpoint() {
    if [ -d "$1" ]; then
        umount "$1" || :
        rmdir "$1"
    fi
}

function cleanup() {
    cleanup_mountpoint "$ISO_MOUNTPOINT"
    cleanup_mountpoint "$ESP_MOUNTPOINT"
    cleanup_mountpoint "$EOSLIVE_MOUNTPOINT"
    if [ -d "$MOUNTS" ]; then
        rmdir "$MOUNTS"
    fi
}

function find_by_type() {
    echo "$1" | grep -i "type=$2" | cut -d' ' -f1
}

while true; do
    case "$1" in
        -o|--os-image)
            shift
            OS_IMAGE="$1"
            shift
            ;;
        -x|--windows-tool)
            shift
            WINDOWS_TOOL="$1"
            WINDOWS_TOOL_PROVIDED=true
            shift
            ;;
        -e|--extra-data)
            shift
            EXTRA_DATA_PATHS+=("$1")
            shift
            ;;
        -l|--latest)
            shift
            FETCH_LATEST=true
            ;;
        -p|--personality)
            shift
            PERSONALITY="$1"
            shift
            ;;
        -r|--product)
            shift
            PRODUCT="$1"
            shift
            ;;
        -n|--ntfs)
            shift
            NTFS=true
            ;;
        -m|--mbr)
            shift
            MBR=true
            ;;
        -w|--writable)
            shift
            WRITABLE=true
            ;;
        --fill)
            shift
            EXPAND=true
            ;;
        -s|--size)
            shift
            EXPAND=true
            SIZE="$1"
            shift
            ;;
        -L|--label)
            shift
            LABEL="$1"
            shift
            ;;
        -f|--force)
            shift
            FORCE=true
            ;;
        -P|--persistent)
            shift
            PERSISTENT=true
            ;;
        -S|--free-space)
            shift
            PERSISTENT_FREE_SPACE="$1"
            shift
            ;;
        --no-bios)
            shift
            BIOS=
            ;;
        --debug)
            shift
            set -x
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        --)
            shift
            break
            ;;
    esac
done

# Be consistent with eos-write-image IMAGE DEVICE
if [ $# -eq 2 ] && [ -z "$OS_IMAGE" ]; then
    OS_IMAGE="$1"
    FETCH_LATEST=true
    shift
fi

if [ $# -ne 1 ] ; then
    if [ $# -lt 1 ] ; then
        echo "Error: missing OUTPUT" >&2
    else
        echo "Error: extra arguments after OUTPUT:" "$@" >&2
    fi
    echo >&2
    usage >&2
    exit 1
fi

OUTPUT="$1"

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

# Check for required tools
declare -A dependencies
dependencies=(
    [dd]='coreutils'
    [mkfs.vfat]='dosfstools'
    [partprobe]='parted'
    [sfdisk]='fdisk'
    [unzip]='unzip'
    [xzcat]='xz-utils'
    [zcat]='gzip'
)

if [ "$NTFS" ]; then
    MKFS_IMAGES="mkfs.ntfs"
    MKFS_ARGS='--quick -L'
    dependencies[$MKFS_IMAGES]='ntfs-3g'
else
    MKFS_IMAGES="mkfs.exfat"
    # By default, mkfs.exfat picks the cluster size (allocation unit) based on
    # the disk size.  On large disks this is 128 KiB, wasting a lot of space on
    # Endless Key with many small files.  Using 4 KiB (the default for very
    # small disks) appears to have some performance impact when reading large

    # files.  As a compromise, force the cluster size to 32 KiB, which is the
    # default for volumes between 256 MiB and 32 GiB.
    MKFS_ARGS='-c 32k -L'
    dependencies[$MKFS_IMAGES]='exfatprogs'
fi

if [ "$EXPAND" ]; then
    dependencies[truncate]='coreutils'
fi

missing_packages=()
for command in "${!dependencies[@]}"; do
    if ! command -v "$command" >/dev/null 2>&1; then
        echo "$command is not installed" >&2
        missing_packages+=("${dependencies[$command]}")
    fi
done
if [ ${#missing_packages[@]} -gt 0 ]; then
    echo "Try 'sudo apt-get install ${missing_packages[*]}'" >&2
    exit 1
fi

if [ -z "$FETCH_LATEST" ]; then
    if [ -z "$WINDOWS_TOOL" ] || [ -z "$OS_IMAGE" ]; then
        echo "--os-image and --windows-tool are required if --latest is not specified" >&2
        usage >&2
        exit 2
    fi
fi

if [ ! "$BIOS" ] && [ "$MBR" ]; then
    echo "--no-bios and --mbr cannot be used together" >&2
    exit 1
fi

if [ -n "$SIZE" ]; then
    SIZE_DESC=$SIZE
else
    if [ -n "$EXPAND" ]; then
        SIZE_DESC="fill disk"
    else
        SIZE_DESC="n/a"
    fi
fi

echo "Summary:"
echo
echo "       Endless OS image: ${OS_IMAGE:-latest $PRODUCT $PERSONALITY image}"
echo "  Installer for Windows: ${WINDOWS_TOOL:-latest release}"
echo "      Image target size: ${SIZE_DESC}"
echo "  Create BIOS partition: ${BIOS:-false}"
echo "                 Target: ${OUTPUT}"
echo

if [ ! -e "$OUTPUT" ]; then
    echo "$OUTPUT does not exist... aborting!" >&2
    exit 1
fi

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

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

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

if [ -n "$FETCH_LATEST" ]; then
    if [ -z "$OS_IMAGE" ]; then
        OS_IMAGE=$($EOS_DOWNLOAD_IMAGE --personality "${PERSONALITY}" --product "${PRODUCT}")
        if [ -z "$LABEL" ]; then
           LABEL="Endless OS ${PERSONALITY}"
        fi
    fi
    if [ -z "$WINDOWS_TOOL" ]; then
        WINDOWS_TOOL=$($EOS_DOWNLOAD_IMAGE --windows-tool)
    fi
fi

if [ -z "$LABEL" ]; then
    LABEL="Endless OS"
fi

check_exists "$OS_IMAGE" "image"

if [[ "$OS_IMAGE" == *.iso ]]; then
    ISO=true
    OS_IMAGE_UNCOMPRESSED="${OS_IMAGE%.iso}.img"
else
    # If image does not have a .gz, .xz or .iso suffix, assume it is the
    # uncompressed .img
    OS_IMAGE_UNCOMPRESSED="${OS_IMAGE%.?z}"
fi
OS_IMAGE_UNCOMPRESSED_BASENAME="$(basename "${OS_IMAGE_UNCOMPRESSED}")"

EXTRACTED_SIGNATURE="${OS_IMAGE_UNCOMPRESSED}.asc"
check_exists "$EXTRACTED_SIGNATURE" "uncompressed image signature"

BOOT_ZIP="${OS_IMAGE_UNCOMPRESSED%.img}.boot.zip"
check_exists "$BOOT_ZIP" "bootloader bundle"
check_exists "$BOOT_ZIP.asc" "bootloader bundle signature"

check_exists "$WINDOWS_TOOL" "Windows tool"

echo
echo "Preparing $OUTPUT..."

# Erase any existing ISO9660 (0x8000 == 8 * 4096) or Joliet
# (0x9000 == 9 * 4096) header on the disk. If you have written ISO image I
# to disk A, overwrite disk A with this tool, write I to disk B, then try
# to boot from B with A still plugged in, the "UUID" from the stale ISO9660
# header on A is still read by the kernel and taken as the UUID for disk A
# as a whole, and its BIOS boot partition too (for good measure). This
# confuses eos-image-boot-setup into trying to mount a partition on disk A
# rather than B.
dd if=/dev/zero of="$OUTPUT" bs=4096 count=2 seek=8
udevadm settle

if [ "$MBR" ]; then
    sfdisk --label dos "$OUTPUT" <<MBR_PARTITIONS
1 : type=7
MBR_PARTITIONS
    udevadm settle
    partprobe "$OUTPUT"
    PARTMAP=$(sfdisk --dump "$OUTPUT")
    DEVICE_IMAGES=$(find_by_type "$PARTMAP" "7")
else
    # https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_type_GUIDs
    PARTITION_SYSTEM_GUID="c12a7328-f81f-11d2-ba4b-00a0c93ec93b"
    PARTITION_BIOS_BOOT_GUID="21686148-6449-6E6F-744E-656564454649"
    PARTITION_BASIC_DATA_GUID="ebd0a0a2-b9e5-4433-87c0-68b6b72699c7"

    # We want the data partition, "eoslive", to occupy all the space on the disk
    # after the UEFI and BIOS boot partitions. But, we also want it to be numbered
    # first: apparently Windows will only mount the partition numbered first,
    # regardless of where it is on the disk.
    #
    # It is important that the offset of the BIOS boot partition matches that
    # used in the Endless OS image builder, since it is embedded in the GRUB
    # image.
    if [ "$BIOS" ]; then
        sfdisk --label gpt "$OUTPUT" <<EFI_PARTITIONS
2 : start=2048, size=62MiB, type=$PARTITION_SYSTEM_GUID
3 : size=1MiB, type=$PARTITION_BIOS_BOOT_GUID
1 : name=eoslive, type=$PARTITION_BASIC_DATA_GUID
EFI_PARTITIONS
    else
        sfdisk --label gpt "$OUTPUT" <<EFI_PARTITIONS
2 : start=2048, size=62MiB, type=$PARTITION_SYSTEM_GUID
1 : name=eoslive, type=$PARTITION_BASIC_DATA_GUID
EFI_PARTITIONS
    fi
    udevadm settle
    partprobe "$OUTPUT"
    PARTMAP=$(sfdisk --dump "$OUTPUT")
    DEVICE_IMAGES=$(find_by_type "$PARTMAP" "$PARTITION_BASIC_DATA_GUID")
    DEVICE_EFI=$(find_by_type "$PARTMAP" "$PARTITION_SYSTEM_GUID")
    if [ "$BIOS" ]; then
        DEVICE_BIOS=$(find_by_type "$PARTMAP" "$PARTITION_BIOS_BOOT_GUID")
    fi
fi

# Give udev a chance to notice the new partitions
udevadm settle

# Below here we start mounting stuff, so register a cleanup function
trap cleanup EXIT

if [ ! "$MBR" ]; then
    mkfs.vfat -n efi "${DEVICE_EFI}"
    partprobe "$OUTPUT"
    udevadm settle
    mkdir -p "$ESP_MOUNTPOINT"
    mount "$DEVICE_EFI" "$ESP_MOUNTPOINT"
    DIR_EFI="$ESP_MOUNTPOINT"
fi

$MKFS_IMAGES $MKFS_ARGS eoslive "${DEVICE_IMAGES}"
partprobe "$OUTPUT"
udevadm settle
mkdir -p "$EOSLIVE_MOUNTPOINT"
mount "$DEVICE_IMAGES" "$EOSLIVE_MOUNTPOINT"

echo
echo "Copying boot files"

DIR_IMAGES_ENDLESS="$EOSLIVE_MOUNTPOINT/endless"
mkdir "$DIR_IMAGES_ENDLESS"
unzip -q -d "${DIR_IMAGES_ENDLESS}" "${BOOT_ZIP}" "grub/*"

if [ -n "$DIR_EFI" ]; then
    unzip -q -d "${DIR_EFI}" "${BOOT_ZIP}" "EFI/*"
fi

if [ ! "$WRITABLE" ]; then
    cp "${BOOT_ZIP}" "${BOOT_ZIP}.asc" "${EXTRACTED_SIGNATURE}" \
        "$DIR_IMAGES_ENDLESS/"
    echo "$OS_IMAGE_UNCOMPRESSED_BASENAME" > "${DIR_IMAGES_ENDLESS}/live"

    for EXTRA_DATA_PATH in "${EXTRA_DATA_PATHS[@]}"; do
        echo "Copying '$EXTRA_DATA_PATH'"
        cp -r "$EXTRA_DATA_PATH" "$EOSLIVE_MOUNTPOINT/"
    done
fi

if [ ! "$WRITABLE" ] || [ "$WINDOWS_TOOL_PROVIDED" ]; then
    echo "Copying Windows-specific files"

    cp "$WINDOWS_TOOL" "$EOSLIVE_MOUNTPOINT/"
    WINDOWS_TOOL_BASENAME="$(basename "$WINDOWS_TOOL")"
    sed 's/$/\r/' <<AUTORUN_INF | iconv -f utf-8 -t utf-16 > "$EOSLIVE_MOUNTPOINT/autorun.inf"
[AutoRun]
label=${LABEL}
icon=${WINDOWS_TOOL_BASENAME}
open=${WINDOWS_TOOL_BASENAME}

[Content]
MusicFiles=false
PictureFiles=false
VideoFiles=false
AUTORUN_INF
fi

echo "Copying image files"
if [ "$ISO" ]; then
    # Loop-mount the ISO and copy the endless.squash and  directory
    mkdir -p "$ISO_MOUNTPOINT"
    mount -o ro,loop "$OS_IMAGE" "$ISO_MOUNTPOINT"
    cp "$ISO_MOUNTPOINT/endless/endless.squash" \
       "$ISO_MOUNTPOINT/endless/${OS_IMAGE_UNCOMPRESSED_BASENAME%.img}.squash.asc" \
       "$DIR_IMAGES_ENDLESS"
else
    "$EOS_WRITE_IMAGE" "$OS_IMAGE" "-" > "${DIR_IMAGES_ENDLESS}/endless.img"
fi

if [ "$EXPAND" ]; then
    IMAGE_BYTES=$(du -b "${DIR_IMAGES_ENDLESS}/endless.img" | cut -f1)
    FREE_SPACE_BYTES=$(df -P -B1 "$DIR_IMAGES_ENDLESS" | awk 'NR==2 {print $4}')

    if [ -n "$SIZE" ]; then
        if [ "$SIZE" -lt "$IMAGE_BYTES" ]; then
            echo "Cannot expand the image to $SIZE bytes, minimum size is $IMAGE_BYTES bytes" >&2
            exit 1
        fi

        # Expand the image to the provided size
        EXTRA_BYTES=$((SIZE - IMAGE_BYTES))

        if [ "$EXTRA_BYTES" -gt "$FREE_SPACE_BYTES" ]; then
            echo "Cannot expand the image to $SIZE bytes, only $FREE_SPACE_BYTES bytes available" >&2
            exit 1
        fi
    else
        # Expand the image to take all the space on the drive
        SIZE=$((IMAGE_BYTES + FREE_SPACE_BYTES))
    fi

    if [ $SIZE -gt "$IMAGE_BYTES" ]; then
        echo "Expanding the image from $IMAGE_BYTES to $SIZE bytes"
        truncate --size "$SIZE" "${DIR_IMAGES_ENDLESS}/endless.img"
    fi
fi

if [ "$PERSISTENT" ]; then
    # We always want at least 16MB for log files, regardless of what the user
    # asked for.
    if (( "${PERSISTENT_FREE_SPACE:-0}" < 16 )); then
        PERSISTENT_FREE_SPACE=16
    fi

    FREE_MB=$(df -P -B1M "$DIR_IMAGES_ENDLESS" | awk 'NR==2 {print $4}')
    # df always rounds up, so leave an extra 1MB to compensate.
    PERSISTENT_SIZE_MB=$(( FREE_MB - (PERSISTENT_FREE_SPACE + 1) ))


    if (( PERSISTENT_SIZE_MB < 1024 )); then
        echo "Not enough free space to create persistent storage" >&2
        cleanup
        exit 1
    fi

    echo "Creating ${PERSISTENT_SIZE_MB}M bytes persistent storage file."
    echo "This will not be fast."

    echo "endless_live_storage_marker" > "${DIR_IMAGES_ENDLESS}/persistent.img"
    truncate -s ${PERSISTENT_SIZE_MB}M "${DIR_IMAGES_ENDLESS}/persistent.img"
fi

echo
echo "Finalizing"

cleanup

if [ "$MBR" ]; then
    # Bootstrap code, which will jump to sector 1
    unzip -q -p "${BOOT_ZIP}" "ntfs/boot.img" | dd of="${OUTPUT}" bs=446 count=1
    udevadm settle  # udev recreates device files after any write to the device
    # The rest of GRUB goes after the MBR and before the first partition.
    unzip -q -p "${BOOT_ZIP}" "ntfs/core.img" | dd of="${OUTPUT}" bs=512 seek=1
else
    # Bootstrap code, built to jump to the offset of BIOS boot partition
    unzip -q -p "${BOOT_ZIP}" "live/boot.img" | dd of="${OUTPUT}" bs=446 count=1
    udevadm settle  # udev recreates device files after any write to the device
    # The rest of GRUB goes into a dedicated BIOS-boot partition
    if [ "$BIOS" ]; then
        unzip -q -p "${BOOT_ZIP}" "live/core.img" | dd of="${DEVICE_BIOS}" bs=512
    fi
fi

udevadm settle
sync
