#!/bin/bash
#  vim:sw=8:noet
# SPDX-License-Identifier:  GPL-2.0+
#
# rootfs installer for ALT Linux
# Current version
VERSION=0.5.7

# Variable
EXT4=0FC63DAF-8483-4772-8E79-3D69D8477DE4
TMPROOT=
[ -n "$TMP" ] || TMP=/tmp
PROGNAME="${0##*/}"
unset ROOTPART_UID

# one Mebibyte = 2^20
MB=1048576

# usage message
usage() {
	printn "
Usage: $PROGNAME --target=TARGET <--rootfs=ROOTFS | --image-in=IMAGE> <--media=DEVICE | --image-out=FILE> [options]

Options:
	--target=TARGET   - target board
	--efi             - add EFI partition, set partition table gpt
	--efi-mbr         - add EFI partition, set partition table msdos
	--vnc=1           - enable vnc mode for alterator-setup
	--vnc=0           - disable vnc mode for alterator-setup

	--rootfs=ROOTFS   - archive rootfs file name
	--image-in=IMAGE  - image file name (img|img.xz)

	--media=DEVICE    - media device file (/dev/[sdX|mmcblkX])
	--image-out=FILE  - out image file

	--addconsole      - Add system console to extlinux.conf

	--image-size=SIZE - size out image file (MiB)
	--resize          - Resize ROOT partition
	--supported       - List of supported hardware
	--version         - Display version and exit
	--stdout          - Don't use log file, put everything to stdout
	--log=FILE        - Log file (default $LOG_FILE)
	-y                - Assumes yes, will not wait for confirmation

Example: $PROGNAME --rootfs=regular-lxqt-20190213-aarch64.tar.xz --target=rpi3 --media=/dev/mmcblk0

Update U-Boot for another Target
Update to a new u-boot from a local host install.

Example: $PROGNAME --target=orangepi_prime --media=/dev/mmcblk0

"
}

pre_exit() {
	local rc=$?
	trap - EXIT
	sync

	unmount_partition

	LC_ALL=C type soc_exit >&/dev/null && soc_exit

	if [ -n "$IMAGE_OUT" -a -n "$MEDIA" ]; then
		kpartx -d -s "$MEDIA" || {
			sleep 10
			kpartx -d -s -v "$MEDIA"
		}
		losetup --detach "$MEDIA"
	fi
	exit $rc
}

final_exit() {
	printn 'Installation Complete!'
	if [ -n "$IMAGE_OUT" ]; then
		printn 'Image created successfully.'
	else
		printn 'Insert into your device and try boot.'
		printn 'Good luck!!!'
	fi
	exit 0
}

# Set some global variables for the command directory, target board directory,
# and valid targets.
DIR="$(dirname $0)"
[ -d "${DIR}/boards.d" ] ||
	DIR=/usr/share/alt-rootfs-installer
[ -d "${DIR}/boards.d" ] ||
	fatal_error "directory boards.d not exist"
BOARDDIR="${DIR}/boards.d"
SOCSDIR="${DIR}/socs.d"

# include logging
. shell-error 2>/dev/null || . $DIR/shell-error 2>/dev/null
. shell-terminfo 2>/dev/null || . $DIR/shell-terminfo 2>/dev/null
. $DIR/log

. ${SOCSDIR}/socs-utils

# missing or wrongs arguments
[ $# -ge 1 ] &&
TMPARGS="$(getopt -n "$0" -o hy -l addconsole,debug,help,\
efi,efi-mbr,target:,rootfs:,media:,image-in:,image-out:,\
image-size:,resize,supported,version,stdout,log:,vnc: -- "$@")"
if [ $? -ne 0 ]; then
	usage
	exit 1
fi

eval set -- "$TMPARGS"
# check the args
while :; do
	case $1 in
		--)
			shift; break
			;;
		--addconsole)
			CONSOLE=1
			;;
		--debug)
			set -x
			;;
		-h|--help)
			usage; exit 0
			;;
		--target)
			shift; TARGET="$1"
			;;
		--efi)
			EFI="1"
			MEDIATABLE=gpt
			;;
		--efi-mbr)
			EFI="1"
			MEDIATABLE=msdos
			;;
		--vnc)
			shift; VNC="$1"
		;;
		--rootfs)
			shift; ROOTFS="$1"
			;;
		--image-in)
			shift; IMAGE="$1"
			;;
		--media)
			shift; MEDIA="$1"
			;;
		--image-out)
			shift; IMAGE_OUT="$1"
			;;
		--image-size)
			shift; IMAGE_SIZE="$1"
			;;
		--resize)
			RESIZE="1"
			;;
		--supported)
			cat "$DIR/SUPPORTED-BOARDS"
			exit 0
			;;
		--version)
			echo "$PROGNAME-$VERSION"
			exit 0
			;;
		--log)
			shift
			LOG_FILE="$(readlink -f "$1")"
			;;
		--stdout)
			LOG_MODE="con"
			;;
		-y)
			NOASK=1
			;;
		*)
			echo "Invalid argument $1"
			usage
			exit 1
			;;
	esac
	shift
done

log_init

# getting full PATH
[ -n "$ROOTFS" ] && ROOTFS_PATH="$(readlink -f "$ROOTFS")"
[ -n "$IMAGE" ] && IMAGE_PATH="$(readlink -f "$IMAGE")"

# checking file existing
if [ -n "$ROOTFS" ]; then
	[ -f "$ROOTFS_PATH" ] || fatal_error "ROOTFS does not exist"
	ROOTFS="$ROOTFS_PATH"
fi
if [ -n "$IMAGE" ]; then
	[ -f "$IMAGE_PATH" ] || fatal_error "IMAGE does not exist"
	IMAGE="$IMAGE_PATH"
fi

# replace old targets
[ "$TARGET" = "rpi2" ] && TARGET="rpi_3_32b"
[ "$TARGET" = "rpi3" ] && TARGET="rpi_3"

# checking the entered value --image-size
if [ -n "${IMAGE_SIZE-}" ]; then
	[[ "$IMAGE_SIZE" =~ [^0-9]+ ]] && fatal_error "bad image size"
fi

# ensure sudo user
if [ "$(whoami)" != "root" ]; then
	fatal_error "this script requires 'sudo' privileges in order to write to disk & mount media"
fi


if [ -z "$IMAGE_OUT" -a -z "$MEDIA" ]; then
	fatal_error "one of the options must be defined: --media or --image-out"
fi

if [ -n "$IMAGE_OUT" -a -n "$MEDIA" ]; then
	fatal_error "At the same time, mutually exclusive options are indicated: --media and --image-out"
fi

if [ -n "$IMAGE" -a -n "$ROOTFS" ]; then
	fatal_error "At the same time, mutually exclusive options are indicated: --image-in and --rootfs"
fi

if [ -n "$IMAGE" -a -n "$IMAGE_OUT" ]; then
	fatal_error "At the same time, mutually exclusive options are indicated: --image-in and --image-out"
fi

if [ -z "$TARGET" -a -z "$ROOTFS" -a -z "$IMAGE" ]; then
	usage
	exit 1
fi

if [ -n "$MEDIA" ] && [[ "$(readlink -f /dev/root)" ==  "$MEDIA"* ]]; then
	fatal_error "$MEDIA is a device contains a host system rootfs"
fi

# rootfs exists
if [ -n "$ROOTFS" -a ! -f "$ROOTFS" ]; then
	fatal_error "$ROOTFS not found! Please choose an existing rootfs"
fi

# image exists
if [ -n "$IMAGE" -a ! -f "$IMAGE" ]; then
	fatal_error "$IMAGE not found! Please choose an existing image"
fi

# device exists
if [ -z "$IMAGE_OUT" -a ! -e "$MEDIA" ]; then
	fatal_error "$MEDIA not found! Please choose an existing device"
fi

# target board exists
if [ -n "$TARGET" -a ! -e "${BOARDDIR}/${TARGET}" ]; then
	printn "Please choose a supported board"
	usage
	exit 1
fi

# check image format
if [ -n "$IMAGE" ]; then
	case "$IMAGE" in
		*.img)
			CAT='cat'
			;;
		*.img.xz)
			CAT='xzcat'
			;;
		*) fatal_error "unknown image format"
			;;
	esac
	# NOTE: We assume there is no arch names containing '-' or '.'
	ARCH="${IMAGE##*-}"
	ARCH="${ARCH%%.*}"
fi

# output current status of tar to /dev/tty every 512 records (5 Mb)
# with default blocking factor 20 (1 record = 20*512 bytes)
TAR='tar --checkpoint=500 --checkpoint-action=ttyout=%s:%{}T\r'
# definition of the rootfs archive format and ARCH
if [ -n "$ROOTFS" ]; then
	case "$ROOTFS" in
		*.tar)
			TAR="$TAR -xf"
			let "TAR_SIZE = $(stat -Lc %s "$ROOTFS") / $MB"
			;;
		*.tar.gz)
			TAR="$TAR -zxf"
			let "TAR_SIZE = $(gzip --list "$ROOTFS" | awk 'END{print $2}') / $MB"
			;;
		*.tar.xz)
			TAR="$TAR -Jxf"
			let "TAR_SIZE = $(xz --robot --list "$ROOTFS" | awk '/^totals/{print $5}') / $MB"
			;;
		*) fatal_error "unknown rootfs archive format"
			;;
	esac
	# NOTE: We assume there is no arch names containing '-' or '.'
	ARCH="${ROOTFS##*-}"
	ARCH="${ARCH%%.*}"
fi

if [ -n "$TARGET" ]; then
	if [ -L "$BOARDDIR/$TARGET" ]; then
		SOCS="$(basename $(readlink -e "$BOARDDIR/$TARGET"))"
		SOCS_ARCH="${SOCS#*-}"
	else
		SOCS="$TARGET".sh
		unset SOCS_ARCH
	fi
else
	SOCS="unknown-soc.sh"
	unset SOCS_ARCH
fi

[ -n "$SOCS" ] || SOCS="unknown-soc.sh"

case "$MEDIA" in
	/dev/mmcblk*|/dev/loop*|/dev/nvme*)
		partsuffix="p"
		;;
	*)
		unset partsuffix
		;;
esac

# Source SoCs specific shell script
. "$SOCSDIR/$SOCS"
soc_init

# Add default system console
if [ -n "$CONSOLE" ] && [ -z "$SYSCON" ]; then
	SYSCON=ttyS0,115200
fi

# Add EFI Partitions
if [ -n "$EFI" -a -n "$ROOTFS" ]; then
	if [ -z "$FIRMPART_NUM" ]; then
		FIRMPART_NUM="${ROOTPART_NUM}"
		let ROOTPART_NUM=${FIRMPART_NUM}+1
	fi
fi

if [ -z "$ROOTFS" -a -z "$IMAGE" ]; then
	find_firmware_partition
fi

change_partition

# Image size calculation
if [ -n "$ROOTFS" -a -n "$IMAGE_OUT" ]; then
	# Add 256 MiB to the image size for bios, firmware, bootloader, etc and increase size by 10%
	# which is reserved by FS, take into account i-node ratio and reserved 32Mb. Add 1 just in case.
	IMAGE_SIZE_MIN="$(awk -vTS="$TAR_SIZE" -vM="$MB" 'BEGIN{ printf ("%d", ((TS+512+32)*(1+256/8192)*1.1+1)) }')"
	# Change image size
	if [ -z "${IMAGE_SIZE-}" ] || [ "$IMAGE_SIZE_MIN" -ge "$IMAGE_SIZE" ]; then
		IMAGE_SIZE="$IMAGE_SIZE_MIN"
	fi
	let "END_SECTOR = $IMAGE_SIZE - 1"
	END_SECTOR="$END_SECTOR"MiB
fi

# Last chance to back out
if [ -n "$IMAGE_OUT" -a -n "$IMAGE_SIZE" ]; then
	printx 'Image out: '; printn "$IMAGE_OUT" green 20
	printx 'Image size: '; printn "$IMAGE_SIZE MiB" green 20
else
	printx 'Selected Media: '; printn "$MEDIA" green 20
fi
# target hardware ARCH
if [ -n "$TARGET" ]; then
	printx 'Target: '; printn "$TARGET" green 20
fi
# rootfs if included
if [ -n "$ROOTFS" ]; then
	printx 'Selected rootfs: '; printn "$ROOTFS" green 20
fi
# image-in if included
if [ -n "$IMAGE" ]; then
	printx 'Selected image: '; printn "$IMAGE" green 20
fi
# addconsole if included
if [ -n "$CONSOLE" ]; then
	printx 'System console: '; printn "$SYSCON" green 20
fi
printx 'Log file: '; printn "$LOG_FILE" green 20

case "$VNC" in
	1)
		printx 'VNC: '; printn "enabled" green 20
	;;
	0)
		printx 'VNC: '; printn "disabled" red 20
	;;
esac

if [ -z "$IMAGE_OUT" ]; then
	if [ -n "$IMAGE" -o -n "$ROOTFS" ]; then
		printn "WARNING! ALL DATA WILL BE DESTROYED" red
	else
		printn "WARNING! BOOTLOADER WILL BE OVERWRITTEN" red
	fi
fi

if [ "$NOASK" != 1 ]; then
	# Ask user before start
	printn "Would you like to continue? [Yes/No] " bold
	read PROCEED
	case "$PROCEED" in
		[Yy][Ee][Ss])
			;;
		*)
			printn "User exit, no rootfs written."
			exit 0
			;;
	esac
fi

trap pre_exit EXIT HUP PIPE INT QUIT TERM

# Check to make image file
if [ -n "$IMAGE_OUT" ]; then
	MEDIA="$(losetup --find)"
	[ -n "$MEDIA" ] || fatal_error "/dev/loop* not allowed"
fi

# Prepare media
[ -z "${MEDIA-}" ] || prepare_media

# Write image file to the media device
if [ -n "$IMAGE" ]; then
	image_to_media
	find_firmware_partition
	change_partition
fi

# Call SoC specific function to write system
[ -z "$ROOTFS" ] || soc_write_rootfs
[ -z "$TARGET" ] || soc_write_bootloader

# Setup boot flag
set_bootflag

# Change system console
if [ -n "$CONSOLE" ]; then
	EXTLINUX_CONF=${TMPROOT-}/boot/extlinux/extlinux.conf
	if [ -n "$EXTLINUX_CONF" ]; then
		sed -e 's/[[:space:]]*console=[a-zA-Z0-9,]*//g' \
			-e "/append/ s/$/ console=tty1 console=$SYSCON/g" \
			-i "$EXTLINUX_CONF"
	fi
fi

# Enable vnc mode for alterator-setup
ALTERATOR_SETUP_CONF=${TMPROOT-}/etc/alterator-setup/config
if [ -f "$ALTERATOR_SETUP_CONF" ]; then
	case "$VNC" in
		1)
			sed '/ALTERATOR_SETUP_VNC=/d' -i "$ALTERATOR_SETUP_CONF"
			echo 'ALTERATOR_SETUP_VNC=1' >> "$ALTERATOR_SETUP_CONF"
		;;
		0)
			sed '/^ALTERATOR_SETUP_VNC=/s/^/#/' -i "$ALTERATOR_SETUP_CONF"
		;;
	esac
fi

final_exit
