#!/bin/sh
#
# Tazusb - SliTaz LiveUSB utility to generate, configure and manipulate
# SliTaz LiveUSB bootable media and/or USB /home partition, such as
# flash keys, SD card or USB harddisk.
#
# Copyright (C) 2014 SliTaz GNU/Linux - GNU gpl v2
#
# Authors: see AUTHORS file
#

VERSION=209

. /lib/libtaz.sh

# i18n
export TEXTDOMAIN='tazusb'

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


#
# Tazusb functions
#


# Print the usage.

usage () {
	_ 'SliTaz Live USB - Version: %s' "$VERSION"
	newline
	boldify "$(_n 'Usage:')"
	echo -n '  '; _ '%s [command] [compression|device]' "$(basename $0)"
	boldify "$(_n 'Commands:')"
	optlist "\
usage		$(_ 'Print this short usage.')
writefs		$(_ 'Write the current filesystem to rootfs.gz. Supported compression: lzma, gzip, none.')
format		$(_ 'Format and label device with ext4, 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.')
"
	newline
}


# Display a list of available partitions.

fdisk_list() {
	separator '-'
	fdisk -l | grep ^/dev/sd*
	separator '-'
}


# We need a USB media to install.

ask_for_device() {
	longline "$(_ "Please specify the target USB device to %s. You can type \
'list' to get a list of devices, type 'exit' or give an empty value to exit." \
"$COMMAND")"
	newline
	_n 'Device to use: '; read answer

	while [ "$answer" = 'list' ]; do
		fdisk_list
		_n 'Device to use: '; read answer
	done

	if [ -z "$answer" -o "$answer" = 'exit' ]; then
		newline
		_ 'No specified device or exit.'
		exit 0
	else
		DEVICE="$answer"
	fi
}


# Verify a device exists before format or install

check_for_device() {
	if [ -z "$(blkid | grep "^$DEVICE:")" ]; then
		newline
		_ 'Unable to find device %s' "$DEVICE"
		exit 0
	fi
}


# gets the UUID and filesystem type

get_part_info() {
	UUID="$(blkid $DEVICE | sed 's|.* UUID="||;s|".*||')"
	FSTYPE="$(blkid $DEVICE | sed 's|.* TYPE="||;s|".*||')"
}


# Get label for device

get_label() {
	_n 'Please specify a label for the partition (TazUSB): '
	read label
	[ -z "$label" ] && label='TazUSB'
}


# Get fs type. Supported fs are ext4, ext3, ext2, fat32

get_fs_type() {
	_n 'Please specify a filesystem type ext2, ext3, ext4 or fat32 (ext3): '
	read fs_type
	case "$fs_type" in
	ext2|ext4|fat32);;
	*) fs_type='ext3'
	esac
}


# We can chose the filesystem type.

ask_for_fs_type() {
	_ 'Please specify the filesystem type to %s.' "$COMMAND"
	_ 'Available formats are ext4, ext3(default), ext2 or fat32.'
	_ 'Press enter to keep the default value.'
	newline
	_n 'File system type: '; read FS_TYPE
	case "$FS_TYPE" in
	ext2|ext4|fat32);;
	*) FS_TYPE='ext3'
	esac
}


check_partition_table() {
	DISK=${DEVICE%[1-99]}
	fdisk -l $DISK | grep -q ^$DEVICE: || fdisk $DISK <<EOT
n
p
1


t
$1
w
EOT
}

 
# Format target device in ext4, ext3, ext2 or fat32.
# Usage: make_fs ext2|ext4|fat32|ext3 (default)

make_fs() {
	local answer

	FS=$(echo $fs_type | tr '[A-Z]' '[a-z]')
	case "$FS" in
		ext4|ext3|ext2)
			check_partition_table 83
			newline; _ 'Processing...'
			_ 'Label: %s' "$label"
			echo "Mkfs: mkfs.$FS -L \"$label\" $DEVICE"
			newline; sleep 2
			mkfs.$FS -L "$label" $DEVICE > $LOG 2>&1
			;;
		fat32)
			if [ -x '/sbin/mkdosfs' ];then
				check_partition_table C
				newline; _ 'Processing...'
				_ 'Label: %s' "$label"
				echo "Mkfs: mkdosfs -F 32 -n \"$label\" $DEVICE"
				newline; sleep 2
				mkdosfs -F 32 -n "$label" $DEVICE
			else
				_ "Can't find %s tool." 'mkdosfs'
				_n 'Would you like to install %s from repository [y/N]? ' 'dosfstools'
				read answer
				case "$answer" in
					y|Y)
						yes | tazpkg get-install dosfstools && make_fs fat32;;
					*)
						exit 1 ;;
				esac
			fi
			;;
		*)
			_ 'Sorry. Filesystem %s is not supported.' "$FS"
			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
		_ 'Unmounting USB target device...'
		umount $DEVICE
	fi
}


# Mount an existing USB device.

mount_target_usb() {
	_ 'Mounting USB target device...'
	mkdir -p $TARGET_ROOT
	mount $DEVICE $TARGET_ROOT 2>/dev/null
}


# Mount SliTaz LiveCD to get needed files.

mount_cdrom() {
	_ 'Mounting CD-ROM 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
		newline
		longline "$(_ 'Unable to mount CD-ROM or to find a filesystem on it (%s).' 'rootfs.gz')"
		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

	_ 'Mounting %s...' "$(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
		longline "$(_ 'Unable to mount ISO or to find a filesystem on it (%s).' 'rootfs.gz')"
		exit 0
	fi
}


# All needed files are in the boot directory of the CD-ROM.

copy_cdrom_files() {
	_n 'Copying needed files from CD-ROM...'
	mkdir -p                         $TARGET_ROOT/boot
	cp /media/cdrom/boot/bzImage*    $TARGET_ROOT/boot
	cp /media/cdrom/boot/rootfs*.gz* $TARGET_ROOT/boot
	cp /media/cdrom/boot/memtest*    $TARGET_ROOT/boot 2>/dev/null
	cp /media/cdrom/boot/*xe*        $TARGET_ROOT/boot 2>/dev/null
	status
}


install_mbr() {
	# MBR
	DISK=${DEVICE%[1-99]}
	if [ -f /usr/share/boot/mbr.bin ]; then
		_n 'Installing a new MBR to %s' "$DISK"
		cat /usr/share/boot/mbr.bin > $DISK
		status
	else
		# Skip MBR install (tazpkg get-install syslinux-extra ? and then cat)
		_ 'No new MBR installed to %s' "$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

	_ 'Installing bootloader: %s' "$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

	# Update DVD autoinstall
	sed -i "s/LABEL=packages-[^,]*/UUID=$UUID/g" $(grep -il append $TARGET_ROOT/boot/$ST/*)

	# Add home=$UUID to kernel line in extlinux or syslinux.cfg
	sed -i -e "s/\(root=.*\)/\1 home=$UUID/" -e '/LABEL md5sum/,/append md5sum/d' \
		$(grep -il append $TARGET_ROOT/boot/$ST/*)

	# Splash screen and help files.
	cp /media/cdrom/boot/isolinux/splash.*	$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
	cp /media/cdrom/boot/isolinux/opts.*	$TARGET_ROOT/boot/$ST
	cp /media/cdrom/boot/isolinux/help.*	$TARGET_ROOT/boot/$ST
	sed -i -e s/'SliTaz GNU\/Linux'/'SliTaz GNU\/Linux LiveUSB'/ \
		-e "s/isolinux/$ST/" $TARGET_ROOT/boot/$ST/$ST.$STE
}


# Let user exit or reboot.

exit_or_reboot() {
	separator
	newline
	longline "$(_n 'Do you want to exit Tazusb or reboot system (Exit/reboot)? ')"
	read answer
	if [ "$answer" = 'reboot' ]; then
		unmount_target_usb
		reboot || reboot -f
	else
		unmount_target_usb
		newline; 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 [ -z "$ISSET" ]; then
		umount $TARGET_ROOT 2>/dev/null
		_n 'Setting %s as bootable...' "$DEVICE"
		fdisk $DISK >/dev/null << EOF
$(fdisk -l $DISK | grep $DISK  | sed "/\\*/!ds#$DISK##;s# .*##;s#.*#&\na#")
a
${DEVICE#$DISK}
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() {
	title "$(_n 'Gen swap')"

	longline "$(_ "Generate a swap file in %s that will be activated on each \
boot to have more memory available (empty value to exit)." '/home/swap')"
	newline
	_n 'Swap file in MB: '
	read size
	if [ -z "$size" ]; then
		_ 'Empty value. Exiting...'
		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
	newline
}


# Clean out old backups to save disk space

clean_usb() {
	title "$(_n 'Clean')"

	longline "$(_n 'Remove old %s backup filesystems to free up disk space.' \
'rootfs.gz.unixtimestamp')"
	newline
	# Locate and interactively remove old filesystems from /home directory
	for file in $(find /home/boot/rootfs.gz.[0-9]*); do
		_n 'Do you wish to remove: %s (Yes/no/exit)? ' "$(basename $file)"
		read answer
		case $answer in
			e|E|"exit"|Exit)
				exit 0 ;;
			y|Y|yes|Yes)
				rm -f $file ;;
			*)
				_ 'No filesystems selected, exiting...' ;;
		esac
	done
}


#
# Tazusb sequence
#

case $COMMAND in
	writefs)
		# Writefs to rootfs.gz
		check_root
		# Compression type (optional): lzma, gzip, none. Default is none
		COMPRESSION="${2:-none}"
		# Full path to rootfs.gz (optional). Default is /rootfs.gz
		ROOTFS_PATH="${3:-/rootfs.gz}"
		# File name
		ROOTFS="$(basename "$ROOTFS_PATH")"

		# Start info
		title 'Write filesystem'

		longline "$(_ "The command writefs will write all the current \
filesystem into a suitable cpio archive (%s) usable on a bootable \
LiveUSB media." "$ROOTFS")"
		newline
		_ 'Archive compression: %s' "$(colorize 36 "$COMPRESSION")"
		newline

		# Clear out tazpkg cache
		rm /var/lib/tazpkg/*.bak /var/cache/tazpkg/* -r -f

		# Optionally remove sound card selection and screen resolution.
		_ 'Do you wish to remove the sound card and screen configs?'
		_n 'Press ENTER to keep or answer (No|yes|exit): '
		read answer
		case $answer in
			e|E|"exit"|Exit)
				exit 0 ;;
			y|Y|yes|Yes)
				_n 'Removing current sound card and screen configurations...'
				rm -f /var/lib/sound-card-driver
				rm -f /var/lib/alsa/asound.state
				rm -f /etc/X11/xorg.conf ;;
			*)
				_n 'Keeping current sound card and screen configurations...' ;;
		esac
		status
		newline

		# Optionally remove i18n settings
		_ 'Do you wish to remove local/keymap settings?'
		_n 'Press ENTER to keep or answer (No|yes|exit): '
		read answer
		case $answer in
			e|E|"exit"|Exit)
				exit 0 ;;
			y|Y|yes|Yes)
				_n 'Removing current locale/keymap settings...'
				echo > /etc/locale.conf
				echo > /etc/keymap.conf ;;
			*)
				_n 'Keeping current locale/keymap settings...'
				grep -qs '^INCLUDE i18n.cfg'        /home/boot/*linux/[se]*linux.c* &&
				sed  -i  's/^INCLUDE i18n.cfg/# &/' /home/boot/*linux/[se]*linux.c* ;;
		esac
		status
		newline

		# Clean-up files by default
		mv -f /var/log/wtmp /tmp/tazusb-wtmp
		touch /var/log/wtmp
		echo > /etc/udev/rules.d/70-persistent-net.rules
		echo > /etc/udev/rules.d/70-persistant-cd.rules

		# Create list of files
		# find / -xdev | sed '/^\/home\//d;/^\/tmp\//d' >/tmp/list
		# for dev in console null tty tty1
		# do
		#	echo /dev/$dev >>/tmp/list
		# done

		for dir in /bin /etc /init /sbin /var /dev /lib /root /usr /opt
		do
			[ -d $dir -o -f $dir ] && find $dir
		done >/tmp/list

		for dir in /home /proc /run /sys /tmp /mnt /media /media/cdrom /media/flash /media/usbdisk
		do
			echo $dir >>/tmp/list
		done
		sed -i '/^\/var\/run\/.*pid$/d' /tmp/list

		for removelog in auth boot messages dmesg daemon utmp slim Xorg tazpanel cups; do
			sed -i "\/var\/log\/$removelog/d" /tmp/list
		done

		# Generate initramfs with specified compression
		case "$COMPRESSION" in
			lzma)
				_n 'Creating %s with lzma compression... ' "$ROOTFS"
				cpio -o -H newc | lzma e -si -so > "$ROOTFS_PATH"
				;;
			gzip)
				_n 'Creating %s with gzip compression... ' "$ROOTFS"
				cpio -o -H newc | gzip -9 > "$ROOTFS_PATH"
				[ -x /usr/bin/advdef ] && advdef -z4 "$ROOTFS_PATH"
				;;
			none|*)
				_n 'Creating %s without compression... ' "$ROOTFS"
				cpio -o -H newc > "$ROOTFS_PATH"
				;;
		esac < /tmp/list

		mv -f /tmp/tazusb-wtmp /var/log/wtmp

		# Get initramfs size
		size=$(du -sh "$ROOTFS_PATH" | cut -f1)

		# If the bootable medium is where it should be, copy across
		if [ "$(ls /home/boot/bzImage*)" ]; then
			longline "$(_ 'Moving %s to media. Remember to unmount for delayed writes!' "$ROOTFS")"

			CONFIG=$(ls /home/boot/*linux/[se]*linux.c*)
			if grep -qs 'gz,' $CONFIG; then

				# Disable loram & many-in-1
				[ -s /home/boot/rootfs.gz ] && mv /home/boot/rootfs.gz /home/boot/rootfs.gz.loram
				cp $CONFIG $CONFIG.org
				sed -i 's|rootfs.*\.gz|rootfs.gz|' $CONFIG
			else
				# Move the old filesystem with the unix timestamp for reference
				if [ -e /home/boot/previous.gz ]; then
					mv /home/boot/previous.gz /home/boot/rootfs.gz.$(date -r /home/boot/previous.gz +%s)
				fi
			
				mv /home/boot/rootfs.gz /home/boot/previous.gz
			fi
			mv "$ROOTFS_PATH" /home/boot/rootfs.gz
			_ '%s is located in %s' 'rootfs.gz' '/home/boot'
		else
			_ '%s is located in %s' "$ROOTFS" "$(dirname "$ROOTFS_PATH")"
		fi

		separator
		_ 'Root filesystem size: %s' "$size"
		separator '-'
		_n 'ENTER to continue...'; read i
		;;


	format)
		# Format a partition.
		check_root

		title 'Format a device'

		DEVICE="$2"
		label="$3"
		fs_type="$4"
		if [ -z "$DEVICE" ]; then
			ask_for_device
			check_for_device
		else
			_ 'Device: %s' "$DEVICE"
		fi
		[ -z "$fs_type" ] && get_fs_type
		get_label
		unmount_target_usb
		make_fs "$fs_type"
		separator
		longline "$(_ 'Device %s is ready to use as LiveUSB and/or /home partition.' "$label ($DEVICE)")"
		;;


	gen-liveusb)
		# Generate a LiveUSB media using files from a LiveCD.
		check_root

		title 'Gen a LiveUSB media'

		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" -o ! -f "$ISO" ]; then
			_ 'Please specify a valid filename on the command line.'
			exit 1
		fi

		title 'Copy ISO file to SliTaz LiveUSB media'

		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
