#!/bin/sh

# (c) 2009-2022 SHIMADA Hirofumi
# based :
#  30-bootmenu.sh (c) 2007 Paul Sokolovsky
#  initramfs-framework Copyright (C) 2011 O.S. Systems Software LTDA.
# Licensed on MIT

# scan_mount_sqs : 
# - Mount main squashfs image, and overlayfs
# - Copy firmware, cocon.cnf to 
# - Mount iso file if squshfs image is inside of iso image
#
# global val :
# - bootparam_sqsfile : squashfs file for mount
# - bootparam_fromiso : iso image for mount (if have fromiso= parameter)
# - bootparam_copytoram : if have this parameeter, copy squashfs image to ramdisk.
# - cocon_main_ready : if mount squashfs correctly, place new root path.

CFGLOC="/.cfg"
COPYTORAMLOC="/.cfg/copytoram"
FIRMWARELOC="/.cfg/firmware"
ISOLOC="/.iso"
MOUNTLOC="/rootfs"
NEWLOC="/.newroot"
OLDLOC="/.oldroot"
SCANLOC="/.cfg/scan"
UNIONLOC="/.cfg/union"
OVERLAY_UPPERLOC="/.cfg/overlay_upper"
OVERLAY_WORKLOC="/.cfg/overlay_work"
PARAM_FILE="$CFGLOC/.bootparam"
MAIN_READY_FILE="$CFGLOC/.main_image_ready"

# mount_tmp_area : allocation randisk area
mount_tmp_area()
{
	mount -t tmpfs none $CFGLOC
	mkdir -p $COPYTORAMLOC
	mkdir -p $FIRMWARELOC
	mkdir -p $UNIONLOC
	mkdir -p $OVERLAY_UPPERLOC
	mkdir -p $OVERLAY_WORKLOC
	mkdir -p $SCANLOC
}

# get_partition_type : get patition type
get_partition_type()
{
	fstype=$(blkid /dev/$1)
	fstype=$(expr "$fstype" : '.*TYPE="\([A-Za-z0-9]*\)".*')
}

# scan_copy_firmware : if found firewall directory, copy to ramdisk
# $1 --- scan directory
# $2 --- device name
scan_copy_firmware()
{
	if test -d "$1/coconfrm" -a ! -e "$CFGLOC/.firmware-loaded.$2";
	then
		# TODO : filter firmware files

		echo " -> copy coconfrm"
		cp -R $1/coconfrm/* $FIRMWARELOC

		# to avoid copy firmware twice, do make state file.
		touch "$CFGLOC/.firmware-loaded.$2"
	fi
}

# $1 --- device name
# $2 --- fstype
copytoram_eject()
{
	# If not set noeject flag, it is safe to eject media.
	if test -n "$bootparam_noeject" -a "$2" = "iso9660";
	then
		echo "--- Eject CD device. ($1) ---"
		fuser -k /dev/$1
		sync
		sync
		sync
		eject -s /dev/$1
		eject /dev/$1

		echo "================================================================"	
		echo "NOTE: Copy-to-RAM mode. It is safe to Eject or Unplug boot media."
		echo "================================================================"	
	fi
}

# copy-to-ram
# $1 ---- squashfs filename
# return : 0 --- ready, 1 --- failed or disabled
copytoram()
{
	# check size : disable for low memory
	sqs_size=$( stat -c %s "$1" )
	memory_free_size_kb=$( free | grep Mem | awk '{print $4}' )
	memory_free_size=$( expr $memory_free_size_kb "*" 1024 )
	padding=33554432  # 32MB = 33554432 byte
	check_size=$( expr $memory_free_size "-" $sqs_size "-" $padding )
	
	echo "free memory size : $memory_free_size byte"
	echo "squashfs size : $sqs_size byte"
	echo "calc size : $check_size"

	if test $check_size -lt 0;
	then
		# It seems do not Copy-on-RAM on this machine.

		echo "==============================================================="
		echo "WARNING: Copy-to-RAM is disbled because memory size is low."
		echo "         DO NOT EJECT OR UNPLUG boot media."
		echo "==============================================================="
		sleep 5

		echo "cocon_copytoram_blocked=1" >> $PARAM_FILE
		return 1
	fi

	# copy squashfs file to ramdisk
	sqs_name="$( basename $1 )"
	echo "--- copy squashfs image to ramdisk ($sqs_name) ---"

	# mount ramdisk for copytoram
	mount -t tmpfs -o size=$( expr $sqs_size "+" $padding ) tmpfs $COPYTORAMLOC
	if test $? -lt 0;
	then
		# It seems do not Copy-on-RAM on this machine.

		echo "==============================================================="
		echo "WARNING: Copy-to-RAM is disbled because memory size is low."
		echo "         DO NOT EJECT OR UNPLUG boot media."
		echo "==============================================================="
		sleep 5

		echo "cocon_copytoram_blocked=2" >> $PARAM_FILE
		return 1
	fi

	cp "$1" "$COPYTORAMLOC/$sqs_name"
	sync
	sync
	sync

	# verify size
	sqs_size_verify=$( stat -c %s $COPYTORAMLOC/$sqs_name )

	if test $sqs_size -ne $sqs_size_verify;
	then
		# Seems copied squashfs file is incoorect.
		rm -f $COPYTORAMLOC/$sqs_name
		echo "==============================================================="
		echo "WARNING: Copy-to-RAM is disbled because copy was failed."
		echo "         DO NOT EJECT OR UNPLUG boot media."
		echo "==============================================================="
		sleep 5

		echo "cocon_copytoram_blocked=3" >> $PARAM_FILE
		return 1
	fi

	echo "cocon_copytoram_enabled=\"$COPYTORAMLOC/$sqs_name\"" >> $PARAM_FILE
	return 0
}

# scan_copy_cfg : if found cocon.cnf and related files. copy to ramdisk
# $1 --- scan directory
# $2 --- device
scan_copy_cfg()
{
	# if contain cocon.cnf/coconnm/coconfrm on media, copy to memory.
	if test -f "$1/cocon.cnf" -a ! -e "$CFGLOC/.cfg-loaded.$2";
	then
		echo " -> copy cocon.cnf.$( stat -c %Y $1/cocon.cnf )"
		# copy cocon.cnf with modify timestamp in filename
		cp $1/cocon.cnf $CFGLOC/cocon.cnf.$( stat -c %Y $1/cocon.cnf )

		# to avoid copy config twice, make state file.
		touch "$CFGLOC/.cfg-loaded.$2"
	fi
}

# scan main squashfs inside .iso file
# $1 --- filesystem type of $SCANLOC
# $2 --- scan device
# return : 0 --- mounted and ready , 1 --- not found or failed
scan_mount_iso()
{
	if test -n "$bootparam_fromiso" -a -n "$bootparam_sqsfile";
	then
		# loopmount ISO image
        	mount -o loop "$SCANLOC/$bootparam_fromiso" $ISOLOC >/dev/null 2>&1

		# check if have squashfs file on ISO
	        if test $? -eq 0 -a -e "$ISOLOC/$bootparam_sqsfile" ;
		then
			echo "scan_mount_iso : found $bootparam_sqsfile on $bootparam_fromiso."

			if test "$1" = "iso9660";
			then
				echo "cocon_need_eject_on_clean=$dev" >> $PARAM_FILE
			fi
			echo "cocon_fromiso_enabled=$ISOLOC" >> $PARAM_FILE

			return 0
		else
			echo "... is not contain $bootparam_sqsfile on $bootparam_fromiso. skip."
			umount -lf $ISOLOC
			return 1
		fi	
	else
		echo "... not found $bootparam_fromiso. skip."
		return 1
	fi
}

# Mount squashfs main image and overlayfs
# $1 --- filename to mount
# $2 --- filesystem type of $1
# $3 --- scan device
# return : 0 --- mounted and ready, 1 --- not found or failed
scan_mount_sqs()
{
	# just search squashfs file
	if test ! -e "$1";
	then
		echo "... $1 is not found."
		return 2
	fi

	sqs_file="$1"

	# TODO : checksum

	if test -n "$bootparam_copytoram";
	then
		# Copy-to-RAM mode.
		copytoram $sqs_file

		if test $? -eq 0;
		then
			copytoram_enabled=1
			sqs_file="$COPYTORAMLOC/$sqs_name"
		fi
	fi

	echo "--- mount squashfs ---"
	mount -o loop -t squashfs "$sqs_file" $UNIONLOC >/dev/null 2>&1
	if test $? -ge 1;
	then
		echo "WARNING : failed to mount squashfs fiile. skip $1."
		umount -lf $UNIONLOC
		rm -f $COPYTORAMLOC/$sqs_name
		umount -lf $ISOLOC
		return 1
	fi

	echo "--- overlay ---"
	mount -t overlay -o lowerdir=$UNIONLOC,upperdir=$OVERLAY_UPPERLOC,workdir=$OVERLAY_WORKLOC overlay $NEWLOC >/dev/null 2>&1
	if test $? -ge 1;
	then
		echo "WARNING : failed to mount unionfs."
		umount -lf $NEWLOC
		umount -lf $UNIONLOC
		rm -f $COPYTORAMLOC/$sqs_name
		umount -lf $ISOLOC
		return 1
	fi

	echo "Main image is ready."
	touch $MAIN_READY_FILE

	if test "$2" = "iso9660" -a -z "$copytoram_enabled";
	then
		echo "cocon_need_eject_on_clean=$3" >> $PARAM_FILE
	fi

	echo "cocon_main_ready=$NEWLOC" >> $PARAM_FILE
	echo "cocon_boot_fs=$2" >> $PARAM_FILE

	return 0
}


# Scan disk devices (main loop)
scan_device()
{
	# set squashfs file to mount if unsetted
	if test -z "$bootparam_sqsfile";
	then
		# Default filename
		bootparam_sqsfile="crusoe.sqs"
		export bootparam_sqsfile
	fi

	echo "Scan start"

	# Scan all available device/partitions
	cat /proc/diskstats | sort -r | while read maj min dev other ;
	do
		if test -z "$maj" -o "$maj" = "major" ;
		then
			continue
		fi

		echo "Found $dev"

		get_partition_type "$dev"
		case "$fstype" in
			"iso9660" | "vfat" | "ntfs" | "ext3" | "ext4" | "exfat" )
				echo "Scanning $dev ($fstype)"

				mount -o ro /dev/$dev $SCANLOC >/dev/null 2>&1
				if test $? -eq 0;
				then
					echo " Mounted /dev/$dev , start search files..."

					# scan firmware
					scan_copy_firmware $SCANLOC $dev

					# scan cocon config file
					scan_copy_cfg $SCANLOC $dev

					# If main image is not ready, search and mount it.
					if test ! -e $MAIN_READY_FILE;
					then

						if test -n "$bootparam_fromiso" -a -e "$SCANLOC/$bootparam_fromiso";
						then
							echo "fromiso mode is enabled"
							# scan ISO image
							scan_mount_iso $fstype $dev
							if test "$?" -eq 1;
							then
								# ISO image was not found.
								# Skip this mount point.
								continue
							fi
						else
							# Normal boot : set squashfs file
							sqs_scanfile="$SCANLOC/$bootparam_sqsfile"
						fi

						# Scan Main squashfs image
						echo "Scan squashfs file. Device=$dev Scanpath=$sqs_scanfile Fs=$fstype"
						scan_mount_sqs "$sqs_scanfile" $fstype $dev
						if test "$?" -eq 1;
						then
							# Main image is not found, or failed to mount. skip this mount point.
							continue
						fi
					fi  # ! -e $MAIN_READY_FILE

					grep "cocon_copytoram_enabled" $PARAM_FILE >/dev/null 2>&1
					is_coptyoram_ready="$?"

					if test ! -e "$MAIN_READY_FILE";
					then
						# unmount device, except main image is ready
						umount -f $SCANLOC >/dev/null 2>&1

						if test -n "$is_coptyoram_ready";
						then
							# eject media if copytoram is ready
							copytoram_eject $dev $fstype
						fi
					fi

				else  # return status of mount -o ro /dev/$dev $SCANLOC
					echo "Couldn't mount /dev/$dev ($fstype). skip."
				fi
				;;

			*)
				echo "Skip $dev ($fstype)"
				;;
		esac   # case of $fstype
	done

	echo "Scan complete"
}

mountsqs_enabled() {
	return 0
}

mountsqs_run() {
	mount_tmp_area

	# set squashfs file to mount if unsetted
	if test -z "$bootparam_sqsfile";
	then
		# Default filename
		bootparam_sqsfile="crusoe.sqs"
		export bootparam_sqsfile
	fi

	i=1
	while test $i -le 20;
	do
		echo "--- scanning root media ($i)---"
		scan_device

		if test -e "$MAIN_READY_FILE";
		then
			echo "Main image is ready ($i). Move to next script."
			return 0
		fi

		echo "failed scan ($i)"

		i=$(expr $i + 1)
		sleep 5
	done

	echo "ERROR : main squashfs image is not found."
	echo "DEBUG : Fallback shell. "
	/bin/sh
}
