#!/bin/bash

# SPDX-FileCopyrightText: 2023 - 2023 ENCRYPTED SUPPORT LP <adrelanos@kicksecure.com>
#
# SPDX-License-Identifier: AGPL-3.0-or-later

# Usage: image-to-iso <raw_image_file> <iso_tmp_dir> <iso_image_file>

# This script generates a bootable ISO image.
# It takes a bootable raw disk image as input and outputs a bootable ISO image.
# The script is designed to support both booting methods:
# - EFI (Unified Extensible Firmware Interface) and,
# - legacy BIOS booting.
#
# Limitations:
# - Only tested with Debian.
# - Only tested with images creates by grml-debootstrap --vmfile --vmefi,
#   see corresponding createimg1 script.
# - Makes some assumptions about the partition layout.
#   First partition needs to be the EFI partition.
#   Last partition needs to be the root file system.
# - Encrypted raw images are unsupported but that point is a bit moot,
#   because to the knowledge of the author there are no raw image creation
#   tools for Debian that can create encrypted raw images.
# - Certain cross builds are impossible. For example, building an ARM64 ISO
#   on AMD64 is unsupported. This is because Debian package grub-efi-arm64
#   is unavailable for AMD64. There is also no package grub-efi-arm64-bin.
#   The build is possible but an ARM64 ISO built on AMD64 will be unbootable.
#   This is not a limitation of this script. This is also not possible to do
#   manually. To work around this, something like an ARM64 build chroot and
#   QEMU might work.

set -x
set -e
set -o pipefail
set -o errtrace
set -o nounset

error_handler() {
  if test -n "${BASH_COMMAND:-}"; then
    # shellcheck disable=SC2039,3028,3054
    true "ERROR: Executed script, function, command executed: '${0}' '${FUNCNAME[1]}' '${BASH_COMMAND}'"
  else
    true "ERROR: Executed script: '${0}'"
  fi
  unmount_cleanup || true "ERROR: Cleanup unmounting failed."
  true "ERROR: Exiting with error."
}

parse_arguments() {
  if [ $# -lt 1 ]; then
    echo "Usage: $0 <raw_image_file> <iso_tmp_dir> <iso_image_file>" >&2
    exit 1
  fi

  if [ "$1" = "check_dependencies" ]; then
    trap "exit 1" err
    check_dependencies
    exit 0
  fi

  if [ $# -lt 3 ]; then
    echo "Usage: $0 <raw_image_file> <iso_tmp_dir> <iso_image_file>" >&2
    exit 1
  fi

  RAW_IMAGE_FILE="$1"
  ISO_TMP_DIR="$2"
  ISO_IMAGE_FILE="$3"

  ## read by binary_iso_wrapper
  export ISO_IMAGE_FILE
}

check_root() {
  if [ ! "$(id -u)" = "0" ]; then
    true "ERROR: Must be run with root privileges (sudo)! sudo $0" >&2
    exit 1
  fi
}

check_dependencies() {
  true "INFO: Ensure necessary commands are available..."

  test -x ./binary_grub-pc_wrapper
  test -x ./binary_grub-efi_wrapper
  test -x ./binary_iso_wrapper
  test -x ./grub-raw-cfg-to-grub-iso-cfg
  command -v mksquashfs >/dev/null
  command -v grub-mkrescue >/dev/null
  command -v xorriso >/dev/null
  command -v kpartx >/dev/null
  command -v mount >/dev/null
  command -v umount >/dev/null
  command -v rmdir >/dev/null

  ## Used by live-build.
  ## Unnecessary dependency.
  ## Avoided in forked live-build scripts.
  #dpkg -l | grep isolinux >/dev/null

  ## Undocumented dependencies by grub-mkrescue.
  ## Also needed by the live-build scripts.
  dpkg -l | grep mtools >/dev/null

  architecture=$(dpkg --print-architecture)

  ## Undocumented dependencies by grub-mkrescue.
  ## Also needed by the live-build scripts.

  case $architecture in
    amd64)
      dpkg -l | grep grub-efi-amd64-signed >/dev/null
      dpkg -l | grep grub-efi-amd64-bin >/dev/null
      dpkg -l | grep grub-efi-ia32-bin >/dev/null
      dpkg -l | grep shim-signed >/dev/null
      ## package: grub-efi-amd64-signed
      test -f /usr/lib/grub/x86_64-efi-signed/gcdx64.efi.signed
      ## package: shim-signed amd64
      test -f /usr/lib/shim/shimx64.efi.signed
      ;;
    i386)
      #dpkg -l | grep grub-efi-amd64-signed >/dev/null
      dpkg -l | grep grub-efi-amd64-bin >/dev/null
      dpkg -l | grep grub-efi-ia32-bin >/dev/null
      dpkg -l | grep shim-signed >/dev/null
      ## package: grub-efi-ia32-signed
      test -f /usr/lib/grub/i386-efi-signed/gcdia32.efi.signed
      ## package: shim-signed i386
      test -f /usr/lib/shim/shimia32.efi.signed
      ;;
    arm64)
      dpkg -l | grep shim-signed >/dev/null
      ## package: shim-signed arm64
      test -f /usr/lib/shim/shimaa64.efi.signed
      ;;
    *)
      true "INFO: Host dependency checking not implemented for the $architecture architecture."
      ;;
  esac
}

setup_directories() {
  BINARY_FOLDER=$(mktemp --directory)

  ## read by binary_grub-efi
  export BINARY_FOLDER
  export ISO_TMP_DIR
}

setup_image_mount() {
  true "INFO: Mount the image file and set up loopback device..."
  DEVINFO=$(kpartx -a -s -v "${RAW_IMAGE_FILE}")

  # grml-debootstrap --arch amd64 example DEVINFO:
  # add map loop1p1 (253:3): 0 204800 linear 7:1 2048
  # add map loop1p2 (253:4): 0 2048 linear 7:1 206848
  # add map loop1p3 (253:5): 0 209504256 linear 7:1 208896

  # grml-debootstrap --arch arm64 example DEVINFO:
  # add map loop0p1 (253:0): 0 18432 linear 7:0 2048
  # add map loop0p2 (253:1): 0 209692672 linear 7:0 20480

  if [ -z "${DEVINFO}" ] ; then
    echo "Error setting up loopback device." >&2
    exit 1
  fi
}

prepare_loopback_info() {
  true "INFO: Extract and prepare loopback device information..."
  local loop_part

  # Extract the first partition (EFI)
  loop_part=$(echo "${DEVINFO}" | head -n 1 | awk '{print $3}')
  # grml-debootstrap --arch arm64 example LOOP_PART:
  # loop0p1

  EFI_TARGET="/dev/mapper/${loop_part}"
  # example: EFI_TARGET
  # /dev/mapper/loop1p1

  # Extract the last partition (Root FS)
  loop_part=$(echo "${DEVINFO}" | tail -n 1 | awk '{print $3}')
  # grml-debootstrap --arch arm64 example LOOP_PART:
  # loop0p3

  ROOTFS_TARGET="/dev/mapper/${loop_part}"
  # example ROOTFS_TARGET:
  # /dev/mapper/loop1p3

  RAW_IMAGE_MNT=$(mktemp -d)
  ## read by binary_...
  export RAW_IMAGE_MNT
}

mount_file_systems() {
  true "INFO: Mount the target file system and EFI partition..."
  mount "${ROOTFS_TARGET}" "${RAW_IMAGE_MNT}"

  ## No longer add /boot/efi folder to squashfs.
  ## The grub.cfg can confuse after installing using Calamares.
  ## If re-enabling this, also re-enable umount below.
  #mount "${EFI_TARGET}" "${RAW_IMAGE_MNT}/boot/efi"
}

create_iso_directories() {
  # Clear and create necessary directories in the temporary ISO directory.
  rm -r -f "${ISO_TMP_DIR}"
  mkdir --parents "${ISO_TMP_DIR}/LiveOS"
}

get_file_names() {
  true "INFO: Retrieve the kernel version from the mounted image..."
  KVERSION=$(cd "${RAW_IMAGE_MNT}/boot" && find vmlinuz-* | tail -1 | sed 's@vmlinuz-@@')

  # Define paths for the kernel file and initial ramdisk.
  kernel_file_name="vmlinuz-${KVERSION}"
  initial_ramdisk_file_name="initrd.img-${KVERSION}"
}

create_squashfs() {
  true "INFO: Creating squashfs..."
  test -r ./squashfs-exclude.txt
  ## Take the contents of the VM image and create the squashfs.img file in the LiveOS folder
  ## within the temporary ISO directory.
  ##
  ## The kernel and initial ramdisk is not needed inside the squashfs.img.
  ## These will later be copied in the ISO's /boot folder in a later function.
  ##
  ## The /LiveOS/squashfs.img path is a default setting used by dracut.
  ##
  ## During development, especially when testing various kernel boot parameters, it may be
  ## unnecessary to repeatedly recreate the squashfs file.
  ##
  ## TODO: Consider porting to squashfs-tools-ng for potential improvements in speed and size reduction.
  mksquashfs "${RAW_IMAGE_MNT}" "${ISO_TMP_DIR}/LiveOS/squashfs.img" -noappend -ef ./squashfs-exclude.txt
}

copy_kernel_and_initial_ramdisk_and_grub_config() {
  true "INFO: Copy kernel and initial ramdisk to the temporary ISO directory..."

  ## Sanity test.
  ## During mksquashfs might be commented out. But if mksquashfs does not exist, the image will boot
  ## but dracut obviously fail to find the root filesystem. This this can be confusing nonetheless
  ## until noticed. Therefore adding a sanity test here.
  test -f "${ISO_TMP_DIR}/LiveOS/squashfs.img"

  architecture=$(dpkg --print-architecture)

  ## Needed by binary_grub-efi_wrapper
  case $architecture in
    amd64)
      test -f "${RAW_IMAGE_MNT}/usr/lib/grub/x86_64-efi-signed/gcdx64.efi.signed"
      test -f "${RAW_IMAGE_MNT}/usr/lib/shim/shimx64.efi.signed"
      ;;
    i386)
      test -f "${RAW_IMAGE_MNT}/usr/lib/grub/i386-efi-signed/gcdia32.efi.signed"
      test -f "${RAW_IMAGE_MNT}/usr/lib/shim/shimia32.efi.signed"
      ;;
    arm64)
      test -f "${RAW_IMAGE_MNT}/usr/lib/shim/shimaa64.efi.signed"
      ;;
    *)
      true "INFO: raw image dependency checking not implemented for the $architecture architecture."
      ;;
  esac

  ## grub + shim
  ## shim: bootx64.efi
  ## grub: grubx64.efi
  ## Creates $BINARY_FOLDER which contains the following files / folders:
  # /boot/grub/i386-efi/*.mod
  # /boot/grub/x86_64-efi/*.mod
  # /boot/grub/efi.img
  # /boot/grub/unicode.pf2
  # /EFI/boot/bootia32.efi
  # /EFI/boot/bootx64.efi
  # /EFI/boot/grubia32.efi
  # /EFI/boot/grubx64.efi
  # efi.img
  ./binary_grub-efi_wrapper

  ## Creates $BINARY_FOLDER which contains the following files / folders:
  # /boot/grub/i386-pc
  # /boot/grub/grub_eltorito
  ./binary_grub-pc_wrapper

  ## Use '/.*' to also copy the hidden '.disk' folder.
  ## But that foldes does only exist after binary_iso_wrapper has been run.
  #mv --verbose "$BINARY_FOLDER"/.* "${ISO_TMP_DIR}/"
  mv --verbose "$BINARY_FOLDER"/* "${ISO_TMP_DIR}/"

  ## Sanity tests.
  test -f "${ISO_TMP_DIR}/boot/grub/efi.img"
  test -f "${ISO_TMP_DIR}/boot/grub/grub_eltorito"

  mkdir --parents "${ISO_TMP_DIR}/boot/grub"
  cp -a "${RAW_IMAGE_MNT}/boot/${kernel_file_name}" "${ISO_TMP_DIR}/boot/${kernel_file_name}"
  cp -a "${RAW_IMAGE_MNT}/boot/${initial_ramdisk_file_name}" "${ISO_TMP_DIR}/boot/${initial_ramdisk_file_name}"

  ## If using grub-mkrescue there is no need to copy the /boot folder except for kernel and initrd because
  ## grub-mkrescue should handle all of that.
  #cp -a -r "${RAW_IMAGE_MNT}/boot" "${ISO_TMP_DIR}/boot"

  ## Create a copy of the grub.cfg from the raw image as a template for later
  ## adjustment to be functional for the ISO.
  cp -a "${RAW_IMAGE_MNT}/boot/grub/grub.cfg" "${ISO_TMP_DIR}/boot/grub/grub-raw-image-temp.cfg"

  ## This is what Debian live-build would be doing:
  #cp -a chroot/usr/lib/grub/x86_64-efi-signed/gcdx64.efi.signed chroot/grub-efi-temp/EFI/boot/grubx64.efi
  #cp -a chroot/usr/lib/shim/shimx64.efi.signed chroot/grub-efi-temp/EFI/boot/bootx64.efi
  #
  #cp -a chroot/usr/lib/grub/i386-efi-signed/gcdia32.efi.signed chroot/grub-efi-temp/EFI/boot/grubia32.efi
  #cp -a chroot/usr/lib/shim/shimia32.efi.signed chroot/grub-efi-temp/EFI/boot/bootia32.efi

  ## On a Debian ISO file '.disk/info' is created using:
  #echo "${TITLE} ${VERSION} \"${DISTRIBUTION}\" - ${STRING} ${DISK_LABEL} Binary ${_DATE}" > binary/.disk/info
  ## Example content:
  #Official Debian GNU/Linux Live 12.4.0 xfce 2023-12-10T17:43:24Z
  ## Existence of that file is important because grub config created by binary_grub-efi contains:
  #search --set=root --file /.disk/info
  ## The actual content seems not crucially important.
  mkdir --parents "${ISO_TMP_DIR}/.disk"
  echo "created by $0" | tee "${ISO_TMP_DIR}/.disk/info" >/dev/null
  test -f "${ISO_TMP_DIR}/.disk/info"
}

handle_fstab() {
  if test -f "${RAW_IMAGE_MNT}/etc/fstab" ; then
    if ! test -f "${RAW_IMAGE_MNT}/etc/fstab_old_image-to-iso" ; then
      mv --verbose "${RAW_IMAGE_MNT}/etc/fstab" "${RAW_IMAGE_MNT}/etc/fstab_old_image-to-iso"
    fi
  fi
}

unmount_cleanup() {
  true "INFO: Unmount partitions and clean up loopback device..."

  ## syntax:
  #~/derivative-maker/help-steps/unmount-helper -- /tmp/user/0/tmp.6UT5XIy8K9

  # XXX: hardcoded path
  if test -x "/home/$SUDO_USER/derivative-maker/help-steps/unmount-helper" ; then
    # XXX: Hacky. Would be better if this was a standalone script.
    umount_tool="sudo -u $SUDO_USER /home/$SUDO_USER/derivative-maker/help-steps/unmount-helper --"
  else
    umount_tool=umount
  fi

  if [ -n "${RAW_IMAGE_MNT:-}" ]; then
    #if findmnt "${RAW_IMAGE_MNT}/boot/efi" &>/dev/null; then
      #$umount_tool "${RAW_IMAGE_MNT}/boot/efi"
    #fi
    if findmnt "${RAW_IMAGE_MNT}" &>/dev/null; then
      $umount_tool "${RAW_IMAGE_MNT}"
    fi
  fi
  if [ -n "${RAW_IMAGE_FILE:-}" ]; then
    kpartx -d -v "${RAW_IMAGE_FILE}"
  fi
  if [ -n "${RAW_IMAGE_MNT:-}" ]; then
    if test -d "${RAW_IMAGE_MNT}" ; then
      rmdir "${RAW_IMAGE_MNT}"
    fi
  fi
}

grub_raw_image_config_to_grub_iso_image_config() {
  true "INFO: Create the GRUB configuration file for the ISO."

  # Notes:
  # - The /boot folder originates from the temporary ISO directory.
  #   It will only be available during very early boot.
  # - Once dracut set up the root system, there might be a different /boot,
  #   which originates from the raw image.
  # - This means kernel and initial ramdisk are duplicated. In theory,
  #   it might be possible to omit adding the /boot folder from squashfs.img.
  # - This might be confusing, but makes it easier to generate a
  #   /boot/grub/grub.cfg file based on the VM image.

  ## Was functional.
  ## Replaced with a more complex implementation below for the ISO to inherit
  ## same grub.cfg from the raw image but adjusted to work with the ISO.
  # cat << EOF > "${ISO_TMP_DIR}/boot/grub/grub.cfg"
  # set default="0"
  # set timeout=10
  # menuentry "Linux" {
  #     linux /boot/${kernel_file_name} root=live:CDLABEL=cdlabel rd.live.image rd.debug rd.live.debug rd.live.overlay.overlayfs=1 console=ttyS0
  #     initrd /boot/${initial_ramdisk_file_name}
  # }
  # EOF

  iso_grub_extra_params=""

  ## Define root device.
  ## Needs to be same used below with grub-mkrescue '-V cdlabel', otherwise the
  ## ISO would be unbootable with dracut.
  ## (initramfs-tools is able to boot such an ISO.)
  iso_grub_extra_params+=" root=live:CDLABEL=cdlabel"

  iso_grub_extra_params+=" rd.live.image"
  iso_grub_extra_params+=" rd.live.overlay.overlayfs=1"

  # Debugging dracut.
  #iso_grub_extra_params+=" rd.debug"
  #iso_grub_extra_params+=" rd.live.debug"

  # console=ttyS0: Show output when booting with qemu with -nographic.
  # console=tty0: Keep normal console output enabled.
  #iso_grub_extra_params+=" console=tty0 console=ttyS0,115200n8"

  # Convert grub.cfg from raw image to grub.cfg usable for ISO image.
  #
  # Keep original grub-raw-image-temp.cfg for later comparison purposes.
  ./grub-raw-cfg-to-grub-iso-cfg "${ISO_TMP_DIR}/boot/grub/grub-raw-image-temp.cfg" "${ISO_TMP_DIR}/boot/grub/grub.cfg" "${iso_grub_extra_params}"

  ## Debian Live ISO has no file /EFI/boot/grub.cfg
  ## /boot/efi folder in raw image has grub.cfg:
  # search.fs_uuid 26ada0c0-1165-4098-884d-aafd2220c2c6 root
  # set prefix=($root)'/boot/grub'
  # configfile $prefix/grub.cfg
  ## This might cause issues because wrong UUID.
}

create_iso_file() {
  true "INFO: Create the ISO using grub-mkrescue based on the temporary ISO directory..."

  ## Was functional,
  ## except for not including shim and therefore likely Secure Boot issues.
  ## grub feature request:
  ## grub-mkrescue ISO Secure Boot / shim support
  ## https://savannah.gnu.org/bugs/?65039
  ##
  ## '-V cdlabel' set label to 'cdlabel'. This must match the 'root=live:CDLABEL=cdlabel' kernel command line.
  ##
  ## XXX: consistent architectures
  ## "grub-mkrescue will automatically include every platform it finds." [1]
  ## [1] https://lists.gnu.org/archive/html/grub-devel/2014-03/msg00009.html
  ## This is an issue for reproducible builds because files copied need to be
  ## predictable and not depending on the host operating system configuration.
  ## Also wastes a bit disk space.
  ## For derivative-maker builds this is solved by installing all of these
  ## packages on the build host.
  ##
  ## During development also useful to edit file:
  ## ~/grml-debootstraptestbin/iso_tmp/boot/grub/grub.cfg
  ## and then:
  ## sudo grub-mkrescue --verbose -V cdlabel -o ~/grml-debootstraptestbin/output.iso ~/grml-debootstraptestbin/iso_tmp
  ##
  #   grub-mkrescue \
  #     --verbose \
  #     -V cdlabel \
  #     -o "${ISO_IMAGE_FILE}" \
  #     "${ISO_TMP_DIR}"

  # grub-mkrescue:
  # xorriso -as mkisofs -graft-points --modification-date=2023121215520700 -b boot/grub/i386-pc/eltorito.img -no-emul-boot -boot-load-size 4 -boot-info-table --grub2-boot-info --grub2-mbr /usr/lib/grub/i386-pc/boot_hybrid.img -hfsplus -apm-block-size 2048 -hfsplus-file-creator-type chrp tbxj /System/Library/CoreServices/.disk_label -hfs-bless-by i /System/Library/CoreServices/boot.efi --efi-boot efi.img -efi-boot-part --efi-boot-image --protective-msdos-label -o /home/user/grml-debootstraptestbin/output.iso -r /tmp/user/0/grub.UQp4UP --sort-weight 0 / --sort-weight 1 /boot -V cdlabel /tmp/tmpdir.

  # Debian 12 .disk/mkisofs
  # xorriso -as mkisofs -R -r -J -joliet-long -l -cache-inodes -iso-level 3 -isohybrid-mbr /usr/lib/ISOLINUX/isohdpfx.bin -partition_offset 16 -A "Debian Live" -p "live-build d70a84f2e9f6808ca24a52cee1ec62898de98db0_2023-09-27T15:30:41+02:00_mod; https://salsa.debian.org/live-team/live-build" -publisher "Debian Live project; https://wiki.debian.org/DebianLive; debian-live@lists.debian.org" -V "d-live 12.2.0 kd amd64" --modification-date=2023100709291900 -b isolinux/isolinux.bin -c isolinux/boot.cat -no-emul-boot -boot-load-size 4 -boot-info-table -eltorito-alt-boot -e boot/grub/efi.img -no-emul-boot -isohybrid-gpt-basdat -isohybrid-apm-hfsplus -o live-image-amd64.hybrid.iso binary

  # Debian live-build from git .disk/mkisofs
  # xorriso -as mkisofs -R -r -J -joliet-long -l -cache-inodes -iso-level 3 --grub2-boot-info --grub2-mbr /usr/lib/grub/i386-pc/boot_hybrid.img -efi-boot-part --efi-boot-image -A "Debian Live" -p "live-build 20230502; https://salsa.debian.org/live-team/live-build" -publisher "Debian Live project; https://wiki.debian.org/DebianLive; debian-live@lists.debian.org" -V "Debian bookworm 20240423-10:45" --modification-date=2024042314455200 -no-emul-boot -boot-load-size 4 -boot-info-table -b boot/grub/grub_eltorito -eltorito-alt-boot -e boot/grub/efi.img -no-emul-boot -isohybrid-gpt-basdat -isohybrid-apm-hfsplus -o live-image-amd64.hybrid.iso binary

  ## Options not used yet but that might make sense to be more similar to Debian:
  ##
  ## Only useful for ISOLINUX, hence not using:
  ## -c boot_catalog

  ## xorriso -as mkisofs -R -r -J -joliet-long -l -cache-inodes -iso-level 3 --grub2-boot-info --grub2-mbr /usr/lib/grub/i386-pc/boot_hybrid.img -efi-boot-part --efi-boot-image -v -A "application_id" -p "preparer" -publisher "publisher" -V "cdlabel" --modification-date=2024042316311900 -no-emul-boot -boot-load-size 4 -boot-info-table -b boot/grub/grub_eltorito -eltorito-alt-boot -e boot/grub/efi.img -no-emul-boot -isohybrid-gpt-basdat -isohybrid-apm-hfsplus -o /home/user/derivative-binary/17.1.3.8/Kicksecure-Xfce-17.1.3.8.iso /home/user/derivative-binary/Kicksecure-Xfce_iso-tmp
  ##
  ## Also creates file:
  ## /.disk/mkisofs
  ./binary_iso_wrapper

  # Fix permissions.
  chown "${SUDO_USER}:${SUDO_USER}" "${ISO_IMAGE_FILE}"
}

cleanup_temporary_files() {
  true "INFO: Cleanup temporary files..."
  # Not deleting for easier debugging.
  #rm -r "${ISO_TMP_DIR}"
}

main() {
  parse_arguments "$@"
  check_root
  check_dependencies
  setup_directories
  setup_image_mount
  prepare_loopback_info
  mount_file_systems
  create_iso_directories
  get_file_names
  create_squashfs
  copy_kernel_and_initial_ramdisk_and_grub_config
  handle_fstab
  unmount_cleanup
  grub_raw_image_config_to_grub_iso_image_config
  create_iso_file
  cleanup_temporary_files
  true "INFO: Done."
}

# shellcheck disable=SC2039,3047
trap 'error_handler $? ${LINENO:-}' ERR

main "$@"
