#!/bin/sh
# vim: set ts=4:
#---help---
# Usage: alpine-make-vm-image [options] [--] <image> [<script> [<script-opts...>]]
#
# This script creates a bootable Alpine Linux disk image for virtual machines.
# If running on Alpine system (detected by file /etc/alpine-release), then it
# also installs needed packages on the host system. On other systems you must
# install them yourself: qemu-img, qemu-nbd, and mkfs utility for the chosen
# ROOTFS. If $APK is not available on the host system, then static apk-tools
# specified by $APK_TOOLS_URI is downloaded and used.
#
# Note that it does not create any partitions (it's really not needed),
# filesystem is created directly on the image.
#
# Arguments:
#   <image>                               Path of disk image to use or create if not exists.
#
#   <script>                              Path of script to execute after installing base system in
#                                         the mounted image and before umounting it.
#
#   <script-opts>                         Arguments to pass to the script.
#
# Options and Environment Variables:
#   -b --branch ALPINE_BRANCH             Alpine branch to install; used only when
#                                         --repositories-file is not specified. Default is
#                                         latest-stable.
#
#   -f --image-format IMAGE_FORMAT        Format of the disk image (see qemu-img --help).
#
#   -s --image-size IMAGE_SIZE            Size of the disk image to create in bytes or with suffix
#                                         (e.g. 1G, 1024M). Default is 2G.
#
#   -i --initfs-features INITFS_FEATURES  List of additional mkinitfs features (basically modules)
#                                         to be included in initramfs (see mkinitfs -L). "base" and
#                                         $ROOTFS is always included, don't specify them here.
#                                         Default is "scsi virtio".
#
#   -k --kernel-flavor KERNEL_FLAVOR      The kernel flavour to install; virt (default), lts
#                                         (Alpine >=3.11), or vanilla (Alpine <3.11).
#
#      --keys-dir KEYS_DIR                Path of directory with Alpine keys to copy into the image.
#                                         Default is /etc/apk/keys. If does not exist, keys for
#                                         x86_64 embedded in this script will be used.
#
#   -m --mirror-uri ALPINE_MIRROR         URI of the Aports mirror to fetch packages; used only
#                                         when --repositories-file is not specified. Default is
#                                         https://nl.alpinelinux.org/alpine.
#
#   -C --no-cleanup (CLEANUP)             Don't umount and disconnect image when done.
#
#   -p --packages PACKAGES                Additional packages to install into the image.
#
#   -r --repositories-file REPOS_FILE     Path of repositories file to copy into the rootfs.
#                                         Default is /etc/apk/repositories. If does not exist,
#                                         repositories file with Alpine's main and community
#                                         repositories on --mirror-uri is created.
#
#      --rootfs ROOTFS                    Filesystem to create on the image. Default is ext4.
#
#   -c --script-chroot (SCRIPT_CHROOT)    Bind <script>'s directory at /mnt inside image and chroot
#                                         into the image before executing <script>.
#
#   -t --serial-console (SERIAL_CONSOLE)  Add configuration for a serial console on ttyS0.
#
#   -h --help                             Show this help message and exit.
#
#   -v --version                          Print version and exit.
#
#   APK                                   APK command to use. Default is "apk".
#
#   APK_OPTS                              Options to pass into apk on each execution.
#                                         Default is "--no-progress".
#
#   APK_TOOLS_URI                         URL of static apk-tools tarball to download if $APK is
#                                         not found on the host system. Default is x86_64 apk-tools
#                                         from https://github.com/alpinelinux/apk-tools/releases.
#
#   APK_TOOLS_SHA256                      SHA-256 checksum of $APK_TOOLS_URI.
#
# Each option can be also provided by environment variable. If both option and
# variable is specified and the option accepts only one argument, then the
# option takes precedence.
#
# https://github.com/alpinelinux/alpine-make-vm-image
#---help---
set -eu

# Some distros (e.g. Arch Linux) does not have /bin or /sbin in PATH.
PATH="$PATH:/usr/sbin:/usr/bin:/sbin:/bin"

readonly PROGNAME='alpine-make-vm-image'
readonly VERSION='0.6.0'
readonly VIRTUAL_PKG=".make-$PROGNAME"

# Alpine APK keys for verification of packages for x86_64.
readonly ALPINE_KEYS='
alpine-devel@lists.alpinelinux.org-4a6a0840.rsa.pub:MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA1yHJxQgsHQREclQu4Ohe\nqxTxd1tHcNnvnQTu/UrTky8wWvgXT+jpveroeWWnzmsYlDI93eLI2ORakxb3gA2O\nQ0Ry4ws8vhaxLQGC74uQR5+/yYrLuTKydFzuPaS1dK19qJPXB8GMdmFOijnXX4SA\njixuHLe1WW7kZVtjL7nufvpXkWBGjsfrvskdNA/5MfxAeBbqPgaq0QMEfxMAn6/R\nL5kNepi/Vr4S39Xvf2DzWkTLEK8pcnjNkt9/aafhWqFVW7m3HCAII6h/qlQNQKSo\nGuH34Q8GsFG30izUENV9avY7hSLq7nggsvknlNBZtFUcmGoQrtx3FmyYsIC8/R+B\nywIDAQAB
alpine-devel@lists.alpinelinux.org-5261cecb.rsa.pub:MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAwlzMkl7b5PBdfMzGdCT0\ncGloRr5xGgVmsdq5EtJvFkFAiN8Ac9MCFy/vAFmS8/7ZaGOXoCDWbYVLTLOO2qtX\nyHRl+7fJVh2N6qrDDFPmdgCi8NaE+3rITWXGrrQ1spJ0B6HIzTDNEjRKnD4xyg4j\ng01FMcJTU6E+V2JBY45CKN9dWr1JDM/nei/Pf0byBJlMp/mSSfjodykmz4Oe13xB\nCa1WTwgFykKYthoLGYrmo+LKIGpMoeEbY1kuUe04UiDe47l6Oggwnl+8XD1MeRWY\nsWgj8sF4dTcSfCMavK4zHRFFQbGp/YFJ/Ww6U9lA3Vq0wyEI6MCMQnoSMFwrbgZw\nwwIDAQAB
'

: ${APK_TOOLS_URI:="https://github.com/alpinelinux/apk-tools/releases/download/v2.10.4/apk-tools-2.10.4-x86_64-linux.tar.gz"}
: ${APK_TOOLS_SHA256:="efe948160317fe78058e207554d0d9195a3dfcc35f77df278d30448d7b3eb892"}

: ${APK:="apk"}
: ${APK_OPTS:="--no-progress"}


# For compatibility with systems that does not have "realpath" command.
if ! command -v realpath >/dev/null; then
	alias realpath='readlink -f'
fi

die() {
	printf '\033[1;31mERROR:\033[0m %s\n' "$@" >&2  # bold red
	exit 1
}

einfo() {
	printf '\n\033[1;36m> %s\033[0m\n' "$@" >&2  # bold cyan
}

# Prints help and exists with the specified status.
help() {
	sed -En '/^#---help---/,/^#---help---/p' "$0" | sed -E 's/^# ?//; 1d;$d;'
	exit ${1:-0}
}

# Cleans the host system. This function is executed before exiting the script.
cleanup() {
	set +eu
	trap '' EXIT HUP INT TERM  # unset trap to avoid loop

	cd /
	if [ -d "$temp_dir" ]; then
		rm -Rf "$temp_dir"
	fi
	if [ "$mount_dir" ]; then
		umount_recursively "$mount_dir" \
			|| die "Failed to unmount $mount_dir; unmount it and disconnect $nbd_dev manually"
		rm -Rf "$mount_dir"
	fi
	if [ "$nbd_dev" ]; then
		qemu-nbd --disconnect "$nbd_dev" \
			|| die "Failed to disconnect $nbd_dev; disconnect it manually"
	fi
	if [ "$INSTALL_HOST_PKGS" = yes ]; then
		_apk del $VIRTUAL_PKG
	fi
}

_apk() {
	"$APK" $APK_OPTS "$@"
}

# Attaches the specified image as a NBD block device and prints its path.
attach_image() {
	local image="$1"
	local format="${2:-}"
	local nbd_dev

	nbd_dev=$(get_available_nbd) || {
		modprobe nbd max_part=0
		sleep 1
		nbd_dev=$(get_available_nbd)
	} || die 'No available nbd device found!'

	qemu-nbd --connect="$nbd_dev" --cache=writeback \
		${format:+--format=$format} "$image" \
		&& echo "$nbd_dev"
}

# Prints UUID of filesystem on the specified block device.
blk_uuid() {
	local dev="$1"
	blkid "$dev" | sed -En 's/.*UUID="([^"]+)".*/\1/p'
}

# Writes Alpine APK keys embedded in this script into directory $1.
dump_alpine_keys() {
	local dest_dir="$1"
	local content file line

	mkdir -p "$dest_dir"
	for line in $ALPINE_KEYS; do
		file=${line%%:*}
		content=${line#*:}

		printf -- "-----BEGIN PUBLIC KEY-----\n$content\n-----END PUBLIC KEY-----\n" \
			> "$dest_dir/$file"
	done
}

# Prints path of available nbdX device, or returns 1 if not any.
get_available_nbd() {
	local dev; for dev in $(find /dev -maxdepth 2 -name 'nbd[0-9]*'); do
		if [ "$(blockdev --getsize64 "$dev")" -eq 0 ]; then
			echo "$dev"; return 0
		fi
	done
	return 1
}

# Prints name of the package needed for creating the specified filesystem.
fs_progs_pkg() {
	local fs="$1"  # filesystem name

	case "$fs" in
		ext4) echo 'e2fsprogs';;
		btrfs) echo 'btrfs-progs';;
		xfs) echo 'xfsprogs';;
	esac
}

# Binds the directory $1 at the mountpoint $2 and sets propagation to private.
mount_bind() {
	mkdir -p "$2"
	mount --bind "$1" "$2"
	mount --make-private "$2"
}

# Prepares chroot at the specified path.
prepare_chroot() {
	local dest="$1"

	mkdir -p "$dest"/proc
	mount -t proc none "$dest"/proc
	mount_bind /dev "$dest"/dev
	mount_bind /sys "$dest"/sys

	install -D -m 644 /etc/resolv.conf "$dest"/etc/resolv.conf
}

# Adds specified services to the runlevel. Current working directory must be
# root of the image.
rc_add() {
	local runlevel="$1"; shift  # runlevel name
	local services="$*"  # names of services

	local svc; for svc in $services; do
		mkdir -p etc/runlevels/$runlevel
		ln -s /etc/init.d/$svc etc/runlevels/$runlevel/$svc
		echo " * service $svc added to runlevel $runlevel"
	done
}

# Installs and configures extlinux.
setup_extlinux() {
	local mnt="$1"  # path of directory where is root device currently mounted
	local root_dev="$2"  # root device
	local modules="$3"  # modules which should be loaded before pivot_root
	local kernel_flavor="$4"  # name of default kernel to boot
	local serial_port="$5"  # serial port number for serial console
	local default_kernel="$kernel_flavor"
	local kernel_opts=''

	[ -z "$serial_port" ] || kernel_opts="console=$serial_port"

	if [ "$kernel_flavor" = 'virt' ]; then
		_apk search --root . --exact --quiet linux-lts | grep -q . \
			&& default_kernel='lts' \
			|| default_kernel='vanilla'
	fi

	sed -Ei \
		-e "s|^[# ]*(root)=.*|\1=$root_dev|" \
		-e "s|^[# ]*(default_kernel_opts)=.*|\1=\"$kernel_opts\"|" \
		-e "s|^[# ]*(modules)=.*|\1=\"$modules\"|" \
		-e "s|^[# ]*(default)=.*|\1=$default_kernel|" \
		-e "s|^[# ]*(serial_port)=.*|\1=$serial_port|" \
		"$mnt"/etc/update-extlinux.conf

	chroot "$mnt" extlinux --install /boot
	chroot "$mnt" update-extlinux --warn-only 2>&1 \
		| grep -Fv 'extlinux: cannot open device /dev' >&2
}

# Configures mkinitfs.
setup_mkinitfs() {
	local mnt="$1"  # path of directory where is root device currently mounted
	local features="$2"  # list of mkinitfs features

	features=$(printf '%s\n' $features | sort | uniq | xargs)

	sed -Ei "s|^[# ]*(features)=.*|\1=\"$features\"|" \
		"$mnt"/etc/mkinitfs/mkinitfs.conf
}

# Unmounts all filesystem under the specified directory tree.
umount_recursively() {
	local mount_point="$1"
	test -n "$mount_point" || return 1

	cat /proc/mounts \
		| cut -d ' ' -f 2 \
		| grep "^$mount_point" \
		| sort -r \
		| xargs umount -rn
}

# Downloads the specified file using wget and checks checksum.
wgets() (
	local url="$1"
	local sha256="$2"
	local dest="${3:-.}"

	cd "$dest" \
		&& wget -T 10 --no-verbose "$url" \
		&& echo "$sha256  ${url##*/}" | sha256sum -c
)


#=============================  M a i n  ==============================#

opts=$(getopt -n $PROGNAME -o b:cCf:hi:k:p:r:s:tv \
	-l branch:,image-format:,image-size:,initfs-features:,kernel-flavor:,keys-dir:,mirror-uri:,no-cleanup,packages:,repositories-file:,rootfs:,script-chroot,serial-console,help,version \
	-- "$@") || help 1 >&2

eval set -- "$opts"
while [ $# -gt 0 ]; do
	n=2
	case "$1" in
		-b | --branch) ALPINE_BRANCH="$2";;
		-f | --image-format) IMAGE_FORMAT="$2";;
		-s | --image-size) IMAGE_SIZE="$2";;
		-i | --initfs-features) INITFS_FEATURES="${INITFS_FEATURES:-} $2";;
		-k | --kernel-flavor) KERNEL_FLAVOR="$2";;
		     --keys-dir) KEYS_DIR="$(realpath "$2")";;
		-m | --mirror-uri) ALPINE_MIRROR="$2";;
		-C | --no-cleanup) CLEANUP='no'; n=1;;
		-p | --packages) PACKAGES="${PACKAGES:-} $2";;
		-r | --repositories-file) REPOS_FILE="$(realpath "$2")";;
		     --rootfs) ROOTFS="$2";;
		-t | --serial-console) SERIAL_CONSOLE='yes'; n=1;;
		-c | --script-chroot) SCRIPT_CHROOT='yes'; n=1;;
		-h | --help) help 0;;
		-V | --version) echo "$PROGNAME $VERSION"; exit 0;;
		--) shift; break;;
	esac
	shift $n
done

: ${ALPINE_BRANCH:="latest-stable"}
: ${ALPINE_MIRROR:="https://nl.alpinelinux.org/alpine"}
: ${CLEANUP:="yes"}
: ${IMAGE_FORMAT:=}
: ${IMAGE_SIZE:="2G"}
: ${INITFS_FEATURES:="scsi virtio"}
: ${KERNEL_FLAVOR:="virt"}
: ${KEYS_DIR:="/etc/apk/keys"}
: ${PACKAGES:=}
: ${REPOS_FILE:="/etc/apk/repositories"}
: ${ROOTFS:="ext4"}
: ${SCRIPT_CHROOT:="no"}
: ${SERIAL_CONSOLE:="no"}

case "$ALPINE_BRANCH" in
	[0-9]*) ALPINE_BRANCH="v$ALPINE_BRANCH";;
esac

if [ -f /etc/alpine-release ]; then
	: ${INSTALL_HOST_PKGS:="yes"}
else
	: ${INSTALL_HOST_PKGS:="no"}
fi

SERIAL_PORT=
[ "$SERIAL_CONSOLE" = 'no' ] || SERIAL_PORT='ttyS0'

[ $# -ne 0 ] || help 1 >&2

IMAGE_FILE="$1"; shift
SCRIPT=
[ $# -eq 0 ] || { SCRIPT=$(realpath "$1"); shift; }

[ "$CLEANUP" = no ] || trap cleanup EXIT HUP INT TERM

#-----------------------------------------------------------------------
if [ "$INSTALL_HOST_PKGS" = yes ]; then
	einfo 'Installing needed packages on host system'

	# We need load btrfs module to avoid the error message:
	# 'failed to open /dev/btrfs-control'
	if ! grep -q -w "$ROOTFS" /proc/filesystems; then
		modprobe $ROOTFS
	fi
	_apk add -t $VIRTUAL_PKG qemu-img $(fs_progs_pkg "$ROOTFS")
fi

#-----------------------------------------------------------------------
temp_dir=''
if ! command -v "$APK" >/dev/null; then
	einfo "$APK not found, downloading static apk-tools"

	temp_dir="$(mktemp -d /tmp/$PROGNAME.XXXXXX)"
	wgets "$APK_TOOLS_URI" "$APK_TOOLS_SHA256" "$temp_dir"
	tar -C "$temp_dir" -xzf "$temp_dir/${APK_TOOLS_URI##*/}"
	APK="$(ls "$temp_dir"/apk-tools-*/apk)"
fi

#-----------------------------------------------------------------------
if [ ! -f "$IMAGE_FILE" ]; then
	einfo "Creating $IMAGE_FORMAT image of size $IMAGE_SIZE"
	qemu-img create ${IMAGE_FORMAT:+-f $IMAGE_FORMAT} "$IMAGE_FILE" "$IMAGE_SIZE"
fi

#-----------------------------------------------------------------------
einfo "Attaching image $IMAGE_FILE as a NBD device"

nbd_dev=$(attach_image "$IMAGE_FILE" "$IMAGE_FORMAT")

#-----------------------------------------------------------------------
einfo "Formatting image to $ROOTFS"

# syslinux 6.0.3 cannot boot from ext4 w/ 64bit feature enabled.
# -E nodiscard / -K - do not attempt to discard blocks at mkfs time (it's
# useless for NBD image and prints confusing error).
[ "$ROOTFS" = ext4 ] && mkfs_args='-O ^64bit -E nodiscard' || mkfs_args='-K'

mkfs.$ROOTFS -L root $mkfs_args "$nbd_dev"

root_uuid=$(blk_uuid "$nbd_dev")
mount_dir=$(mktemp -d /tmp/$PROGNAME.XXXXXX)

#-----------------------------------------------------------------------
einfo "Mounting image at $mount_dir"

mount "$nbd_dev" "$mount_dir"

#-----------------------------------------------------------------------
einfo 'Installing base system'

cd "$mount_dir"

mkdir -p etc/apk/keys
if [ -f "$REPOS_FILE" ]; then
	install -m 644 "$REPOS_FILE" etc/apk/repositories
else
	cat > etc/apk/repositories <<-EOF
		$ALPINE_MIRROR/$ALPINE_BRANCH/main
		$ALPINE_MIRROR/$ALPINE_BRANCH/community
	EOF
fi
if [ -d "$KEYS_DIR" ]; then
	cp "$KEYS_DIR"/* etc/apk/keys/
else
	dump_alpine_keys etc/apk/keys/
fi

_apk add --root . --update-cache --initdb alpine-base
prepare_chroot .

#-----------------------------------------------------------------------
einfo "Installing and configuring mkinitfs"

_apk add --root . mkinitfs
setup_mkinitfs . "base $ROOTFS $INITFS_FEATURES"

#-----------------------------------------------------------------------
einfo "Installing kernel linux-$KERNEL_FLAVOR"

_apk add --root . linux-$KERNEL_FLAVOR

#-----------------------------------------------------------------------
einfo 'Setting up extlinux bootloader'

_apk add --root . --no-scripts syslinux
setup_extlinux . "UUID=$root_uuid" "$ROOTFS" "$KERNEL_FLAVOR" "$SERIAL_PORT"

cat > etc/fstab <<-EOF
	# <fs>		<mountpoint>	<type>	<opts>		<dump/pass>
	UUID=$root_uuid	/		$ROOTFS	noatime		0 1
EOF

#-----------------------------------------------------------------------
if [ "$SERIAL_PORT" ]; then
	echo "$SERIAL_PORT" >> "$mount_dir"/etc/securetty
	sed -Ei "s|^[# ]*($SERIAL_PORT:.*)|\1|" "$mount_dir"/etc/inittab
fi

#-----------------------------------------------------------------------
einfo 'Enabling base system services'

rc_add sysinit devfs dmesg mdev hwdrivers
[ -e etc/init.d/cgroups ] && rc_add sysinit cgroups ||:  # since v3.8

rc_add boot modules hwclock swap hostname sysctl bootmisc syslog
rc_add shutdown killprocs savecache mount-ro

#-----------------------------------------------------------------------
if [ "$PACKAGES" ]; then
	einfo 'Installing additional packages'
	_apk add --root . $PACKAGES
fi

#-----------------------------------------------------------------------
if [ "$SCRIPT" ]; then
	script_name="${SCRIPT##*/}"

	if [ "$SCRIPT_CHROOT" = 'no' ]; then
		einfo "Executing script: $script_name $*"
		"$SCRIPT" "$@" || die 'Script failed'
	else
		einfo "Executing script in chroot: $script_name $*"
		mount_bind "${SCRIPT%/*}" mnt/
		chroot . /bin/sh -c "cd /mnt && ./$script_name \"\$@\"" -- "$@" \
			|| die 'Script failed'
	fi
fi

#-----------------------------------------------------------------------
einfo 'Completed'

cd - >/dev/null
ls -lh "$IMAGE_FILE"