#! /bin/sh
set -e

# Copyright (C) 2010, 2011 Canonical Ltd.
#           (C) 2019-2020 Steve McIntyre <93sam@debian.org>
# Author: Colin Watson <cjwatson@ubuntu.com>
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.

# Make an EFI boot image.

usage () {
	echo "usage: $0 [OPTIONS]"
	echo "Options:"
	echo "-o <output directory>    Set the output directory. Must be specified."
	echo "-g <grub platform name>  Set the grub platform name. Must be specified."
	echo "-e <efi platform name>   Set the EFI platform name. Must be specified."
	echo "-n <netboot prefix>      Set the netboot prefix for Grub. Must be specified."
	echo "-s <y|n>                 Whether to set up for EFI Secure Boot. Must be specified."
	echo "-b <y|n>                 Whether to create a bundled EFI bootloader. Optional."
	echo "-d <dtb directory>       Set the input directory for DTB files to be included. Optional."
}

outdir=""
platform=""
efi_name=""
netboot_prefix=""
efi_signed=""
efi_bundled=""
dtb_dir=""

while getopts ":o:g:e:n:s:b:d:" o; do
    case "${o}" in
        o)
            outdir="${OPTARG}"
            ;;
        g)
            platform="${OPTARG}"
            ;;
        e)
            efi_name="${OPTARG}"
            ;;
        n)
            netboot_prefix="${OPTARG}"
            ;;
        s)
            efi_signed="${OPTARG}"
            ;;
        b)
            efi_bundled="${OPTARG}"
            ;;
        d)
            dtb_dir="${OPTARG}"
            ;;
    esac
done
shift $((OPTIND-1))

# Kali: force efi bundle for platforms where it's supported
case "${platform}" in
    x86_64-efi|arm64-efi) efi_bundled=y ;;
esac

# Validate inputs
if [ "$outdir"x = ""x ] || [ "$platform"x = ""x ] || \
	[ "$efi_name"x = ""x ] || [ "$netboot_prefix"x = ""x ] || \
	[ "$efi_signed"x = ""x ]; then
	usage
	exit 1
fi

memdisk_img=
workdir=

cleanup () {
	[ -z "$memdisk_img" ] || rm -f "$memdisk_img"
	[ -z "$workdir" ] || rm -rf "$workdir"
}
trap cleanup EXIT HUP INT QUIT TERM

copy_files_into_image () {
	mmd -i "$outdir/efi.img" ::efi
	mmd -i "$outdir/efi.img" ::efi/boot
	mcopy -i "$outdir/efi.img" "$workdir/boot$efi_name.efi" \
		"::efi/boot/boot$efi_name.efi"

	if [ $efi_signed = y ] || [ "$efi_bundled" = y ]; then
		# In this case, also add the grub binary
		mcopy -i "$outdir/efi.img" "$workdir/grub$efi_name.efi" \
			"::efi/boot/grub$efi_name.efi"

		# And finally the skeleton grub.cfg - the grub
		# binaries will have prefix set to "/EFI/<vendor>", so
		# put this there to redirect to the correct location.
		mmd -i "$outdir/efi.img" ::efi/kali
		mcopy -i "$outdir/efi.img" "$workdir/boot/grub/grub.cfg" \
			"::efi/kali/grub.cfg"
	fi

	# Copy in any DTB files we might have/want
	if [ "$dtb_dir"x != ""x ]; then
		mmd -i "$outdir/efi.img" ::dtb
		mcopy -s -i "$outdir/efi.img" "$dtb_dir/." ::dtb
	fi
}

count_free_clusters () {
	image=$1
	clusters_free=$(fsck.msdos -v $1 | \
	awk '
		/files.*cluster/ {
			used_free=$4
			split(used_free, sizes, "/")
			print (sizes[2] - sizes[1])
		}
	')
	echo $clusters_free
}

rm -rf "$outdir"
mkdir -p "$outdir"

memdisk_img="$(mktemp efi-image.XXXXXX)"
workdir="$(mktemp -d efi-image.XXXXXX)"

# Skeleton configuration file which finds the real boot disk.
mkdir -p "$workdir/boot/grub"
cat >"$workdir/boot/grub/grub.cfg" <<EOF
search --file --set=root /.disk/info
set prefix=(\$root)/boot/grub
source \$prefix/$platform/grub.cfg
EOF

mkdir -p "$outdir/boot/grub/$platform"
(for i in /usr/lib/grub/$platform/part_*.mod; do
    i=`echo $i | sed 's?^.*/??g;s?\.mod$??g;'`
	echo "insmod $i"
 done; \
 echo "source /boot/grub/grub.cfg") >"$outdir/boot/grub/$platform/grub.cfg"

if [ $efi_signed = y ]; then
	# Just copy existing shim and Grub binaries into place.
	# First, the binaries we use for disc or CD boot.
	echo "$0: Using pre-signed grub-efi and shim binaries for $efi_name"
	cp /usr/lib/shim/shim$efi_name.efi.signed \
		$workdir/boot$efi_name.efi
	cp /usr/lib/grub/$platform-signed/grub$efi_name.efi.signed \
		$workdir/grub$efi_name.efi

	# Now the binaries for netboot. These are *not* actually used
	# in generating the image here, but will be picked up later on
	# in the d-i build and copied into the output netboot
	# tree. Hence the different output directory.
	cp /usr/lib/shim/shim$efi_name.efi.signed \
		$outdir/bootnet$efi_name.efi
	# Copy the pre-built signed grub-netboot binary that is
	# configured just for d-i, with the right prefix baked in
	# (debian-installer/$arch/grub)". We can't change the prefix,
	# hence the special build. (See #928750)
	cp /usr/lib/grub/$platform-signed/grubnet$efi_name-installer.efi.signed \
		$outdir/grub$efi_name.efi

elif [ "$efi_bundled" = y ]; then
	# "<package>: <location>" for unsigned files:
	# shim-unsigned: /usr/lib/shim/shimx64.efi
	# grub-efi-amd64-bin: /usr/lib/grub/x86_64-efi/monolithic/grubx64.efi
	# grub-efi-amd64-bin: /usr/lib/grub/x86_64-efi/monolithic/grubnetx64-installer.efi

	# Just copy existing shim and Grub binaries into place.
	# First, the binaries we use for disc or CD boot.
	echo "$0: Using unsigned grub-efi and shim binaries for $efi_name"
	cp /usr/lib/shim/shim$efi_name.efi \
		$workdir/boot$efi_name.efi
	cp /usr/lib/grub/$platform/monolithic/grub$efi_name.efi \
		$workdir/grub$efi_name.efi

	# Now the binaries for netboot. These are *not* actually used
	# in generating the image here, but will be picked up later on
	# in the d-i build and copied into the output netboot
	# tree. Hence the different output directory.
	cp /usr/lib/shim/shim$efi_name.efi \
		$outdir/bootnet$efi_name.efi
	# Copy the pre-built signed grub-netboot binary that is
	# configured just for d-i, with the right prefix baked in
	# (debian-installer/$arch/grub)". We can't change the prefix,
	# hence the special build. (See #928750)
	cp /usr/lib/grub/$platform/monolithic/grubnet$efi_name-installer.efi \
		$outdir/grub$efi_name.efi

else
	# Build the core image for disc/CD boot
	echo "$0: Building non-signed grub-efi binaries for $efi_name"
	find "$workdir" -newermt "@$SOURCE_DATE_EPOCH" -print0 \
		| xargs -0r touch --no-dereference --date="@$SOURCE_DATE_EPOCH"
	(cd "$workdir"; find boot -print0 | LC_ALL=C sort -z | tar --null -T - -cf -) >"$memdisk_img"
	grub-mkimage -O "$platform" -m "$memdisk_img" \
		-o "$workdir/boot$efi_name.efi" -p '(memdisk)/boot/grub' \
		search iso9660 configfile normal memdisk tar part_msdos part_gpt fat

	# Now the image for netboot. This is *not* used to generate
	# the image here, but will be picked up later on in the d-i
	# build and copied into the output netboot tree. Hence the
	# different output directory.
	grub-mkimage -O "$platform" \
		-o "$outdir/bootnet$efi_name.efi" -p "$netboot_prefix/grub" \
		search configfile normal efinet tftp net
fi

# Stuff it into a FAT filesystem, making it as small as
# possible. Count up the sizes of all the files and directories we need to fit in there.
#
# 32KiB headroom seems to be enough; (x+31)/32*32 rounds up to multiple of 32.
size=0
block_size=1024
for file in $(find $workdir $dtb_dir -type f -name *.efi -o -name grub.cfg -o -name *.dtb); do
	this_size=$(( ($(stat -c %s $file) / $block_size + 1 + 32 + 31) / 32 * 32 ))
	size=$(($size + $this_size))
done
for dir in $(find $workdir $dtb_dir -type d); do
	size=$(($size + 1)) # Is 1 block enough for each directory?
done

# Specify a deterministic volume ID; we use DOS_VOLUME_ID here for
# consistency.
echo "$0: Running first pass of mkfs.msdos with size $size blocks"
mkfs.msdos --invariant -v -i deb00001 -C "$outdir/efi.img" $size
copy_files_into_image

# The worst-case calculations above are potentially giving
# significantly too-large EFI filesystems. Stop guessing. Look at what
# we have using fsck.msdos, then recalculate directly how much we've
# used and try again with a more accurate size.
cluster_size=$(fsck.msdos -v "$outdir/efi.img" | awk '/bytes per cluster/ {print $1}')
cluster_blocks=$(($cluster_size / $block_size))
clusters_free=$(count_free_clusters "$outdir/efi.img")

echo "$0: First pass has $clusters_free free clusters"
# Add a little bit of free space, 32kB for luck?
size=$(($size + 32 - ($cluster_blocks * $clusters_free)))

echo "$0: Running second pass of mkfs.msdos with size $size blocks"
rm -f "$outdir/efi.img"
mkfs.msdos --invariant -v -i deb00001 -C "$outdir/efi.img" $size
copy_files_into_image
clusters_free=$(count_free_clusters "$outdir/efi.img")
echo "$0: Second pass has $clusters_free free clusters"
grub-cpmodules "$outdir" "$platform"

exit 0
