#!/bin/sh
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Zygmunt Krynicki
set -eu

# Project directory or place of invocation.
GARDEN_PROJECT_DIR="$(pwd)"
GARDEN_ITSELF="$0"

GARDEN_PREFIX=@PREFIX@
GARDEN_INCLUDE=@PREFIX@/include
GARDEN_DOC=@PREFIX@/share/doc/image-garden
# This hacks supports both installed and uninstalled version equally.
if [ "$GARDEN_PREFIX" = @'PREFIX'@ ]; then
	GARDEN_ITSELF="$(readlink -f "$0")"
	GARDEN_PREFIX="$(dirname "$GARDEN_ITSELF")"
	GARDEN_INCLUDE="$GARDEN_PREFIX"
	GARDEN_DOC="$GARDEN_PREFIX"
fi
: "${XDG_CACHE_HOME:=$HOME/.cache}"

if [ -n "${SNAP_USER_COMMON:-}" ]; then
	: "${GARDEN_CACHE_DIR:=$SNAP_USER_COMMON/cache}"
else
	: "${GARDEN_CACHE_DIR:=$XDG_CACHE_HOME/image-garden}"
fi
: "${GARDEN_DL_DIR:=$GARDEN_CACHE_DIR/dl}"

# Set XDG_RUNTIME_DIR to just /run as a fallback. This happens in CI/CD systems
# that run a container runtime and don't use systemd or its PAM module during
# the login process.
: "${XDG_RUNTIME_DIR:=/run}"

download_with_sha512sums() {
	set -e
	url="${1:-}"
	sha512sums_url="${2:-"$(dirname "$url")/SHA512SUMS"}"
	directory_prefix="${3:-"$GARDEN_DL_DIR"}"
	filename="${4:-$(basename "$url")}"

	if [ -z "$url" ]; then
		echo "cannot download file without an URL" 1>&2
		return 1
	fi

	for attempt in first retry; do
		if [ ! -f "$directory_prefix"/"$filename".sha512sums ]; then
			rm -vf "$directory_prefix"/"$filename".interrupted
			trap 'touch "$directory_prefix"/"$filename".interrupted' INT

			echo "Downloading SHA512SUMS from $sha512sums_url to $directory_prefix/$filename.sha512sums.partial"
			wget --output-document="$directory_prefix/$filename".sha512sums.partial "$sha512sums_url"

			trap - INT

			# Wget returns 0 on SIGINT. On interrupt remove the interrupted file and
			# exit with an unsuccesful exit code. Remove the partial file as it is
			# our only protection against corruptions of the main content and the
			# checksum files are small.
			if [ -f "$directory_prefix"/"$filename".interrupted ]; then
				rm -vf "$directory_prefix"/"$filename".interrupted "$directory_prefix"/"$filename".sha512sums.partial
				return 1
			fi

			mv "$directory_prefix"/"$filename".sha512sums.partial "$directory_prefix"/"$filename".sha512sums
		fi

		if [ ! -f "$directory_prefix"/"$filename" ]; then
			rm -vf "$directory_prefix"/"$filename".interrupted
			trap 'touch "$directory_prefix"/"$filename".interrupted' INT

			echo "Downloading $filename from $url to $directory_prefix/$filename"
			# shellcheck disable=SC2046
			wget $(test -n "${CI:-}" && echo "--no-verbose") --continue --output-document="$directory_prefix"/"$filename".partial "$url"

			trap - INT

			# Wget returns 0 on SIGINT. On interrupt remove the interrupted file
			# and exit with an unsuccesful exit code. Leave the partial file around
			# so that we can resume the download later.
			if [ -f "$directory_prefix"/"$filename".interrupted ]; then
				rm -vf "$directory_prefix"/"$filename".interrupted
				return 1
			fi

			mv "$directory_prefix"/"$filename".partial "$directory_prefix"/"$filename"
		fi

		echo "Verifying sha512 checksum of $directory_prefix/$filename"
		grep -F "$filename" "$directory_prefix"/"$filename".sha512sums
		if grep -F "$filename" "$directory_prefix"/"$filename".sha512sums | (cd "$directory_prefix" && sha512sum --check); then
			return 0
		else
			# On failure unlink both the files so that we re-download the
			# checksum file. This can help if we have a stale checksum from the
			# past that will never match the current download.
			rm -vf "$directory_prefix"/"$filename" "$directory_prefix"/"$filename".sha512sums
			if [ "$attempt" = first ]; then
				continue
			fi
			return 1
		fi
	done
}

download_with_sha256sums() {
	set -e
	url="${1:-}"
	sha256sums_url="${2:-"$(dirname "$url")/SHA256SUMS"}"
	directory_prefix="${3:-"$GARDEN_DL_DIR"}"
	filename="${4:-$(basename "$url")}"

	if [ -z "$url" ]; then
		echo "cannot download file without an URL" 1>&2
		return 1
	fi

	for attempt in first retry; do
		if [ ! -f "$directory_prefix"/"$filename".sha256sums ]; then
			rm -vf "$directory_prefix"/"$filename".interrupted
			trap 'touch "$directory_prefix"/"$filename".interrupted' INT

			echo "Downloading SHA256SUMS from $sha256sums_url to $directory_prefix/$filename.sha256sums.partial"
			wget --output-document="$directory_prefix/$filename".sha256sums.partial "$sha256sums_url"

			trap - INT

			# Wget returns 0 on SIGINT. On interrupt remove the interrupted file and
			# exit with an unsuccesful exit code. Remove the partial file as it is
			# our only protection against corruptions of the main content and the
			# checksum files are small.
			if [ -f "$directory_prefix"/"$filename".interrupted ]; then
				rm -vf "$directory_prefix"/"$filename".interrupted "$directory_prefix"/"$filename".sha256sums.partial
				return 1
			fi

			mv "$directory_prefix"/"$filename".sha256sums.partial "$directory_prefix"/"$filename".sha256sums
		fi

		if [ ! -f "$directory_prefix"/"$filename" ]; then
			rm -vf "$directory_prefix"/"$filename".interrupted
			trap 'touch "$directory_prefix"/"$filename".interrupted' INT

			echo "Downloading $filename from $url to $directory_prefix/$filename.partial"
			# shellcheck disable=SC2046
			wget $(test -n "${CI:-}" && echo "--no-verbose") --continue --output-document="$directory_prefix"/"$filename".partial "$url"

			trap - INT

			# Wget returns 0 on SIGINT. On interrupt remove the interrupted file
			# and exit with an unsuccesful exit code. Leave the partial file around
			# so that we can resume the download later.
			if [ -f "$directory_prefix"/"$filename".interrupted ]; then
				rm -vf "$directory_prefix"/"$filename".interrupted
				return 1
			fi

			mv -v "$directory_prefix"/"$filename".partial "$directory_prefix"/"$filename"
		fi

		echo "Verifying sha256 checksum of $directory_prefix/$filename"
		grep -F "$filename" "$directory_prefix"/"$filename".sha256sums
		if grep -F "$filename" "$directory_prefix"/"$filename".sha256sums | (cd "$directory_prefix" && sha256sum --check); then
			return 0
		else
			# On failure unlink both the files so that we re-download the
			# checksum file. This can help if we have a stale checksum from the
			# past that will never match the current download.
			rm -vf "$directory_prefix"/"$filename" "$directory_prefix"/"$filename".sha256sums
			if [ "$attempt" = first ]; then
				continue
			fi
			return 1
		fi
	done
}

download_with_sha256sum() {
	set -e
	url="${1:-}"
	sha256sum="${2:-}"
	directory_prefix="${3:-"$GARDEN_DL_DIR"}"
	filename="${4:-$(basename "$url")}"

	if [ -z "$url" ]; then
		echo "cannot download file without an URL" 1>&2
		return 1
	fi

	if [ -z "$sha256sum" ]; then
		echo "cannot download file without a sha256sum" 1>&2
		return 1
	fi

	for attempt in first retry; do
		if [ ! -f "$directory_prefix"/"$filename" ]; then
			rm -vf "$directory_prefix"/"$filename".interrupted
			trap 'touch "$directory_prefix"/"$filename".interrupted' INT

			echo "Downloading $filename from $url to $directory_prefix/$filename.partial"
			# shellcheck disable=SC2046
			wget $(test -n "${CI:-}" && echo "--no-verbose") --continue --output-document="$directory_prefix"/"$filename".partial "$url"

			trap - INT

			# Wget returns 0 on SIGINT. On interrupt remove the interrupted file
			# and exit with an unsuccesful exit code. Leave the partial file around
			# so that we can resume the download later.
			if [ -f "$directory_prefix"/"$filename".interrupted ]; then
				rm -vf "$directory_prefix"/"$filename".interrupted
				return 1
			fi

			mv -v "$directory_prefix"/"$filename".partial "$directory_prefix"/"$filename"
		fi

		echo "Verifying sha256 checksum of $directory_prefix/$filename (expecting $sha256sum)"
		if [ "$sha256sum" = "$(sha256sum "$directory_prefix"/"$filename" | awk '{ print $1 }')" ]; then
			return 0
		else
			rm -vf "$directory_prefix"/"$filename"
			if [ "$attempt" = first ]; then
				continue
			fi
			return 1
		fi
	done
}

download() {
	set -e
	url="${1:-}"
	directory_prefix="${2:-"$GARDEN_DL_DIR"}"
	filename="${3:-$(basename "$url")}"

	if [ -z "$url" ]; then
		echo "cannot download file without an URL" 1>&2
		return 1
	fi

	if [ ! -f "$directory_prefix"/"$filename" ]; then
		trap 'touch "$directory_prefix"/"$filename".interrupted' INT

		echo "Downloading $filename from $url to $directory_prefix/$filename.partial"
		# shellcheck disable=SC2046
		wget $(test -n "${CI:-}" && echo "--no-verbose") --continue --output-document="$directory_prefix/$filename".partial "$url"

		trap - INT

		# Wget returns 0 on SIGINT. On interrupt remove both the interrupted file
		# and the partial file. Without the checksum to protect us against corruption,
		# it's better to be safe than sorry.
		if [ -f "$directory_prefix"/"$filename".interrupted ]; then
			rm -vf "$directory_prefix"/"$filename".interrupted "$directory_prefix"/"$filename".partial
			return 1
		fi

		mv -v "$directory_prefix"/"$filename".partial "$directory_prefix"/"$filename"
	fi
}

kvm_check() {
	if [ -e /dev/kvm ] && [ ! -w /dev/kvm ]; then
		cat <<__KVM_WARNING__

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

Your system supports hardware-assisted virtualization because /dev/kvm exists,
but you do not have write access to the device. Lack of hardware-assisted
virtualization severely impacts performance and is strongly not recommended for
production use.

Please ensure you have write access for much better performance.
__KVM_WARNING__
		if [ -n "${CI-}" ]; then
			cat <<__ADVICE_CI__
This seems to be a CI/CD environment. You may simply grant everyone access to
/dev/kvm. This is not recommended unless the environment is ephemeral (e.g.
GitHub or GitLab runner).

$ sudo chmod 666 /dev/kvm

__ADVICE_CI__
		else
			cat <<'__ADVICE__'
On a typical system /dev/kvm is owned by the "kvm" group. You can add yourself
to the group with:

$ sudo usermod -aG kvm $LOGNAME

You will have to log out and log back in for the group membership to take effect.
__ADVICE__
		fi
	fi
}

case "${1:-}" in
'' | -h | -\? | --help)
	echo "Usage: image-garden make GOAL"
	echo ""
	echo "       image-garden allocate SYSTEM"
	echo "       image-garden rebase SYSTEM.ARCH"
	echo "       image-garden discard ADDRESS"
	echo "       image-garden version"
	echo "       image-garden news"
	echo "       image-garden download-with-sha512sums URL SHA512SUMS DIRECTORY_PREFIX [FILENAME]"
	echo "       image-garden download-with-sha256sums URL SHA256SUMS DIRECTORY_PREFIX [FILENAME]"
	echo "       image-garden download-with-sha256sum  URL SHA256SUM  DIRECTORY_PREFIX [FILENAME]"
	echo "       image-garden download URL DIRECTORY_PREFIX [FILENAME]"
	echo ""
	echo "Typically GOAL is 'SYSTEM.qcow2' or 'SYSTEM.run'."
	echo "Use 'image-garden make list-systems' to see all defined systems."
	;;
--version | version)
	echo "image-garden @VERSION@"
	;;
news)
	for NEWS_SUFFIX in "" ".snappy" ".docker"; do
		NEWS_FILE="$GARDEN_DOC"/NEWS"$NEWS_SUFFIX"
		if [ ! -f "$NEWS_FILE" ] && [ -n "${SNAP:-}" ]; then
			NEWS_FILE="${SNAP:-}"/"$GARDEN_DOC"/NEWS"$NEWS_SUFFIX"
		fi

		if [ ! -f "$NEWS_FILE" ]; then
			if [ -n "$NEWS_SUFFIX" ]; then
				continue
			fi

			echo "Cannot find the NEWS file $NEWS_FILE" >&2
			exit 1
		fi

		LAST=${2:-1}
		IFS=''
		while read -r L; do
			if [ "$(echo "$L" | cut -b 1-11)" = "Changes in " ]; then
				LAST=$((LAST - 1))
				if [ "$LAST" -lt 0 ]; then
					break
				fi
				echo "*** $L ***"
			else
				echo "$L"
			fi
		done <"$NEWS_FILE"
	done
	;;
rebase)
	shift
	if [ $# -ne 1 ]; then
		echo "Usage: image-garden rebase SYSTEM.ARCH"
		echo
		echo "Rebase the system image to the current location of the GARDEN_DL_DIR."
		exit 0
	fi
	IMAGE="$1"
	IMAGE_PATH="$1"
	shift
	if [ -d "$GARDEN_PROJECT_DIR"/.image-garden ]; then
		IMAGE_PATH="$GARDEN_PROJECT_DIR"/.image-garden/"$IMAGE"
	fi
	if [ ! -f "$IMAGE_PATH" ]; then
		echo "image-garden: cannot find image at $IMAGE_PATH" >&2
		exit 1
	fi
	if [ ! -d "$GARDEN_DL_DIR" ]; then
		echo "image-garden: cannot rebase without downloaded image directory: $GARDEN_DL_DIR" >&2
		exit 1
	fi
	# I don't quite want to depend on jq if getting the value of the quoted string can be done with somewhat more crude tools.
	BACKING_IMG="$(qemu-img info --output=json "$IMAGE_PATH" | grep full-backing-filename | cut -d : -f 2 | sed -e 's/^ \+"//' -e 's/",\?$//')"
	NEW_BACKING_IMG="$(find "$GARDEN_DL_DIR" -name "$(basename "$BACKING_IMG")" -print -quit)"
	echo "image-garden: re-basing $IMAGE_PATH on top of $NEW_BACKING_IMG"
	exec qemu-img rebase -b "$NEW_BACKING_IMG" -F qcow2 "$IMAGE_PATH"
	;;
make)
	shift

	kvm_check

	exec make \
		--warn-undefined-variables \
		GARDEN_ITSELF="$GARDEN_ITSELF" \
		GARDEN_PROJECT_DIR="$GARDEN_PROJECT_DIR" \
		-f "$GARDEN_INCLUDE"/image-garden.mk \
		-C "$(test -d "$GARDEN_PROJECT_DIR"/.image-garden && echo "$GARDEN_PROJECT_DIR"/.image-garden || echo .)" \
		"$@"
	;;

allocate)
	shift
	case "${XDG_SESSION_TYPE:-}" in
	wayland | x11)
		nice="nice"
		;;
	esac
	while [ $# -gt 0 ]; do
		case "$1" in
		--premade)
			premade=1
			shift
			;;
		--nice)
			nice="nice"
			shift
			;;
		-*)
			echo "image-garden: unknown option: $1" >&2
			shift
			;;
		*)
			break
			;;
		esac
	done
	SPREAD_SYSTEM="$1"
	SPREAD_SYSTEM_PROPER="$(echo "$SPREAD_SYSTEM" | sed -e 's/x86-64/x86_64/g')"
	shift

	kvm_check

	if [ "${premade:=0}" -eq 0 ]; then
		echo "Making $SPREAD_SYSTEM_PROPER..." >&2
		# Lock on a per-system lock file since parallel-invocations of make should
		# not attempt to download the same file over and over again.
		${nice:-} flock "$(test -d .image-garden && echo .image-garden/ || echo ./)""$SPREAD_SYSTEM_PROPER".lock "$GARDEN_ITSELF" make "$SPREAD_SYSTEM_PROPER".run "$SPREAD_SYSTEM_PROPER".qcow2 >&2
	fi

	# Lock on spread.yaml as we need to prevent port numbers from clashing.
	# This artificially limits virtual machine construction but it is better
	# than having qemu racing to allocate the same port.
	exec 4<>spread.yaml
	echo "Looking available TCP port for $SPREAD_SYSTEM_PROPER..." >&2
	while ! flock --exclusive -n 4; do
		sleep 1
	done

	# Find an unused port to give to qemu as the local version of port 22 in the
	# virtual machine. This is not really very good but it's better than nothing.
	mkdir -p "${XDG_RUNTIME_DIR}"/image-garden
	MAYBE_PORT=
	while test -z "$MAYBE_PORT" || nc -w 1 127.0.0.1 "$MAYBE_PORT" </dev/null >/dev/null; do
		MAYBE_PORT=$(shuf -i 5000-9999 -n 1)
	done
	PORT="$MAYBE_PORT"

	# Disable display unless the user wants to have something custom.
	export QEMU_DISPLAY_OPTION="${QEMU_DISPLAY_OPTION:=-display none}"
	# Forward tcp port $PORT to port 22 in the virtual machine so that spread can connect.
	export QEMU_NETDEV_USER_EXTRA=",hostfwd=tcp:127.0.0.1:$PORT-:22"
	# Run qemu redirecting serial port to a file, disabling parallel port,
	# disabling QEMU monitor and daemonizing with a pidfile.
	parent=$$                                                              # this shell
	parent_parent="$(cut -f 4 -d ' ' </proc/"$parent"/stat)"               # bash from spread
	parent_parent_parent="$(cut -f 4 -d ' ' </proc/"$parent_parent"/stat)" # spread, hopefully
	parent_parent_parent_comm="$(cut -f 3 -d ' ' </proc/"$parent_parent_parent"/comm)"
	echo "Starting $SPREAD_SYSTEM_PROPER..." >&2
	# shellcheck disable=SC2086
	if (if [ -d .image-garden ]; then cd .image-garden; fi && ${nice:-} ./"$SPREAD_SYSTEM_PROPER".run \
		${QEMU_SNAPSHOT_OPTION=-snapshot} \
		-parallel none \
		-monitor none \
		-serial file:"$SPREAD_SYSTEM_PROPER"."$PORT".serial.log \
		-pidfile "${XDG_RUNTIME_DIR}"/image-garden/port."$PORT".pid \
		-daemonize \
		"$@" \
		>"$SPREAD_SYSTEM_PROPER"."$PORT".stdout.log \
		2>"$SPREAD_SYSTEM_PROPER"."$PORT".stderr.log \
		4<&-); then
		# Print the address where spread can connect. If spread is the parent
		# process of the bash process then implement ADDRESS directly.
		# Otherwise fall back to legacy mode where we just print the value.
		# Parent of parent - bash invoked by spread (unconditionally)
		# Parent of parent of parent - possibly spread
		if echo "$parent_parent_parent_comm" | grep -Eq '^spread(-[a-z][a-z0-9-]*)?$'; then
			echo "<ADDRESS localhost:$PORT>"
		else
			echo "localhost:$PORT"
		fi
		exit
	else
		rm -f "${XDG_RUNTIME_DIR}"/image-garden/port."$PORT".pid
		# Parent of parent - bash invoked by spread (unconditionally)
		# Parent of parent of parent - possibly spread
		if [ "$parent_parent_parent_comm" = spread ]; then
			echo "<FATAL $(tail -n 1 "$(test -d .image-garden && echo .image-garden/ || echo ./)""$SPREAD_SYSTEM_PROPER"."$PORT".stderr.log)>"
			exit 213
		else
			echo "Allocation failed, see $(test -d .image-garden && echo .image-garden/ || echo ./)$SPREAD_SYSTEM_PROPER.$PORT.stderr.log" >&2
			exit 1
		fi
	fi
	;;

discard)
	shift
	SPREAD_SYSTEM_ADDRESS="$1"
	shift

	case "$SPREAD_SYSTEM_ADDRESS" in
	localhost:*)
		PORT="$(echo "$SPREAD_SYSTEM_ADDRESS" | sed -e 's/localhost://')"
		if [ -f "${XDG_RUNTIME_DIR}"/image-garden/port."$PORT".pid ]; then
			PID="$(cat "${XDG_RUNTIME_DIR}"/image-garden/port."$PORT".pid)"
			if [ -n "$PID" ]; then
				kill "$PID" || true
			fi
			rm -f "${XDG_RUNTIME_DIR}"/image-garden/port."$PORT".pid
		fi
		;;
	*)
		exit 1
		;;
	esac
	;;

download-with-sha512sums)
	shift
	download_with_sha512sums "$@"
	;;

download-with-sha256sums)
	shift
	download_with_sha256sums "$@"
	;;

download-with-sha256sum)
	shift
	download_with_sha256sum "$@"
	;;

download)
	shift
	download "$@"
	;;

*)
	echo "image-garden: unknown command: $1" >&2
	exit 1
	;;
esac
