#!/usr/bin/env bash
#
# Copyright (C) 2024 Laurent Jourden <laurent85@enarel.fr>
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
# Make a live usb drive with persistence
# using an Archuseriso iso image.

set -e -u

appname="${0##*/}"
bootmodes=()
bcachefs_options=('--compression=lz4')
btrfs_compress=""
btrfs_subvol_home=""
btrfs_subvol_root=""
copytoram_size=""
cow_files_settings=()
cow_flags="defaults"
cow_label=""
cow_size=""
encryption="no"
esp_files_settings=()
esp_label=""
esp_size=""
ext4_journal="yes"
img_label=""
install_dir=""
iso_label=""
iso_file=""
iso_ize=0
medium_version=""
partition_table="gpt"
raw_write="no"
rootfs="ext4"
usb_device=""
ia32_uefi_default_bootloader=""
x64_uefi_default_bootloader=""
WD="${PWD}"
workdir="$(mktemp -u auiwork.XXXXXXXX)"
valid_medium_version="v13"

_usage() {
    echo
    echo "${appname}: make a bootable usb drive with persistence."
    echo
    echo 'Synopsis:'
    echo "${appname} [options] <archuseriso iso image> <usb device>"
    echo
    echo 'Help:'
    echo "${appname} --help"
    exit "${1}"
}

_help() {
    echo
    echo "${appname}: make a bootable usb drive with persistence."
    echo
    echo 'Synopsis:'
    echo "${appname} [options] <archuseriso iso image> <usb device>"
    echo
    echo 'Options:'
    echo '-h, --help                  Command line help'
    echo '--bcachefs-compress=algo    Bcachefs compression algorithm: gzip, lz4 (default), none, zstd'
    echo '--btrfs-compress=algo       Btrfs compression algorithm: lzo (default), none, zlib, zstd'
    echo '--encrypt                   Disk encryption, encrypt the persistent partition'
    echo '--ext4-no-journal           Disable ext4 journal'
    echo '--gpt                       GPT partition table layout (default)'
    echo '--mbr                       MBR partition table layout'
    echo '--raw-write                 Raw write to USB drive (dd like mode), no persistence'
    echo '--rootfs=fstype             Persistent filesystem type: bcachefs, btrfs, ext4 (default), f2fs'
    echo '--size-esp=integer[g|G]     EFI System Partition size in GiB (ESP)'
    echo '--size-cow=integer[g|G]     Copy on write partition size in GiB (persistence)'
    echo
    echo 'Example:'
    echo "sudo ${appname} aui-xfce-linux_5_7_10-0724-x64.iso /dev/sdc"
    echo
    echo 'Custom partitioning of a 64 GiB USB flash drive, unallocated free space left for other usage:'
    echo "sudo ${appname} --size-esp=1G --size-cow=40G aui-xfce-linux_5_7_10-i3-0724-x64.iso /dev/sdc"
    echo
    exit "${1}"
}

_msg_info() {
    local _msg="${1}"
    printf '[%s] INFO: %s\n' "${appname}" "${_msg}"
}

_cleanup() {
    local _workdirs
    _workdirs=('overlay/boot' 'overlay' 'squashfs' 'home' 'iso' 'usbro' 'usbesp' 'usbrw')
    for _workdir in ${_workdirs[*]}; do
        if mountpoint -q -- "${WD}/${workdir}/${_workdir}"; then
            umount -- "${WD}/${workdir}/${_workdir}"
        fi
    done
    for _workdir in ${_workdirs[*]}; do
        if [[ -d "${WD}/${workdir}/${_workdir}" ]]; then
            rmdir -- "${WD}/${workdir}/${_workdir}"
        fi
    done
    if [[ -f "${WD}/${workdir}/grub-embed.cfg" ]]; then
        rm -- "${WD}/${workdir}/grub-embed.cfg"
    fi
    if [[ -d "${WD}/${workdir}" ]]; then
        rmdir -- "${WD}/${workdir}"
    fi
    if [[ -e "/dev/mapper/${crypt_mapper:-auicrypt}" ]]; then
        cryptsetup close "${crypt_mapper:-auicrypt}"
    fi
}

_unmount() {
    echo
    _msg_info "Unmount working directories, may take some time"
    _cleanup
    _msg_info "Done!"
}

_luks_format() {
    echo
    _msg_info "Disk encryption setup, type in a passphrase of your choice"
    if ! cryptsetup --type=luks2 --label="${crypt_label:-AUI_LUKS}" --uuid="${crypt_uuid:=$(uuidgen)}" -q luksFormat -- "${usb_device}3" > /dev/null; then
        echo 'Disk encryption setup failed, exiting!'
        exit 1
    fi
    _msg_info "Done!"
    echo
    _msg_info "Type in your passphrase to unlock partition"
    if ! cryptsetup -- open "${usb_device}3" "${crypt_mapper:-auicrypt}"; then
        echo 'Error: Could not unlock partition! Exiting.'
        exit 1
    fi
    _msg_info "Done!"
    echo
    cow_device="/dev/mapper/${crypt_mapper:-auicrypt}"
}

_encryption_setup_mount_overlay() {
    echo
    _msg_info "Disk encryption setup"
    mkdir -p -- "${WD}/${workdir}/"{overlay,squashfs}
    mount -o ro -- "${WD}/${workdir}/iso/${install_dir}/x86_64/airootfs.sfs" "${WD}/${workdir}/squashfs"
    mkdir -p -- "${WD}/${workdir}/usbrw/persistent_${img_label}/x86_64/workdir"
    mount -t overlay airootper -o lowerdir="${WD}/${workdir}/squashfs",upperdir="${WD}/${workdir}/usbrw/persistent_${img_label}/x86_64/upperdir",workdir="${WD}/${workdir}/usbrw/persistent_${img_label}/x86_64/workdir" -- "${WD}/${workdir}/overlay"
    mount -o bind -- "${WD}/${workdir}/usbesp" "${WD}/${workdir}/overlay/boot"
}

_encryption_setup_unmount_overlay() {
    umount -- "${WD}/${workdir}/overlay/boot" "${WD}/${workdir}/overlay" "/${WD}/${workdir}/squashfs"
    sleep 3
    rmdir -- "${WD}/${workdir}/overlay" "/${WD}/${workdir}/squashfs"
}

_luks_encryption_setup() {
    _encryption_setup_mount_overlay
    sed -i -- "s|block|& encrypt|" "${WD}/${workdir}/overlay/etc/mkinitcpio.conf"
    _msg_info "initramfs update"
    arch-chroot -- "${WD}/${workdir}/overlay" mkinitcpio -P &> /dev/null
    _msg_info "Done!"
    if [[ "${bootmodes[*]}" =~ 'grub.esp' ]]; then
        sed -i -- "s|: persistence|& \& encryption|;
                   s|overlay|cryptdevice=UUID=${crypt_uuid}:${crypt_mapper:-auicrypt} &|" \
                  "${WD}/${workdir}/overlay/boot/grub/grub.cfg"
    fi
    if [[ "${bootmodes[*]}" =~ 'uefi-x64.systemd-boot.esp' ]]; then
        sed -i -- "s|persistence$|& \& encryption|;
                   s|overlay|cryptdevice=UUID=${crypt_uuid}:${crypt_mapper:-auicrypt} &|" \
                  "${WD}/${workdir}/overlay/boot/loader/entries/"*.conf
    fi
    if [[ "${bootmodes[*]}" =~ 'bios.syslinux.mbr' ]]; then
        sed -i -- "s|persistence$|& \& encryption|;
                   s|overlay|cryptdevice=UUID=${crypt_uuid}:${crypt_mapper:-auicrypt} &|" \
                  "${WD}/${workdir}/overlay/boot/syslinux/archiso_sys-linux.cfg"
    fi
    _encryption_setup_unmount_overlay
    _msg_info "Done!"
}

_check() {
    if [[ $# -ne 2 ]]; then
        echo 'Error: Invalid arguments!'
        _usage 1
    fi
    if [[ ${EUID} -ne 0 ]]; then
        echo 'This script must be run as root!'
        exit 1
    fi
    iso_file="${1}"
    usb_device="${2}"
    # command arguments reverse order compatibility
    if [[ $(file -- "${usb_device}" 2> /dev/null) =~ 'MBR boot sector' ]]; then
        # reverse order compatibility
        iso_file="${2}"
        usb_device="${1}"
    fi
    if [[ ! $(LC_ALL=C stat -c %F -- "${usb_device}" 2> /dev/null) =~ 'block special file' ]]; then
        echo "Error: ${usb_device} is not a block device!"
        _usage 1
    fi
    case "$(LC_ALL=C stat -c %t -- "${usb_device}" 2> /dev/null)" in
        '8')
            # SCSI disk devices (0-15)
            ;;
        '65')
            # SCSI disk devices (16-31)
            ;;
        '66')
            # SCSI disk devices (32-47)
            ;;
        '67')
            # SCSI disk devices (48-63)
            ;;
        '68')
            # SCSI disk devices (64-79)
            ;;
        '69')
            # SCSI disk devices (80-95)
            ;;
        '70')
            # SCSI disk devices (96-111)
            ;;
        '71')
            # SCSI disk devices (112-127)
            ;;
        '128')
            # SCSI disk devices (128-143)
            ;;
        '129')
            # SCSI disk devices (144-159)
            ;;
        '130')
            # SCSI disk devices (160-175)
            ;;
        '131')
            # SCSI disk devices (176-191)
            ;;
        '132')
            # SCSI disk devices (192-207)
            ;;
        '133')
            # SCSI disk devices (208-223)
            ;;
        '134')
            # SCSI disk devices (234-239)
            ;;
        '135')
            # SCSI disk devices (240-255)
            ;;
        *)
            echo "Error: ${usb_device} is not a scsi disk device!"
            _usage 1
    esac
    if ! [[ $(lsblk -dnro rm      -- "${usb_device}" 2> /dev/null) -eq 1 || \
            $(lsblk -dnro hotplug -- "${usb_device}" 2> /dev/null) -eq 1      ]]; then
        echo "Error: ${usb_device} is not a removable block device!"
        _usage 1
    fi
    if [[ ! "$(lsblk -dnro tran -- "${usb_device}" 2> /dev/null)" == 'usb' ]]; then
        echo "Error: ${usb_device} is not a usb device!"
        _usage 1
    fi
    if grep -q -- "${usb_device}" /proc/self/mountinfo > /dev/null; then
        echo "Error: ${usb_device} appears in active mounts, unmount drive partitions before proceeding!"
        exit 1
    fi
    if [[ ! -f "${iso_file}" ]]; then
        echo "file ${iso_file} not found!"
        _usage 1
    fi
    if [[ ! $(file -- "${iso_file}" 2> /dev/null) =~ 'MBR boot sector' ]]; then
        echo "Error: ${iso_file} is not an iso image!"
        _usage 1
    fi
    if [[ "${raw_write}" == "yes" ]]; then
        echo
        echo 'raw_write option set. Ignoring other options!'
        cow_size=""
        encryption="no"
        esp_size=""
    fi
    # Set efi boot partition size in MiB
    if [[ -n "${esp_size}" ]]; then
        if ! [[ "${esp_size}" =~ ^[1-9][0-9]?+$ ]]; then
            echo "FAT partition size error: Invalid size argument (GiB): ${esp_size}"
            _usage 1
        fi
        esp_size=$(( esp_size * 1024 ))
    else
        esp_size=512
    fi
    # Set persistent partition size in MiB, free space left by default
    if [[ -n "${cow_size}" ]]; then
        if ! [[ "${cow_size}" =~ ^[1-9][0-9]?+$ ]]; then
            echo "Persistent partition size error: Invalid size argument (GiB): ${cow_size}"
            _usage 1
        fi
        cow_size=$(( cow_size * 1024 ))
    fi
    if [[ "${rootfs}" == "bcachefs" ]]; then
        if ! pacman -Q bcachefs-tools &> /dev/null; then
            echo 'bcachefs-tools package not installed, aborting!'
            exit 0
        fi
    fi
    if [[ "${rootfs}" == "btrfs" ]]; then
        if ! pacman -Q btrfs-progs &> /dev/null; then
            echo 'btrfs-progs package not installed, aborting!'
            exit 0
        fi
    fi
    if [[ "${rootfs}" == "f2fs" ]]; then
        if ! pacman -Q f2fs-tools &> /dev/null; then
            echo 'f2fs-tools package not installed, aborting!'
            exit 0
        fi
    fi
}

_init() {
    local _esp_size _cow_size
    drivesize=$(blockdev --getsize64 "${usb_device}")
    FREESPACE=1073741824 # 1 GiB
    iso_size=$(stat -c %s -- "${iso_file}") # Bytes
    # USB Medium type isohybrid / FAT
    if [[ "${raw_write}" == "yes" ]]; then
        esp_size=0
        FREESPACE=0
    fi
    # Check usb drive capacity
    if [[ ! ${drivesize} -gt $(( iso_size + FREESPACE )) ]]; then
        echo 'Storage capacity error!'
        exit 1
    fi
    # check partitions size don't exceed drive's capacity
    _esp_size=$(( esp_size * 1048576 )) # Bytes
    if [[ -n "${cow_size}" ]]; then
        _cow_size=$(( cow_size * 1048576 )) # Bytes
    else
        _cow_size=0
    fi
    if [[ ! ${drivesize} -gt $(( iso_size + _esp_size + _cow_size )) ]]; then
        echo "Size settings error: exceeds drive storage capacity!"
        exit 1
    fi
    # ISO size in MiB
    iso_size=$(( iso_size / 1048576 ))
}

_confirm_write() {
    # Confim write
    echo
    _msg_info "ISO file:   ${iso_file}"
    _msg_info "USB device: ${usb_device}"
    echo
    read -r -n1 -p "Confirm write to $(lsblk -dno model,size -- "${usb_device}") (N/y)? "
    echo
    if [[ ! "${REPLY}" =~ ^[Yy]$ ]]; then
        echo 'Operation canceled by user!'
        exit 0
    fi
}

_raw_write() {
    echo
    _msg_info "Raw copy to usb drive (dd like mode), no persistence"
    if ! cp -v -- "${iso_file}" "${usb_device}"; then
        echo 'Write failed!'
        exit 1
    fi
    _msg_info "Done!"
}

_usb_prepare() {
    # Check & prepare working directory
    for mountpoint in "${workdir}" "${workdir}/"{iso,overlay,squashfs} "${workdir}/usb"{ro,esp,rw}; do
        if findmnt -nr -- "${mountpoint}" > /dev/null; then
            echo "Error: ${mountpoint} appears in active mounts, unmount before proceeding!"
            exit 1
        fi
    done
    for mountpoint in "${workdir}" "${workdir}/"{iso,overlay,squashfs} "${workdir}/usb"{ro,esp,rw}; do
        if [[ -e "${WD}/${mountpoint}" ]]; then
            echo "Error: ${mountpoint} exists in working directory! Delete or rename before proceeding!"
            exit 1
        fi
    done
    if [[ -e "/dev/mapper/${crypt_mapper:-auicrypt}" ]]; then
        echo "Error: cryptsetup mapping /dev/mapper/${crypt_mapper:-auicrypt} exists! cannot proceed."
        exit 1
    fi
    mkdir -p -- "${WD}/${workdir}/iso" "${WD}/${workdir}/usb"{ro,esp,rw}

    # Mount iso
    echo
    _msg_info 'Mount iso'
    _msg_info "mountpoint: ${WD}/${workdir}/iso"
    if ! mount -o ro -- "${iso_file}" "${WD}/${workdir}/iso"; then
        echo "Error: mount iso failed!"
        rmdir -- "${WD}/${workdir}/usb"{ro,esp,rw} "${WD}/${workdir}/iso" "${WD}/${workdir}"
        exit 1
    fi

    # check iso
    if [[ ! -e "${WD}/${workdir}/iso/.drive/mediumdef.sh" ]]; then
        echo "Error: ${iso_file} is not a compatible iso image!"
        echo
        _usage 1
    fi

    # source medium def variables
    source "${WD}/${workdir}/iso/.drive/mediumdef.sh"
    if [[ ! "${medium_version}" == "${valid_medium_version}" ]]; then
        echo -e "Error: ${medium_version} gen iso image, gen version mismatch. ${appname} requires gen ${valid_medium_version}!"
        echo
        _usage 1
    fi

    if [[ -f "${WD}/${workdir}/iso/${install_dir}/${arch}/airootfs.sfs" ]]; then
        iso_size=$(( $(du -bs "${WD}/${workdir}/iso/${install_dir}/${arch}" | awk '{print $1}') / 1048576 * 108 / 100 ))
        copytoram_size=$(( $(stat -c %s "${WD}/${workdir}/iso/${install_dir}/${arch}/airootfs.sfs") + 1024 * 64 ))
    fi

    if [[ ! -d "${WD}/${workdir}/iso/EFI/" ]]; then
        echo "Error: missing EFI directory, ${iso_file} is not a compatible iso image!"
        echo
        _usage 1
    fi

    _msg_info "Done!"
}

_partitions() {
    local _part_type_linux _part_type_msft _partition_table_layout

    # GPT default
    if ! [[ "${partition_table}" == 'mbr' ]]; then
        # GPT partition table
        _partition_table_layout="gpt"
        # GPT Partition type : Linux Filesystem
        _part_type_linux="0FC63DAF-8483-4772-8E79-3D69D8477DE4"
        # GPT Partition type : Microsoft basic data
        _part_type_msft="EBD0A0A2-B9E5-4433-87C0-68B6B72699C7"
    else
        # MBR partition table
        _partition_table_layout="dos"
        # MBR Partition type : Linux Filesystem
        _part_type_linux="83"
        # MBR Partition type : Microsoft basic data
        _part_type_msft="07"
    fi

    # Wipe drive
    udevadm lock --device="${usb_device}" -- wipefs --all --force -- "${usb_device}"* > /dev/null
    sleep 3
    partprobe -- "${usb_device}"
    echo

    # New gpt partition table
    _msg_info "Drive partitions"
    _msg_info "New ${partition_table} partition table"
    if ! echo "label: ${_partition_table_layout}" | udevadm lock --device="${usb_device}" -- \
            sfdisk -W always -- "${usb_device}" > /dev/null; then
        echo 'Failed to create new gpt partition table!'
        exit 1
    fi
    sleep 3
    partprobe -- "${usb_device}"

    # partitions
    _msg_info "Partition #1 in MiB: ${iso_size}"
    _msg_info "Partition #2 in MiB: ${esp_size}"
    _msg_info "Partition #3 in MiB: ${cow_size:-All Free Space}"
    if ! echo -e ",${iso_size}M,${_part_type_linux},\n\
                  ,${esp_size}M,${_part_type_msft},\n\
                  ,${cow_size:+${cow_size}M},${_part_type_linux},\n" | \
            udevadm lock --device="${usb_device}" -- sfdisk --append -W always -- "${usb_device}" > /dev/null; then
        echo 'Failed to create partitions!'
        exit 1
    fi
    sleep 3

    partprobe -- "${usb_device}"
    _msg_info "Done!"
}

_format_bcachefs() {
    local 
    _msg_info "partition #3: type Bcachefs, label ${cow_label}"
    if ! udevadm lock --device="${cow_device}" -- \
            mkfs.bcachefs -L "${cow_label}" ${bcachefs_options[*]} "${cow_device}" > /dev/null; then
        echo 'Failed to format partition!'
        exit 1
    fi
    if bcachefs unlock -c "${cow_device}" > /dev/null; then
        _msg_info "Type in your passphrase to unlock partition"
        bcachefs unlock -k session "${cow_device}"
    fi
}

_format_btrfs() {
    _msg_info "partition #3: type Btrfs, label ${cow_label}"
    if ! udevadm lock --device="${cow_device}" -- \
            mkfs.btrfs -L "${cow_label}" -- "${cow_device}" > /dev/null; then
        echo 'Failed to format partition!'
        exit 1
    fi
}

_format_f2fs() {
    _msg_info "partition #3: type F2FS, label ${cow_label}"
    if ! udevadm lock --device="${cow_device}" -- \
            mkfs.f2fs -l "${cow_label}" -O encrypt,extra_attr,compression -- "${cow_device}" > /dev/null; then
        echo 'Failed to format partition!'
        exit 1
    fi
}

_format_ext4() {
    _msg_info "partition #3: type Ext4, label ${cow_label}"
    if ! udevadm lock --device="${cow_device}" -- \
            mkfs.ext4 -L "${cow_label}" -O encrypt -q -- "${cow_device}" > /dev/null; then
        echo 'Failed to format partition!'
        exit 1
    fi
    # disable ext4 journal
    if [[ "${ext4_journal}" == "no" ]]; then
        _msg_info "Ext4 partitions: disable journal"
        tune2fs -O '^has_journal' -- "${usb_device}1" &> /dev/null
        tune2fs -O '^has_journal' -- "${cow_device}" &> /dev/null
    fi
}

_format() {
    echo
    _msg_info "Format partitions"
    _msg_info "partition #1: type Ext4, label ${img_label}"
    if ! udevadm lock --device="${usb_device}1" -- \
            mkfs.ext4 -L "${img_label}" -O encrypt -m 0 -q -- "${usb_device}1" > /dev/null; then
        echo 'Failed to format partition!'
        exit 1
    fi
    _msg_info "partition #2: type FAT, label ${esp_label}"
    if ! udevadm lock --device="${usb_device}2" -- \
            mkfs.fat -F32 -n "${esp_label}" -- "${usb_device}2" &> /dev/null; then
        echo 'Failed to format partition!'
        exit 1
    fi
    cow_device="${usb_device}3"
    if [[ "${encryption}" == "yes" ]]; then
        case "${rootfs}" in
            'btrfs'|'f2fs'|'ext4')
                _luks_format ;;
            'bcachefs')
                bcachefs_options+=('--encrypted');;
        esac
    fi
    # Format persistent partition, filesystem type
    _format_"${rootfs}"
    _msg_info "Done!"
}

_btrfs_subvolumes() {
    # rootfs and home subvolumes creation for persistence
    btrfs subvolume create -- "${WD}/${workdir}/usbrw/${btrfs_subvol_root}" > /dev/null
    btrfs subvolume create -- "${WD}/${workdir}/usbrw/${btrfs_subvol_home}" > /dev/null
}

_mount() {
    # Mount usb device
    echo
    _msg_info "Mount partitions"
    _msg_info "device: ${usb_device}1, mountpoint: ${WD}/${workdir}/usbro"
    mount -- "${usb_device}1" "${WD}/${workdir}/usbro"
    install -m 0755 -d -- "${WD}/${workdir}/usbro/${install_dir}"
    _msg_info "device: ${usb_device}2, mountpoint: ${WD}/${workdir}/usbesp"
    mount -- "${usb_device}2" "${WD}/${workdir}/usbesp"
    _msg_info "device: ${cow_device}, mountpoint: ${WD}/${workdir}/usbrw"
    mount -- "${cow_device}" "${WD}/${workdir}/usbrw"
    if [[ "${rootfs}" == "btrfs" ]]; then
        _btrfs_subvolumes
        umount "${WD}/${workdir}/usbrw"
        sleep 3
        # Mount rootfs subvolume
        mount -o "${cow_flags}" -- "${cow_device}" "${WD}/${workdir}/usbrw"
    fi
    _msg_info "Done!"
}

_copy() {
    # Copy iso to partition #1
    echo
    _msg_info "Copy live image to ${WD}/${workdir}/usbro"
    cp -r -- "${WD}/${workdir}/iso/${install_dir}/${arch}" "${WD}/${workdir}/usbro/${install_dir}"
    _msg_info "Done!"
}

# Make GRUB standalone EFI binary
_make_grub_efi_binary() {
    local _grub_efi_arch="${1}"
    local _grub_efi_name="${2}"
    local _grubmodules=()

    IFS='' read -r -d '' grubembedcfg <<'EOF' || true
regexp --set=1:archiso_bootdevice '^\(([^)]+)\)\/?[Ee][Ff][Ii]\/?' "$cmdpath"
if ! [ -d "$cmdpath" ]; then
     # On some firmware, GRUB has a wrong cmdpath when booted from an optical disc.
     # https://gitlab.archlinux.org/archlinux/archiso/-/issues/183
     if regexp '^\(([^)]+)\)\/?[Ee][Ff][Ii]\/[Bb][Oo][Oo][Tt]\/?$' "$cmdpath"; then
         cmdpath="${archiso_bootdevice}/EFI/BOOT"
     fi
     if regexp '^\(([^)]+)\)\/?[Ee][Ff][Ii]\/[Gg][Rr][Uu][Bb]\/?$' "$cmdpath"; then
         cmdpath="${archiso_bootdevice}/EFI/grub"
     fi
fi
configfile "(${archiso_bootdevice})/grub/grub.cfg"
EOF

    printf '%s\n' "$grubembedcfg" > "${WD}/${workdir}/grub-embed.cfg"

    # Create EFI binary
    # Module list from https://bugs.archlinux.org/task/71382#comment202911
    _grubmodules=(all_video at_keyboard boot btrfs cat chain configfile echo efifwsetup efinet ext2 f2fs fat font  \
                  gfxmenu gfxterm gzio halt hfsplus iso9660 jpeg keylayouts linux loadenv loopback lsefi lsefimmap \
                  minicmd normal part_apple part_gpt part_msdos png read reboot regexp search search_fs_file       \
                  search_fs_uuid search_label serial sleep tpm usb usbserial_common usbserial_ftdi                 \
                  usbserial_pl2303 usbserial_usbdebug video xfs zstd)
     grub-mkstandalone -O "${_grub_efi_arch}"                                  \
                       --modules="${_grubmodules[*]}"                          \
                       --locales="en@quot"                                     \
                       --themes=""                                             \
                       --sbat=/usr/share/grub/sbat.csv                         \
                       --disable-shim-lock                                     \
                       -o "${WD}/${workdir}/usbesp/EFI/grub/${_grub_efi_name}" \
                       "boot/grub/grub.cfg=${WD}/${workdir}/grub-embed.cfg"
}

_esp_bootloader_cfg_uefi-ia32.grub.esp() {
    # Copy new GRUB cfg
    cp -rT -- "${WD}/${workdir}/iso/.drive/liveusb/grub" "${WD}/${workdir}/usbesp/grub"

    _make_grub_efi_binary "i386-efi" "grubia32.efi"
}

_esp_bootloader_cfg_uefi-x64.grub.esp() {
    # Copy new GRUB cfg
    cp -rT -- "${WD}/${workdir}/iso/.drive/liveusb/grub" "${WD}/${workdir}/usbesp/grub"

    _make_grub_efi_binary "x86_64-efi" "grubx64.efi"
}

_esp_bootloader_cfg_uefi-x64.refind.esp() {
    # Copy new rEFInd config
    cp -- "${WD}/${workdir}/iso/.drive/liveusb/refind/refind.conf" "${WD}/${workdir}/usbesp/EFI/refind/"
    esp_files_settings+=('EFI/refind/refind.conf')
}

_esp_bootloader_cfg_uefi-x64.systemd-boot.esp() {
    # Copy new systemd-boot config
    cp -r -- "${WD}/${workdir}/iso/.drive/liveusb/loader" "${WD}/${workdir}/usbesp/"
}

_esp_default_bootloader_uefi-ia32.grub.esp() {
    # Move ia32 EFI binary to ESP /EFI/BOOT
    mv -- "${WD}/${workdir}/usbesp/EFI/grub/grubia32.efi" "${WD}/${workdir}/usbesp/EFI/BOOT/BOOTIA32.EFI"
}

_esp_default_bootloader_uefi-x64.grub.esp() {
    # Move EFI binary to ESP /EFI/BOOT
    mv -- "${WD}/${workdir}/usbesp/EFI/grub/grubx64.efi" "${WD}/${workdir}/usbesp/EFI/BOOT/BOOTx64.EFI"
}

_esp_default_bootloader_uefi-x64.refind.esp() {
    mv -- "${WD}/${workdir}/usbesp/EFI/refind/refind.conf" "${WD}/${workdir}/usbesp/EFI/BOOT/"
}

# Make ESP for usb device
_esp() {
    echo
    _msg_info "ESP setup"
    _msg_info "Copy to ${WD}/${workdir}/usbesp"
    cp -LrT -- "${WD}/${workdir}/iso/.drive/liveusb/esp" "${WD}/${workdir}/usbesp"
    cp -r -- "${WD}/${workdir}/iso/.drive/liveusb/syslinux" "${WD}/${workdir}/usbesp/"

    # EFI boot loader setup
    for bootmode in "${bootmodes[@]}"; do
        if typeset -f "_esp_bootloader_cfg_${bootmode}" &> /dev/null; then
            _esp_bootloader_cfg_${bootmode}
        fi
    done

    for mediumfile in "${esp_files_settings[@]}"; do
        if [[ -f "${WD}/${workdir}/usbesp/${mediumfile}" ]]; then
            sed -i -- "s|%ARCHISO_LABEL%|${iso_label}|g;
                       s|%COW_LABEL%|${cow_label}|g;
                       s|%COW_FLAGS%|${cow_flags}|g;
                       s|%ESP_LABEL%|${esp_label}|g;
                       s|%IMG_LABEL%|${img_label}|g;
                       s|--label ${iso_label}$|--label ${esp_label}|g;
                       s|=${iso_label}|=${img_label}|g" \
                      "${WD}/${workdir}/usbesp/${mediumfile}"
            if [[ -n "${copytoram_size}" ]]; then
                sed -i -- "s|copytoram|copytoram_size=${copytoram_size} &|" "${WD}/${workdir}/usbesp/${mediumfile}"
            fi
        fi
    done

    # Set default EFI boot loader
    if typeset -f "_esp_default_bootloader_${ia32_uefi_default_bootloader}" &> /dev/null; then
        _esp_default_bootloader_${ia32_uefi_default_bootloader}
    fi
    if typeset -f "_esp_default_bootloader_${x64_uefi_default_bootloader}" &> /dev/null; then
        _esp_default_bootloader_${x64_uefi_default_bootloader}
    fi
    _msg_info "Done!"
}

_persistence_btrfs_home() {
    _msg_info "Copy home files to btrfs home subvolume"
    # copy squashfs home to btrfs home subvolume
    mkdir -p -- "${WD}/${workdir}/"{squashfs,home}
    mount -o ro -- "${WD}/${workdir}/iso/${install_dir}/x86_64/airootfs.sfs" "${WD}/${workdir}/squashfs"
    mount -o subvol="${btrfs_subvol_home}" -- "${cow_device}" "${WD}/${workdir}/home"
    cp -aT -- "${WD}/${workdir}/squashfs/home/" "${WD}/${workdir}/home/"
    umount "${WD}/${workdir}/"{squashfs,home}
    sleep 3
}

_persistence() {
    echo
    _msg_info "Persistence setup"
    _msg_info "Copy persistence config to ${WD}/${workdir}/usbrw"
    cp -a -- "${WD}/${workdir}/iso/.drive/liveusb/persistent_${iso_label}" "${WD}/${workdir}/usbrw/persistent_${img_label}"
    for mediumfile in "${cow_files_settings[@]}"; do
        if [[ -f "${WD}/${workdir}/usbrw/${mediumfile}" ]]; then
            sed -i -- "s|%ESP_LABEL%|${esp_label}|g;
                       s|%COW_LABEL%|${cow_label}|g;
                       s|%BTRFS_SUBVOL_HOME%|${btrfs_subvol_home}|g" \
                      "${WD}/${workdir}/usbrw/${mediumfile}"
            if [[ "${rootfs}" != "btrfs" && "${mediumfile}" =~ 'fstab' ]]; then
                # remove unused btrfs line in fstab
                sed -i -- '/btrfs/d' "${WD}/${workdir}/usbrw/${mediumfile}"
            fi
        fi
    done
    if [[ "${rootfs}" == "btrfs" ]]; then
        _persistence_btrfs_home
    fi
    _msg_info "Done!"

    if [[ "${encryption}" == "yes" && ! "${rootfs}" == "bcachefs" ]]; then
        _luks_encryption_setup
    fi
}

_biosbootloader_gpt(){
    if ! dd bs=440 count=1 conv=notrunc if=/usr/lib/syslinux/bios/gptmbr.bin of="${usb_device}" 2> /dev/null; then
        echo 'Bootloader installation failed!'
        exit 1
    fi
    if ! sfdisk --lock --part-attrs "${usb_device}" 2 LegacyBIOSBootable &> /dev/null; then
        echo 'Bootloader installation failed!'
        exit 1
    fi
}

_biosbootloader_mbr(){
    if ! dd bs=440 count=1 conv=notrunc if=/usr/lib/syslinux/bios/mbr.bin of="${usb_device}" 2> /dev/null; then
        echo 'Bootloader installation failed!'
        exit 1
    fi
    if ! sfdisk --lock --activate "${usb_device}" 2 &> /dev/null; then
        echo 'Bootloader installation failed!'
        exit 1
    fi
}

_biosbootloader() {
    echo
    _msg_info "Install boot loader"
    if ! syslinux --directory syslinux --install -- "${usb_device}2" &> /dev/null; then
        echo 'Bootloader installation failed!'
        exit 1
    fi
    if ! [[ "${partition_table}" == 'mbr' ]]; then
        _biosbootloader_gpt
    else
        _biosbootloader_mbr
    fi
    _msg_info "Done!"
}

# arguments
OPTS=$(getopt -o 'h' --long 'btrfs-compress:,encrypt,ext4-no-journal,gpt,help' \
                     --long 'mbr,raw-write,rootfs:,size-esp:,size-cow:'        \
                     -n "${appname}" -- "$@")
[[ $? -eq 0 ]] || _usage 1
eval set -- "${OPTS}"
unset OPTS
[[ $# -eq 1 ]] && _usage 1

while true; do
    case "$1" in
        '-h'|'--help')
            _help 0 ;;
        '--bcachefs-compress')
            if [[ "${2}" == 'gzip' || \
                  "${2}" == 'lz4'  || \
                  "${2}" == 'none' || \
                  "${2}" == 'zstd' ]]; then
                 bcachefs_options+=("--compression=${2}")
            fi
            shift 2 ;;
        '--btrfs-compress')
            if [[ "${2}" == 'lzo'  || \
                  "${2}" == 'no'   || \
                  "${2}" == 'none' || \
                  "${2}" == 'zlib' || \
                  "${2}" == 'zstd' ]]; then
                btrfs_compress="${2}"
            fi
            shift 2 ;;
        '--encrypt')
            encryption="yes"
            shift ;;
        '--ext4-no-journal')
            ext4_journal="no"
            shift ;;
        '--gpt')
            partition_table="gpt"
            shift;;
        '--mbr')
            partition_table="mbr"
            shift;;
        '--raw-write')
            raw_write="yes"
            shift ;;
        '--rootfs')
            case "${2}" in
                'btrfs')
                    rootfs="${2}"
                    btrfs_subvol_home="/home"
                    btrfs_subvol_root="/rootfs"
                    cow_flags="subvol=${btrfs_subvol_root},compress=${btrfs_compress:-lzo}" ;;
                'bcachefs'|'f2fs'|'ext4')
                    rootfs="${2}" ;;
                *)
                    echo "${2} is not a valid filesystem type option"
                    _usage 1 ;;
            esac
            shift 2 ;;
        '--size-esp')
            esp_size="${2/[gG]}"
            shift 2 ;;
        '--size-cow')
            cow_size="${2/[gG]}"
            shift 2 ;;
        '--')
            shift
            break ;;
    esac
done

trap _cleanup EXIT
_check "$@"
_init
_confirm_write
if [[ "${raw_write}" == "yes" ]]; then
    _raw_write
else
    _usb_prepare
    _partitions
    _format
    _mount
    _copy
    _esp
    _persistence
    _unmount
    _biosbootloader
fi
echo
_msg_info "Success!"
echo
echo "IMPORTANT NOTICE REGARDING PERSISTENCE !!!"
echo "- The root account and the live user account have no password !!!"
echo "- You must set a root and a user password by yourself after the first boot !!!"

# vim:ts=4:sw=4:et:
