#!/usr/bin/env perl
# vim: softtabstop=2 shiftwidth=2 expandtab

use strict;
use warnings;

our $VERSION = '3.1.0';

use Getopt::Long qw(:config no_ignore_case auto_version);
use Pod::Usage   qw(pod2usage);
use File::Basename;
use File::Temp qw(tempfile tempdir);
use File::Copy;
use File::stat;
use File::Path qw(make_path remove_tree);
use File::Glob qw(:globally :nocase);
use Sort::Versions;
use Config;
use bigint qw(hex);

use Pod::Usage qw(pod2usage);

use Data::Dumper;
$Data::Dumper::Indent   = 1;
$Data::Dumper::Sortkeys = 1;
$Data::Dumper::Purity   = 1;

use Sort::Versions;
use YAML::PP;
use boolean;

use Storable qw( dclone );

use constant REFARRAY => ref [];

sub versionedKernel;
sub latestKernel;
sub createInitramfs;
sub createUEFIBundle;
sub execute;
sub safeCopy;
sub has_value;
sub cleanupMount;
sub enabled;
sub maxRevision;
sub groupKernels;
sub pruneVersions;
sub purgeFiles;
sub verboseUnlink;

BEGIN {
  $SIG{INT}  = \&cleanupMount;
  $SIG{TERM} = \&cleanupMount;
}

my ( %runConf, %config );

$runConf{config}  = "/etc/zfsbootmenu/config.yaml";
$runConf{bootdir} = "/boot";

GetOptions(
  "version|v=s" => \$runConf{version},
  "kernel|k=s"  => \$runConf{kernel},
  "kver|K=s"    => \$runConf{kernel_version},
  "prefix|p=s"  => \$runConf{kernel_prefix},
  "bootdir|b=s" => \$runConf{bootdir},
  "confd|C=s"   => \$runConf{confd},
  "cmdline|l=s" => \$runConf{cmdline},
  "config|c=s"  => \$runConf{config},
  "enable"      => \$runConf{enable},
  "disable"     => \$runConf{disable},
  "initcpio|i!" => \$runConf{usecpio},
  "hookd|H=s@"  => \$runConf{cpio_hookd},
  "debug|d"     => \$runConf{debug},
  "showver|V"   => \$runConf{showver},
  "help|h"      => sub {
    pod2usage( -verbose => 2 );
    exit;
  },
) or exit 1;

if ( defined $runConf{showver} and $runConf{showver} ) {
  printf "%s\n", $VERSION;
  exit 0;
}

if ( -r $runConf{config} ) {
  eval {
    local $SIG{'__DIE__'};
    my $yaml = YAML::PP->new( boolean => 'boolean' )->load_file( $runConf{config} );
    %config = %$yaml;
  } or do {
    my $error = <<"EOF";
Unable to parse configuration $runConf{config} as YAML.
EOF
    print $error;

    warn $@ if $@;
    exit 1;
  };
} else {
  printf "Configuration %s does not exist or is unreadable\n", $runConf{config};
  exit 1;
}

if ( $runConf{disable} ) {
  $runConf{enable} = false;
}

if ( defined $runConf{enable} ) {

  $config{Global}{ManageImages} = boolean( $runConf{enable} );

  my $yaml = YAML::PP->new(
    boolean => 'boolean',
    header  => 0,
  );

  $yaml->dump_file( $runConf{config}, \%config );
  my $state = $runConf{enable} ? "true" : "false";
  printf "ManageImages set to '%s' in %s\n", $state, $runConf{config};
  exit;
}

unless ( $config{Global}{ManageImages} ) {
  print "ManageImages not enabled, no action taken\n";
  exit;
}

unless ( defined $runConf{usecpio} ) {

  # If usecpio wasn't set by cmdline, try loading from the config
  if ( defined $config{Global}{InitCPIO} ) {
    $runConf{usecpio} = $config{Global}{InitCPIO};
  } else {
    my @output = execute(qw(sh -c "command -v dracut"));
    my $status = pop(@output);
    if ( $status eq 0 ) {
      print "No initramfs generator specified; using dracut\n";
      $runConf{usecpio} = false;
    } else {
      print "No initramfs generator specified; using mkinitcpio\n";
      $runConf{usecpio} = true;
    }
  }
}

unless ( defined $runConf{confd} ) {

  # Set initramfs configuration from defaults or config file

  # Defaults and config key depend on initramfs generator
  my $ckey;
  if ( $runConf{usecpio} ) {
    $runConf{confd} = "/etc/zfsbootmenu/mkinitcpio.conf";
    $ckey = "InitCPIOConfig";
  } else {
    $runConf{confd} = "/etc/zfsbootmenu/dracut.conf.d";
    $ckey = "DracutConfDir";
  }

  # Replace the default if a configuration option exists
  if ( defined $config{Global}{$ckey} ) {
    $runConf{confd} = $config{Global}{$ckey};
  }
}

if ( $runConf{usecpio} and not defined $runConf{cpio_hookd} ) {

  # With initcpio mode, load hookdirs from config when not specified on cmdline
  my @hooks;

  if ( defined $config{Global}{InitCPIOHookDirs} ) {
    if ( ref $config{Global}{InitCPIOHookDirs} eq REFARRAY ) {
      foreach my $hookd ( @{ $config{Global}{InitCPIOHookDirs} } ) {
        push( @hooks, $hookd );
      }
    } else {
      push( @hooks, $config{Global}{InitCPIOHookDirs} );
    }
  }

  $runConf{cpio_hookd} = \@hooks;
}

# Ensure our bootloader partition is mounted
$runConf{umount_on_exit} = 0;
if ( has_value $config{Global}{BootMountPoint} ) {
  my $mounted = 0;

  my @output = execute(qq(mountpoint $config{Global}{BootMountPoint}));
  my $status = pop(@output);
  unless ( $status eq 0 ) {
    print "Mounting $config{Global}{BootMountPoint}\n";
    my @output = execute(qq(mount $config{Global}{BootMountPoint}));
    my $status = pop(@output);
    if ( $status eq 0 ) {
      $runConf{umount_on_exit} = 1;
    } else {
      foreach my $line (@output) {
        print $line;
      }
      printf "Unable to mount %s\n", $config{Global}{BootMountPoint};
      exit $status;
    }
  }
}

if ( has_value $config{Global}{PreHooksDir} and -d $config{Global}{PreHooksDir} ) {
  while ( my $hook = <$config{Global}{PreHooksDir}/*> ) {
    next unless -x $hook;
    Log("Processing hook: $hook");
    my @output = execute(qq($hook));
    Log( \@output );
  }
}

# Create a temp directory
# It is automatically purged on program exit
my $dir     = File::Temp->newdir();
my $tempdir = $dir->dirname;

# Config file may provide some default values for command-line args
if ( has_value $config{Kernel}{Path} and !has_value $runConf{kernel} ) {
  $runConf{kernel} = $config{Kernel}{Path};
}
if ( has_value $config{Kernel}{Prefix} and !has_value $runConf{kernel_prefix} ) {
  $runConf{kernel_prefix} = $config{Kernel}{Prefix};
}

if ( has_value $config{Kernel}{Version} and !has_value $runConf{kernel_version} ) {
  $runConf{kernel_version} = $config{Kernel}{Version};
  $runConf{kernel_version} =~ s/%current\b/%{current}/i;
}

if ( has_value $config{Global}{Version} and !has_value $runConf{version} ) {
  $runConf{version} = $config{Global}{Version};
}

if ( has_value $config{Kernel}{CommandLine} and !has_value $runConf{cmdline} ) {
  $runConf{cmdline} = $config{Kernel}{CommandLine};
}

if ( has_value $runConf{version} ) {
  $runConf{version} =~ s/%current\b/%{current}/i;
  $runConf{version} =~ s/%\{current\}/$VERSION/i;
} else {
  $runConf{version} = $VERSION;
}

# Map "%current" kernel version to output of `uname r`
if ( has_value $runConf{kernel_version} and $runConf{kernel_version} =~ /%\{current\}/i ) {
  my @uname  = execute(qw(uname -r));
  my $status = pop(@uname);
  unless ( $status eq 0 and scalar @uname ) {
    print "Cannot determine current kernel version\n";
    exit $status;
  }
  chomp @uname;
  $runConf{kernel_version} =~ s/%\{current\}/$uname[0]/i;
}

if ( has_value $runConf{kernel} ) {

  # Make sure the provided kernel file exists
  unless ( -f $runConf{kernel} ) {
    printf "The provided kernel %s was not found, unable to continue\n", $runConf{kernel};
    exit 1;
  }
} else {

  # Try to determine a kernel file when one was not provided
  if ( has_value $runConf{kernel_version} ) {
    my $exactVersion;
    ( $runConf{kernel}, $exactVersion ) = versionedKernel $runConf{kernel_version};

    # Make sure a kernel was found
    unless ( has_value $runConf{kernel} ) {
      print "Unable to find file for kernel version $runConf{kernel_version}\n";
      exit 1;
    }

    # If the kernel version was not exact, allow it to be determined later
    unless ($exactVersion) {
      undef $runConf{kernel_version};
    }

  } else {
    $runConf{kernel} = latestKernel;
    unless ( has_value $runConf{kernel} ) {
      print "Unable to find latest kernel; specify version or path manually\n";
      exit 1;
    }
  }
}

# Try to determine kernel_prefix or kernel_version if necessary
unless ( has_value $runConf{kernel_version} ) {

  # Kernel version comes from either file name or internal strings
  $runConf{kernel_version} = kernelVersion( $runConf{kernel} );
  unless ( has_value $runConf{kernel_version} ) {
    printf "Unable to determine kernel version from %s\n", $runConf{kernel};
    exit 1;
  }
}

unless ( has_value $runConf{kernel_prefix} ) {

  # Prefix is basename of file, less any "-<version>" suffix
  $runConf{kernel_prefix} = basename( $runConf{kernel} );
  $runConf{kernel_prefix} =~ s/-\Q$runConf{kernel_version}\E$//;
  unless ( has_value $runConf{kernel_prefix} ) {
    printf "Unable to determine kernel prefix from %s\n", $runConf{kernel};
    exit 1;
  }
}

printf "Creating ZFSBootMenu %s from kernel %s\n", $VERSION, $runConf{kernel};

my $spl_hostid = "/sys/module/spl/parameters/spl_hostid";
if ( -f $spl_hostid ) {
  open PROC, $spl_hostid;
  $runConf{hostid}{module} = sprintf( "%08x", <PROC> );
  close PROC;
} else {
  $runConf{hostid}{module} = "00000000";
}

my $etc_hostid = "/etc/hostid";
if ( $runConf{hostid}{module} ne "00000000" and -f $etc_hostid ) {
  open SPL, '<:raw', $etc_hostid;
  read SPL, my $hostid, 4;
  close SPL;

  $runConf{hostid}{etc} = sprintf( "%08x", unpack( 'L<4', $hostid ) );

  if ( $runConf{hostid}{module} ne $runConf{hostid}{etc} ) {
    print "SPL ($runConf{hostid}{module}) and system ($runConf{hostid}{etc}) hostids do not match!\n";
  }
}

# Create the initramfs as long as some output will consume it
my $initramfs;
if ( enabled $config{EFI} or enabled $config{Components} ) {
  $initramfs = createInitramfs( $tempdir, $runConf{kernel_version} );
}

# Create a unified kernel/initramfs/command line EFI file
if ( enabled $config{EFI} ) {
  my $unified_efi = createUEFIBundle( $tempdir, $runConf{kernel}, $initramfs );

  my $efi_target;

  my $efi_prefix = sprintf( "%s/%s", $config{EFI}{ImageDir}, $runConf{kernel_prefix} );
  Log("Setting \$efi_prefix: $efi_prefix");

  my $efi_versions = int $config{EFI}{Versions};

  make_path $config{EFI}{ImageDir};

  if ( $efi_versions > 0 ) {
    Log("EFI.Versions is $efi_versions");

    # Find UEFI bundles and group by apparent version
    my @efi        = glob sprintf( "%s-*.EFI", $efi_prefix );
    my $efi_groups = groupKernels( \@efi, $efi_prefix, ".EFI" );
    Log($efi_groups);

    # Determine the revision to use for this image
    my $revision = maxRevision( $efi_groups->{ $runConf{version} }, ".EFI" ) + 1;
    $efi_target = sprintf( "%s-%s_%s.EFI", $efi_prefix, $runConf{version}, $revision );
    Log("Setting \$efi_target: $efi_target");

    # Attempt to copy the file, clean up if it does not
    unless ( safeCopy( $unified_efi, $efi_target, 0 ) ) {
      verboseUnlink( $efi_target, "Failed to create $efi_target" );
      exit 1;
    }

    # Prune the old versions
    pruneVersions( $efi_groups, $runConf{version}, $efi_versions );
  } else {
    $efi_target = sprintf( "%s.EFI", $efi_prefix );

    # Copy to a placeholder location to ensure success
    my ( $efi_fh, $efi_tempfile ) = tempfile( "zbm.XXXXXX", DIR => $config{EFI}{ImageDir}, UNLINK => 0 );
    close $efi_fh;

    unless ( safeCopy( $unified_efi, $efi_tempfile, 0 ) ) {
      verboseUnlink( $efi_tempfile, "Failed to create $efi_target" );
      exit 1;
    }

    # Roll backups
    my $efi_backup = sprintf( "%s-backup.EFI", $efi_prefix );
    if ( -f $efi_target and rename( $efi_target, $efi_backup ) ) {
      printf "Created backup %s -> %s\n", $efi_target, $efi_backup;
    }

    unless ( rename( $efi_tempfile, $efi_target ) ) {
      verboseUnlink( $efi_tempfile, "Failed to create $efi_target" );
      exit 1;
    }
  }

  printf "Created new UEFI image %s\n", $efi_target;
}

# Create a separate kernel / initramfs. Used by syslinux/extlinux/grub.
if ( enabled $config{Components} ) {
  my ( $kernel_target, $initramfs_target );

  my $component_prefix   = sprintf( "%s/%s", $config{Components}{ImageDir}, $runConf{kernel_prefix} );
  my $component_versions = int $config{Components}{Versions};

  make_path $config{Components}{ImageDir};

  if ( $component_versions > 0 ) {

    # Find ZBM kernels and group by apparent version
    my @kernels     = glob( sprintf( "%s-*", $component_prefix ) );
    my $kern_groups = groupKernels( \@kernels, $component_prefix );

    my $revision = maxRevision( $kern_groups->{ $runConf{version} } ) + 1;
    $kernel_target = sprintf( "%s-%s_%s", $component_prefix, $runConf{version}, $revision );
    $initramfs_target =
      sprintf( "%s/initramfs-%s_%s.img", $config{Components}{ImageDir}, $runConf{version}, $revision );

    unless ( safeCopy( $initramfs, $initramfs_target, 0 ) ) {
      verboseUnlink( $initramfs_target, "Failed to create $initramfs_target" );
      exit 1;
    }

    unless ( safeCopy( $runConf{kernel}, $kernel_target, 0 ) ) {
      verboseUnlink( $kernel_target,    "Failed to create $kernel_target" );
      verboseUnlink( $initramfs_target, "" );
      exit 1;
    }

    # Prune old versions of the kernel
    pruneVersions( $kern_groups, $runConf{version}, $component_versions );

    # Map each kernel to initramfs and prune those too
    keys %$kern_groups;
    while ( my ( $kver, $image ) = each %$kern_groups ) {
      foreach (@$image) {
        s/\Q$component_prefix\E/$config{Components}{ImageDir}\/initramfs/;
        s/$/.img/;
      }
    }
    pruneVersions( $kern_groups, $runConf{version}, $component_versions );
  } else {
    $kernel_target    = sprintf( "%s-bootmenu",               $component_prefix );
    $initramfs_target = sprintf( "%s/initramfs-bootmenu.img", $config{Components}{ImageDir} );

    # Copy to a placeholder location to ensure success
    my ( $init_fh, $init_tempfile ) = tempfile( "init.XXXXXX", DIR => $config{Components}{ImageDir}, UNLINK => 0 );
    close $init_fh;

    unless ( safeCopy( $initramfs, $init_tempfile, 0 ) ) {
      verboseUnlink( $init_tempfile, "Failed to create $initramfs_target" );
      exit 1;
    }

    my ( $kern_fh, $kern_tempfile ) = tempfile( "kern.XXXXXX", DIR => $config{Components}{ImageDir}, UNLINK => 0 );
    close $kern_fh;

    unless ( safeCopy( $runConf{kernel}, $kern_tempfile, 0 ) ) {
      verboseUnlink( $kern_tempfile, "Failed to create $kernel_target" );
      verboseUnlink( $init_tempfile, "" );
      exit 1;
    }

    # Roll backups
    my $kernel_backup = sprintf( "%s-backup", $kernel_target );
    if ( -f $kernel_target and rename( $kernel_target, $kernel_backup ) ) {
      printf "Created backup %s -> %s\n", $kernel_target, $kernel_backup;
    }

    my $initramfs_backup = sprintf( "%s/initramfs-bootmenu-backup.img", $config{Components}{ImageDir} );
    if ( -f $initramfs_target and rename( $initramfs_target, $initramfs_backup ) ) {
      printf "Created backup %s -> %s\n", $initramfs_target, $initramfs_backup;
    }

    unless ( rename( $init_tempfile, $initramfs_target ) ) {
      verboseUnlink( $init_tempfile, "Failed to create $initramfs_target" );
      verboseUnlink( $kern_tempfile, "" );
      exit 1;
    }

    unless ( rename( $kern_tempfile, $kernel_target ) ) {
      verboseUnlink( $kern_tempfile, "Failed to create $kernel_target" );
      exit 1;
    }
  }

  printf "Created initramfs image %s\n", $initramfs_target;
  printf "Created kernel image %s\n",    $kernel_target;
}

if ( has_value $config{Global}{PostHooksDir} and -d $config{Global}{PostHooksDir} ) {
  while ( my $hook = <$config{Global}{PostHooksDir}/*> ) {
    next unless -x $hook;
    Log("Processing hook: $hook");
    my @output = execute(qq($hook));
    Log( \@output );
  }
}

END {
  cleanupMount;
}

# Finds specifically versioned kernel in /boot
sub versionedKernel {
  my ( $kver, ) = @_;

  foreach my $prefix (qw(vmlinuz linux vmlinux kernel)) {
    my $pattern = join( '/', ( $runConf{bootdir}, join( '-', ( $prefix, $kver ) ) ) );

    # Try an exact match first
    if ( -f $pattern ) {
      return $pattern, true;
    }

    # Otherwise, try to glob
    my @kernels = sort versioncmp glob($pattern);

    next unless @kernels;
    return pop @kernels, false;
  }

  return;
}

# Finds the latest kernel in /boot, if possible
sub latestKernel {
  my @prefixes = ( "vmlinuz*", "vmlinux*", "linux*", "kernel*" );

  for my $prefix (@prefixes) {
    my $glob = join( '/', ( $runConf{bootdir}, $prefix ) );
    my %kernels;

    for my $kernel ( glob($glob) ) {
      my $version = kernelVersion($kernel);
      next unless defined($version);
      Log("Identified version $version for kernel $kernel");
      $kernels{$version} = $kernel;
    }

    next unless ( keys %kernels );

    for ( sort { versioncmp( $b, $a ) } keys %kernels ) {
      Log("Latest kernel: $_");
      return $kernels{$_};
    }
  }

  return;
}

# Attempts to determine a version for the given kernel, by
#
# a. Identifying the first version-looking string in the file, or
# b. Identifying a version-like part in the name of the file
#
# If one of these exists and not the other, that value is used; if both exist,
# the name-derived value is used if that version string can be matched
# somewhere in the file contents, otherwise the version is undefined.
sub kernelVersion {
  my $kernel = shift;

  my ( $filever, $namever );

  # Consider an unreadable file to have no version
  unless ( -r $kernel ) {
    Log("Unable to read path $kernel, assuming no version");
    return;
  }

  # Read version from the file name, if possible
  basename($kernel) =~ m/-([0-9]+\.[0-9]+\.[0-9]+.*)/;
  if ( defined $1 ) {
    $namever = $1;
  }

  # Read strings in the kernel to recover a version, if possible
  my @output = execute(qq(strings $kernel));
  my $status = pop(@output);
  if ( $status eq 0 ) {
    for (@output) {

      # Versions are any three dot-separated numbers followed by non space
      next unless (/([0-9]+\.[0-9]+\.[0-9]+\S+)/);

      my $ver = $1;

      # First version match is always the file version
      $filever = $ver unless ( has_value $filever );

      # When there is no version from the file name, we have a match
      last unless ( has_value $namever );

      # A version that equals the file version supersedes the first match
      if ( $namever eq $ver ) {
        $filever = $ver;
        last;
      }
    }
  }

  # Kernel is unusable if no version could be detected
  return if ( !has_value $filever and !has_value $namever );

  # If only one is defined, that's the version
  unless ( has_value $filever ) {
    Log("No version found in kernel strings, using $namever from path $kernel");
    return $namever;
  }

  unless ( has_value $namever ) {
    Log("No version found in path $kernel, using $filever from kernel strings");
    return $filever;
  }

  # Warn if the two alternatives do not agree
  if ( $namever ne $filever ) {
    my $warning = <<"EOF";
WARNING: ignoring inconsistent versions in kernel $kernel:
  Path suggests version $namever.
  Kernel strings suggest version $filever.
  To use this kernel, explicitly specify the path and version.
EOF
    print $warning;
    return;
  }

  return $namever;
}

# Given a sections size, calculate where the next section should be placed,
# while respecting the stub alignment value
sub increaseBundleOffset {
  my ( $step, $offset, $alignment ) = @_;
  $offset += int( ( $step + $alignment - 1 ) / $alignment * $alignment );
  Log( "New offset is: " . hex($offset) );
  return $offset;
}

# Adds the commands necessary to put another section into the EFI bundle,
# and then calculates where the bundle offset has been moved to
sub addBundleSection {
  my ( $cmds, $secname, $filename, $offset, $alignment ) = @_;

  my $hex_offset = sprintf( "0x%X", $offset );
  push( @$cmds, ( "--add-section", "$secname=\"$filename\"" ), qw(--change-section-vma), ("$secname=\"$hex_offset\""),
  );

  my $sb = stat($filename);
  return increaseBundleOffset( $sb->size, $offset, $alignment );

}

# Creates a UEFI bundle from an initramfs and kernel
# Returns the path to the bundle or dies with an error

sub createUEFIBundle {
  my ( $imagedir, $kernel, $initramfs ) = @_;

  my $output_file = join( '/', $imagedir, "zfsbootmenu.EFI" );

  unless ( -f $kernel and -f $initramfs ) {
    print "Cannot find kernel or initramfs to create UEFI bundle\n";
    exit 1;
  }

  my $uefi_stub;

  if ( has_value $config{EFI}{Stub} ) {
    $uefi_stub = $config{EFI}{Stub};
    unless ( -f $uefi_stub ) {
      print "UEFI stub loader '$uefi_stub' does not exist\n";
      exit 1;
    }
  } else {

    my @uefi_stub_defaults;
    if ( $Config{archname} =~ m/x86_64/ ) {
      push( @uefi_stub_defaults, '/usr/lib/systemd/boot/efi/linuxx64.efi.stub' );
    } elsif ( $Config{archname} =~ m/aarch64/ ) {
      push( @uefi_stub_defaults, '/usr/lib/systemd/boot/efi/linuxaa64.efi.stub' );
    }

    foreach my $stubloc (@uefi_stub_defaults) {
      if ( -f $stubloc ) {
        $uefi_stub = $stubloc;
        last;
      }
    }

    unless ( defined $uefi_stub and -f $uefi_stub ) {
      print "Unable to find UEFI stub loader at default locations:\n";
      foreach my $stubloc (@uefi_stub_defaults) {
        print " $stubloc\n";
      }
      exit 1;
    }
  }

  my ( $uki_alignment, $uki_offset );

  # Determine stub alignment, most likely 4096
  my @cmd = qw(objdump -p);
  push( @cmd, $uefi_stub );

  my @output = execute(@cmd);
  my $status = pop(@output);
  if ( $status eq 0 ) {
    foreach my $line (@output) {
      if ( $line =~ m/SectionAlignment\s+(\d+)/ ) {
        Log( "Alignment is: " . hex($1) );
        $uki_alignment = hex($1);
      }
    }
  } else {
    print "Unable to determine stub alignment!\n";
    exit 1;
  }

  # Determine initial UKI offset value by grabbing the size and VMA of
  # the last section of the EFI stub.
  @cmd = qw(objdump -w -h);
  push( @cmd, $uefi_stub );

  @output = execute(@cmd);
  $status = pop(@output);
  if ( $status eq 0 ) {
    my @sizes = split( /\s+/, @output[ scalar @output - 1 ] );

    my $size = "0x" . $sizes[3];
    my $vma  = "0x" . $sizes[4];
    my $sum  = hex($size) + hex($vma);

    $uki_offset = increaseBundleOffset( $sum, 0, $uki_alignment );
    Log( "Initial offset is: " . hex($uki_offset) );
  } else {
    print "Unable to determine initial stub offset!\n";
    exit 1;
  }

  @cmd = qw(objcopy);

  my ( $hex_offset, $sb );

  if ( -f "/etc/os-release" ) {
    $uki_offset = addBundleSection( \@cmd, ".osrel", "/etc/os-release", $uki_offset, $uki_alignment );
  }

  if ( has_value $runConf{cmdline} ) {
    my $cmdline = join( '/', $imagedir, "cmdline.txt" );

    open( my $fh, '>', $cmdline );
    print $fh $runConf{cmdline};
    close($fh);

    $uki_offset = addBundleSection( \@cmd, ".cmdline", $cmdline, $uki_offset, $uki_alignment );
  }

  if ( has_value $config{EFI}{SplashImage} and -f $config{EFI}{SplashImage} ) {

    # only supported with systemd-boot's efistub,
    # but gummiboot doesn't care if the section exists
    $uki_offset = addBundleSection( \@cmd, ".splash", $config{EFI}{SplashImage}, $uki_offset, $uki_alignment );
  }

  if ( has_value $config{EFI}{DeviceTree} ) {
    my $dtb_try;

    if ( $config{EFI}{DeviceTree} =~ m,^/, ) {
      if ( $config{EFI}{DeviceTree} =~ m/\Q%{kernel}/ ) {

        # Possibly a templated file path
        $dtb_try = $config{EFI}{DeviceTree};
        $dtb_try =~ s/\Q%{kernel}/$runConf{kernel_version}/;
      } else {

        # Possibly an absolute path to a file
        $dtb_try = $config{EFI}{DeviceTree};
      }
    } else {

      # Possibly a partial path
      $dtb_try = sprintf( "/boot/dtbs/dtbs-%s/%s", $runConf{kernel_version}, $config{EFI}{DeviceTree} );
    }

    if ( has_value $dtb_try ) {
      if ( -f $dtb_try ) {
        $uki_offset = addBundleSection( \@cmd, ".dtb", $dtb_try, $uki_offset, $uki_alignment );
      } else {
        printf "EFI.DeviceTree key set to '%s', file not found at '%s'\n", $config{EFI}{DeviceTree}, $dtb_try;
        exit 1;
      }
    } else {
      printf "EFI.DeviceTree key set to '%s', but no DeviceTree file could be found\n", $config{EFI}{DeviceTree};
      exit 1;
    }
  }

  $uki_offset = addBundleSection( \@cmd, ".initrd", $initramfs, $uki_offset, $uki_alignment );

  # Add the kernel last, so that it can decompress without overflowing other sections
  $uki_offset = addBundleSection( \@cmd, ".linux", $kernel, $uki_offset, $uki_alignment );

  push( @cmd, ( $uefi_stub, $output_file ) );

  my $command = join( ' ', @cmd );
  Log("Executing: $command");

  @output = execute(@cmd);
  $status = pop(@output);
  if ( $status eq 0 ) {
    foreach my $line (@output) {
      Log($line);
    }
    return $output_file;
  } else {
    foreach my $line (@output) {
      print $line;
    }
    print "Failed to create $output_file\n";
    exit $status;
  }
}

# Creates an initramfs and returns its path, or dies with an error
sub createInitramfs {
  my ( $imagedir, $kver ) = @_;

  my $output_file = join( '/', $imagedir, "zfsbootmenu.img" );

  my @cmd;
  my $flagsKey;

  if ( $runConf{usecpio} ) {
    push( @cmd, ( qw(mkinitcpio --config), $runConf{confd} ) );
    push( @cmd, qw(-v) ) if $runConf{debug};

    # Add hook directories as appropriate
    if ( defined $runConf{cpio_hookd} ) {
      foreach my $hookd ( @{ $runConf{cpio_hookd} } ) {
        push( @cmd, ( "--hookdir", $hookd ) );
      }
    }

    $flagsKey = "InitCPIOFlags";
  } else {
    push( @cmd, ( qw(dracut -f --confdir), $runConf{confd} ) );
    push( @cmd, qw(-q) ) unless $runConf{debug};

    $flagsKey = "DracutFlags";
  }

  # Load custom flag additions from configuration file
  if ( defined $config{Global}{$flagsKey} ) {
    if ( ref $config{Global}{$flagsKey} eq REFARRAY ) {
      foreach my $flag ( @{ $config{Global}{$flagsKey} } ) {
        push( @cmd, $flag );
      }
    } else {
      push( @cmd, $config{Global}{$flagsKey} );
    }
  }

  # Specify kernel version and ouptut location
  if ( $runConf{usecpio} ) {
    push( @cmd, ( qw(-A zfsbootmenu), "--generate", $output_file, "--kernel", $kver ) );
  } else {
    push( @cmd, ( $output_file, $kver ) );
  }

  my $command = join( ' ', @cmd );
  Log("Executing: $command");

  my @output = execute(@cmd);
  my $status = pop(@output);
  if ( $status eq 0 ) {
    foreach my $line (@output) {
      Log($line);
    }
    return $output_file;
  } else {
    foreach my $line (@output) {
      print $line;
    }
    print "Failed to create $output_file\n";
    exit $status;
  }
}

sub execute {
  ( @_ = qx{@_ 2>&1}, $? >> 8 );
}

sub safeCopy {
  my ( $source, $dest, $savetime ) = @_;

  my $preserve = ( defined $savetime ) ? boolean($savetime) : true;
  Log("safeCopy called with: $source, $dest, $preserve");

  unless ( copy( $source, $dest ) ) {
    printf "Unable to copy %s to %s: %s\n", $source, $dest, $!;
    return 0;
  }

  if ($preserve) {

    # Copy the access and mod times if possible
    my $sb = stat $source;
    utime( $sb->atime, $sb->mtime, $dest );
  }

  return 1;
}

sub has_value {
  my $item = shift;
  return ( defined $item and length $item );
}

sub enabled {
  my $section = shift;
  return ( defined $section->{Enabled} and $section->{Enabled} );
}

sub cleanupMount {
  my $signal = shift;

  if ( $runConf{umount_on_exit} ) {
    print "Unmounting $config{Global}{BootMountPoint}\n";
    execute(qq(umount $config{Global}{BootMountPoint}));
  }

  if ( defined $signal ) {
    print "$0 terminating on signal $signal\n";
    exit 1;
  }
}

sub maxRevision {
  my ( $files, $suffix ) = @_;
  my $revision = 0;

  $suffix = "" unless ( defined $suffix );

  foreach my $file (@$files) {
    if ( $file =~ /_(\d+)\Q$suffix\E$/ ) {
      $revision = $1 if ( $1 > $revision );
    }
  }

  Log("maxRevision discovered: $revision");
  return $revision;
}

sub groupKernels {
  my ( $kernels, $prefix, $suffix ) = @_;
  my %groups;

  $suffix = "" unless ( defined $suffix );

  foreach my $kernel (@$kernels) {
    next unless ( $kernel =~ /^\Q$prefix\E-(.+)_\d+\Q$suffix\E$/ );
    push( @{ $groups{$1} }, $kernel );
  }

  return \%groups;
}

sub pruneVersions {
  my ( $versions, $current, $keep ) = @_;
  my $old_version;

  Log("pruneVersions called with: $current, $keep");
  Log($versions);

  $keep = 0 unless ( defined $keep and $keep gt 0 );

  # Keep revisions current version
  purgeFiles( $versions->{$current}, $keep );

  # Sort the versions and remove the current
  my @old_versions = sort versioncmp keys %$versions;

  my $index = 0;
  foreach my $key (@old_versions) {
    if ( $key eq $current ) {
      splice( @old_versions, $index, 1 );
    } else {
      $index++;
    }
  }

  # Purge all of the too-old revisions
  while ( scalar @old_versions > $keep ) {
    $old_version = shift @old_versions;
    purgeFiles( $versions->{$old_version} );
  }

  # Purge all but the remaining revision of the leftover versions
  foreach $old_version (@old_versions) {
    purgeFiles( $versions->{$old_version}, 1 );
  }
}

sub purgeFiles {
  my ( $files, $keep ) = @_;

  return unless ( defined $files );

  $keep = 0 unless ( defined $keep and $keep gt 0 );

  if ( $keep gt 0 ) {
    my @sorted_files = sort versioncmp @$files;
    while ( scalar @sorted_files > $keep ) {
      my $file = shift @sorted_files;
      verboseUnlink $file;
    }
  } else {
    foreach my $file (@$files) {
      verboseUnlink $file;
    }
  }
}

sub verboseUnlink {
  my ( $file, $message ) = @_;

  return unless ( -f $file );

  # If a message is defined, display regardless of unlink success
  if ( defined $message and ( $message ne "" ) ) {
    print "$message\n";
  }

  if ( unlink $file ) {

    # Print a default success message if none was defined
    print "Removed file $file\n" unless ( defined $message );
  } else {
    print "ERROR: unable to remove $file: $!\n";
  }
}

sub Log {
  my $entry = shift;

  return unless $runConf{debug};
  chomp($entry);
  unless ( ref($entry) ) {
    print STDERR "## $entry\n";
  } elsif ( ref $entry eq REFARRAY ) {
    foreach my $line ( @{$entry} ) {
      chomp $line;
      print STDERR "## $line\n";
    }
  } else {
    print STDERR Dumper($entry);
  }
}

__END__

=for comment
KEEP IN SYNC WITH docs/man/generate-zbm.8.rst

=head1 NAME

B<generate-zbm> - ZFSBootMenu initramfs generator

=head1 SYNOPSIS

B<generate-zbm> [options]

=head1 OPTIONS

Where noted, command-line options supersede options in the B<generate-zbm>(5) configuration file.

=over 4

=item B<--version|-v> I<zbm-version>

Override the ZFSBootMenu version in output files; supersedes I<Global.Version>

=item B<--kernel|-k> I<kernel-path>

Manually specify a specific kernel; supersedes I<Kernel.Path>

=item B<--kver|-K> I<kernel-version>

Manually specify a specific kernel version; supersedes I<Kernel.Version>

=item B<--prefix|-p> I<image-prefix>

Manually specify the output image prefix; supersedes I<Kernel.Prefix>

=item B<--initcpio|-i>

Force the use of mkinitcpio instead of dracut.

=item B<--no-initcpio|-i>

Force the use of dracut instead of mkinitcpio.

=item B<--confd|-C> I<config-path>

Specify initramfs configuration path

=over 4

=item For dracut: supersedes I<Global.DracutConfDir>

=item For mkinitcpio: supersedes I<Global.InitCPIOConfig>

=back

=item B<--hookd|-H> I<hookd-path>

Specify mkinitcpio hook directory; supersedes I<Global.InitCPIOHookDirs>

May be specified more than once. Ignored when using dracut.

=item B<--cmdline|-l> I<options>

Override the kernel command line; supersedes I<Kernel.CommandLine>

=item B<--bootdir|-b> I<boot-path>

Specify the path to search for kernel files; default: I</boot>

=item B<--config|-c> I<conf-file>

Specify the path to a configuration file; default: I</etc/zfsbootmenu/config.yaml>

=item B<--enable>

Set the I<Global.ManageImages> option to true, enabling image generation.

=item B<--disable>

Set the I<Global.ManageImages> option to false, disabling image generation.

=item B<--debug|d>

Enable debug output

=item B<--showver|V>

Print ZFSBootMenu version and quit.

=back

=head1 SEE ALSO

B<generate-zbm>(5) B<zfsbootmenu>(7)

=head1 AUTHOR

ZFSBootMenu Team L<https://github.com/zbm-dev/zfsbootmenu>

=cut
