#!/bin/bash

# simplechroot

#
# Copyright (C) 2018 Caleb Woodbine <calebwoodbine.public@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#

mountNode=false
_verbose=false
version=1.1.1

if [ -z "$args" ]
then
	args="/bin/bash"
fi

if (($# == 0))
then
	echo "Error: No args provided. Please run 'simplechroot -h' for help."
	exit 1
fi

function msg() {
# print if _verbose is true
if [ "$_verbose" = true ]
then
	echo "::> $@"
fi
}

function parse_rootDir() {
# resolve dir or node
rootDir="$_rootDir"
if [ ! -e "$rootDir" ]
then
	echo "Error: directory or node not found."
	trap "" EXIT
	exit 1
fi

if [[ "$_rootDir" = "/dev"* ]]
then
	_rootDir_basename=$(basename $_rootDir)
	if ! lsblk | grep -q "$_rootDir_basename" || [[ "$_rootDir_basename" = "loop"* ]] || ( [[ ! "${_rootDir_basename#${_rootDir_basename%?}}" = [[:digit:]] ]] && [[ ! $(dirname $_rootDir) = "/dev/mapper" ]] && [[ ! $(dirname $_rootDir) = "/dev/dm"* ]] )
	then
		echo "Error: please enter a valid root partition disk node (must be a partition)."
		exit 1
	fi
	mkdir -p /tmp/simplechroot-mount-root
	if mount | grep -q /tmp/simplechroot-mount-root || ! mount "$_rootDir" /tmp/simplechroot-mount-root
	then
		echo "Cannot mount '$_rootDir'."
		exit 1
	fi
	msg "Status: mounted '$_rootDir' to '/tmp/simplechroot-mount-root'."
	rootDir="/tmp/simplechroot-mount-root"
	mountNode=true

elif [ -d "$_rootDir" ] || [[ "$_rootDir" = "/dev"* ]]
then
	rootDir=$(readlink -f "$rootDir")
	return
else
	echo "'$_rootDir' is not a valid directory or node."
	exit 1
fi
}

function setupChroot() {
# bind mount points

msg "Setting up '$_rootDir'"
for point in dev proc sys run
do
	if ! mount --bind "/$point" "$rootDir/$point"
	then
		echo "Failed to mount '/$point' to '$rootDir/$point'; Unmounting all."
		unmountChroot
		exit
	fi
done
}

function unmountChroot() {
# unmount all binds
msg "Unmounting partitions"
for point in dev proc sys run
do
	if mount | grep -q "$rootDir/$point"
	then
		umount "$rootDir/$point" || echo "Failed to unmount '$rootDir/$point'"
	fi
done

if [ ! -z "$bootDir" ] && mount | grep -q "$rootDir"/boot
then
	umount -f "$rootDir"/boot
fi

if [ "$mountNode" = true ]
then
	umount "$rootDir"
fi
}

function validate_fsDir() {
# check if $_rootDir contains a standard Linux filesystem
fsUse="$1"
for i in bin/bash boot dev etc home lib mnt proc root run sbin sys tmp usr var
do
	if [ -e "$fsUse/$i" ]
	then
		continue
	else
		[[ ! "$NOECHO" = true ]] && echo "Cannot find standard Linux filesystem folders or programs."
		return 1
	fi
done

if [ "$reinstallgrub" = true ]
then
	if [[ -e "$rootDir/usr/sbin/grub-install" ]] || [[ -e "$rootDir/usr/sbin/grub-mkconfig" ]]
	then
		grubver=""

	elif [[ -e "$rootDir/usr/sbin/grub2-install" ]] || [[ -e "$rootDir/usr/sbin/grub2-mkconfig" ]]
	then
		grubver="2"

	else
		echo "Cannot find grub install programs on destination root."
		exit 1
	fi
fi

unset i
return
}

function parse_bootDir() {
# parse given boot node
if [[ "$bootDir" = "/dev"* ]]
then
	_bootDir_basename=$(basename $bootDir)
	if ! lsblk | grep -q "$_bootDir_basename" || [[ "$_bootDir_basename" = "loop"* ]] || ( [[ ! "${_bootDir_basename#${_bootDir_basename%?}}" = [[:digit:]] ]] && [[ ! $(dirname $bootDir) = "/dev/mapper" ]]  && [[ ! $(dirname $bootDir) = "/dev/dm"* ]] )
	then
		echo "Error: please enter a valid boot partition disk node (must be a partition)."
		exit 1
	fi
	return
fi

echo "Error: arg '-b' must be a node in '/dev'."
exit 1
}

function mount_bootDir() {
# mount boot node when given
if [[ "$bootDir" = "$_rootDir" ]]
then
	return
fi

if ! mount "$bootDir" "$rootDir"/boot
then
	echo "Error: failed to mount '$bootDir' to '$rootDir/boot'."
	exit
fi
msg "Status: mounted '$bootDir' to '$rootDir/boot'."
}

function listPartitions() {
# list partitions that are mountable
partitionList=""
partitionUUIDs=""
partCount=0
echo "Listing partitions which should be usable"
echo "-----------------------------------------"
if grep -q "/mnt" /etc/mtab
then
	mountPoint="/tmp/simplechroot-mount-test"
	mkdir -p "$mountPoint"
else
	mountPoint="/mnt"
fi
while read o
do
	o_cut=$(echo $o | cut -d ' ' -f1)
	o_basename=$(basename $o_cut)
	o_blkuuid=$(lsblk -no UUID $o_cut)
	if echo "sr0" | grep -q "$o_basename" || ! lsblk | grep -q "$o_basename" || [[ "$o_basename" = "loop"* ]] || ( [[ ! "${o_basename#${o_basename%?}}" = [[:digit:]] ]] && [[ ! $(dirname $o_cut) = "/dev/mapper" ]] && [[ ! $(dirname $o_cut) = "/dev/dm"* ]] ) || echo "swap LVM2_member ntfs hfs hfs+ apfs" | grep -q "$(echo $o | cut -d ' ' -f2)" || echo "$partitionList" | grep -q "$o_cut" || echo "$partitionUUIDs" | grep -q "$o_blkuuid"
	then
		continue
	fi
	partitionList="${partitionList} $o_cut"
	partitionUUIDs="${partitionUUIDs} ${o_blkuuid}"
	if grep -q "$o_cut" /etc/mtab
	then
		o_cut_mountpoint=$(findmnt -nr -o target -S $o_cut)
		if NOECHO=true validate_fsDir "$o_cut_mountpoint"
		then
			if [ -e "$o_cut_mountpoint/boot/grub/grub.cfg" ] || [ -e "$o_cut_mountpoint/boot/grub2/grub.cfg" ]
			then
				echo "[${partCount}-mr] $o_cut    (probably an installed Linux system with grub, already mounted)"
			else
				echo "[${partCount}-mr] $o_cut    (probably an installed Linux system, already mounted)"
			fi
		else
			echo "[${partCount}-m.] $o_cut    (already mounted)"
		fi
		partCount=$((partCount += 1))
		continue
	fi
	if ! mount "$o_cut" "$mountPoint" > /dev/null 2>&1
	then
		echo "[${partCount}-!!] $o_cut    (was unable to mount)"
		partCount=$((partCount += 1))
		continue
	fi
	if NOECHO=true validate_fsDir "$mountPoint"
	then
		if [ -e "$mountPoint/boot/grub/grub.cfg" ] || [ -e "$mountPoint/boot/grub2/grub.cfg" ]
		then
			echo "[${partCount}-rb] $o_cut    (probably an installed Linux system with grub)"
		else
			echo "[${partCount}-r.] $o_cut    (probably an installed Linux system)"
		fi
		partCount=$((partCount += 1))
		umount -f "$mountPoint"
		continue
	fi
	if [ -e "$mountPoint/grub/grub.cfg" ] || [ -e "$mountPoint/grub2/grub.cfg" ]
	then
		echo "[${partCount}-b.] $o_cut    (probably a Linux boot partition)"
		partCount=$((partCount += 1))
		umount -f "$mountPoint"
		continue
	fi
	echo "[${partCount}-??] $o_cut    (neither an installed Linux system or boot partition, could be a data partition such as a home partition)"
	if ! umount -f "$mountPoint"
	then
		echo "[${partCount}-!!] $o_cut    (was unable to unmount)"
		partCount=$((partCount += 1))
		exit 1
	fi
done < <(lsblk -rpf | sed -n '1!p' && realpath /dev/mapper/* | sed '/control/d')
echo "-----------------------------------------"
}

# parse args
while [ $# -gt 0 ]
do
	case "$1" in
		-r|--root)
			# root directory or node
			shift
			_rootDir="$1"
			;;

		-b|--boot)
			# boot directory or node
			shift
			bootDir="$1"
			;;

		-rg|--reinstall-grub)
			# grub reinstallation
			reinstallgrub=true
			if [ -z "$bootDir" ] || [ -z "$_rootDir" ]
			then
				echo "Error: args '-r' and '-b' must be used. It is required for reinstalling grub. If the system has a single partition, just specify in '-r' and '-b'"
				exit 1
			fi
			bootDir_base=$(lsblk -no pkname --path $bootDir)
			;;

		-c|--comand)
			# commands to run in chroot
			shift
			sendcommands=true
			args=$@
			;;

		-v|--verbose)
			# output progress
			_verbose=true
			;;

		-l|--list-partitions)
			listPartitions
			exit
			;;

		*)
			# help menu
			echo "simplechroot ($version)"
			echo
			echo "Usage examples:"
			echo "	simplechroot -r /dev/sda2 -b /dev/sda1"
			echo "	simplechroot -r /mnt/rootMount"
			echo "	simplechroot -r /mnt/mapper/lvmPartition"
			echo "	simplechroot -r /mnt/rootMount -b /dev/sda1"
			echo "	simplechroot -r /dev/sda2 -b /dev/sda1 --reinstall-grub"
			echo "	simplechroot -r /mnt/rootMount -b /dev/sda1 -c 'ls /bin/.'"
			echo
			echo "==========================================================="
			echo "| -l               | list partitions and LVM partitions   |"
			echo "| -r               | mounted root partition or devfs node |"
			echo "| -b               | boot partition devfs node            |"
			echo "| -v               | make output verbose                  |"
			echo "| -c               | send a command to the chroot         |"
			echo "| --reinstall-grub | reinstall grub (requires -b and -r)  |"
			echo "==========================================================="
			echo
			echo "To find your disks, use 'simplechroot -l' or your disk manager of choice."
			echo
			exit
			;;
	esac
	shift
done

if [ ! "$(id -u)" = 0 ]
then
	echo "Error: you must be root to use this script."
	exit 1
fi

if [ -z "$_rootDir" ]
then
	echo "Error: please enter root directory or node using the '-r' arg."
	exit 1
fi

if [ ! -z "$sendcommands" ] && [ ! -z "$reinstallgrub" ]
then
	echo "Error: please either use '-c' or '--reinstall-grub'"
	exit 1
fi

trap unmountChroot EXIT

# main process
parse_rootDir
validate_fsDir "$rootDir" || exit 1
if [ ! -z "$bootDir" ]
then
	parse_bootDir
	mount_bootDir
fi
setupChroot
if [ "$reinstallgrub" = true ]
then
	msg "Reinstalling grub${grubver} on '$bootDir_base'."
	args=("grub${grubver}-install $bootDir_base --recheck" "grub${grubver}-mkconfig -o /boot/grub${grubver}/grub.cfg")
fi
msg "=======CHROOT START======="
for i in "${args[@]}"
do
	chroot "$rootDir" ${i[@]}
done
msg "========CHROOT END========"