#!/bin/bash -e
# Helper for creating autopkgtest/ROCM-capable QEMU images
#
# Author:  Christian Kastner <ckk@kvr.at>
# License: MIT

# Option defaults
authorized_keys_file=
mirror="http://deb.debian.org/debian"
release="unstable"

function usage() {
	cat >&2 <<-EOF

	Create a QEMU VM for suitable for running ROCm package autopkgtests.

	IMAGE, if not specified, defaults to RELEASE.img.

	This script will build an EFI-bootable image IMAGE, with packages
	firmware-amd-graphics and openssh-server preinstalled, and (optionally) other
	minor customizations.

	The image will have a regular (non-system) user 'user', who will be in all
	the necessary system groups.

	You will be able to share a host directory with the guest.

	Synopsis:
	  $0 -h

	  $0 [-a AKFILE] [-m MIRROR] [-r RELEASE] [IMAGE]

	Options:
	  -h          Show this help
	  -a AKFILE   File to copy to /root/.ssh/authorized_keys and to
	              /home/rocm/.ssh/authorized_keys in the guest. 
	  -m MIRROR   Download packages from here. This will be also used within
	              the image. Ideally, this points to an APT cache on the host
	              machine, such as apt-cacher-ng or approx.
	  -r RELEASE  Release (default: '$release')
	              If RELEASE is 'experimental', APT sources for both 'unstable'
	              and 'experimental' will be added to the image.
	              If RELEASE contains a dash, APT sources for the "basename"
	              will be added, e.g.: 'bookworm-backports' will include
	              sources for 'bookworm' and 'bookworm-backports'.

	Examples:

	  # Configure the system for GPU pass-through

	  \$ rocm-qemu-setup -u <user>

	  # Creates unstable.img based on unstable

	  \$ sudo $0

	  # Same image as above, but named unstable-autopkgtest-amd64.img

	  \$ sudo $0 unstable-autopkgtest-amd64.img

	  # Create unstable.img using the mirror at 10.1.2.3:9999, which ideally
	  # is a local APT cache

	  \$ sudo $0 -m http://10.1.2.3:9999/debian

	  # Creates bookworm.img

	  \$ sudo $0 -r bookworm

	  # Copies this file to {/root,/home/rocm}/.ssh/authorized_keys

	  \$ sudo $0 -a ~/.ssh/id_rsa.pub

	EOF
	exit 0
}

while getopts "a:hm:r:" OPTNAME
do
	case "$OPTNAME" in
	a)	authorized_keys_file="$OPTARG";;
	h)  usage;;
	m)	mirror="$OPTARG";;
	r)  release="$OPTARG";;
	?)	usage;;
	esac
done
shift $((OPTIND - 1))

if [ "$UID" -ne 0 ]
then
	echo "This script must be run as root." >&2
	exit 1
fi

imagename="${1:-$release.img}"

# bookworm           -> bookworm
# bookworm-backports -> bookworm
# experimental       -> unstable
release_base=$(echo "$release" | sed -r 's/-.+$//')
[ "$release_base" != "experimental" ] || release_base="unstable"

script="/usr/share/debootstrap/scripts/$release_base"
if grep -q archive.ubuntu.com "$script"
then
	distro="ubuntu"
else
	distro="debian"
fi

# /etc/apt/sources.list within the VM
if [ "$distro" = "ubuntu" ]
then
		AUTOPKGTEST_APT_SOURCES=$(cat <<- EOF
		deb     $mirror $release          main restricted universe multiverse
		deb-src $mirror $release          main restricted universe multiverse
		deb     $mirror $release-updates  main restricted universe multiverse
		deb-src $mirror $release-updates  main restricted universe multiverse
		deb     $mirror $release-security main restricted universe multiverse
		deb-src $mirror $release-security main restricted universe multiverse
EOF
)
else
	if [ "$release" != "$release_base" ]
	then
		AUTOPKGTEST_APT_SOURCES=$(cat <<- EOF
		deb     $mirror $release_base main contrib non-free-firmware
		deb-src $mirror $release_base main contrib non-free-firmware

	deb     $mirror $release main contrib non-free-firmware
	deb-src $mirror $release main contrib non-free-firmware
EOF
)
	else
		AUTOPKGTEST_APT_SOURCES=$(cat <<- EOF
		deb     $mirror $release main contrib non-free-firmware
		deb-src $mirror $release main contrib non-free-firmware
EOF
)
	fi
fi
export AUTOPKGTEST_APT_SOURCES

# Image customization script
#
# Unfortunately the modscript cannot be stuffed in a variable; it must be an
# executable file
modscript_file="$(mktemp || exit 1)"
chmod +x "$modscript_file"
cleanup() {
	rm -f "$modscript_file"
}
trap cleanup EXIT

# The script generated below will be executed by autopkgtest-build-qemu. The
# one and only argument to this script is the path to where the guest's root fs
# is temporarily mounted on the host.
cat > "$modscript_file" << EOF
[ -n "\$1" ] || exit 1
VMROOT="\$1"

# This user is created by autopkgtest-build-qemu
chroot "\$VMROOT" usermod -a -G audio,render,video user

chroot "\$VMROOT" apt-get install --quiet -y openssh-server firmware-amd-graphics pciutils

if [ -n "$authorized_keys_file" ]
then
	mkdir --mode=0700 "\$VMROOT/root/.ssh"
	install -m 0600 "$authorized_keys_file" "\$VMROOT/root/.ssh/authorized_keys"

	mkdir --mode=0700 "\$VMROOT/home/user/.ssh"
	install -m 0600 "$authorized_keys_file" "\$VMROOT/home/user/.ssh/authorized_keys"
	chroot "\$VMROOT" chown -R "user": "/home/user/.ssh"
fi

# Mount point for a shared folder, if the VM is started with one
/bin/echo "9p\n9pnet\n9pnet_virtio" >> "\$VMROOT/etc/initramfs-tools/modules"
mkdir -m 755 "\$VMROOT/shared"
echo "rocm-guest /shared 9p trans=virtio,version=9p2000.L,auto,nofail 0 0" >> "\$VMROOT/etc/fstab"
chroot "\$VMROOT" update-initramfs -u

echo "Setting hostname"
echo "$release-guest" > "\$VMROOT/etc/hosts"
echo "$release-guest" > "\$VMROOT/etc/hostname"
sed -i -r 's/^GRUB_TIMEOUT=[0-9]+/GRUB_TIMEOUT=2/' "\$VMROOT/etc/default/grub"
# Ubuntu sets this
sed -i 's/^GRUB_TIMEOUT_STYLE=hidden/GRUB_TIMEOUT_STYLE=menu/' "\$VMROOT/etc/default/grub"
chroot "\$VMROOT" update-grub
EOF

# Finally: build the image
# This invocation will create files inside the chroot that are expected to be
# world-readable, e.g. /etc/apt/sources.list
umask 0022
autopkgtest-build-qemu \
	--boot=efi \
	--architecture=amd64 \
	--mirror "$mirror" \
	--size=50G \
	--script "$modscript_file" \
	"$release_base" "$imagename"
