#!/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:
#   -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; vanilla or virt.
#                                         Default is virt.
#
#      --keys-dir KEYS_DIR                Path of directory with Alpine keys to copy into the image.
#                                         Default is /etc/apk/keys.
#
#   -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 image.
#                                         Default is /etc/apk/repositories.
#
#      --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>.
#
#   -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

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

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

: ${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'
}

# 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
}

# Installs tools in the host system needed for creating the specified filesystem.
install_fs_tools() {
	local fs="$1"  # filesystem name

	# We need load btrfs module early to avoid the error message:
	# 'failed to open /dev/btrfs-control'
	if ! grep -q -w "$fs" /proc/filesystems; then
		modprobe $fs
	fi

	local pkg
	case "$fs" in
		ext4) pkg='e2fsprogs';;
		btrfs) pkg='btrfs-progs';;
		xfs) pkg='xfsprogs';;
		*) die "Unsupported filesystem: $fs";;
	esac

	_apk add -t $VIRTUAL_PKG $pkg
}

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

	mkdir -p "$dest"/proc "$dest"/dev "$dest"/sys
	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 default_kernel="${4:-"vanilla"}"  # name of default kernel to boot
	local kernel_opts="${5:-}"  # default kernel options

	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|" \
		"$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 cCf:hi:k:p:r:s:v \
	-l image-format:,image-size:,initfs-features:,kernel-flavor:,keys-dir:,no-cleanup,packages:,repositories-file:,rootfs:,script-chroot,help,version \
	-- "$@") || help 1 >&2

eval set -- "$opts"
while [ $# -gt 0 ]; do
	n=2
	case "$1" in
		-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")";;
		-C | --no-cleanup) CLEANUP='no'; n=1;;
		-p | --packages) PACKAGES="${PACKAGES:-} $2";;
		-r | --repositories-file) REPOS_FILE="$(realpath "$2")";;
		     --rootfs) ROOTFS="$2";;
		-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

: ${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"}

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

[ $# -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'

	_apk add -t $VIRTUAL_PKG qemu-img
	install_fs_tools "$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
install -m 644 "$REPOS_FILE" etc/apk/repositories
cp "$KEYS_DIR"/* etc/apk/keys/

_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#virt}

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

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

rc_add sysinit devfs dmesg mdev hwdrivers
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 . sh -c "cd /mnt && ./$script_name \"\$@\"" -- "$@" \
			|| die 'Script failed'
	fi
fi

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

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