mkdiskimage

#!/bin/bash

readonly UDC_DOWNLOAD_URL="https://igeldownloadprod-bydsc8hmbsaegvdy.a01.azurefd.net/files/IGEL_UNIVERSAL_DESKTOP_CONVERTER/UDC3_Stick_V10/UDC3_10.06.220.zip"
readonly DEPENDENCIES=(
    "wget"
    "unzip"
    "osirrox"
    "xorriso"
    "mksquashfs"
    "unsquashfs"
    "fdisk"
    "igelfs-cli"
)

# Log message to stderr
#   $1: message content
#   $2: status level, optional (default: "INFO")
function log() {
    local message="$1"
    local status="INFO"
    if [[ -n "$2" ]]; then
        status="$2"
    fi
    echo "[${status}] ${message}" >&2
}

# Check dependencies are available
function check_dependencies() {
    for prog in "${DEPENDENCIES[@]}"; do
        if ! command -v "${prog}" > /dev/null; then
            log "Dependency '${prog}' not available" "ERROR"
            exit 1
        fi
    done
}

# Download and extract UDC archive
#   $1: path to extract ISO
function download_udc() {
    local iso_extract_path="$1"

    local udc_zip_path
    local udc_extract_path
    local udc_iso_path
    udc_zip_path="$(mktemp)"
    udc_extract_path="$(mktemp --directory)"

    log "Downloading UDC archive from '${UDC_DOWNLOAD_URL}' to '${udc_zip_path}'..."
    wget -O "${udc_zip_path}" "${UDC_DOWNLOAD_URL}"
    log "Extracting ISO from archive '${udc_zip_path}' to '${udc_extract_path}'..."
    unzip -j "${udc_zip_path}" "preparestick/udc*.iso" -d "${udc_extract_path}"
    udc_iso_path=("${udc_extract_path}"/udc*.iso)
    log "Extracting ISO '${udc_iso_path[0]}' to '${iso_extract_path}'..."
    osirrox \
        -indev "${udc_iso_path[0]}" \
        -extract "/" "${iso_extract_path}"
    rm -rf "${udc_zip_path}" "${udc_extract_path}"
}

# Copy required EFI binaries
#   $1: path to EFI image
#   $2: path to target directory
function copy_efi_files() {
    local boot_image="$1"
    local esp_path="$2"

    local boot_image_mountpoint
    boot_image_mountpoint="$(mktemp --directory)"

    sudo mount -o ro "${boot_image}" "${boot_image_mountpoint}"
    log "Copying EFI binaries from '${boot_image_mountpoint}/efi' to '${esp_path}'..."
    cp -R "${boot_image_mountpoint}/efi" "${esp_path}/"  # EFI files
    sudo umount "${boot_image_mountpoint}"
    rmdir "${boot_image_mountpoint}"
}

# Copy required boot files from ISO, e.g. kernel, GRUB configuration, etc.
#   $1: path to ISO files
#   $2: path to target directory
#   $3: path to overlay directory, optional
function copy_boot_files() {
    local iso_path="$1"
    local esp_path="$2"
    local esp_overlay="$3"

    log "Copying boot files from '${iso_path}' to '${esp_path}'..."
    mkdir -p "${esp_path}/boot/grub"
    cp "${iso_path}/boot/bzImage" "${esp_path}/boot/"  # Kernel
    cp -R "${iso_path}"/boot/grub/{fonts,igel.conf} "${esp_path}/boot/grub/"  # GRUB
    log "Creating 'igel-ud-converter' file hint for GRUB..."
    touch "${esp_path}/boot/igel-ud-converter"  # Required for GRUB to find device
    if [[ -d "${esp_overlay}" ]]; then
        log "Patching EFI partition from '${esp_overlay}' to '${esp_path}'..."
        cp -R "${esp_overlay}/." "${esp_path}/"
    fi
}

# Prepare ISO directory
#   $1: path to ISO files
#   $2: boot ID, optional (default: IGEL_UDC_TO)
function prepare_iso() {
    local iso_path="$1"
    local boot_id="IGEL_UDC_TO"
    if [[ -n "$2" ]]; then
        boot_id="$2"
    fi

    log "Removing unnecessary files..."
    find "${iso_path}" -not -name "ddimage.bin" -delete
    log "Creating boot_id '${boot_id}' file hint for initramfs..."
    touch "${iso_path}/.${boot_id}"  # Required for init of initramfs
}

# Modify system partition of ddimage.bin
#   $1: path to ddimage.bin
#   $2: path to overlay directory, optional
#   $3: path to custom SquashFS, optional
function patch_system() {
    local ddimage_path="$1"
    local overlay_path="$2"
    local custom_squashfs="$3"

    local squashfs_extract
    local squashfs_extract_custom
    local squashfs_output
    local squashfs_root
    squashfs_extract="$(mktemp --directory)"
    squashfs_extract_custom="$(mktemp --directory)"
    squashfs_output="$(mktemp)"
    rm "${squashfs_output}"  # Prevent appending to file

    extract_ddimage "${ddimage_path}" "${squashfs_extract}"
    rm "${ddimage_path}"
    squashfs_root="${squashfs_extract}"
    if [[ -f "${custom_squashfs}" ]]; then
        log "Extracting custom SquashFS image to '${squashfs_extract_custom}'..."
        sudo unsquashfs -dest "${squashfs_extract_custom}" "${custom_squashfs}"
        log "Creating required 'igfimage' directory..."
        sudo mkdir "${squashfs_extract_custom}/igfimage"
        log "Copying kernel modules from '${squashfs_extract}/lib/modules' to '${squashfs_extract_custom}'..."
        sudo cp -R "${squashfs_extract}/lib/modules" "${squashfs_extract_custom}/lib/"
        squashfs_root="${squashfs_extract_custom}"
    fi
    if [[ -d "${overlay_path}" ]]; then
        log "Patching system partition from '${overlay_path}' to '${squashfs_root}'..."
        sudo cp -R --remove-destination "${overlay_path}/." "${squashfs_root}/"
    fi
    log "Creating SquashFS image from '${squashfs_root}' to '${squashfs_output}'..."
    sudo mksquashfs "${squashfs_root}" "${squashfs_output}" -comp zstd
    rebuild_ddimage "${ddimage_path}" "${squashfs_output}"
    sudo rm -rf \
        "${squashfs_extract}" \
        "${squashfs_extract_custom}" \
        "${squashfs_output}"
}

# Extract SquashFS contents from ddimage.bin
#   $1: path to ddimage.bin
#   $2: path to extract SquashFS
function extract_ddimage() {
    local ddimage_path="$1"
    local extract_path="$2"

    local ddimage_extract
    ddimage_extract="$(mktemp --directory)"

    log "Extracting system partition from '${ddimage_path}' to '${ddimage_extract}'..."
    igelfs-cli extract --partitions 1 "${ddimage_extract}" "${ddimage_path}"
    log "Extracting SquashFS image '${ddimage_extract}/1/payload.squashfs' to '${extract_path}'..."
    sudo unsquashfs -dest "${extract_path}" "${ddimage_extract}/1/payload.squashfs"
    sudo rm -rf "${ddimage_extract}"
}

# Create new ddimage.bin with specified SquashFS
#   $1: path to ddimage.bin
#   $2: path to SquashFS
function rebuild_ddimage() {
    local ddimage_path="$1"
    local squashfs_path="$2"

    log "Rebuilding ddimage.bin from '${squashfs_path}' to '${ddimage_path}'..."
    igelfs-cli new \
        "$(( $(igelfs-cli info --section-count "${squashfs_path}") + 1 ))" \
        "${ddimage_path}"
    igelfs-cli add "${squashfs_path}" 1 "${ddimage_path}"
}

# Create disk image
#   $1: path to disk image
#   $2: size of image, optional (default: 256M)
function create_disk_image() {
    local disk_path="$1"
    local disk_size="256M"
    if [[ -n "$2" ]]; then
        disk_size="$2"
    fi

    log "Creating disk image of size '${disk_size}' at '${disk_path}'..."
    fallocate -l "${disk_size}" "${disk_path}"
}

# Attach disk image as loop device
#   $1: path to disk image
#   output: path to loop device
function attach_loop_device() {
    local disk_path="$1"

    log "Attaching '${disk_path}' as loop device..."
    sudo losetup --partscan --find --show "${disk_path}"
}

# Detach loop device
#   $1: path to loop device
function detach_loop_device() {
    local device_path="$1"

    log "Detaching loop device '${device_path}'..."
    sudo losetup --detach "${device_path}"
}

# Create ISO image
#   $1: path to ISO directory
#   $2: path to output image
function create_iso_image() {
    local iso_path="$1"
    local iso_output="$2"

    log "Creating ISO image from '${iso_path}' to '${iso_output}'..."
    xorriso -no_rc -as mkisofs -output "${iso_output}" "${iso_path}"
}

# Partition loop device
#   $1: path to disk
#   $2: size of ISO partition
#   $3: size of ESP partition, optional (default: 100% free)
function partition_disk() {
    local disk_path="$1"
    local iso_size="$2"
    local esp_size="$3"

    log "Partitioning disk '${disk_path}'..."
    log " - ISO size: ${iso_size}"
    if [[ -n "${esp_size}" ]]; then
        log " - ESP size: ${esp_size}"
        esp_size="+${esp_size}"
    else
        log " - ESP size: 100% free"
    fi

    (
        echo "g"                 # Create a new empty GPT partition table
        echo "n"                 # Add a new partition
        echo "1"                 # Partition number
        echo                     # First sector
        echo "+${iso_size}"      # Last sector
        echo "n"
        echo "2"
        echo
        echo "${esp_size}"
        echo "t"                 # Partition type
        echo "2"                 # Partition number
        echo "1"                 # EFI System
        echo "w"                 # Write changes
    ) | sudo fdisk "${disk_path}"
}

# ISO Partition
#   $1: path to loop device
#   $2: path to ISO file
function create_iso_partition() {
    local device="$1"
    local iso_path="$2"

    log "Writing '${iso_path}' to '${device}p1'..."
    sudo dd if="${iso_path}" of="${device}p1" status=progress
}

# EFI System Partition
#   $1: path to loop device
#   $2: path to EFI directory
#   $3: filesystem label, optional (default: "ESP")
function create_efi_partition() {
    local device="$1"
    local esp_path="$2"
    local label="ESP"
    if [[ -n "$3" ]]; then
        label="$3"
    fi

    local esp_mountpoint
    esp_mountpoint="$(mktemp --directory)"

    log "Formatting EFI System Partition as FAT32 on '${device}p2' with label '${label}'..."
    sudo mkfs.fat -F32 -n "${label}" "${device}p2"
    sudo mount "${device}p2" "${esp_mountpoint}"
    log "Copying EFI files from '${esp_path}' to '${esp_mountpoint}'..."
    sudo cp -R "${esp_path}/." "${esp_mountpoint}/"
    sudo umount "${esp_mountpoint}"
    rmdir "${esp_mountpoint}"
}

# Parse command-line arguments
#   $@: arguments to parse
function parse_command_line() {
    while getopts "s:l:e:r:h" arg; do
        case "${arg}" in
            s) DISK_SIZE="${OPTARG}" ;;
            l) ESP_LABEL="${OPTARG}" ;;
            e) ESP_OVERLAY="${OPTARG}" ;;
            r) ROOT_OVERLAY="${OPTARG}" ;;
            h | *)
                usage
                exit 1
                ;;
        esac
    done
    DISK_PATH="${*:${OPTIND}:1}"
    CUSTOM_SQUASHFS="${*:${OPTIND}+1:1}"
    if [[ -z "${DISK_PATH}" ]]; then
        usage
        exit 1
    fi
}

# Output usage information
function usage() {
    cat <<- END
	$0: [-s SIZE] [-l LABEL] [-e ESP_OVERLAY] [-r ROOT_OVERLAY] PATH [SQUASHFS]

	-h              display usage information
	-s              size of disk image (default: "256M")
	-l              label of EFI System Partition (default: "ESP")
	-e              directory to copy to EFI System Partition
	-r              directory to copy to SquashFS
	PATH            path to disk image
	SQUASHFS        path to custom SquashFS
	END
}

# Create patched disk image from UDC installation archive
function main() {
    local iso_path
    local esp_path
    local iso_output
    local loop_device
    local iso_size
    local iso_size_mb
    iso_path="$(mktemp --directory)"
    esp_path="$(mktemp --directory)"
    iso_output="$(mktemp)"

    check_dependencies
    parse_command_line "$@"

    download_udc "${iso_path}"

    # EFI System Partition
    copy_efi_files "${iso_path}/igel_efi.img" "${esp_path}"
    copy_boot_files "${iso_path}" "${esp_path}" "${ESP_OVERLAY}"

    # ISO & ddimage.bin
    prepare_iso "${iso_path}"
    patch_system "${iso_path}/boot/ddimage.bin" "${ROOT_OVERLAY}" "${CUSTOM_SQUASHFS}"

    # Disk image
    create_disk_image "${DISK_PATH}" "${DISK_SIZE}"
    if ! loop_device="$(attach_loop_device "${DISK_PATH}")"; then
        log "Attaching loop device failed" "ERROR"
        exit 1
    fi

    create_iso_image "${iso_path}" "${iso_output}"
    iso_size="$(stat -c %s "${iso_output}")"
    iso_size_mb=$(( ("${iso_size}" + (1024 ** 2) - 1) / 1024 ** 2 ))

    partition_disk "${loop_device}" "${iso_size_mb}M"
    create_iso_partition "${loop_device}" "${iso_output}"
    create_efi_partition "${loop_device}" "${esp_path}" "${ESP_LABEL}"

    detach_loop_device "${loop_device}"

    log "Cleaning up..."
    rm -rf "${iso_path}" "${esp_path}" "${iso_output}"

    log "Disk image created at '${DISK_PATH}'" "SUCCESS"
}

main "$@"