#!/bin/sh

# Copyright (C) 2014-2015 Red Hat Inc.
# SPDX-License-Identifier:  GPL-2.0+

# Automate Media Creation for Fedora ARM
# Current version
VERSION=2.18

# usage message
usage() {
    echo "
Usage: $(basename ${0}) <options>

	--image=IMAGE   - xz compressed image file name
	--media=DEVICE  - media device file (/dev/[sdX|mmcblkX])
Optional
	--addconsole    - Add system console kernel parameter for the target
	--addkey        - /path/to/ssh-public-key
	--args          - Optional kernel parameters listed in quotes
	--norootpass    - Remove the root password
	--relabel       - SELinux relabel root filesystem on first boot
	--resizefs      - Resize root filesystem to fill media device
	--sysrq         - Enable System Request debugging of the kernel
	--target=TARGET - target board for uboot
	-y              - Assumes yes, will not wait for confirmation
Help
	--supported     - List of supported hardware
	--version       - Display version and exit

Example: $(basename ${0}) --image=Fedora-Rawhide.xz --target=Bananapi --media=/dev/mmcblk0

"
}

# Set some global variables for the command directory, target board directory,
# and valid targets.
DIR=$(dirname $0)
if [ -d "/usr/share/arm-image-installer/boards.d" ]; then
	BOARDDIR="/usr/share/arm-image-installer/boards.d"
	DOC_DIR="/usr/share/doc/arm-image-installer/"
else
	DIR=$(dirname $0)
	BOARDDIR="${DIR}/boards.d"
	DOC_DIR="${DIR}/"
fi

# check the args
while [ $# -gt 0 ]; do
	case $1 in
		--debug)
			set -x
			;;
		-h|--help)
			usage
			exit 0
			;;
		--target*)
			if echo $1 | grep '=' >/dev/null ; then
				TARGET=$(echo $1 | sed 's/^--target=//')
			else
				TARGET=$2
				shift
			fi
			;;
		--image*)
			if echo $1 | grep '=' >/dev/null ; then
				IMAGE=$(echo $1 | sed 's/^--image=//')
			else
				IMAGE=$2
				shift
			fi
			;;
		--media*)
			if echo $1 | grep '=' >/dev/null ; then
				MEDIA=$(echo $1 | sed 's/^--media=//')
			else
				MEDIA=$2
				shift
			fi
			;;
		--addkey*)
			if echo $1 | grep '=' >/dev/null ; then
				SSH_KEY=$(echo $1 | sed 's/^--addkey=//')
			else
				SSH_KEY=$2
				shift
			fi
			;;
		--selinux*)
			if echo $1 | grep '=' >/dev/null ; then
				SELINUX=$(echo $1 | sed 's/^--selinux=//')
			else
				SELINUX=$2
				shift
			fi
			;;
		--args*)
			if echo $1 | grep '=' >/dev/null ; then
				OPT_ARGS=$(echo $1 | sed 's/^--args=//')
			else
				OPT_ARGS=$2
				shift
			fi
			;;
		--norootpass)
			NOROOTPASS=1
			;;
		--sysrq)
			SYSRQ=1
			;;
		--resizefs)
			RESIZEFS=1
			;;
		--addconsole)
			CONSOLE=1
			;;
		--blacklistvc4)
			FIX_RPI=1
			;;
		--supported)
			cat $DOC_DIR/SUPPORTED-BOARDS
			exit 0
			;;
		--relabel)
			RELABEL=1
			;;
		--version)
			echo "$(basename ${0})-"$VERSION""
			exit 0
			;;
		-y)
			NOASK=1
			;;
		*)
			echo "$(basename ${0}): Error - ${1}"
			usage
			exit 1
			;;
	esac
	shift
done

contains() {
	string="$1"
	substring="$2"
	if test "${string#*$substring}" != "$string"; then
		return 0    # $substring is in $string
	else
		return 1    # $substring is not in $string
	fi
}

# ensure sudo user
if [ "$(whoami)" != "root" ]; then
	echo "Error: This script requires 'sudo' privileges in order to write to disk & mount media."
	exit 1
fi

# check to make sure populated
if [ "$MEDIA" = "" ]; then
	usage
	exit 1
fi

if [ "$MEDIA" = "/dev/sda" ]; then
	echo ""
	echo " ***********************************************************"
	echo " ** WARNING: You have requested the image be written to sda."
	echo " ** /dev/sda is usually the root filesystem of the host. "
	echo " ***********************************************************"
	echo " ** Do you wish to continue? (type 'yes' to continue)"
	echo " ***********************************************************"
	# wait for agreement
	read -p " = Continue? " AGREE
	if [ "$(echo ${AGREE} | tr [:lower:] [:upper:])" != "YES" ]; then
		echo "User exit, no image written."
		exit 0
	fi
fi

# change cubietruck target to uppercase
if [ "$TARGET" = "cubietruck" ]; then
	TARGET="Cubietruck"
fi

# check for boards
if [ "$TARGET" != "" -a ! -e "${BOARDDIR}/${TARGET}" ]; then
	echo "Error: You must choose a supported board or none at all."
	usage
	exit 1
fi

# image exists
if [ ! -f "$IMAGE" ] && [ "$IMAGE" != "" ]; then
	echo "Error: $IMAGE not found! Please choose an existing image."
	exit 1
fi

# device exists
if [ ! -e "$MEDIA" ]; then
	echo "Error: $MEDIA not found! Please choose an existing device."
	exit 1
fi

# Last chance to back out
echo ""
echo "====================================================="

# check to see if host system uses LVM and VG named fedora for root
# if writing the aarch64 server image rename to fedora-server
if [ -b /dev/fedora/root ] && [ "$(echo $IMAGE|awk '/aarch64/ && /Server/')" != "" ]; then
	echo "**************************************************"
	echo "= NOTE: This host system uses the same VG name as "
	echo "= the AArch64 disk image. To avoid issues, the VG "
	echo "= on the image will be renamed to 'fedora-server'."
	echo "**************************************************"
	RENAME_LVM="1"
fi

# Image if included
if [ "$IMAGE" != "" ]; then
	echo "= Selected Image:                                 "
	echo "= $IMAGE"
fi
echo "= Selected Media : $MEDIA"
# target hardware platform
if [ "$TARGET" != "" ]; then
	echo "= U-Boot Target : $TARGET"
fi
# SE Linux On/Off
if [ "$SELINUX" != "" ]; then
	echo "= SELINUX = $SELINUX"
fi
if [ "$RELABEL" != "" ]; then
	echo "= SELinux relabel will be completed on first boot."
fi
# Remove root password
if [ "$NOROOTPASS" != "" ]; then
	echo "= Root Password will be removed."
fi
# Enable System Request debugging of the kernel
if [ "$SYSRQ" != "" ]; then
	echo "= System Request debugging of the kernel will be enabled."
fi
# Resize root filesystem to fill media device
if [ "$RESIZEFS" != "" ]; then
	echo "= Root partition will be resized"
fi
# Console to be added
if [ "$CONSOLE" != "" ]; then
	echo "= Console for $TARGET will be added."
fi
# User ssh key
if [ "$SSH_KEY" != "" ]; then
	echo "= SSH Public Key $SSH_KEY will be added."
fi
echo "====================================================="
echo " "
echo "*****************************************************"
echo "*****************************************************"
echo "******** WARNING! ALL DATA WILL BE DESTROYED ********"
echo "*****************************************************"
echo "*****************************************************"
if [ "$NOASK" != 1 ]; then
	echo " "
	echo " Type 'YES' to proceed, anything else to exit now "
	echo " "
	# wait for agreement
	read -p "= Proceed? " PROCEED
	if [ "$(echo ${PROCEED} | tr [:lower:] [:upper:])" != "YES" ]; then
		echo "User exit, no image written."
		exit 0
	fi
fi
# umount before starting
umount $MEDIA* &> /dev/null
if [ "$(fdisk -l $MEDIA | grep LVM)" != "" ]   ; then
	vgchange -a n $(pvs 2> /dev/null | grep $MEDIA | awk '{print $2}') &> /dev/null
fi

# Write the disk image to media
if [ "$IMAGE" != "" ]; then
	echo "= Writing: "
	echo "= $IMAGE "
	echo "= To: $MEDIA ...."
	xzcat $IMAGE | dd of=$MEDIA bs=4M status=progress; sync; sleep 3
	echo "= Writing image complete!"
fi
# check to see how many partitions on the image
partprobe "$MEDIA"

get_lvm_name () {
	if [ "$(fdisk -l $MEDIA | grep LVM)" != "" ]   ; then
		LVM_NAME=$(pvs 2> /dev/null | grep $MEDIA | awk '{print $2}')
	fi
}

add_kernel_parameter () {
	if [ -f /tmp/boot/extlinux/extlinux.conf ]; then
		sed -i "s|append|& $1 |" /tmp/boot/extlinux/extlinux.conf
	elif [ -f /tmp/fw/EFI/fedora/grub.cfg ]; then
		sed -i "s|GRUB_CMDLINE_LINUX=\"|& $1 |" ${PREFIX}/etc/default/grub
		sed -i "s|kernelopts=|& $1 |" /tmp/fw/EFI/fedora/grubenv
	fi
}

if [ -e "$MEDIA"p5 ]; then
	export FIRMPART="${MEDIA}p1"
	BOOTPART="${MEDIA}p2"
	ROOTPART="${MEDIA}p5"
	PARTNUM=5
elif [ -e "$MEDIA"p4 ]; then
	export FIRMPART="${MEDIA}p1"
	BOOTPART="${MEDIA}p2"
	ROOTPART="${MEDIA}p4"
	PARTNUM=4
elif [ -e "$MEDIA"p3 ]; then
	export FIRMPART="${MEDIA}p1"
	BOOTPART="${MEDIA}p2"
	ROOTPART="${MEDIA}p3"
	PARTNUM=3
elif [ -e "$MEDIA"5 ]; then
	export FIRMPART="${MEDIA}1"
	BOOTPART="${MEDIA}2"
	ROOTPART="${MEDIA}5"
	PARTNUM=5
elif [ -e "$MEDIA"4 ]; then
	export FIRMPART="${MEDIA}1"
	BOOTPART="${MEDIA}2"
	ROOTPART="${MEDIA}4"
	PARTNUM=4
else
	export FIRMPART="${MEDIA}1"
	BOOTPART="${MEDIA}2"
	ROOTPART="${MEDIA}3"
	PARTNUM=3
fi

# resize root filesystem before mounting
if [ "$RESIZEFS" != "" ]; then
	echo "= Resizing $MEDIA ...."
	sync
	count=0
	if [ "$PARTNUM" = "4" ] || [ "$PARTNUM" = "5" ]; then
		echo ", +" | sfdisk -N "4" "$MEDIA"
		while [ $? != '0' ]; do
			sleep 5
			echo ", +" | sfdisk -N "4" "$MEDIA"
			((count++))
			if [ $count -gt 5 ]; then
				echo "= Partition Resize Failed."
				continue
			fi
		done
		partprobe "$MEDIA"
		if [ "$PARTNUM" = "5" ]; then
			echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
			while [ $? != '0' ]; do
				sleep 5
				echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
				((count++))
				if [ $count -gt 5 ]; then
					echo "= Partition Resize Failed."
					continue
				fi
			done
		fi
		partprobe "$MEDIA"
	fi

	if [ "$PARTNUM" = "3" ]; then
		if [ "$RENAME_LVM" = "1" ]; then
		# rename VG if host system uses conflicting names
			LVM_UUID=$(pvs -o +vguuid | grep $MEDIA | awk '{print $7}')
			vgrename $LVM_UUID fedora-server
			LVM_NAME=fedora-server
		fi

		get_lvm_name
		vgchange -a n $LVM_NAME &> /dev/null

		echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
		while [ $? != '0' ]; do
			sleep 5
			echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
			((count++))
			if [ $count -gt 5 ]; then
				echo "= Partition Resize Failed."
				continue
			fi
		done
		sleep 5
		partprobe "$MEDIA"
		if [ "$LVM_NAME" != "" ]; then
			vgchange -a y $LVM_NAME &> /dev/null
			pvresize "$ROOTPART"
			lvextend -l +100%FREE /dev/$LVM_NAME/root
			ROOTPART="/dev/$LVM_NAME/root"
		fi
	fi

	FS_TYPE=$(blkid "$ROOTPART" -s TYPE -o value)
	if [ "$FS_TYPE" = "xfs" ] && [ "$LVM_NAME" != "" ]; then
		mkdir /tmp/root &> /dev/null
		mount "$ROOTPART" /tmp/root &> /dev/null
		xfs_growfs /tmp/root
	elif [ "$FS_TYPE" = "btrfs" ]; then
		mkdir /tmp/root &> /dev/null
		mount "$ROOTPART" /tmp/root &> /dev/null
		btrfs filesystem resize max /tmp/root
		BTRFS=1
	elif [ "$FS_TYPE" = "ext4" ]; then
		fsck.ext4 -fy "$ROOTPART"
		resize2fs "$ROOTPART"
	fi
fi

sleep 5

# make temp mount points
mkdir /tmp/{boot,root,fw} &> /dev/null
mount "$BOOTPART" /tmp/boot &> /dev/null
if [ $? -ne 0 ]; then
	echo "Error: mount $BOOTPART /tmp/boot failed"
	exit 1
fi
mount "$FIRMPART" /tmp/fw
if [ $? -ne 0 ]; then
	echo "Error: mount $FIRMPART /tmp/fw failed"
	exit 1
fi

if [ "$RESIZEFS" = "" ]; then
	get_lvm_name
	vgchange -a y $LVM_NAME &> /dev/null
fi

if [ "$(grep /tmp/root /proc/mounts)" = "" ]; then
	if [ "$LVM_NAME" != "" ]; then
		mount "/dev/$LVM_NAME/root" /tmp/root &> /dev/null
	else
        	mount "$ROOTPART" /tmp/root
	fi
fi

if [ "$(echo $IMAGE | grep -i iot)" != "" ]; then
	IOT_IMAGE="1"
	OSTREE_ROOT_HOME="/tmp/root/ostree/deploy/fedora-iot/var/roothome"
	OSTREE_PREFIX="/tmp/root/ostree/deploy/fedora-iot/deploy/*/"
fi

# fix up grub.cfg to reflect the new vg name
if [ "$RENAME_LVM" != "" ]; then
	sed -i 's|/dev/mapper/fedora-root|/dev/mapper/fedora--server-root|g; s|rd.lvm.lv=fedora/root|rd.lvm.lv=fedora-server/root|g' /tmp/fw/EFI/fedora/grub.cfg
	sed -i 's|/dev/mapper/fedora-root|/dev/mapper/fedora--server-root|g; s|rd.lvm.lv=fedora/root|rd.lvm.lv=fedora-server/root|g' /tmp/fw/EFI/fedora/grubenv
fi

if [ "$IOT_IMAGE" = "1" ]; then
	# dd doesnt support wildcards, echo to expand
	PREFIX=$(echo $OSTREE_PREFIX)
elif [ "$BTRFS" = "1" ]; then
	PREFIX="/tmp/root/root/"
else
	PREFIX=/tmp/root
fi

# determine uboot and write to disk
if [ "$TARGET" != "" ]; then
	if [[ "$TARGET" == rpi[2,3,4] ]] || [ "$TARGET" = "olpc_xo175" ]; then
		. "${BOARDDIR}/${TARGET}"
	elif [ -d "${PREFIX}/usr/share/uboot/${TARGET}" ]; then
		. "${BOARDDIR}/${TARGET}"
	else
		echo "= No U-Boot files found for $TARGET."
	fi
else
	echo "= No U-boot will be written."
	TARGET="board"
fi

# turn off selinux
if [ "$SELINUX" != "" ]; then
	if [ "$(echo ${SELINUX} | tr [:lower:] [:upper:])" = "OFF" ]; then
		echo "= Turning SELinux off ..."
		sed -i 's/SELINUX=enforcing/SELINUX=permissive/' ${PREFIX}/etc/selinux/config
		# turn on selinux
	elif [ "$(echo ${SELINUX} | tr [:lower:] [:upper:])" = "ON" ]; then
		echo "= Turning SELinux on ..."
		sed -i 's/SELINUX=permissive/SELINUX=enforcing/' ${PREFIX}/etc/selinux/config
	fi
fi
# Remove root password
if [ "$NOROOTPASS" = "1" ]; then
	echo "= Removing the root password."
	sed -i 's/root:x:/root::/' ${PREFIX}/etc/passwd
fi
# Enable System Request debugging of the kernel
if [ "$SYSRQ" != "" ]; then
	echo "= Enabling System Request debugging of the kernel."
	cat >> ${PREFIX}/etc/sysctl.d/arm-image-installer-sysrq.conf <<-EOH
		# Controls the System Request debugging functionality of the kernel
		kernel.sysrq = 1
	EOH
fi
# Add ssh key to the image
if [ "$SSH_KEY" != "" ]; then
	if [ -f $SSH_KEY ]; then
		echo "= Adding SSH key to authorized keys."
		if [ "$IOT_IMAGE" = "1" ]; then
			mkdir $OSTREE_ROOT_HOME/.ssh/ &> /dev/null
			cat $SSH_KEY >> $OSTREE_ROOT_HOME/.ssh/authorized_keys
			chmod -R u=rwX,o=,g= $OSTREE_ROOT_HOME/.ssh/
		else
			mkdir ${PREFIX}/root/.ssh/ &> /dev/null
			cat $SSH_KEY >> ${PREFIX}/root/.ssh/authorized_keys
			chmod -R u=rwX,o=,g= ${PREFIX}/root/.ssh/
		fi
	else
		echo "= SSH key $SSH_KEY : Not Found!"
		echo "= WARNING: No SSH Key Added."
	fi
fi

# Add console
if [ "$CONSOLE" = "1" ]; then
	if [ "$SYSCON" = "" ]; then
		SYSCON="ttyS0,115200"
		echo "= No console listed for $TARGET, adding default $SYSCON ."
	fi

	echo "= Adding console $SYSCON to kernel parameters ..."
	add_kernel_parameter "console=$SYSCON console=tty0"

	if [[ "$TARGET" == rpi[2,3,4] ]]; then
		sed -i "s|# enable_uart=1|enable_uart=1|" /tmp/fw/config.txt
	fi
fi

# fix up rpi2/3
if [ "$FIX_RPI" != "" ]; then
	echo "= Blacklisting the VC4 Driver for the Raspberry Pi 2/3"
	echo blacklist vc4 > ${PREFIX}/etc/modprobe.d/blacklist-vc4.conf
	add_kernel_parameter "rd.driver.blacklist=vc4"
fi

# check if host system has selinux disabled, if it does autorelabel is required
if [ "$(getenforce)" = "Disabled" ]; then
	echo "= NOTE: System Relabel required on first boot."
	RELABEL="1"
fi

# touch /.autorelabel
if [ "$RELABEL" != "" ]; then
	if [ "$IOT_IMAGE" = "1" ]; then
		echo "= SELinux relabel not supported on IoT images."
	else
		echo "= Touch /.autorelabel on rootfs."
		touch ${PREFIX}/.autorelabel
	fi
fi
# add option kernel parameters
if [ "$OPT_ARGS" != "" ] ; then
	echo "= Adding optional kernel parameters for $TARGET : "
	echo "= Parameter: $OPT_ARGS"
	add_kernel_parameter "$OPT_ARGS"
fi

sync

umount /tmp/root $BOOTPART $FIRMPART &> /dev/null
if [ "$LVM_NAME" != "" ]; then
	vgchange -a n $LVM_NAME &> /dev/null
fi
rmdir  /tmp/root /tmp/boot /tmp/fw &> /dev/null

if [ "$URL" != "" ]; then
	echo
	echo "= NOTE"
	echo "= Additional instructions for $TARGET can be found at:"
	echo "= $URL "
	echo
fi

echo ""
echo "= Installation Complete! Insert into the "$TARGET" and boot."
exit 0
