#!/bin/sh
# SPDX-License-Identifier: MIT

# shellcheck disable=SC3060,SC3028,SC3043,SC2015,SC2064,SC3057

# Helper script to rollback to previous install version
# from mkswu scripts (pre_init for env, post_boot for actual switch)

usage() {
	echo "Usage: $0 [options] [action]"
	echo
	echo "Armadillo Base OS control utility"
	echo
	echo "The default action is status."
	echo
	echo "Actions:"
	echo "    status: print and check current active partition"
	echo "    reboot: reboot after making sure no update is in progress"
	echo "    update: install update from network (/etc/swupdate.watch)"
	echo "    rollback: force a rollback to the other version"
	echo "    rollback-force <dev>: rollback to specified target device"
	echo "    rollback-clone: copy current system on backup partition"
	echo "    mount-old: mount previous baseos on /target"
	echo "    make-installer: make install disk image from current system"
	echo "    podman-storage: display or change podman main storage (disk or tmpfs)"
	echo "    podman-rw: manually run podman with its read only storage as read-write"
	echo "    container-clear: remove all containers and related data"
	echo "    reset-default: delete user settings and related files"
	echo "    restore-kernel: remove kernel pinning and reinstall version from repo"
	echo "    certificates: manage /etc/swupdate.pem certificates"
	echo "    mount: probe type & mount command"
	echo "    umount: recursive umount command"
	echo "    installer-setting: configure the production Armadillo's settings"
	echo "    check-secure: display security settings"
	echo "    usb-filter: manage USB devices allowed to connect"
	echo "    rtos: manage rtos firmware"
	echo
	echo "Options:"
	echo "    -n/--dry-run: do not actually rollback, just print what would be done"
	echo "    -v/--verbose: add trace messages"
	echo "    -q/--quiet: do not display info messages"
	echo "                repeat to disable warnings as well"
	echo "    -V/--version: show version and exit"
}

trace() {
	[ "${debug:-2}" -lt 3 ] && return
	printf "%s\n" "$@"
}

info() {
	[ "${debug:-2}" -lt 2 ] && return
	printf "%s\n" "$@"
}

warning() {
	[ "${debug:-2}" -lt 1 ] && return
	printf "WARNING: %s\n" "$@" >&2
}

# shellcheck disable=SC2317 # redefining error is seen as unreachable
error() {
	# redefine error to allow soft errors on cleanup
	error() {
		printf "While cleaning up, another error happened\n" >&2
		printf "ERROR: %s\n" "$@" >&2
		return
	}
	printf "ERROR: %s\n" "$@" >&2
	exit 1
}

prompt_yesno() {
	local default="$1" YN=""
	while true; do
		[ -z "$noprompt" ] && read -r YN
		[ -z "$YN" ] && YN="$default"
		case "$YN" in
		[Yy]|[Yy][Ee][Ss]) return 0;;
		[Nn]|[Nn][Oo]) return 1;;
		esac
		echo "Please answer y, n, or empty value (default $default)"
	done
}

has_function() {
	type "$1" | grep -q "is a function"
}

get_appdev() {
	[ -n "$appdev" ] && return

	appdev="$(findmnt -nr --nofsroot -o SOURCE /var/tmp)"
	[ -n "$appdev" ] || error "Could not find appfs"
}

fw_printenv_nowarn() {
	FILTER="Cannot read environment, using default|Cannot read default environment from file|Environment WRONG|Environment OK" \
		filter command fw_printenv "$@"
}

fw_setenv_nowarn() {
	FILTER="Cannot read environment, using default|Environment WRONG|Environment OK" \
		filter command fw_setenv "$@"
}

filter() {
	local output ret

	output="$(mktemp /tmp/cmd_output.XXXXXX)" \
		|| error "Could not create tmpfile for $1"

	"$@" > "$output" 2>&1
	ret=$?

	if [ -n "$FILTER" ]; then
		grep -vE "$FILTER" < "$output"
	else
		cat "$output"
	fi
	rm -f "$output"
	return "$ret"
}

###################################
# encryption helpers              #
# from mkswu/scripts/common.sh    #
###################################

luks_unlock() {
	# modifies dev if unlocked
	local target="$1"
	[ -n "$dev" ] || error "\$dev must be set"

	if [ -e "/dev/mapper/$target" ]; then
		# already unlocked, use it
		dev="/dev/mapper/$target"
		return
	fi

	# no cryptsetup command = assume not luks, skip
	command -v cryptsetup > /dev/null \
		|| return 0

	# not luks? nothing to do!
	cryptsetup isLuks "$dev" \
		|| return 0

	command -v caam-decrypt > /dev/null \
		|| return 0

	local index offset
	case "$dev" in
	*mmcblk*p*)
		# keys are stored in $rootdev as follow
		# 0MB        <GPT header and partition table>
		# 9MB        key for part 1
		# 9MB+4k     key for part 2
		# 9MB+(n*4k) key for part n+1
		# 10MB       first partition
		index=${dev##*p}
		index=$((index-1))
		offset="$(((9*1024 + index*4)*1024))"
		;;
	*) error "LUKS only supported on mmcblk*p* partitions" ;;
	esac

	mkdir -p /run/caam
	local KEYFILE=/run/caam/lukskey
	# use unshared tmpfs to not leak key too much
	# key is:
	# - 112 bytes of caam black key
	# - 16 bytes of iv followed by rest of key
	unshare -m sh -c "mount -t tmpfs tmpfs /run/caam \
		&& dd if=$rootdev of=$KEYFILE.mmc bs=4k count=1 status=none \
			iflag=skip_bytes skip=$offset \
		&& dd if=$KEYFILE.mmc of=$KEYFILE.bb bs=112 count=1 status=none \
		&& dd if=$KEYFILE.mmc of=$KEYFILE.enc bs=4k status=none \
			iflag=skip_bytes skip=112 \
		&& caam-decrypt $KEYFILE.bb AES-256-CBC $KEYFILE.enc \
			$KEYFILE.luks >/dev/null 2>&1 \
		&& cryptsetup luksOpen --key-file $KEYFILE.luks \
			--allow-discards $dev $target >/dev/null 2>&1" \
		|| return 0

	dev="/dev/mapper/$target"
}

luks_format() {
	# modifies dev with new target
	local target="$1"
	[ -n "$dev" ] || error "\$dev must be set"

	command -v cryptsetup > /dev/null \
		|| error "cryptsetup must be installed in current rootfs"
	command -v caam-decrypt > /dev/null \
		|| error "caam-decrypt must be installed in current rootfs"

	local index offset
	case "$dev" in
	*mmcblk*p*)
		index=${dev##*p}
		index=$((index-1))
		offset="$(((9*1024 + index*4)*1024))"
		;;
	*) error "LUKS only supported on mmcblk*p* partitions" ;;
	esac

	mkdir -p /run/caam
	local KEYFILE=/run/caam/lukskey
	# lower iter-time to speed PBKDF phase up,
	# since our key is random PBKDF does not help
	# also, we don't need a 16MB header so make it as small as possible (1MB)
	# by limiting the maximum number of luks keys (3 here, same size with less)
	# key size is 112
	unshare -m sh -c "mount -t tmpfs tmpfs /run/caam \
		&& caam-keygen create ${KEYFILE##*/} ccm -s 32 \
		&& dd if=/dev/random of=$KEYFILE.luks bs=$((4096-112-16)) count=1 status=none \
		&& dd if=/dev/random of=$KEYFILE.iv bs=16 count=1 status=none \
		&& cat $KEYFILE.iv $KEYFILE.luks > $KEYFILE.toenc \
		&& caam-encrypt $KEYFILE.bb AES-256-CBC $KEYFILE.toenc $KEYFILE.enc \
		&& cat $KEYFILE.bb $KEYFILE.iv $KEYFILE.enc > $KEYFILE.mmc \
		&& { if ! [ \$(stat -c %s $KEYFILE.mmc) = 4096 ]; then \
			echo \"Bad key size \$(stat -c %s $KEYFILE.mmc)\"; false; \
		fi; } \
		&& cryptsetup luksFormat -q --key-file $KEYFILE.luks \
			--pbkdf pbkdf2 --iter-time 1 \
			--luks2-keyslots-size=768k \
			$dev > /dev/null \
		&& cryptsetup luksOpen --key-file $KEYFILE.luks \
			--allow-discards $dev $target \
		&& dd if=$KEYFILE.mmc of=$rootdev bs=4k count=1 status=none \
			oflag=seek_bytes seek=$offset" \
		|| error "Could not create luks partition on $dev"

	dev="/dev/mapper/$target"
}

# robustly lock a file while allowing cleanup
LOCKFILE=/var/lock/swupdate.lock
LOCKSEM=0
_lock_update() {
	local try="$1"

	LOCKSEM=$((LOCKSEM+1))
	if [ "$LOCKSEM" != 1 ]; then
		# lock already taken
		return 0
	fi

	local FD_INO FILE_INO

	# fd cannot be put in variable without eval, keep this simple.
	if [ -e "/proc/$$/fd/8" ]; then
		info "fd 8 already taken, overwriting it: $(ls -l /proc/$$/fd/8)"
	fi

	while :; do
		exec 8<>"$LOCKFILE" || error "Could not open $LOCKFILE"
		if ! flock -n 8; then
			if [ -n "$try" ]; then
				exec 8<&-
				LOCKSEM=$((LOCKSEM-1))
				return 1
			fi
			info "swupdate currently locked, waiting until lock is free"
			flock 8 || error "Could not lock $LOCKFILE"
		fi

		# Check if file has changed since opening
		# if it has, drop lock and try again
		FD_INO=$(stat -c "%i" -L "/proc/self/fd/8")
		FILE_INO=$(stat -c "%i" "$LOCKFILE" 2>/dev/null)

		[ "$FD_INO" = "$FILE_INO" ] && break

		exec 8<&-
	done
	[ -e "/run/swupdate_rebooting" ] && error "reboot in progress!!"

	return 0
}

lock_update() {
	_lock_update
}

try_lock_update() {
	_lock_update "try"
}

unlock_update() {
	LOCKSEM=$((LOCKSEM-1))
	if [ "$LOCKSEM" = 0 ]; then
		rm -f "$LOCKFILE"
		exec 8<&-
	fi
}

is_mountpoint() {
	local dir="$1"

	# normalize path to match what we have in mountinfo
	dir="$(realpath "$1" 2>/dev/null)" || return 1

	# busybox 'mountpoint' stats target and checks for device change, so
	# bind mounts like /var/lib/containers/overlay are not properly detected
	# as mountpoint by it.
	# util-linux mountpoint parses /proc/self/mountinfo correctly though so we
	# could use it if installed, but it is simpler to always reuse our
	# implementation instead
	! awk -v dir="$dir" '$5 == dir { exit 1 }' < /proc/self/mountinfo
}

umount_if_mountpoint() {
	local dir="$1"

	# nothing to do if not a mountpoint
	is_mountpoint "$dir" || return 0

	# findmnt outputs in tree order, so umounting from
	# last line first should always work
	findmnt -nr -o TARGET -R "$dir" | tac | xargs -r umount --
}

probe_disk_fstype() {
	local src="$1" fstype

	# store in a variable to return failure if empty
	fstype="$(blkid "$src" | LANG=C sed -ne 's/.*\bTYPE="\([^"]*\)\".*/\1/p')"

	if [ -z "$fstype" ]; then
		return 1
	fi
	echo "$fstype"
}

probe_and_mount() {
	local src fstype skip=0

	# find the first non-option argument
	for arg in "$@"; do
		if [ "$skip" -gt 0 ]; then
			skip=$((skip-1))
			continue
		fi
		if [ "$skip" = "-1" ]; then
			src="$arg"
			break
		fi
		case "$arg" in
		--)
			skip=-1
			;;
		-*)
			skip=1
			;;
		*)
			src="$arg"
			break
			;;
		esac
	done

	fstype="$(probe_disk_fstype "$src")"
	mount "$@" ${fstype:+-t "$fstype"}
}

mount_rootfs_btrfs() {
	# don't mount with discard here: this can be used for loop devices
	# and does not really matter in practice
	mount -t btrfs -o "compress-force=zstd,noatime$opts" \
		"$src" "$dst" >/dev/null 2>&1
}

mount_rootfs_ext4() {
	mount -t ext4 -o "noatime$opts" "$src" "$dst" >/dev/null 2>&1
}

mount_rootfs() {
	local src="$1" dst="$2" fstype="$fstype" opts="$opts"

	[ -n "$fstype" ] \
		|| fstype="$(probe_disk_fstype "$src")" \
		|| fstype="$(findmnt -nr -o FSTYPE /live/rootfs)" \
		|| { echo "Could not get rootfs fstype" >&2; return 1; }

	case "$fstype" in
	btrfs)
		mount_rootfs_btrfs || mount_rootfs_ext4
		;;
	ext4)
		mount_rootfs_ext4 || mount_rootfs_btrfs
		;;
	*)
		echo "Unsupported filesystem for rootfs: $fstype" >&2
		return 1
		;;
	esac
}

###################################
# rollback/status                 #
###################################

dev_from_string() {
	local dev="$1"

	# we still need to query currently booted part for index case
	# and warning on full path case
	dev_from_current

	case "$dev" in
	rootfs_[01])
		target_ab="${dev#rootfs_}"
		target_rootpart="${rootpart%[0-9]}$((target_ab+1))"
		return
		;;
	0|1)
		target_ab="$dev"
		target_rootpart="${rootpart%[0-9]}$((target_ab+1))"
		return
		;;
	/dev/mmcblk*) target_rootpart="$dev";;
	mmcblk*) target_rootpart="/dev/$dev";;
	*) error "Unrecognized device $dev";;
	esac
	target_ab=${target_rootpart#*mmcblk[012]p}
	target_ab=$((target_ab-1))
	case "$target_ab" in
	0|1) ab=$((!target_ab));;
	*) error "Could not parse $target_rootpart for rollback target";;
	esac
	if [ "${rootpart%[0-9]}" != "${target_rootpart%[0-9]}" ]; then
		warning "Not booted on rootdev device, rollback requires more than just rebooting"
		rootdev="${target_rootpart%[0-9]}"
		[ "${rootdev#/dev/mmcblk}" = "$rootdev" ] \
			|| rootdev="${rootdev%p}"
	fi
	trace "Found $target_rootpart on $rootdev"
}

dev_from_current() {
	[ -n "$ab" ] && return

	# get current root block device (mmc or sd card or..),
	# and set ab to booted ab and target_ab to the rollback target ab
	rootpart=$(swupdate -g)

	[ -e "$rootpart" ] || rootpart="$(findfs "$rootpart")"
	[ -e "$rootpart" ] || error "Could not find what partition linux booted from"

	if [ "${rootpart##*[a-z]}" = "1" ]; then
		ab=0
		target_ab=1
	else
		ab=1
		target_ab=0
	fi
	rootdev="${rootpart%[0-9]}"
	[ "${rootdev#/dev/mmcblk}" = "$rootdev" ] \
		|| rootdev="${rootdev%p}"


	target_rootpart="${rootpart%[0-9]}$((target_ab+1))"
}

###################################
# atlog-write                     #
###################################

get_hostname() {
	# $HOSTNAME variable cannot be relied upon (either shell does not set it or early boot)
	case "$HOSTNAME" in
	""|"(none)") ;;
	*) return;;
	esac
	HOSTNAME=$(cat /etc/hostname 2>/dev/null)
	[ -n "$HOSTNAME" ] || HOSTNAME=armadillo
}


ctrl_atlog_write() {
	local atlog="/var/at-log/atlog" dev
	local progname="abos-ctrl"

	if [ "$#" -ge 2 ]; then
		progname="$1"
		shift
	fi

	if ! mountpoint -q /var/at-log; then
		atlog=/var/log/swupdate/atlog
	fi

	# if /var/log is encrypted also prefer /var/log
	dev=$(findmnt -nr -o SOURCE /var/log)
	if [ -n "$dev" ] && [ "$(lsblk -n -o type "$dev")" = "crypt" ]; then
		atlog=/var/log/swupdate/atlog
	fi

	if [ "$atlog" = /var/log/swupdate/atlog ]; then
		if ! mountpoint -q /var/log; then
			warning "at-log was not mounted, did not log $*"
			return
		fi
		[ -d "/var/log/swupdate" ] || mkdir /var/log/swupdate || return
	fi

	# rotate file if it got too big
	if [ "$(stat -c %s "$atlog" 2>/dev/null || echo 0)" -gt $((3*1024*1024)) ]; then
		mv -v "$atlog" "$atlog.1" \
			|| warning "Could not rotate atlog"
	fi

	get_hostname
	echo "$(date "+%b %_d %H:%M:%S") $HOSTNAME NOTICE $progname: $*" >> "$atlog"
	sync "$atlog"
}

###################################
# abos-ctrl main switch           #
###################################

set_mode() {
	if [ -n "$mode" ]; then
		[ -n "$break_unknown_opt" ] && return 1
		error "$arg cannot be used with $mode_arg"
	fi

	# shellcheck disable=SC3060 # busybox sh/mksh support this
	mode="${1//-/_}"
	mode_arg="$arg"
	if [ -f "$scripts_dir/$mode.sh" ]; then
		. "$scripts_dir/$mode.sh" || exit
	fi
}

check_root_user() {
	# only status can run as non-root user
	case "$1" in
	status) return;;
	esac

	USER=${USER:-"$(id -un)"}
	[ "$USER" = "root" ] || error "abos-ctrl $1 must be run as root"
}

main() {
	local rootdev rootpart target_rootpart target_ab="" ab=""
	local mode="" mode_arg="" break_unknown_opt=""
	local scripts_dir=/usr/libexec/abos-ctrl
	local debug=2 dryrun=""

	# for tests
	if [ -e "$(dirname "$0")/abos-ctrl.d" ]; then
		scripts_dir="$(realpath "$(dirname "$0")/abos-ctrl.d")"
	fi

	while [ "$#" -ge 1 ]; do
		local arg="$1"
		case "$arg" in
		status|"")
			set_mode status || break
			break_unknown_opt=1
			shift
			break
			;;
		atlog-write|make-installer|rollback|rollback-clone| \
		container-clear|reset-default|installer-setting| \
		usb-filter|rtos)
			set_mode "$arg" || break
			break_unknown_opt=1
			;;
		mount-old|restore-kernel|reboot|update|check-secure)
			set_mode "$arg" || break
			;;
		cert|certs|certificate|certificates)
			set_mode certificates || break
			break_unknown_opt=1
			;;
		rollback-force)
			# compat, go through rollback --force
			[ "$#" -ge 2 ] || error "$arg requires an argument";
			set -- "$@" rollback --force "$2"
			shift
			;;
		mount)
			shift
			probe_and_mount "$@"
			exit
			;;
		umount)
			# ugly wrapper, don't bother with mode etc...
			[ "$#" -ge 2 ] || error "$arg requires an argument";
			shift
			local dir fail=0
			for dir in "$@"; do
				umount_if_mountpoint "$dir" || fail="$?"
			done
			exit "$fail"
			;;
		internal)
			# run whatever we ask next. This is used for unshare wrapping
			# where we cannot easily run a function
			shift
			# probably need to get something from a script...
			if [ -f "$scripts_dir/$1.sh" ]; then
				. "$scripts_dir/$1.sh" || exit
			fi
			shift
			"$@"
			exit
			;;
		podman-storage)
			set_mode "$arg" || break
			wrapper="${wrapper:-"$0 podman-storage"}"
			break_unknown_opt=1
			shift
			break # does its own option parsing
			;;
		podman-rw)
			set_mode "$arg" || break
			wrapper="${wrapper:-"$0 podman-rw"}"
			# stop parsing when we encounter podman options
			break_unknown_opt=1
			;;
		"-n"|"--dry-run")
			# shellcheck disable=SC2034 ## used in other scripts
			dryrun=1
			;;
		"-q"|"--quiet")
			debug=$((debug-1))
			;;
		"-v"|"--verbose")
			debug=$((debug+1))
			;;
		"-V"|"--version")
			apk list -I abos-base
			exit 0
			;;
		"--help")
			if has_function "${mode}_help"; then
				"${mode}_help"
			else
				usage
			fi
			exit 0
			;;
		"--")
			break
			;;
		*)
			[ -n "$break_unknown_opt" ] && break
			warning "Invalid option $arg"
			usage >&2
			exit 1
			;;
		esac
		shift
	done

	[ -z "$mode" ] && set_mode status
	[ -n "$break_unknown_opt" ] \
		|| [ "$#" = 0 ] \
		|| error "$mode takes no argument, but got this left: $*"

	check_root_user "$mode"

	"ctrl_$mode" "$@"
}

main "$@"
