#!/bin/bash
# SPDX-License-Identifier: GPL-3.0-or-later

##########################################################################
### Utility for downloading Talos artifacts and manipulating the cache ###
##########################################################################

# Safety first
set -o errexit
set -o noglob
set -o nounset
set -o errtrace

# Arguments
schematic=
release=
aarch64=
bootid=
action=
jsonform=
verbose=
quiet=

# Talos artifacts are not cached by default
taloscache=

# The default source server for downloading artifacts
talosserver="https://factory.altlinux.space"

# The default Talos schematic without any customizations
talosschematic=376567988ad370138ad8b2698212367b8edcb69b5fd68c80be1f2ec7d603b4ba

# Internal IPv4-address of the netboot server (required)
ipv4_address=

# Full path to this script
readonly scriptname="$(realpath -- "$0")"

# Short name of this program
readonly progname="${scriptname##*/}"

# The default path to the program configuration
readonly defconf=/etc/sysconfig/netboot-setup

# Supplemental sources location
readonly libdir=/usr/libexec/netboot-setup

# Program data location
readonly datadir=/srv/talos

# The "End Of Line" character
readonly EOL='
'


# Displays the help message for this program
#
show_help()
{
	local A="" X=""

	# Determining the default platform
	[ -n "$aarch64" ] && A=" (default)" || X=" (default)"

	cat <<-EOF
	Usage: $progname --get [--] [<schematic>] <release>
	   or: $progname --sel [--] <boot-id>
	   or: $progname --del [--] <boot-id>
	   or: $progname <other action and options>

	Actions and options:
	  -A, --active, --show     Show the last selected (active) <boot-id>.
	                           This action is used unless otherwise specified.
	  -a, --arm64, --aarch64   Select the ARM64 platform suffix${A}.
	  -c, --clean, --cleanup   Clear Talos cache completely ($datadir).
	  -d, --del, --delete      Delete the specified <boot-id> files.
	  -g, --get, --download    Download the specified build into the cache
	                           and set it as the currently active <boot-id>
	                           for the corresponding machine platform.
	  -j, --json               Use JSON output format with --show/--list.
	  -l, --lst, --list        List all loaded <boot-id>'s in the cache.
	  -q, --quiet              Silent execution. Useful if $progname is
	                           run in a script or as some kind of backend.
	  -s, --sel, --select      Select the specified <boot-id> as active.
	  -S, --server=<server>    Specify an alternative download server.
	                           This option can only be used with the '-g'.
	  -x, --amd64, --x86_64    Select the AMD64 platform suffix${X}.
	      --cache              Allow loading of kernel and initrd into cache.
	      --no-cache           Prevent loading of kernel and initrd into cache.
	  -v, --verbose            Output diagnostics for each processed action.
	  -V, --version            Show the version of this program and exit.
	  -h, --help               Show this help message and exit.

	Please, report bugs to https://bugzilla.altlinux.org/
	EOF
	exit 0
}

# Checks and parses command line arguments
#
parse_cmdline()
{
	local msg="With the action '%s' you must specify %s."

	local l_opts="arm64,aarch64,active,json,show,clean,cleanup,delete"
	      l_opts="$l_opts,del,get,download,lst,list,sel,select,server:"
	      l_opts="$l_opts,amd64,x86_64,cache,no-cache,quiet,verbose"
	      l_opts="$l_opts,version,help"
	local s_opts="+AacdgjlqS:sxvVh"

	l_opts=$(getopt -n "$progname" -o "$s_opts" -l "$l_opts" -- "$@") ||
		show_usage
	eval set -- "$l_opts"
	while [ "$#" != 0 ]; do
		case "$1" in
		-a|--arm64|--aarch64)
			aarch64=1
			;;
		-x|--amd64|--x86_64)
			aarch64=
			;;
		-A|--active|--show)
			set_action show
			;;
		-c|--clean|--cleanup)
			set_action cleanup
			;;
		-d|--del|--delete)
			set_action delete
			;;
		-g|--get|--download)
			set_action download
			;;
		-l|--lst|--list)
			set_action list
			;;
		-s|--sel|--select)
			set_action select
			;;
		-S|--server)
			check_arg "$1" "${2-}" "a source server"
			talosserver="$2"
			shift
			;;
		-j|--json)
			jsonform=1
			;;
		-q|--quiet)
			quiet=-q
			;;
		-v|--verbose)
			verbose=-v
			;;
		--cache)
			taloscache=1
			;;
		--no-cache)
			taloscache=
			;;
		-V|--version)
			show_version
			;;
		-h|--help)
			show_help
			;;
		--)	shift
			break
			;;
		-*)	show_usage "Unsupported option: '%s'." "$1"
			;;
		*)	break
			;;
		esac
		shift
	done

	s_opts=0
	l_opts=0

	case "$action" in
	"delete"|"select")
		s_opts=1
		l_opts=1
		;;

	"download")
		s_opts=1
		l_opts=2
		;;

	"")	set_action show
		;;
	esac

	if [ "$#" -gt "$l_opts" ]; then
		show_usage "Too many arguments specified."
	elif [ "$#" -lt "$s_opts" ]; then
		show_usage "Not enough arguments."
	fi

	case "$action" in
	"delete"|"select")
		[ -n "$1" ] ||
			show_usage "$msg" "$action" "the <boot-id>"
		is_number "$1" ||
			show_usage "The <boot-id> must be an integer."
		[ "${#1}" = 8 ] ||
			show_usage "Invalid <boot-id>: '%s'." "$1"
		bootid="$1"
		;;

	"download")
		if [ "$#" = 1 ]; then
			[ -n "$1" ] ||
				show_usage "$msg" "$action" "<release>"
			release="$1"
		else
			[ -n "$2" ] ||
				show_usage "$msg" "$action" "<release>"
			schematic="$1"
			release="$2"
		fi
		schematic="${schematic:-$talosschematic}"
		;;
	esac
}

# Checking the specified $bootid record in the cache,
# setting up the $platform and other variables
#
check_bootid()
{
	local v="$bootid"

	if [ -z "${1-}" ]; then
		local schematic=
		local release=
		local server=
		local cached=
	fi

	# Checking the list and the record
	msg "Checking cache..."
	[ -s "$datadir/LIST" ] ||
		fatal "The cache is empty."
	[ -s "$datadir/$v/META" ] && grep -qs -E "^$v\s+" "$datadir"/LIST ||
		fatal "The specified images were not found."
	platform=

	# Loading metadata
	. "$datadir/$v"/META

	# Consistency check
	[ -n "$server" ] &&
	[ -n "$schematic" ] &&
	[ -n "$release" ] &&
	[ -n "$platform" ] &&
	[ "$v" = "$bootid" ] &&
	in_array "$platform" amd64 arm64 &&
	[ -s "$datadir/$v/script-$platform.ipxe" ] ||
		fatal "An invalid record found in the cache: '%s'." "$v"
	return 0
}

# Show the last selected (current) <boot-id>
#
talos_show()
{
	local s add=
	local server=
	local cached=
	local release=
	local schematic=
	local platform=amd64

	[ -z "$aarch64" ] ||
		platform=arm64
	bootid="$(head -n1 -- "$datadir/CURRENT-$platform" 2>/dev/null ||:)"
	[ -n "$bootid" ] ||
		exit 0
	check_bootid --show
	s="$datadir/script-$platform.ipxe"
	cached="$(test -n "$cached" && echo true || echo false)"

	if [ -n "$jsonform" ]; then
		if [ -L "$s" ]; then
			add="$(readlink -- "$s")"
			add=",$EOL  \"symlink\": \"$add\""
		fi

		cat <<-EOF
		{
		  "bootid": "$bootid",
		  "cached": $cached,
		  "server": "$server",
		  "release": "$release",
		  "schematic": "$schematic",
		  "platform": "$platform"$add
		}
		EOF
	else
		cat <<-EOF
		bootid:		$bootid
		cached:		$cached
		server:		$server
		release:	$release
		schematic:	$schematic
		platform:	$platform

		[iPXE script]
		EOF
		cat -- "$datadir/$bootid/script-$platform.ipxe"

		if [ -L "$s" ]; then
			printf "[Symlink]\n"
			readlink -- "$s"
		fi
	fi
}

# Selects the specified <boot-id> as active
#
talos_select()
{
	local platform=

	check_bootid
	printf "%s\n" "$bootid" >"$datadir/CURRENT-$platform"
	ln -snf $verbose -- "$bootid/script-$platform.ipxe" "$datadir"/
	talos_show
}

# Deletes the specified <boot-id> files
#
talos_delete()
{
	local a v platform=
	local list="$datadir/LIST"

	check_bootid
	msg "Deleting record '%s'..." "$bootid"
	v="$(head -n1 -- "$datadir"/CURRENT-amd64 2>/dev/null ||:)"

	# Removing links
	if [ "$v" = "$bootid" ] && [ "$platform" = amd64 ]; then
		rm -f $verbose -- "$datadir"/script-amd64.ipxe
		rm -f $verbose -- "$datadir"/CURRENT-amd64
		aarch64=
		v=X
	else
		v=
		a="$(head -n1 -- "$datadir"/CURRENT-arm64 2>/dev/null ||:)"

		if [ "$a" = "$bootid" ] && [ "$platform" = arm64 ]; then
			rm -f $verbose -- "$datadir"/script-arm64.ipxe
			rm -f $verbose -- "$datadir"/CURRENT-arm64
			aarch64=1
			v=A
		fi
	fi

	# Removing the record
	rm -rf $verbose -- "$datadir/$bootid"
	sed -i -E "/^$bootid\s+.*$/d" "$list"

	# Replacing links
	if [ -n "$v" ]; then
		bootid="`sed -n -E "s/^([0-9]+)\s+$v\s+.*$/\1/p" \
					"$list" |head -n1`"
		[ -z "$bootid" ] || talos_select
	fi
}

# List all <boot-id>'s loaded into the cache
#
talos_list()
{
	# First, checking the list
	[ -s "$datadir/LIST" ] ||
		fatal "The cache is empty."
	[ -z "$jsonform" ] ||
		talos_list_json
	[ -n "$jsonform" ] ||
		talos_list_plain
	return 0
}

# List <boot-id>'s using JSON format. We should see something like this:
#
# [
#   {
#     "bootid": "25042701",
#     "platform": "amd64",
#     "schematic": "376567988ad370138ad8b2698212367b8edcb69b5fd68c80be1f2ec7d603b4ba",
#     "release": "v1.10.0-alpha.3"
#   }
# ]
#
talos_list_json()
{
	local IFS p fmt start="  "

	printf "[\n"

	fmt="%s{\"bootid\": \"%s\", \"platform\": \"%s\","
	fmt="$fmt \"schematic\": \"%s\", \"release\": \"%s\"}"

	while IFS='	' read -r bootid p schematic release; do
		[ "$p" = A ] && p="arm64" || p="amd64"
		printf "$fmt" "$start" "$bootid" "$p" "$schematic" "$release"
		start=",$EOL  "
	done <"$datadir"/LIST

	printf "\n]\n"
}

# List <boot-id>'s using plain text. We should see something like this:
#
# Boot-ID   Arch   Talos schematic                                                   Talos release
# ========  =====  ================================================================  =============
# 25042701  amd64  376567988ad370138ad8b2698212367b8edcb69b5fd68c80be1f2ec7d603b4ba  v1.10.0-alpha.3
#
talos_list_plain()
{
	local IFS p x="========"
	local fmt="%-8s  %-5s  %-64s  %s\n"

	printf "$fmt" "Boot-ID"  "Arch"  "Talos schematic"  "Talos release"
	printf "$fmt" "========" "=====" "$x$x$x$x$x$x$x$x" "============="

	while IFS='	' read -r bootid p schematic release; do
		[ "$p" = A ] && p="arm64" || p="amd64"
		printf "$fmt" "$bootid" "$p" "$schematic" "$release"
	done <"$datadir"/LIST
}

# Clears the Talos cache completely
#
talos_cleanup()
{
	local dname

	# Deleting files and symlinks
	rm -f $verbose -- "$datadir"/script-amd64.ipxe
	rm -f $verbose -- "$datadir"/script-arm64.ipxe
	rm -f $verbose -- "$datadir"/CURRENT-amd64
	rm -f $verbose -- "$datadir"/CURRENT-arm64
	rm -f $verbose -- "$datadir"/LIST

	# Removing all sub-directories with data
	find "$datadir" -maxdepth 1 -type d -name '????????' -printf '%f\n' |
	while read -r dname; do
		rm -rf $verbose -- "$datadir/$dname"
	done

	# Deleting the working directory
	rm -rf $verbose -- "$datadir"/.download
}

# Downloads Talos artifacts from the Image Factory server into the cache
#
talos_download()
{
	local v p d url
	local list="$datadir/LIST"
	local wd="$datadir/.download"

	# Checking configuration
	msg "Checking configuration..."
	[ -n "$talosserver" ] && [ -n "$schematic" ] && [ -n "$ip4_address" ] ||
		fatal "talosserver, schematic and ip4_address are required."

	# Determining the letter depending on the platform
	[ -n "$aarch64" ] && p="A" || p="X"

	# Have the requested images been downloaded previously?
	if [ -s "$list" ]; then
		v="$(str2regex "$release")"
		d="$(str2regex "$schematic")"
		d="`sed -n -E "s/^([0-9]+)\s+$p\s+$d\s+$v$/\1/p" \
						"$list" |tail -n1`"
		[ -z "$d" ] || [ ! -s "$datadir/$d/META" ] ||
			fatal "Images already loaded: '%s'." "$d"

		# Deleting invalid records from the list
		if [ -n "$d" ]; then
			sed -i -E "/^$d\s+.*$/d" "$list"
			rm -rf $verbose -- "$datadir/$d"
		fi
	fi

	# Preparing the working directory
	trap 'rm -rf $verbose -- "$datadir"/.download' EXIT
	mkdir -p $verbose -- "$wd"

	# Determining the next <boot-id>
	v="$(env LC_TIME=C date +'%y%m%d')"
	d="$(find "$datadir" -maxdepth 1 -type d -name "${v}[0-9][0-9]" \
				-printf '%f\n' |sort -nr |head -n1)"
	[ -z "$d" ] && bootid="${v}01" || bootid="$((1 + $d))"
	rm -rf $verbose -- "$datadir/$bootid"

	# Determining the platform-specific suffix
	[ -n "$aarch64" ] && d="arm64" || d="amd64"

	# Setting up downloading options
	local opts="${quiet:+-s -S }$verbose"

	# Source iPXE script
	msg "Downloading iPXE script..."
	url="$talosserver/pxe/$schematic/$release"
	curl $opts -o "$wd/metal-$d" -- "$url/metal-$d" >&2

	# Reading kernel boot parameters
	v="$(sed -n -E 's/^kernel //p' "$wd/metal-$d" |head -n1)"
	[ -n "$v" ] && grep -qs -E '^initrd ' "$wd/metal-$d" ||
		fatal "An invalid iPXE script was received."
	v="${v#* }"

	# Caching (or not) kernel and initrd
	url="$talosserver/image/$schematic/$release"
	if [ -z "$taloscache" ]; then
		msg "Creating a new iPXE script..."
		cat >"$wd/script-$d.ipxe" <<-EOF
		#!ipxe

		imgfree
		kernel $url/kernel-$d initrd=initramfs-$d.xz $v
		initrd $url/initramfs-$d.xz
		boot

		EOF
	else
		# Linux kernel
		msg "Downloading Linux kernel..."
		curl $opts -o "$wd/vmlinuz" -- "$url/kernel-$d" >&2

		# Initramfs
		msg "Downloading initramfs image..."
		curl $opts -o "$wd/initrd.img" -- "$url/initramfs-$d.xz" >&2

		# Creating a new iPXE script
		msg "Committing changes to cache..."
		cat >"$wd/script-$d.ipxe" <<-EOF
		#!ipxe

		imgfree
		kernel http://$ip4_address/talos/$bootid/vmlinuz initrd=initrd.img $v
		initrd http://$ip4_address/talos/$bootid/initrd.img
		boot

		EOF
	fi

	# Finishing an action and removing the exit handler
	printf "%s\t%s\t%s\t%s\n" "$bootid" "$p" \
		"$schematic" "$release" >>"$list"
	mv -f $verbose -- "$wd" "$datadir/$bootid"
	trap - EXIT

	# Changing links
	printf "%s\n" "$bootid" >"$datadir/CURRENT-$d"
	ln -snf $verbose -- "$bootid/script-$d.ipxe" "$datadir"/

	# Saving metadata
	cat >"$datadir/$bootid"/META <<-EOF
	server=$talosserver
	schematic=$schematic
	release=$release
	cached=$taloscache
	bootid=$bootid
	platform=$d
	EOF

	printf "%s\n" "$bootid"
}


# Entry point
. "$libdir"/common.sh

# Catch all unexpected errors
trap 'unexpected_error "${BASH_SOURCE[0]##*/}" "$LINENO"' ERR

# Use default configuration (required)
[ -s "$defconf" ] && . "$defconf" ||
	fatal "The netboot server must be configured first."
umask 0022

# Determining the current platform
[ "$(uname -m)" != aarch64 ] || aarch64=1

# Parsing command line arguments
parse_cmdline "$@"

# Security check
check_regular_user

# Executing the specified action
"talos_$action"

