From 9301fc2728e31afab4e917af78856e1331cda335 Mon Sep 17 00:00:00 2001
From: MichaIng <micha@dietpi.com>
Date: Wed, 22 Jan 2025 21:06:44 +0100
Subject: [PATCH] v9.10

- DietPi-Imager | Allow generating images from the host's root drive, by mounting its rootfs and in case bootfs R/O
---
 .build/images/dietpi-imager | 108 ++++++++++++++++++++++++------------
 1 file changed, 74 insertions(+), 34 deletions(-)

diff --git a/.build/images/dietpi-imager b/.build/images/dietpi-imager
index af0885f60c..133ff4f77a 100755
--- a/.build/images/dietpi-imager
+++ b/.build/images/dietpi-imager
@@ -40,7 +40,7 @@
 	readonly G_PROGRAM_NAME='DietPi-Imager'
 	G_CHECK_ROOT_USER
 	G_CHECK_ROOTFS_RW
-	readonly FP_ORIGIN=$PWD # Store origin dir
+	FP_ORIGIN=$PWD # Store origin dir
 	G_INIT
 	G_EXEC cd "$FP_ORIGIN" # Process everything in origin dir instead of /tmp/$G_PROGRAM_NAME
 	# Import DietPi-Globals ---------------------------------------------------------------
@@ -58,6 +58,7 @@
 	##########################################
 	# Process inputs
 	##########################################
+	HOST_ROOT_DRIVE=$(lsblk -npo PKNAME "$G_ROOTFS_DEV")
 	SOURCE_TYPE='Drive'
 	FP_SOURCE_IMG=
 	PART_TABLE_TYPE=
@@ -116,12 +117,21 @@
 		[[ -d $FP_MNT_TMP ]] && G_EXEC rmdir "$FP_MNT_TMP"
 		Delete_Loopback
 		[[ -e 'tmpiso' ]] && G_EXEC rm -R tmpiso
+		if [[ $FP_SOURCE == "$HOST_ROOT_DRIVE" ]]
+		then
+			/var/lib/dietpi/services/fs_partition_resize.sh
+			findmnt -M /boot > /dev/null && G_EXEC mount -o remount,rw /boot
+			G_EXEC mount -o remount,rw /
+			/boot/dietpi/dietpi-services start
+		fi
 	}
 
 	Run_fsck()
 	{
 		if [[ $ROOT_FS_TYPE == 'ext4' ]]
 		then
+			# If root filesystem is R/O, e2fsck can run but returns 2 stating a reboot is required
+			[[ $FP_SOURCE == "$HOST_ROOT_DRIVE" ]] && G_EXEC_POST_FUNC(){ (( $exit_code == 2 )) && exit_code=0; }
 			G_EXEC_OUTPUT=1 G_EXEC e2fsck -fyD "$FP_ROOT_DEV"
 
 		elif [[ $ROOT_FS_TYPE == 'f2fs' ]]
@@ -162,37 +172,64 @@
 		if [[ $SOURCE_TYPE == 'Drive' ]]
 		then
 			# Detect drives with a partition table, containing a partition with and ext4, F2FS or Btrfs filesystem, excluding the hosts root filesystem drive
-			mapfile -t G_WHIP_MENU_ARRAY < <(mawk -v root="$(lsblk -npo PKNAME "$G_ROOTFS_DEV")" '$1!=root && $2~/^(ext4|f2fs|btrfs)$/ {print $1"\n"$1"_details"}' < <(lsblk -rnpo PKNAME,FSTYPE) | sort -u)
-
-			if [[ ! ${G_WHIP_MENU_ARRAY[0]} ]]
-			then
-				G_DIETPI-NOTIFY 1 'No drive with an ext4, F2FS or Btrfs formatted partition found, aborting...'
-				G_DIETPI-NOTIFY 2 'NB: This is the list of available block devices:'
-				lsblk -npo NAME,SIZE,MAJ:MIN,FSTYPE,MOUNTPOINT,MODEL
-				read -rp 'Press any key to return to menu...'
-				return 1
-			fi
+			mapfile -t G_WHIP_MENU_ARRAY < <(mawk -v root="$HOST_ROOT_DRIVE" '$1!=root && $2~/^(ext4|f2fs|btrfs)$/ {print $1"\n"$1"_details"}' < <(lsblk -rnpo PKNAME,FSTYPE) | sort -u)
 
 			# Visually separate dev name and size and add model and serial
 			for ((i=1;i<${#G_WHIP_MENU_ARRAY[@]};i+=2)); do G_WHIP_MENU_ARRAY[$i]=": $(lsblk -drno SIZE,MODEL,SERIAL "${G_WHIP_MENU_ARRAY[$i-1]}")"; done
 
+			# Add this hosts root filesystem
+			G_WHIP_MENU_ARRAY+=("$HOST_ROOT_DRIVE" ": $(lsblk -drno SIZE "$HOST_ROOT_DRIVE") This host's root drive!")
+
 			G_WHIP_DEFAULT_ITEM=$FP_SOURCE
 			G_WHIP_MENU 'Please select the drive you wish to create the image from:
 \nNB: All mounted partitions of the selected drive will be unmounted.' || return 0
 			FP_SOURCE=$G_WHIP_RETURNED_VALUE
 			FP_ROOT_DEV=
 
+			if [[ $FP_SOURCE == "$HOST_ROOT_DRIVE" ]]
+			then
+				G_WHIP_BUTTON_CANCEL_TEXT='Abort'
+				G_WHIP_YESNO '[WARNING] You selected this host'\''s root drive as source!
+\nThis is generally possible, but the root (and boot) filesystem(s) needs to be remounted read-only (R/O) to assure a consistent image. This may fail if certain system services are active, we hence stop services like during DietPi updates, dietpi-software installs and dietpi-backup runs. If remounting the root filesystem still fails, image generation is aborted.
+\nAdditionally, the image needs to be stored to an external drive. If you did not navigate to one before running this script, you are asked to select one.
+\nDo you want to continue?' || exit 0
+				/boot/dietpi/dietpi-services stop
+				[[ $(findmnt -M /boot) ]] && until mount -o remount,ro /boot
+				do
+					G_WHIP_BUTTON_OK_TEXT='Retry'
+					G_WHIP_BUTTON_CANCEL_TEXT='Abort'
+					G_WHIP_YESNO '[WARNING] Mounting the /boot filesystem R/O failed
+\nThis may be due to running system services or for other reasons. Please assure that no process is accessing the /boot filesystem in writing mode, then retry.' || Error_Exit 'Mounting the boot filesystem R/O failed'
+				done
+				until mount -o remount,ro /
+				do
+					G_WHIP_BUTTON_OK_TEXT='Retry'
+					G_WHIP_BUTTON_CANCEL_TEXT='Abort'
+					G_WHIP_YESNO '[WARNING] Mounting the root filesystem R/O failed
+\nThis may be due to running system services or for other reasons. Please assure that no process is accessing the root filesystem in writing mode, then retry.' || Error_Exit 'Mounting the root filesystem R/O failed'
+				done
+				until [[ $(findmnt -no TARGET -T .) != '/' ]]
+				do
+					/boot/dietpi/dietpi-drive_manager 1 || Error_Exit 'No target drive selected'
+					read -r FP_ORIGIN < /tmp/dietpi-drive_manager_selmnt
+					[[ $FP_ORIGIN == '/' ]] || break
+					G_WHIP_MSG '[FAILED] You cannot select this host'\''s root filesystem to store its own image. Please select a different mount'
+					G_EXEC cd "$FP_ORIGIN"
+					rm /tmp/dietpi-drive_manager_selmnt
+				done
+			fi
+
 			G_DIETPI-NOTIFY 2 "Unmounting all filesystems below selected $FP_SOURCE ..."
 			local mountpoint
 			for i in "$FP_SOURCE"?*
 			do
 				mountpoint=$(findmnt -no TARGET "$i")
-				[[ $mountpoint ]] && G_EXEC umount -R "$mountpoint"
+				[[ $mountpoint && $mountpoint != '/' && $mountpoint != '/boot' ]] && G_EXEC umount -R "$mountpoint"
 			done
 		else
 			# Open DietPi-Explorer for image file selection
 			/boot/dietpi/dietpi-explorer 1 || { G_DIETPI-NOTIFY 1 'No image file selected, aborting...'; read -rp 'Press any key to return to menu...'; return 1; }
-			FP_SOURCE_IMG=$(</tmp/.dietpi-explorer_selected_location)
+			read -r FP_SOURCE_IMG < /tmp/.dietpi-explorer_selected_location
 			FP_ROOT_DEV=
 			rm /tmp/.dietpi-explorer_selected_location
 			[[ -f $FP_SOURCE_IMG ]] || { G_DIETPI-NOTIFY 1 "Selected image file ($FP_SOURCE_IMG) does not exist, aborting..."; read -rp 'Press any key to return to menu...'; return 1; }
@@ -360,32 +397,35 @@
 
 		Run_fsck
 
-		# Remount image for any required edits
-		G_EXEC mkdir "$FP_MNT_TMP"
-		G_EXEC mount "$FP_ROOT_DEV" "$FP_MNT_TMP"
-		# - Remove bash history and DHCP leases, which are stored on shutdown, hence cannot be removed via DietPi-PREP
-		G_EXEC rm -f "$FP_MNT_TMP/"{root,home/*}/.bash_history "$FP_MNT_TMP/var/lib/dhcp/"*.leases
-		# - Enable DietPi-FS_partition_resize to have both, old and new image being resized on next boot automatically
-		if [[ $SKIP_FIRSTBOOT_RESIZE != 1 && -f $FP_MNT_TMP/etc/systemd/system/dietpi-fs_partition_resize.service && ! -L $FP_MNT_TMP/etc/systemd/system/local-fs.target.wants/dietpi-fs_partition_resize.service ]]
+		# Remount image for any required edits, if not host rootfs
+		if [[ $FP_ORIGIN != "$HOST_ROOT_DRIVE" ]]
 		then
-			[[ -d $FP_MNT_TMP/etc/systemd/system/local-fs.target.wants ]] || G_EXEC mkdir "$FP_MNT_TMP/etc/systemd/system/local-fs.target.wants"
-			G_EXEC ln -s {'/etc/systemd/system',"$FP_MNT_TMP/etc/systemd/system/local-fs.target.wants"}'/dietpi-fs_partition_resize.service'
-		fi
-		if [[ $MOUNT_IT == 'On' ]] && G_WHIP_MSG "The ${SOURCE_TYPE,,} has been mounted to allow you reviewing or editing its content:
+			G_EXEC mkdir "$FP_MNT_TMP"
+			G_EXEC mount "$FP_ROOT_DEV" "$FP_MNT_TMP"
+			# - Remove bash history and DHCP leases, which are stored on shutdown, hence cannot be removed via DietPi-PREP
+			G_EXEC rm -f "$FP_MNT_TMP/"{root,home/*}/.bash_history "$FP_MNT_TMP/var/lib/dhcp/"*.leases
+			# - Enable DietPi-FS_partition_resize to have both, old and new image being resized on next boot automatically
+			if [[ $SKIP_FIRSTBOOT_RESIZE != 1 && -f $FP_MNT_TMP/etc/systemd/system/dietpi-fs_partition_resize.service && ! -L $FP_MNT_TMP/etc/systemd/system/local-fs.target.wants/dietpi-fs_partition_resize.service ]]
+			then
+				[[ -d $FP_MNT_TMP/etc/systemd/system/local-fs.target.wants ]] || G_EXEC mkdir "$FP_MNT_TMP/etc/systemd/system/local-fs.target.wants"
+				G_EXEC ln -s {'/etc/systemd/system',"$FP_MNT_TMP/etc/systemd/system/local-fs.target.wants"}'/dietpi-fs_partition_resize.service'
+			fi
+			if [[ $MOUNT_IT == 'On' ]] && G_WHIP_MSG "The ${SOURCE_TYPE,,} has been mounted to allow you reviewing or editing its content:
 - $FP_ROOT_DEV > $FP_MNT_TMP
 \nAn interactive bash subshell will open.
 \nPlease use the \"exit\" command when you are finished, to return to $G_PROGRAM_NAME."
-		then
-			# Prevent dietpi-login call in subshell
-			local reallow_dietpi_login=1
-			[[ $G_DIETPI_LOGIN ]] && reallow_dietpi_login=0
-			export G_DIETPI_LOGIN=1
-			G_EXEC cd "$FP_MNT_TMP"
-			bash &> /dev/tty < /dev/tty
-			G_EXEC cd "$FP_ORIGIN"
-			(( $reallow_dietpi_login )) && unset -v G_DIETPI_LOGIN
+			then
+				# Prevent dietpi-login call in subshell
+				local reallow_dietpi_login=1
+				[[ $G_DIETPI_LOGIN ]] && reallow_dietpi_login=0
+				export G_DIETPI_LOGIN=1
+				G_EXEC cd "$FP_MNT_TMP"
+				bash &> /dev/tty < /dev/tty
+				G_EXEC cd "$FP_ORIGIN"
+				(( $reallow_dietpi_login )) && unset -v G_DIETPI_LOGIN
+			fi
+			Unmount_tmp
 		fi
-		Unmount_tmp
 		# Failsafe
 		G_EXEC partprobe "$FP_SOURCE"
 		G_EXEC partx -u "$FP_SOURCE"