#!/bin/bash

# remaster-squashfs-editor
# Copyright (C) 2019-2020 Richard Nelson <unixabg@gmail.com>
#
# This program comes with ABSOLUTELY NO WARRANTY; for details see COPYING.
# This is free software, and you are welcome to redistribute it
# under certain conditions; see COPYING for details.

#set -x
############################
# Dependencies
############################
if [[ $EUID -ne 0 ]]; then
	echo "This script must be run as root"
	exit 1
fi

if [ ! -e /usr/bin/mksquashfs ]; then
	echo "No /usr/bin/mksquashfs file found!"
	echo "Exiting ..."
	echo "On Debian based systems, mksquashfs can be installed with:"
	echo "  apt install squashfs-tools"
	exit 1
fi

# Source in config
echo "Loading remaster-iso.conf ..."
if [ -f remaster-iso.conf ]; then
	echo "Found ./remaster-iso.conf"
	echo "Sourcing in ./remaster-iso.conf"
	. remaster-iso.conf
elif [ -f /usr/share/remaster-iso/remaster-iso.conf ]; then
	echo "Found /usr/share/remaster-iso/remaster-iso.conf"
	echo "Sourcing in /usr/share/remaster-iso/remaster-iso.conf"
	. /usr/share/remaster-iso/remaster-iso.conf
else
	echo "No remaster-iso.conf file found!"
	echo "Exiting ..."
	exit 1
fi

# Sanity check for work area
if [ ! -d "${_ISOExtractPath}/${_ISOLivePath}" ]; then
	echo "E: No iso-extract/live folder of ${_ISOExtractPath}/${_ISOLivePath} !"
	echo "Exiting ..."
	exit 1
fi

Cleanup_mounts () {
	echo "Unmounting the areas we created."
	for mount_point in "${_MOUNTLIST}"; do
		umount -l $mount_point
		rmdir $mount_point
	done
	# The upperdir and work dir need to be cleaned up also.
	rm -rf ./psu_overlay_rw ./psu_overlay_work
}

show_menus() {
	clear
	echo "#################################"
	echo "    remaster-squashfs-editor"
	echo "    remaster-iso version ${_VER}"
	echo "#################################"
	echo "  (C)hroot - Chroot in to the filesystem.squashfs + psu-*.squashfs stack."
	echo "  (J)oin   - Join the partial squashfs update files to new single psu-DATE.squashfs"
	echo "  (N)ew    - New master filesystem.squashfs file which joins all *.squashfs file to new single filesystem.squashfs"
	echo "  E(x)it   - Exit the program with no action"
}

############################
# Main logic - infinite loop
############################
while true
do
	show_menus
	read -r -p "Enter choice [C , J , N , X] " _RUNTYPE
	case $_RUNTYPE in
		[cC])	echo "I: Chroot option selected!"
			_RUNTYPE="C"
			break
			;;
		[jJ])	echo "I: Join option selected!"
			_RUNTYPE="J"
			break
			;;
		[nN])	echo "I: New option selected!"
			_RUNTYPE="N"
			break
			;;
		[xX])	echo "I: Exiting the program with no action taken."
			exit 0
			;;
		*)	echo -e "E: Invalid selection ..." && sleep 2
	esac
done

echo "I: change directory to target live folder"
cd "${_ISOExtractPath}/${_ISOLivePath}"

echo "I: strting mount list and points operations"
# Let's start with mount list and mount points
_MOUNTDIR=""
_MOUNTLIST="./filesystem_squashfs ./psu_overlay"
_LSTACK="./filesystem_squashfs"
for squash_file in $(ls ./psu-*.squashfs); do
	# Setup mount for the squashfs images.
	# Note that we stack here on the listing provided by ls command.
	_MOUNTDIR="$(basename $squash_file .squashfs)_squashfs"
	echo "I: found $squash_file ... setting up mount point of ${_MOUNTDIR}"
	if [ ! -d "${_MOUNTDIR}" ]; then
		mkdir "${_MOUNTDIR}"
	fi
	echo "I: mounting $squash_file on ${_MOUNTDIR}"
	mount -o loop "./$squash_file" "./${_MOUNTDIR}"
	# The lowerdir stacks the list from right to left, hence the filesystem.squashfs
	# is the last one in the list, so we pre-append names to the list.
	_LSTACK="./${_MOUNTDIR}:${_LSTACK}"
	_MOUNTLIST="./${_MOUNTDIR} ${_MOUNTLIST}"
done
echo "${_LSTACK}"
echo "${_MOUNTLIST}"

# Setup mount for the filesystem.squashfs mount.
if [ ! -d "filesystem_squashfs" ]; then
	mkdir filesystem_squashfs
fi
mount -o loop ./filesystem.squashfs ./filesystem_squashfs

# Setup mount for the overlay upperdir which is rw which we will squash later.
if [ ! -d "psu_overlay_rw" ]; then
	mkdir psu_overlay_rw
fi
# Setup mount for the overlay work folder which appears to be required.
if [ ! -d "psu_overlay_work" ]; then
	mkdir psu_overlay_work
fi
# Setup mount for the overlay where we stack ro squashfs list.
if [ ! -d "psu_overlay" ]; then
	mkdir psu_overlay
fi


if [ ${_RUNTYPE} = "C" ]
then
	#mount -t overlay overlay -olowerdir=./psu-2_squashfs:./psu-1_squashfs:filesystem_squashfs,upperdir=./psu_overlay_rw,workdir=./psu_overlay_work ./psu_overlay
	mount -t overlay overlay -olowerdir="${_LSTACK}",upperdir=./psu_overlay_rw,workdir=./psu_overlay_work ./psu_overlay
	echo "You selected to run a chroot"
	echo "Setup the chroot to overlay mount."
	echo "Mounting bindings for dev, proc, sys, pts in overlay."
	mount --bind /dev ./psu_overlay/dev
	mount --bind /proc ./psu_overlay/proc
	mount --bind /sys ./psu_overlay/sys
	mount --bind /dev/pts ./psu_overlay/dev/pts

	echo "Backup the overlay/etc/hosts."
	mv ./psu_overlay/etc/hosts ./hosts.bak
	echo "Copying /etc/hosts to overlay/etc/hosts."
	cp /etc/hosts ./psu_overlay/etc/hosts

	echo "Backup the overlay/etc/resolv.conf."
	mv ./psu_overlay/etc/resolv.conf ./resolv.conf.bak
	echo "Copying /etc/resolv.conf to the overlay/etc/resolv.conf."
	cp /etc/resolv.conf ./psu_overlay/etc/resolv.conf

	echo "Starting chroot in the overlay directory."
	chroot ./psu_overlay /bin/bash

	echo "Exited the chroot so time to clean up."
	umount -l ./psu_overlay/dev
	umount -l ./psu_overlay/proc
	umount -l ./psu_overlay/sys

	echo "Restore original overlay/etc/hosts."
	mv ./hosts.bak ./psu_overlay/etc/hosts
	echo "Restore overlay/etc/resolv.conf."
	mv ./resolv.conf.bak ./psu_overlay/etc/resolv.conf
	echo "Remove overlay/root/.bash_history."
	rm ./psu_overlay/root/.bash_history

	# Prompt to save changes or not.
	_DEFAULTYESNO="N"
	echo "#############################################################"
	echo " (Y)es save my chroot modifications."
	echo " (N)o do not save my chroot modifications."
	echo "Select to save your chroot modifications (default is ${_DEFAULTYESNO}):
	"
	read _YESNO

	_YESNO="${_YESNO:-${_DEFAULTYESNO}}"

	if [[ ${_YESNO} =~ ^[Yy]$ ]]
	then
		_DATE=$(date +%Y%m%d-%H%M%S)
		echo "Now making the updated squashfs ${_DATE}."
		mksquashfs ./psu_overlay_rw psu-${_DATE}.squashfs

		# Since we use sed to put text at the top of the file make sure the file exists.
		if [ ! -e psu-changelog.txt ]
		then
			echo "Created changelog." >> psu-changelog.txt
		fi

		# Add entry to changelog for the psu file.
		echo "Please provide a small changelog entry for the psu-${_DATE}.squashfs :"
		read _NOTES
		sed -i "1s/^/${_DATE} - ${_NOTES}.\n/" psu-changelog.txt
	else
		echo "No chroot modifications saved to a psu file."
	fi

	Cleanup_mounts

elif [ ${_RUNTYPE} = "J" ] || [ ${_RUNTYPE} = "N" ]
then
	if [ ${_RUNTYPE} = "J" ]
	then
		# Modify the _LSTACK to not include the filesystem_squashfs mount so we just stack the updates together
		_LSTACK=$(echo ${_LSTACK} | sed s@:./filesystem_squashfs@@g)
		_ACTIONTYPE="join psu squashfs files to new single psu-DATE.squashfs"
	else
		_ACTIONTYPE="create a new single filesystem.squashfs"
	fi

	echo "${_LSTACK}"
	#mount -t overlay overlay -olowerdir=./psu-2_squashfs:./psu-1_squashfs,upperdir=./psu_overlay_rw,workdir=./psu_overlay_work ./psu_overlay
	mount -t overlay overlay -olowerdir="${_LSTACK}",upperdir=./psu_overlay_rw,workdir=./psu_overlay_work ./psu_overlay

	# Prompt to join squashfs files or not.
	_DEFAULTYESNO="N"
	echo "#############################################################"
	echo " (Y)es, ${_ACTIONTYPE}."
	echo " (N)o, do not ${_ACTIONTYPE}."
	echo "Select to ${_ACTIONTYPE} (default is ${_DEFAULTYESNO}):
	"
	read _YESNO

	_YESNO="${_YESNO:-${_DEFAULTYESNO}}"

	if [[ ${_YESNO} =~ ^[Yy]$ ]]
	then
		# Since we use sed to put text at the top of the file make sure the file exists.
		if [ ! -e psu-changelog.txt ]
		then
			echo "Created changelog." >> psu-changelog.txt
		fi

		_DATE=$(date +%Y%m%d-%H%M%S)

		if [ ${_RUNTYPE} = "N" ]
		then
			if [ ! -d "new_squashfs" ]; then
				mkdir new_squashfs
			fi
			echo "Now making a new and updated filesystem.squashfs."
			mksquashfs ./psu_overlay ./new_squashfs/filesystem.squashfs
			# Add entry to changelog for the joining of the *.squashfs files to a new filesystem.squashfs.
			sed -i "1s/^/******* ${_DATE} - New filesystem.squashfs generated. ********\n/" psu-changelog.txt
		else
			echo "Now making the updated squashfs ${_DATE} from overlay."
			mksquashfs ./psu_overlay psu-${_DATE}.squashfs
			# Add entry to changelog for the joining of the psu files.
			sed -i "1s/^/******* ${_DATE} - Joined squashfs files. ********\n/" psu-changelog.txt
		fi
	else
		echo "No psu files will be joined."
	fi

	Cleanup_mounts

	# Remove old directories and move old partials out.
	if [[ ${_YESNO} =~ ^[Yy]$ ]]
	then
		if [ ! -d "psu_overlay_rw-OLD-${_DATE}" ]; then
			mkdir psu-OOS-${_DATE}
		fi
		if [ ${_RUNTYPE} = "N" ]
		then
			rsync -av --remove-source-files ./psu-*.squashfs ./psu-OOS-${_DATE}/
			rsync -av --remove-source-files ./filesystem.squashfs ./psu-OOS-${_DATE}/
			rsync -av --remove-source-files ./new_squashfs/filesystem.squashfs ./
			rmdir ./new_squashfs
		else
			rsync -av --remove-source-files --exclude psu-${_DATE}.squashfs ./psu-*.squashfs ./psu-OOS-${_DATE}/
		fi
	else
		echo "No rsync of current partials to an Out Of Service folder will occur."
	fi
else
	echo "You entered an unknowon choice!"
	echo "Aborting."
	Cleanup_mounts
	exit 1
fi

echo "All done goodbye!"
exit 0
