imgctl

#!/bin/bash
# Utility to create bootable disk images and/or devices.

# Extra modules for your GRUB install, normally can be empty. For example,
#
# "luks cryptodisk" if you put your ISOs on an encrypted partition
# "part_gpt" if you create a GPT partition table
# "part_msdos" if you're testing from a loopback device, since released
#     grub-install cannot detect loop device partitions properly.
#     (experimental branch can do this, it's unnecessary there.)
GRUB_MODULES="${GRUB_MODULES:-}"
GRUB_TARGET="${GRUB_TARGET:-x86_64-efi}"
QEMU_TARGET="${GRUB_TARGET/-*/}"

GRUB_OPTS="--no-floppy $GRUB_OPTS"

FDISK=${FDISK:-}

QEMU_OPTS="-enable-kvm $QEMU_OPTS"

###############################################################################

USAGE="Usage: $0 {--help|<CMD>} <ARGS>"

# provide access to some admin utilities
PATH="$PATH:/sbin:/usr/sbin"

# set variable defaults
FDISK_DEFAULT=fdisk
if [ -n "$DISPLAY" ] && which gparted >/dev/null; then FDISK_DEFAULT=gparted; fi
FDISK=${FDISK:-$FDISK_DEFAULT}

ERR_NO_2_LOOPDEV="need at least 2 loop devices to infer loop.max_part"
ERR_LOOP_BUILTIN="your loop module is a kernel builtin.
You need to add loop.max_part=xx to your kernel parameters[1], where xx>=1
(recommended >=15), then restart your computer. You can do this by editing
GRUB_CMDLINE_LINUX in /etc/default/grub, then run sudo update-grub."
ASK_FDISK_BEHAVE="I will now run $FDISK twice:
the 1st time, please only modify partitions (no new filesystems);
the 2nd time, please only modify filesystems (no partition changes).
Proceed?"

require_accept() {
	read -p "$1 [y/N] " x
	if [ "$x" != "y" ]; then return 1; fi
}

allow_reject() {
	read -p "$1 [Y/n] " x
	if [ "$x" = "n" ]; then return 1; fi
}

fdisk_can_mkfs() {
	case "$FDISK" in
	gparted )
		return 0;;
	fdisk | cfdisk | sfdisk )
		return 1;;
	* )
		return 0;;
	esac
}

infer_loop_max_part() {
	if [ ! -e /dev/loop1 ]; then echo "$ERR_NO_2_LOOPDEV"; exit 1; fi
	expr $((0x$(stat -c %T /dev/loop1))) - $((0x$(stat -c %T /dev/loop0))) - 1 || true
}

ensure_loop_part() {
	if ! modinfo loop >/dev/null; then LOOP_BUILTIN=true; else LOOP_BUILTIN=false; fi

	if $LOOP_BUILTIN; then
		#LOOP_MAX_PART=$(modprobe -c | grep 'options loop max_part' | cut -d= -f2)
		LOOP_MAX_PART=$(infer_loop_max_part)
	elif [ -f /sys/modules/loop/parameters/max_part ]; then
		# at time of writing, this feature is not yet released
		LOOP_MAX_PART=$(< /sys/modules/loop/parameters/max_part)
	else
		LOOP_MAX_PART=$(infer_loop_max_part)
	fi

	LOOP_MAX_PART=$(expr "0$LOOP_MAX_PART" + 0 || true)
	echo "detected loop.max_part=$LOOP_MAX_PART"
	if [ "$LOOP_MAX_PART" -ge 15 ]; then
		return
	elif [ "$LOOP_MAX_PART" -ge 1 ]; then
		allow_reject "this is a bit low; set this to 63?" || return
	fi

	if $LOOP_BUILTIN; then
		echo "$ERR_LOOP_BUILTIN"
		return 1
	else
		echo "reloading kernel module loop with max_part=63"
		sudo modprobe -r loop || return 1
		sudo modprobe loop max_part=63 || return 1
	fi
}

loopdev_for_image() {
	sudo losetup -j "$1" | head -n1 | cut -d: -f1
}

blockdev_for_dir() {
	df -P "$1" | tail -n1 | cut '-d ' -f1
}

device_ancestor() {
	CDEV=$1
	while [ -n "$CDEV" -a "$CDEV" != "${CDEV%?}/" ]; do
		PARENTS="$PARENTS $CDEV"
		CDEV=${CDEV%?}
	done
	ls -1 $PARENTS 2>/dev/null | head -n1
}

is_mounted() {
	mount | cut '-d ' -f1 | grep -qFx "$1"
}

is_loopdev() {
	ls -1 /dev/loop* | grep -qFx "$1"
}

CMD="$1"
if [ -n "$1" ]; then shift; fi

case "$CMD" in
c | create )
	IMAGE=${1}
	FSIZE=${2:-17179869184}

	set -o errexit
	touch "$IMAGE"
	ensure_loop_part

	dd if=/dev/zero of="$IMAGE" bs=1 count=0 seek="$FSIZE" 2>/dev/null

	if fdisk_can_mkfs; then require_accept "$ASK_FDISK_BEHAVE"; fi
	sudo $FDISK "$IMAGE" || true # gparted has weird exit codes sometimes

	if ! fdisk_can_mkfs; then return; fi
	cleanup() { sleep 1; sudo losetup -d "$DEV"; }
	DEV=$(loopdev_for_image "$IMAGE")
	DEV=${DEV:-$(sudo losetup --show -f "$IMAGE")}
	trap cleanup EXIT INT TERM KILL
	sudo blockdev --rereadpt "$DEV"
	sudo $FDISK "$DEV"
	;;
m | mount )
	IMAGE=${1}
	MOUNTPT=${2}
	PARTNUM=${3:-1}

	set -o errexit
	head -c0 "$IMAGE"
	mkdir -p "$MOUNTPT"
	ensure_loop_part

	DEV=$(loopdev_for_image "$IMAGE")
	DEV=${DEV:-$(sudo losetup --show -f "$IMAGE")}
	if is_mounted "${DEV}p$PARTNUM"; then return; fi
	sudo mount "${DEV}p$PARTNUM" "$MOUNTPT"
	;;
u | umount )
	TARGET=${1}

	set -o errexit
	if [ -f "$TARGET" ]; then
		TESTDEV=$(loopdev_for_image "$TARGET")
	elif [ -d "$TARGET" ]; then
		mountpoint -q "$TARGET" && TESTDEV=$(blockdev_for_dir "$TARGET")
	elif [ -b "$TARGET" ]; then
		TESTDEV=$TARGET
	fi
	if ! is_loopdev "$TESTDEV"; then echo >&2 "could not find a loop device for: $TARGET"; return 1; fi

	LOOPDEV=$(device_ancestor "$TESTDEV")
	if [ "$LOOPDEV" = "$TESTDEV" ]; then
		echo "unmount all under $LOOPDEV"
		for i in "$LOOPDEV"*; do
			if is_mounted "$i"; then sudo umount "$i"; fi
		done
		sudo losetup -d "$LOOPDEV"
		sudo blockdev --rereadpt "$LOOPDEV"
	elif is_mounted "$TESTDEV"; then
		sudo umount "$TESTDEV"
	fi
	;;
i | install )
	BOOT=${1}
	DEV=${2:-$(device_ancestor "$(blockdev_for_dir "$1")")} || exit 1

	set -o errexit
	#mkdir -p "$BOOT"
	mountpoint -x "$DEV" >/dev/null

	# part_msdos needs to be explicitly given since released grub-install cannot
	# detect loop device partitions properly. (experimental branch can do this)
	test "${DEV#/dev/loop}" != "${DEV}" && GRUB_MODULES="$GRUB_MODULES part_msdos part_gpt"

	require_accept "BOOT=$BOOT; DEV=$DEV; install?"
	set -x
	sudo /usr/sbin/grub-install --modules="$GRUB_MODULES" $GRUB_OPTS \
	  --no-uefi-secure-boot --efi-directory="$BOOT" --removable \
	  --target="${GRUB_TARGET}" --boot-directory="$BOOT" "$DEV"
	set +x
	if [ "${GRUB_TARGET%-efi}" != "$GRUB_TARGET" ]; then
		echo >&2 "WARNING: installing *without* UEFI secure boot since the latter causes us problems like launchpad #1851311"
	fi
	sudo cp ./isodetect.cfg "$BOOT/grub/"
	sudo cp -i ./grub.cfg "$BOOT/grub/"
	if ! mountpoint -q "$BOOT"; then
		if mountpoint -q "$BOOT/.." && [ "$(basename "$BOOT")" = "boot" ]; then
			sed -i -re 's,(\s)/grub/,\1/boot/grub/,g' "$BOOT/grub/grub.cfg"
		else
			echo >&2 "non-standard layout detected: you probably need to manually fix the paths in "
			echo >&2 "$BOOT/grub/grub.cfg - and even then things might not work "
			echo >&2 "due to the path assumptions in isodetect.cfg"
		fi
	fi
	;;
t | test )
	TARGET="${1}"
	sync # ensure recent changes synced to disk
	if [ -b "$TARGET" ]; then sudo chown $(id -u) "$TARGET"*; fi
	if [ "${GRUB_TARGET%-efi}" != "$GRUB_TARGET" ]; then
		if ! [ -f /usr/share/ovmf/OVMF.fd ]; then
			echo >&2 "You need to install the ovmf package to test an EFI target"
		fi
		QEMU_OPTS="-bios /usr/share/ovmf/OVMF.fd $QEMU_OPTS"
	fi
	set -x
	qemu-system-"${QEMU_TARGET}" $QEMU_OPTS "$TARGET"
	set +x
	;;
h | help | -h | --help )
	cat >&2 <<-EOF
	$USAGE

	Commands (initials also work, e.g. i for install):

	create <image> [<size>]
	    Create a virtual disk <image> (raw format) of size <size>.
	    default <size>: 16 GiB

	mount <image> <mount_dir> <partnum>
	    Mount (the <partnum>th partition of a virtual disk <image>) onto <mount_dir>.
	    (You may use this to mount the boot partition for other commands.)

	umount <image|device|mount_dir>
	    Free the disk loop device for <image|device|mount_dir>, and unmount all its partitions.
	    Or, unmount the disk partition for <device|mount_dir>, but *do not free* the loop device.

	install <boot_dir> [<device>]
	    Use grub-install onto <device>, into <boot_dir>.
	    default <device>: ancestor of mountpoint device for <boot_dir>
	      e.g. <boot_dir>/../ mounted at /dev/sda1 -> <device> = /dev/sda
	      e.g. <boot_dir> mounted at /dev/loop0p3 -> <device> = /dev/loop0

	test <image|device>
	    Use qemu-system-${QEMU_TARGET} to test <image|device>.

	EOF
	;;
* )
	echo >&2 "$USAGE"
	exit 2
	;;
esac