#!/bin/sh
# Tazusb - SliTaz LiveUSB
#
# Tazusb is an utility to generate, configure and manipulate SliTaz LiveUSB
# bootable media and/or USB /home partition, such as flash keys, SD card or 
# USB harddisk. Authors : see AUTHORS
#
VERSION=3.0

COMMAND=$1
TARGET_ROOT=/media/flash
DRIVE_NAME=`cat /proc/sys/dev/cdrom/info | grep "drive name" | cut -f 3`
CDROM=/dev/$DRIVE_NAME
LOG=/tmp/$(basename $0).log

#
# Tazusb functions
#

# Print the usage.
usage ()
{
	echo -e "\nSliTaz Live USB - Version: $VERSION\n
\033[1mUsage: \033[0m `basename $0` [command] [compression|device]
\033[1mCommands: \033[0m\n
  usage     	Print this short usage.
  writefs 	Write the current filesystem to rootfs.gz.
  		tazSupported compression: lzma. gzip, none.
  format	Format and label device with ext3, ext2 or fat32 filesystem 
                (for LiveUSB or /home). Default is ext3.
  gen-liveusb 	Generate a bootable LiveUSB using files from the LiveCD.
  gen-swap      Create or recreate and activate additional swap memory.
  gen-iso2usb   Generate a bootable LiveUSB using files from ISO file.
  clean         Remove old backup filesystems to free disk space.\n"
}

# Status function.
status()
{
	local CHECK=$?
	echo -en "\\033[70G[ "
	if [ $CHECK = 0 ]; then
		echo -en "\\033[1;33mOK"
	else
		echo -en "\\033[1;31mFailed"
	fi
	echo -e "\\033[0;39m ]"
}

# Exit if user is not root.
check_root()
{
	if test $(id -u) != 0 ; then
	   echo -e "\nThis program requires being run as root.\n"
	   exit 0
	fi
}

# Display a list of available partitions.
fdisk_list()
{
	echo "----"
	fdisk -l | grep ^/dev/sd*
	echo "----"
}

# We need a USB media to install.
ask_for_device()
{
	echo -n "\
Please specify the target USB device to $COMMAND. You can type 'list' to 
get a list of devices, type 'exit' or give an empty value to exit.

Device to use : "; read anser
	while [ "$anser" == "list" ]; do
		fdisk_list
		echo -n "Device to use : "; read anser
	done
	if [ "$anser" = "" -o "$anser" = "exit" ]; then
		echo -e "\nNo specified device or exit.\n"
		exit 0
	else
		DEVICE=$anser
	fi
}

# Verify a device exists before format or install
check_for_device()
{
	IFDEV=`fdisk -l | grep $DEVICE`
	if [ -z "$IFDEV" ]; then
		echo -e "\nUnable to find device: $DEVICE\n"
		exit 0
	fi
}

# gets the UUID and filesystem type
get_part_info()
{
	UUID=`blkid -s UUID -o value $DEVICE`
	FSTYPE=`blkid -s TYPE -o value $DEVICE`
}

# Format target device and label partition.
mkfs_ext3()
{
	echo -n "Please specify a label for the partition (TazUSB): "
	read label
	
	if [ -z $label ]; then
		label=TazUSB
	fi
	
	echo "Label  : $label"
	echo "Mkfs   : mkfs.ext3 -L \"$label\" $DEVICE"
	echo "" && sleep 2
	mkfs.ext3 -L "$label" $DEVICE
	
}

# Get label for device
get_label()
{
	echo -n "Please specify a label for the partition (TazUSB): "
	read label
	
	if [ -z $label ]; then
		label=TazUSB
	fi
}

# Get fs type. Supported fs are ext3, ext2, fat32
get_fs_type()
{
	echo -n "Please specify a filesystem type ext2, ext3 or fat32 (ext3): "
	read  fs_type
	
	if [ -z $fs_type ]; then
		fs_type=ext3
	fi
}

# We can chose the filesystem type.
ask_for_fs_type()
{
	echo -n "\
Please specify the filesystem type to $COMMAND. 
Available formats are ext3(default), ext2 or fat32. 
Press enter to keep the default value.

File system type : "; read anser
	if [ "$anser" = "" ]; then
		FS_TYPE=ext3
	else
		FS_TYPE=$anser
	fi
}

# Format target device in ext3, ext2 or fat32.
# Usage: make_fs ext2|fat32|ext3 (default)
make_fs()
{
	local answer
	
	FS=$(echo $fs_type | tr [A-Z] [a-z])
	case "$FS" in
		ext3|ext2)
			echo -e "\nProcessing..."
			echo "Label  : $label"
			echo "Mkfs   : mkfs.$FS -L \"$label\" $DEVICE"
			echo "" && sleep 2
			mkfs.$@ -L "$label" $DEVICE > $LOG 2>&1
			;;
		[Ff]at32)
			if [ -x /sbin/mkdosfs ];then
				echo -e "\nProcessing..."
				echo "Label  : $label"
				echo "Mkfs   : mkdosfs -F 32 -n \"$label\" $DEVICE"
				echo "" && sleep 2
				mkdosfs -F 32 -n "$label" $DEVICE  
			else
				echo -ne "Can't find mkdosfs tool.\nWould you like to install dosfstools from repository [y/N] ? "; read answer
				case $answer in
					y|Y)
						yes | tazpkg get-install dosfstools  && make_fs fat32;;
					*)
						exit 1 ;;
				esac
			fi
			;;
		*)
			echo "Sorry. Filesystem $FS is not supported."
			exit 
	esac
}

# Mount an existing USB device.
unmount_target_usb()
{
	# If mount point is in use, unmount
	if mount | grep -q $TARGET_ROOT; then
		umount $TARGET_ROOT
	fi
	
	# Device could be mounted elsewhere, so unmount
	if mount | grep -q $DEVICE; then
		echo "Unmounting USB target device..."
		umount $DEVICE
	fi
}

# Mount an existing USB device.
mount_target_usb()
{
	echo "Mounting USB target device..."	
	mkdir -p $TARGET_ROOT
	mount $DEVICE $TARGET_ROOT 2>/dev/null
}

# Mount SliTaz LiveCD to get needed files.
mount_cdrom()
{
	echo "Mounting cdrom device..."
	
	if mount | grep /media/cdrom; then
		umount /media/cdrom
	fi
	
	mkdir -p /media/cdrom
	mount -r -t iso9660 $CDROM /media/cdrom 2>/dev/null
	
	if [ ! -f /media/cdrom/boot/rootfs.gz -a \
	     ! -f /media/cdrom/boot/rootfs1.gz ]; then
		echo -e "\nUnable to mount cdrom or to find a filesystem on it (rootfs.gz).\n"
		exit 0
	fi
}

# Mount SliTaz ISO to get needed files.
mount_iso()
{
	if mount | grep /media/cdrom ; then
		umount /media/cdrom
	fi
	
	test -d /media/cdrom || mkdir -p /media/cdrom
	
	# Add support to other distros.
    if [ ! -f /etc/slitaz-version ]; then
         OPTIONS="-o loop"
    else
         OPTIONS=""
    fi
	
	echo  "Mounting `basename $ISO`..."
	mount $OPTIONS $ISO /media/cdrom 2>/dev/null
	
	if [ ! -f /media/cdrom/boot/rootfs.gz -a \
	     ! -f /media/cdrom/boot/rootfs1.gz ]; then
		echo -e "\nUnable to mount iso or to find a filesystem on it (rootfs.gz).\n"
		exit 0
	fi
}

# All needed files are in the boot directory of the cdrom.
copy_cdrom_files()
{
	echo -n "Copying needed files from cdrom..."
	mkdir -p $TARGET_ROOT/boot
	cp /media/cdrom/boot/bzImage $TARGET_ROOT/boot
	rem=0
	for i in $(ls /media/cdrom/boot/rootfs*.gz | sort -r); do
		[ $rem -ne 0 ] && 
			dd if=/dev/zero bs=1 count=$((4 - $rem)) 2> /dev/null
		cat $i
		rem=$(stat -c %s $i)
		rem=$(($rem % 4))
	done > $TARGET_ROOT/boot/rootfs.gz
	status
}

install_mbr()
{
	# MBR
	DISK=${DEVICE%[1-99]}
	if [ -f /usr/share/boot/mbr.bin ]; then
		echo -n "Installing a new MBR to: $DISK"
		cat /usr/share/boot/mbr.bin > $DISK
		status
	else
		# Skip MBR install (tazpkg get-install syslinux-extra ? and then cat)
		echo "No new MBR installed to: $DISK"
	fi
}

# ext/syslinux install
install_boot()
{
	# Decide if we're installing syslinux or extlinux
	if [ "$FSTYPE" = "vfat" ]; then
		ST=syslinux
		STC="syslinux --install -d /boot/syslinux/ $DEVICE"
		STE=cfg
	else
		ST=extlinux
		STC="extlinux --install $TARGET_ROOT/boot/$ST"
		STE=conf
	fi
	
	echo "Installing bootloader: $ST"
	mkdir -p $TARGET_ROOT/boot/$ST
	$STC
	
	# Use existing isolinux.cfg for extlinux.conf or syslinux.cfg
	cp /media/cdrom/boot/isolinux/isolinux.cfg $TARGET_ROOT/boot/$ST/$ST.$STE
	sed -i "s/isolinux.msg/$ST.msg/"	 $TARGET_ROOT/boot/$ST/$ST.$STE
	
	# Add home= to append in extlinux or syslinux.cfg
	sed -i -e "s/\(append.*\)/\1 home=$UUID/" $(grep -l append $TARGET_ROOT/boot/$ST/*)
	
	# Splash screen and help files.
	cp /media/cdrom/boot/isolinux/isolinux.msg $TARGET_ROOT/boot/$ST/$ST.msg
	sed -i s/'SliTaz GNU\/Linux'/'SliTaz GNU\/Linux LiveUSB'/ $TARGET_ROOT/boot/$ST/$ST.msg
	cp /media/cdrom/boot/isolinux/splash.lss $TARGET_ROOT/boot/$ST
	cp /media/cdrom/boot/isolinux/*.txt $TARGET_ROOT/boot/$ST
	cp /media/cdrom/boot/isolinux/*.cfg $TARGET_ROOT/boot/$ST
	cp /media/cdrom/boot/isolinux/*.kbd $TARGET_ROOT/boot/$ST
	cp /media/cdrom/boot/isolinux/*.c32 $TARGET_ROOT/boot/$ST
	
	# Modifing all cfg files.
	for cfg in $TARGET_ROOT/boot/$ST/*.cfg
	do
		sed -i s/isolinux.msg/$ST.msg/ $cfg
	done

	# Modifing include file if exists.
	if [ -f $TARGET_ROOT/boot/$ST/common.inc ]; then
		sed -i -e "s/isolinux.msg/$ST.msg/" $TARGET_ROOT/boot/$ST/common.inc
	fi

	# Un-meta-ize a multi-in-one flavor
	if grep -qs "label slitaz" $TARGET_ROOT/boot/$ST/common.cfg ; then
		sed -i "s/isolinux/$ST/;s/label slitaz/label multi/" 	 $TARGET_ROOT/boot/$ST/common.cfg
		sed -i 's/\(.*\), flavors.*/                 \1/' \
			$TARGET_ROOT/boot/$ST/$ST.msg
		for i in $TARGET_ROOT/boot/$ST/$ST.$STE \
			 $TARGET_ROOT/boot/$ST/??.$STE \
			 $TARGET_ROOT/boot/$ST/??_??.$STE; do
				sed '/label /!d;/label /{s/label .*/\nlabel slitaz/;NN;s/rootfs..gz.*gz /rootfs.gz /;N;q}' \
					< $i >> $i
		done
	fi
	
}

# Let user exit or reboot.
exit_or_reboot()
{
	echo "==============================================================================="
	echo ""
	echo -n "Do you want to exit Tazusb or reboot system (Exit/reboot) ? "
	read anser
	if [ "$anser" == "reboot" ]; then
		umount $TARGET_ROOT
		umount /media/cdrom
		reboot || reboot -f
	else
		umount $TARGET_ROOT
		umount /media/cdrom
		echo ""
		exit 0
	fi
}

set_bootable()
{
	# As the boot flag is toggable, need to check it before hand
	DISK=${DEVICE%[1-99]}
	ISSET=`fdisk -l $DISK | grep $DEVICE | grep "\*"`
	
	# If not set, set bootable
	if [ "$ISSET" == "" ]; then
		umount $TARGET_ROOT 2>/dev/null
		echo -n "Setting $DEVICE as bootable..."
		fdisk $DISK >/dev/null << EOF
a
1
w
EOF
		status
	fi
}

# Generate a virtual swap file in /home/swap. SliTaz boot scripts 
# will activate it, useful for low memory systems
gen_swap_file()
{
	echo ""
	echo -en "\033[1mGen swap\033[0m
===============================================================================
Generate a swap file in /home/swap that will be activated on each boot to have
more memory available (Empty value to exit).

Swap file in Mb : "
	read size
	if [ -z "$size" ]; then
		echo -e "\nEmpty value. Exiting...\n"
		exit 0
	fi
	cd /home
	# Sanity check
	if [ -f swap ]; then
		swapoff swap 2>/dev/null
	fi
	# DD to gen a virtual disk.
	dd if=/dev/zero of=swap bs=1M count=$size
	# Make swap filesystem.
	mkswap swap
	swapon swap
	echo ""
}

# Clean out old backups to save disk space
clean_usb() 
{
	echo ""
	echo -en "\033[1mClean\033[0m
===============================================================================
Remove old rootfs.gz.unixtimestamp backup filesystems to free up disk space.\n\n"
	# Locate and interactively remove old filesystems from /home directory
	for file in `find /home/boot/rootfs.gz.[0-9]*`
	do
		echo -n "Do you wish to remove: `basename $file` (Yes/no/exit) ? "
		read answer
		case $answer in
			e|E|"exit"|Exit)
				exit 0 ;;
			y|Y|yes|Yes)
				rm -f $file ;;
			*)
			echo -en "No filesystems selected, exiting...\n" ;;
		esac
	done
}

#
# Tazusb sequence
#

case $COMMAND in
	writefs)
		# Writefs to rootfs.gz
		check_root
		if [ -z $2 ]; then
			COMPRESSION=none
		else
			COMPRESSION=$2
		fi
		# Start info
		echo ""
		echo -e "\033[1mWrite filesystem\033[0m
===============================================================================
The command writefs will write all the current filesystem into a suitable cpio 
archive (rootfs.gz) usable on a bootable LiveUSB media. 

Archive compression: $COMPRESSION"
		echo ""
		
		# Clear out tazpkg cache
		rm /var/cache/tazpkg/* -r -f

		# Optionally remove sound card selection
		echo -n "Do you wish to remove the sound card selection (No/yes/exit) ? "
		read anser
		case $anser in
			e|E|"exit"|Exit)
				exit 0 ;;
			y|Y|yes|Yes)
				echo -n "Removing current sound card selection..."
				rm -f /var/lib/sound-card-driver
				rm -f /var/lib/alsa/asound.state ;;
			*)
				echo -n "Keeping current sound card selection..." ;;
		esac
		status
		# Optionally remove screen resolution
		echo -n "Do you wish to remove the screen resolution (No/yes/exit) ? "
		read anser
		case $anser in
			e|E|"exit"|Exit)
				exit 0 ;;
			y|Y|yes|Yes)
				echo -n "Removing current screen resolution..."
				rm -f /etc/X11/screen.conf ;;
			*)
				echo -n "Keeping current screen resolution..." ;;
		esac
		status
		
		# Create list of files
		find /bin /etc /init /sbin /var /dev /lib /root /usr >/tmp/list

		for dir in /home /proc /sys /tmp /mnt /media /media/cdrom /media/flash /media/usbdisk
		do
			echo $dir >>/tmp/list
		done

		# Generate initramfs with specified compression
		if [ "$COMPRESSION" = "lzma" ]; then
			echo -n "Creating rootfs.gz with lzma compression... "
			cat /tmp/list | cpio -o -H newc | lzma e -si -so > /rootfs.gz

		elif [ "$COMPRESSION" = "gzip" ]; then
			echo -n "Creating rootfs.gz with gzip compression... "
			cat /tmp/list | cpio -o -H newc | gzip -9 > /rootfs.gz

		else
			echo -n "Creating rootfs.gz without compression... "
			cat /tmp/list | cpio -o -H newc > /rootfs.gz
		fi

		# Get initramfs size
		size=`du -sh /rootfs.gz | cut -f 1`

		# If the bootable medium is where it should be, copy across
		if (test -e /home/boot/bzImage); then
			echo "Moving rootfs.gz to media. Remember to unmount for delayed writes!"

			# Move the old filesystem with the unix timestamp for reference
			if (test -e /home/boot/previous.gz); then
				mv /home/boot/previous.gz /home/boot/rootfs.gz.$(date +%s)
			fi
			
			mv /home/boot/rootfs.gz /home/boot/previous.gz
			mv /rootfs.gz /home/boot/.
		else
			echo "rootfs.gz is located in /"
		fi

		echo "==============================================================================="
		echo "Root filesystem size: $size"
		echo ""
		echo -en "----\nENTER to continue..."; read i
		;;
	format)
		# Format a partition.
		check_root
		echo ""
		echo -e "\033[1mFormat a device\033[0m"
		echo "==============================================================================="
		DEVICE=$2
		label=$3
		fs_type=$4
		if [ -z $DEVICE ]; then
			ask_for_device
			check_for_device
		else
			echo "Device : $DEVICE"
		fi
		test -z $fs_type &&	get_fs_type
		get_label
		unmount_target_usb 
		make_fs "$fs_type"
		# mkfs_ext3
		echo "==============================================================================="
		echo "Device $label ($DEVICE) is ready to use as LiveUSB and/or /home partition."
		echo ""
		;;
	gen-liveusb)
		# Generate a LiveUSB media using files from a LiveCD.
		check_root
		echo ""
		echo -e "\033[1mGen a LiveUSB media\033[0m"
		echo "==============================================================================="
		DEVICE=$2
		if [ -z $DEVICE ]; then
			ask_for_device
		fi
		
		check_for_device
		mount_cdrom
		get_part_info
		unmount_target_usb
		install_mbr
		set_bootable
		mount_target_usb
		copy_cdrom_files
		install_boot
		exit_or_reboot
		;;
	gen-swap)
		check_root
		gen_swap_file
		;;
	gen-iso2usb|iso2usb)
		check_root

		# Check if file exists
		ISO=$2
		if [ -z $ISO ] || [ ! -f $ISO ]; then
			echo -e "\nPlease specify a valid filename on the command line.\n"
			exit 1
		fi
		echo ""
		echo -e "\033[1mCopy ISO file to SliTaz LiveUSB media\033[0m"
		echo "==============================================================================="
		echo ""
		DEVICE=$3
		if [ -z $DEVICE ]; then
			ask_for_device
		fi
		check_for_device
		mount_iso
		get_part_info
		unmount_target_usb
		install_mbr
		set_bootable
		mount_target_usb
		copy_cdrom_files
		install_boot
		umount /media/cdrom
		losetup -d /dev/loop0
		exit_or_reboot
		;;	
	clean)
		check_root
		clean_usb
		;;	
	usage|*)
		# Display usage by default.
		usage
		;;
esac

exit 0
