#!/bin/dash
########  Click image files to mount & unmount  ########
# Originally by: Terry Becker	aka: SunBurnt
# Rewritten by jamesbond 2012, 2014, 2016, 2018, 2019, 2021 for Fatdog64 (License: GPLv3)
# Support for LUKS encrypted images/savefiles by JakeSFR 2021 (License: GPLv3)
# Supported filetypes: iso, sfs, 2fs/ext2, 3fs/ext3, 4fs/ext4, btrfs, img (filesystem and diskimage), initrd/initrd.gz
# parameter: $1 - filename

### configuration
MKINITRD_NAME=repack-initrd



########### identification helpers #############

# $extension
is_supported() {
	is_readonly_image || is_readwrite_image || is_initrd
}

# $extension
is_readonly_image() {
	case $extension in
		iso|sfs) return 0
	esac
	return 1
}

# $extension
is_readwrite_image() {
	case $extension in
		2fs|3fs|4fs|ext2|ext3|ext4|img|btrfs) return 0
	esac
	return 1
}

# $extension
is_initrd() {
	case $extension in
		initrd|initrd.?z|initrd.?z?|initrd-*) return 0
	esac
	return 1
}

# $filesystem $extension
is_luks() {
	test "$filesystem" = crypto_LUKS
}




############ generic helpers #################

# $1-fullpath
get_filesystem() {
	local fs=$(guess_fstype "$1")
	[ "$fs" = "unknown" ] && file "$1" | grep -q "DOS/MBR" && fs="diskimage"
	echo "$fs"
}

# $1-fullpath
find_loop_device() {
	local p pp
	for p in /sys/block/loop*; do
		if [ -e $p/loop ]; then
			read pp < $p/loop/backing_file
			case "$pp" in
				*"$1") echo /dev/${p##*/}; break ;;
			esac
		fi
	done
}

# $1-mountpoint
remove_mountpoint() {
	# workaround a nasty aufs bug
	# only remove mountpoint if there is exactly one aufs root
	# otherwise aufs will *hard-lockup*
	if [ $(ls -d /sys/fs/aufs/* | wc -l) -eq 1 ]; then
		rmdir "$1"
	fi
	return 0 # always return success
}

# $1-mountpoint $2-fullpath $filesystem $filename
unmount_ops() {
	local mountpoint="$1" fullpath="$2"
	case "$filesystem" in
		diskimage)
			if Xdialog --title "Confirm" --yesno "You have to un-mount all partition(s) used by $filename first.\n\nContinue?\n" 0 0;
			then
				kpartx -d "$fullpath" > /dev/null
			else
				exit
			fi
			;;
		*)
			umount -d /mnt/$mountpoint && remove_mountpoint /mnt/$mountpoint
			;;
	esac
}

# $1-mountpoint $2-fullpath $3-filepath $filesystem
is_mounted() {
	# preserve variables from getting modified by luks_get_modified_mountpoint
	local mountpoint="$1" fullpath="$2" filepath="$3"

	case "$filesystem" in
		diskimage)
			test "$(find_loop_device "$fullpath")"
			;;
		crypto_LUKS)
			if ls /dev/mapper | grep -qw "^${cryptpoint}$"; then
				luks_get_modified_mountpoint "$mountpoint" "$filepath"
				grep -q $mountpoint /proc/mounts
			else
				return 1
			fi
			;;
		*)
			grep -q $mountpoint /proc/mounts
			;;
	esac
}




############ initrd helpers #################

# create a simple script to re-pack initrd
# $1-location of original initrd, $2-compression type
create_mkinitrd() {
	local COMPRESSOR 
	touch $MKINITRD_NAME; chmod +x $MKINITRD_NAME
	case $2 in
		gzip) COMPRESSOR="gzip -9" ;;
		bzip2) COMPRESSOR="bzip2" ;;
		xz) COMPRESSOR="xz --check=crc32 --lzma2=dict=512KiB" ;;
		zstd) COMPRESSOR="zstd -T0 -19" ;;
		none) COMPRESSOR=cat;;
	esac
	cat >> $MKINITRD_NAME << EOF
#!/bin/dash
[ \$(id -u) -ne 0 ] && exec gtksu "Rebuild initrd" \$(readlink -f "\$0")
SOURCE="\$(dirname "\$0")"; cd "\$SOURCE"
exec Xdialog --no-buttons --title "$MKINITRD_NAME" --infobox "Working, please wait ..." 0 0 1000000 &
XPID=\$!
find . | grep -v $MKINITRD_NAME | cpio -o -H newc | $COMPRESSOR > "$1"
kill \$XPID
su \$USER -c "rox -D \"\$SOURCE\""; rm -rf "\$SOURCE"
exec Xdialog --title "$MKINITRD_NAME" --infobox "Repacking done, $1 is now updated." 0 0 10000
EOF
}

# $fullpath $filename $MKINITRD_NAME
expand_initrd() {
	local tmpdir TYPE
	
	tmpdir=$(mktemp -dp /tmp initrd-XXXXXX)
	chgrp $USER $tmpdir
	chmod a+rwx $tmpdir
	cd $tmpdir
	if case $(file -b --mime-type "$fullpath") in
		*gzip) TYPE="gzip" && zcat "$fullpath" | cpio -i ;;
		*xz) TYPE="xz" && xzcat "$fullpath" | cpio -i ;;
		*bzip2) TYPE="bzip2" && bzcat "$fullpath" | cpio -i ;;
		*zstd) TYPE="zstd" && zstdcat "$fullpath" | cpio -i ;;
		*) TYPE="none" && cat "$fullpath" | cpio -i ;;
	esac; then
		if [ -w "$fullpath" ]; then
			create_mkinitrd "$fullpath" $TYPE &&
			disposition="
When done, you can rebuild $filename by running $MKINITRD_NAME.\n
If you decide otherwise, simply remove $tmpdir."
		else
			disposition="
When done, simply remove $tmpdir.\n\n
NOTE: you can't rebuild $filename because it isn't writable."
		fi
		su $USER -c "rox $tmpdir"
		Xdialog --left --title "Success" --infobox "
$filename is an initrd archive, and it has been extracted to $tmpdir.\n
Do what you want to do with it.\n
$disposition" 0 0 10000 
	else
		Xdialog --title "Error!" --infobox "Failed to extract $filename." 0 0 10000 
	fi
}




############ LUKS helpers #################

# NOTE: needs to be called again, after running open_luks
# in: __ORIGINAL__ $1-mountpoint, $__ORIGINAL__ 2-filepath
# out: $mountpoint $mountpoint $orig_mountpoint $orig_filepath
luks_get_modified_mountpoint() {
	orig_mountpoint="$1"
	orig_filepath="$2"
	
	# now the corresponding dm-* becomes a new $mountpoint and /dev/$mountpoint becomes a new filepath
	mountpoint=$(grep -H "^${1}$" /sys/devices/virtual/block/dm-*/dm/name | grep -wo 'dm-[0-9]\+' | head -n 1)
	filepath="/dev/$mountpoint"
}

# NOTE: has to be called with the original $filepath
# __ORIGINAL__ $1-filepath, $cryptpoint $filename and 
luks_open() {
	local filepath="$1"
	local password checkbox keyfile
	if ! ls /dev/mapper | grep -qw "^${cryptpoint}$"; then
		result=$(Xdialog --stdout --title "Request" --check "Use keyfile instead of a password" false \
					--password --input "Enter password for $filename:" 0 0) || exit

		password="$(echo "$result" | head -n 1)"
		checkbox="$(echo "$result" | tail -n 1)" 

		if [ "$checkbox" = 'checked' ]; then
			keyfile=$(Xdialog --stdout --title "Request" --backtitle "Select a keyfile for $filename:" \
						--fselect "$HOME" 0 0) || exit
			cryptsetup -d "$keyfile" open "$filepath" "$cryptpoint"
		else
			echo "$password" | cryptsetup open "$filepath" "$cryptpoint"
		fi

		if [ $? -ne 0 ]; then
			Xdialog --title "Error!" --infobox "Failed to decrypt $filename." 0 0 10000
			exit
		fi
	fi
}

# $1-cryptpoint $filesystem
luks_close() {
	if [ "$filesystem" = 'crypto_LUKS' ]; then
		cryptsetup close "$1"
	fi
}




############### actual tasks ##################

do_mount() {
	# with LUKS we have to deal with 2 filesystems:
	# crypto_LUKS and whatever there is underneath
	local actual_filesystem="$filesystem"
	
	# try to open luks device
	if is_luks; then
		luks_open "$filepath"
		luks_get_modified_mountpoint "$mountpoint" "$filepath"
		actual_filesystem=$(get_filesystem "$filepath")
	fi

	# ensure we understand the filesystem. Otherwise bail out.
	case $actual_filesystem in
		diskimage)
			# check if it is already mounted elsewhere
			elsewhere=$(find_loop_device "$fullpath")
			if [ "$elsewhere" ]; then
				Xdialog --title "Error!" --infobox "$filename is already mounted elsewhere! This is not supposed to happen." 0 0 10000
				exit
			else
				kpartx -a "$fullpath"
				Xdialog --title "Done" --infobox "$filename is a disk image containing multiple partition(s).\n\n
Access the partitions using the drive icons.\n" 0 0 10000
			fi
			;;

		unknown|"")
			luks_close "$cryptpoint"
			Xdialog --title "Error!" --infobox "Cannot mount $filename.\nIt contains unknown filesystem." 0 0 10000
			;;

		*)
			# prepare to mount
			mountpoint="/mnt/$mountpoint"
			mkdir -p "$mountpoint" 2> /dev/null

			# check if it is already mounted elsewhere
			# (does not work with LUKS containers)
			elsewhere=$(find_loop_device "$fullpath")

			if [ "$elsewhere" ] && ! [ "$filesystem" = 'crypto_LUKS' ]; then
				# yes, do bind mount
				elsewhere="$(awk -v loopdev="$elsewhere" '$1 == loopdev {print $2}' /proc/mounts)"
				mount -o bind "$elsewhere" "$mountpoint"

			else
				# no, do loop mount
				if type fatdog-drive-icon-mount-helper.sh > /dev/null; then
					# use fatdog specific helper if available - this makes sure ntfs-3g is mounted with the correct options etc
					fatdog-drive-icon-mount-helper.sh "$filepath" "$mountpoint"
				else
					# else, use standard mount command.
					# mount -o loop "$filepath" "$mountpoint"
					# implicit '-o loop' is not needed anymore and prevents mounting dm-* devices
					mount "$filepath" "$mountpoint"
				fi
			fi &&

			su $USER -c "rox $mountpoint" ||
			{
				luks_close "$cryptpoint"
				Xdialog --title "Error!" --infobox "Failed to mount $filename." 0 0 10000
			}
			;;
	esac
}

do_unmount() {
	#if ! grep $mountpoint /proc/mounts && busybox losetup -a | grep $fullpath; then
	#	load_sfs.sh	# in losetup but not in /proc/mounts, probably mounted by SFS loader
	#	exit
	#fi

	# adjust 4mountpoint and $filepath for luks
	if is_luks; then
		luks_get_modified_mountpoint "$mountpoint" "$filepath"
	fi
	
	# already mounted, will try to unmount
	if unmount_ops "$mountpoint" "$fullpath"; then
		luks_close "$cryptpoint"
		su $USER -c "rox -D /mnt/$mountpoint"
		Xdialog --title "Success!" --infobox "$filename has been un-mounted." 0 0 10000 
	else
		if Xdialog --title "Error!" --yesno "Failed to un-mount $filename, it is currently used by the following process:
$(fuser -m /mnt/$mountpoint | xargs ps -o cmd -p | sed 's/CMD//')\n\n
Do you want me to stop them and try again?" 0 0; then
			[ $filesystem != "diskimage" ] && fuser -m -k /mnt/$mountpoint
			if ! is_luks; then
				exec $0 "$filepath"
			else
				exec $0 "$orig_filepath"
			fi
		fi
	fi
}




############### main ##################
! [ -e "$1" ] && exit 1

### need root
[ $(id -u) -ne 0 ] && exec gtksu "File Mounter" "$0" $(readlink -f "$1")

### get image file and generate mountpoint
filepath=$1
fullpath=$(readlink -f "$filepath")
filename=${filepath##*/}
extension=${filename##*.}
[ "$extension" = "$filename" ] && extension=
case  "$filename" in
	initrd|initrd.?z|initrd.?z?|initrd-*) extension=$filename ;;
esac
extension=$(echo $extension | tr "A-Z" "a-z")
mountpoint=$(echo $fullpath | tr "'\\\\/ .;:?*[]{}()<>&|\"\`\t" +)
cryptpoint="$mountpoint"	# use the same name for /dev/mapper/... device
filesystem=$(get_filesystem "$fullpath")
#echo $filepath $fullpath $filename $extension $mountpoint $filesystem

# only supported extensions
if ! is_supported; then
	Xdialog --title "Error!" --infobox "Unsupported file type: $extension" 0 0 10000
	exit
fi

# mount only if it is not already mounted
if ! is_mounted "$mountpoint" "$fullpath" "$filepath"; then

	if is_initrd; then
		expand_initrd
	else
		if is_readwrite_image; then
			# for writable images, touch the image before mounting
			# if in savefile:ram mode, this will copy it from pup_save to pup_rw
			# without this they operate as read-only (aufs quirks)
			touch "$fullpath"
		fi
		do_mount
	fi

else
	do_unmount
fi
