#!/bin/sh

# update-ukis v0.1
# Update Unified Kernel Images
# Copyright (c) 2025 Raphaël Halimi <raphael.halimi@gmail.com>

#
# Variables
#

UNITS="KMG"
OVERHEAD=4096
ESP_MIN_SIZE="10485760"
VERBOSE=0
DEBUG=0

# Architecture
case "$(dpkg --print-architecture)" in
  amd64) EFI_ARCH="x64" ;;
  arm64) EFI_ARCH="aa64" ;;
  *) return ;;
esac

# EFI stub
EFI_STUB="/usr/lib/systemd/boot/efi/linux$EFI_ARCH.efi.stub"


#
# Functions
#

# Print functions
print_error () { printf "ERROR: %s\n" "$@" >&2 ; }
print_verbose () { [ $VERBOSE -eq 1 ] && printf "%s\n" "$@" ; }
print_debug () { [ $DEBUG -eq 1 ] && printf "%s\n" "$@" >&2 ; }
print_option () { printf "  %-20s%s\n" "$1" "$2" ; }

# Exit on fatal error
die () { print_error "${1:-Fatal error}, aborting." ; exit 1 ; }

# Get file size
get_size () { stat -c%s "$1" ; }

# Convert bytes to KB, MB...
convert_unit () {
  local VAL COUNTER
  VAL="$1"
  [ -n "$VAL" ] || return 1
  printf %d "$VAL" > /dev/null 2>&1 || return 1
  COUNTER=0
  while  [ $VAL -gt 1024 ] ; do
    VAL=$((VAL/1024))
    COUNTER=$((COUNTER+1))
  done
  printf "%i %sB" "$VAL" "$(printf %s "$UNITS" | cut -c $COUNTER)"
}

# Check ESP space for files to install
# This should really be done by kernel-install
check_esp_space () {

  # Get available space
  AVAIL=$(($(df --output=avail "$ESP" | tail -n 1)*1024))

  # Start with the stub and add some overhead
  NEEDED=$(($(get_size "$EFI_STUB")+OVERHEAD))

  # Add kernel and initrd
  for FILE in "$KERNEL" "${KERNEL%/*}/initrd.img-$VERSION" ; do
    NEEDED=$((NEEDED+$(get_size "$FILE")))
  done

  # Check if ESP has enough space left
  if [ $((NEEDED+ESP_MIN_SIZE)) -gt $AVAIL ] ; then
    print_error "No space left on ESP to copy UKI version ${1#*-} ($(convert_unit $NEEDED) needed, $(convert_unit $AVAIL) free)"
    return 1
  fi
}

# Cleanup entry
# Don't use "bootctl unlink", too much output and tries to remove stuff twice
cleanup_entry () {
  local DIR
  DIR="$ESP/$MACHINE_ID/$VERSION"
  if [ -n "$(find "$ESP/loader/entries" -type f -name "$ENTRY*")" ] ; then
    print_verbose "Cleaning Type #1 entry '$ENTRY'..."
    rm -f $RM_OPTS "$ESP/loader/entries/$ENTRY"* >&2
  fi
  if [ -d "$DIR" ] ; then
    print_verbose "Cleaning Type #1 files for '$ENTRY'..."
    rm -f $RM_OPTS "$DIR/linux" "$DIR/initrd.img-$VERSION" >&2
    rmdir $RM_OPTS "$DIR" >&2
    if [ -z "$(find "${DIR%/*}" -mindepth 1 -maxdepth 1)" ] ; then
      rmdir $RM_OPTS "${DIR%/*}" >&2
    fi
  fi
  if [ -n "$(find "$ESP/EFI/Linux" -type f -name "$ENTRY*")" ] ; then
    print_verbose "Cleaning Type #2 (UKI) entry '$ENTRY'..."
    rm -f $RM_OPTS "$ESP/EFI/Linux/$ENTRY"* >&2
  fi
}

# Print usage
print_usage () {
  printf "Usage: %s [OPTION]...\n" "${0##*/}"
  printf "\nOPTIONS:\n"
  print_option "-v" "Verbose mode"
  print_option "-d" "Debug mode"
  print_option "-h" "Print this help message"
}


#
# Options
#

# Command-line options
while getopts "vdh" OPTION ; do
  case $OPTION in
    v) VERBOSE=1 ;;
    d) VERBOSE=1 ; DEBUG=1 ;;
    h) print_usage ; exit 0 ;;
    *) print_usage ; exit 1 ;;
  esac
done ; shift $((OPTIND-1))


#
# Checks
#

# Check root
[ $(id -u) -eq 0 ] || die "this script must be run as root"

# Check dependencies
for BINARY in bootctl linux-version ; do
  command -v $BINARY >/dev/null || die "binary '$BINARY' not found"
done

# Check mandatory variables
ESP="$(bootctl --quiet --print-esp-path 2>/dev/null)"
[ -n "$ESP" ] || die "ESP not found"
MACHINE_ID="$(cat /etc/machine-id)"
[ -n "$MACHINE_ID" ] || die "Machine ID not found"

# Add options to rm
[ $DEBUG -eq 1 ] && RM_OPTS="-v"


#
# Main
#

# List installed kernels
VERSIONS="$(linux-version list)"

# Update UKIs for currently installed kernels
for VERSION in $VERSIONS ; do
  ENTRY="$MACHINE_ID-$VERSION"
  KERNEL="/boot/vmlinuz-$VERSION"
  cleanup_entry
  if check_esp_space "$KERNEL" ; then
    print_verbose "Generating UKI '$ESP/EFI/Linux/$ENTRY.efi'..."
    kernel-install add "$VERSION" "$KERNEL"
  fi
done

exit 0
