#!/bin/bash
# sudo mkdir -p /mnt/Original_OS/
# sudo mount -U 86209506-1709-4a97-8518-5cda2bf0b794 "/mnt/Original_OS/"
# sudo chroot /mnt/Original_OS/var/squashfs/ /bin/bash -c "cd /tmp/ ; tmmount -g"
# Modify the below information to correlate with the software you are designing.
PROGNAME="redit"
PROGRAMMER="theemahn"
BUILDDATE="03/15/2023"
VERSION="2.1.0"
WEBSITE="os-builder.com"
AUTHOR="TheeMahn"
EMAIL="<$AUTHOR@$WEBSITE>"
APPNAME="Ram Editor"

if [[ -f "ultimate-common" ]]; then
	source ultimate-common
elif [[ -f "/usr/share/ultimate_edition/ultimate-common" ]]; then
	source /usr/share/ultimate_edition/ultimate-common
else
	echo "No Ultimate Edition common source. Please install ultimate-edition-common."
	exit 1;
fi
if ! [[ "${SUPPRESS}" ]]; then
	VersionDump "$@"
	Timer "Start" "${APPNAME}"
fi

QuestionReboot () {
	declare -i TESTIT;
	declare -i RESULTS;
	if [[ -f "/etc/default/grub" ]]; then
		TESTIT=1
		RESULTS=1
		if [[ -f "/etc/default/grub" ]]; then
			TIMEOUT=$(grep -i "GRUB_TIMEOUT=" "/etc/default/grub" | sed "s/GRUB_TIMEOUT=//g")
			TESTIT=$(grep -i "GRUB_DEFAULT=" "/etc/default/grub")
			TEMPRESULTS=$(grep -i "GRUB_DEFAULT=" "/etc/default/grub" | sed "s/GRUB_DEFAULT=//g")
			TSTYLE=$(grep -i "GRUB_TIMEOUT_STYLE=" "/etc/default/grub" | sed "s/GRUB_TIMEOUT_STYLE=//g")
			RESULTS=$(echo "${TEMPRESULTS}" | sed 's/"//g')
			TRESULT=$(echo "${TSTYLE}" | sed 's/"//g')
		else
			Error "Critical: /etc/default/grub not found. Exiting."
			exit 1;
		fi
		if [[ "${TSTYLE}" ]]; then
			Encapsulate "GRUB_TIMEOUT_STYLE=${TSTYLE}"
			Encapsulate "Timeout in seconds: ${TIMEOUT}"
			if ! [[ "${TSTYLE}" == "menu" ]]; then
				Encapsulate "Setting menu to on."
				sed -i "s/^GRUB_TIMEOUT_STYLE=.*/GRUB_TIMEOUT_STYLE=menu/g" "/etc/default/grub"
				UpdateGrub
			fi
		else
			Encapsulate "GRUB_TIMEOUT_STYLE has not been set. Setting to menu."
			echo "GRUB_TIMEOUT_STYLE=menu" | sudo tee -a "/etc/default/grub" 1>/dev/null;
			Encapsulate "Timeout in seconds: ${TIMEOUT}"
			UpdateGrub
		fi
		if [[ "${TESTIT}" == 0 ]]; then
			if [[ -f "/original_os" && ! "${RAMSESSION}" ]]; then
				Center "Original O/S is current."
				if ! [[ "${SUPPRESS}" && "${RESULTS}" == 0 ]]; then
					read -p "${BRIGHT}${bldblu}▒ ${BRIGHT}${txtgrn}Do you wish to reboot now to use your Ram Session? [y/N]: " ANSWER
					echo -en "\033[1A\033[2K"
					case "${ANSWER}" in
						y) Center "Rebooting...";
						sudo reboot now;;
						*) 	FullBar;;
					esac
				fi
			fi
		fi
	else
		Encapsulate "Default is not boot from RAM, exiting..."
		exit 0;
	fi
}
############################
# Only run if user is root #
############################

CheckRoot "$@"

##################################################################################
# Check if the user is trying to run this script from outside of the RAM Session #
##################################################################################

if [ ! -e /ram_session ]
then
	Error "This script cannot be run from outside of the RAM Session."
	exit 0
fi

##########################################################################
# Check if the user is trying to run this script from inside of a chroot #
##########################################################################

ischroot
if [ $? -eq 0 ]
then
	Error "This script cannot be run from inside of a chroot"
	exit 0
fi

####################
# Global Variables #
####################

#Path to the file that contains all the functions for this script
RLIB="/bin/rlib"
#Figure out $DEST
DEST=$(cat /var/lib/ram_booter/conf | grep -v '^#' | grep DEST= | sed 's/DEST=//g')
#Remove leading slash from $DEST
DEST="${DEST#/}"
COMPRESSION=$(cat /var/lib/ram_booter/conf | grep -v '^#' | grep -i compression= | sed 's/COMPRESSION=//g')
#The mountpoint of the Original OS device
#Note: If you change it here, be sure to change it in the rupdate script as well
ORIG_OS='/mnt/Original_OS/'

##########################################################
# Source the file with all the functions for this script #
##########################################################

if [[ -e $RLIB ]]
then
        . $RLIB

	#Check if there was a problem
	if [[ "$?" != "0" ]]
	then
		echo
		echo "The library that comes with RAM Booter ($RLIB) failed to be sourced"
		echo "Is it broken?"
		exit 1
	fi
else
        echo "The library that comes with RAM Booter ($RLIB) was not found!"
        exit 1
fi

#####################
# The usage message #
#####################

Usage()
{
	echo "Usage:"
	echo "$0 -s, --squashfs"
	echo "	Skip chroot and recreate the squashfs image"
	echo "$0 -h, --help"
	echo "	This help message"
}

###############################
# Check args passed to script #
###############################

if [[ "$#" -gt 1 ]]
then
	echo 'Invalid number of arguments'
	Usage
	exit 1
fi

case "$1" in
	-s|--squashfs)
		#Ask user if he wants to recreate filesystem.squashfs
		#This basically just lets the user skip the chroot completely
		Encapsulate "Are you sure you want to recreate the squashfs image?"
		Encapsulate "This may take some time to complete"
		Encapsulate "You may use the computer normally during this process - you will NOT disrupt it"
		echo
		read -p "Your choice [y/N]: " answer

		#Convert answer to lowercase
		answer=$(toLower $answer)

		case $answer in
			y|yes)
				#Create squashfs image
				CreateSquashfs
				;;
			*)
				echo
				echo "You chose to NOT recreate the squashfs image"
				exit 0
				;;
		esac

		exit 0
		;;
	-h|--help)
		Usage
		exit 0
		;;
	"")
		#No args is fine
		;;
	*)
		echo "'$1' is not a valid argument"
		Usage
		exit 1
		;;
esac


############
# Mount FS #
############

#Mount the RAM Session, and bind all the pseudo filesystems to it
MOUNT_RS

#####################################################
# Tell user about any user home dirs we had to bind #
#####################################################

if mount | grep bind | grep -q $"${ORIG_OS%/}/${DEST%/}/home/[^ \t]*"
then
	echo "The following encrypted home directories were automatically mounted,"
	echo "since you are currently logged into them:"
	echo
	mount | grep bind | grep $"${ORIG_OS%/}/${DEST%/}/home/[^ \t]*" | sed 's/on.*//g'
	echo
fi

############################
# Let user into the chroot #
############################
if [[ -f "/usr/share/ultimate_edition/extras/both/tmmount" ]]; then
	Encapsulate "Installing tmmount into the chroot environment."
	sudo cp "/usr/share/ultimate_edition/extras/both/tmmount" "$ORIG_OS/$DEST/bin/"
else
	Error "Error tmmount, not found. Continuing."
fi
sudo chroot $ORIG_OS/$DEST/ /bin/bash -c "cd /tmp/ ; tmmount -g"

######################
# Some quick cleanup #
######################

#Clean apt cache quietly
sudo chroot $ORIG_OS/$DEST/ /bin/bash -c "apt-get -y clean" &>/dev/null

#Update fake /boot, just in case any packages being installed in the
#RAM Session want /boot to be consistent with the running system
#Note: This is usually only necessary during an update, but we do it
#	anyway in case the user did updates in the redit instead of
#	using rupdate
sudo mount -o bind $ORIG_OS/$DEST/ $ORIG_OS/$DEST/mnt/ &>/dev/null &&
sudo rsync --delete -a $ORIG_OS/$DEST/boot/ $ORIG_OS/$DEST/mnt/boot/ --exclude=/IMPORTANT_README &>/dev/null &&
sudo umount $ORIG_OS/$DEST/mnt/ &>/dev/null

##########################################################################
# Check if /etc/kernel/postinst.d/zd_warn has set a flag that shows that #
# a kernel update occurred. If it has, the user should get a warning     #
##########################################################################

if [[ -e $ORIG_OS/$DEST/WARN ]]
then
	echo "WARNING: A KERNEL UPDATE JUST OCCURRED, AND THE RAM SESSION GRUB ENTRY"
	echo "HAS BEEN MODIFIED TO USE THE NEW KERNEL. YOU MUST RECREATE THE SQUASHFS"
	echo "IMAGE BEFORE YOU REBOOT OR YOUR RAM SESSION WILL NOT HAVE THE MODULES"
	echo "TO USE THE NEW KERNEL, AND MAY FAIL TO BOOT!"
	rm $ORIG_OS/$DEST/WARN
	echo
fi

#Unmount everything
UMOUNT_RS

#If there was a problem unmounting stuff, tell user to reboot
if [[ "$?" != "0" ]]
then
	Error "There was an error unmounting ${ORIG_OS}"
	Encapsulate "A reboot should fix that"
	Encapsulate "After the reboot, if you need to save your changes by recreating"
	Encapsulate "the squashfs image, just run \"sudo redit -s\""
	exit 1
fi

#################################
# Recreating the SquashFS Image #
#################################

#Ask user if he wants to recreate filesystem.squashfs
FullBar
Encapsulate "Would you like to save your changes by recreating the squashfs image?"
Encapsulate "This may take some time to complete"
Encapsulate "You may use the computer normally during this process - you will NOT disrupt it"
Encapsulate ""
read -p "${BRIGHT}${bldblu}▒ ${BRIGHT}${txtgrn}Your choice [y/N]: " answer
echo -en "\033[1A\033[2K"

#Convert answer to lowercase
answer=$(toLower $answer)

case $answer in
	y|yes)
		#Create squashfs image
		CreateSquashfs
		;;
	*)
		echo
		echo "You chose to NOT recreate the squashfs image"
		echo "If you change your mind, run 'sudo redit -s'"
		exit 0
		;;
esac

