#!/bin/bash

#
# Copyright (c) 2012, AVG Technologies CZ, s.r.o., www.avg.com
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.


DIALOG=dialog
DIALOG_OK=0
DIALOG_CANCEL=1
DIALOG_HELP=2
DIALOG_EXTRA=3
DIALOG_PSCR=131
DIALOG_ESC=255
ARL_ERROR_LOG_FILE="/tmp/arl_error.log"
VALUE=""
ARL_USBDIR="${0%/*}"
USB_DEVICES_FILE="$ARL_USBDIR/usb_devs_file"

error()
{
	# $1 - error message

	$DIALOG --colors --backtitle "AVG Rescue CD" --msgbox "Error:\n\n\Z1${1}\Zn" 10 70
}

success()
{
	# $1 - success message

	$DIALOG --colors --backtitle "AVG Rescue CD" --msgbox "\Z2$1\Zn" 8 70
}

infobox()
{
	# $1 - title
	# $2 - message

	$DIALOG --backtitle "AVG Rescue CD" --title "$1" --infobox "\n$2" 5 70
}

cmd()
{
	# 1 - command
	# 2 - infobox message (optional)
	#   - cmd_INFOBOX_TITLE (optional)

	local cmd="$1"
	local info_msg="$2"
	local rv

	[ -n "$info_msg" ] && infobox "$cmd_INFOBOX_TITLE" "$info_msg"

	VALUE="`eval $cmd 2>&1`"
	rv=$?
	if [ $rv -ne 0 ] 
	then
		echo -e "$cmd  returned $rv:\n\t$VALUE" >> $ARL_ERROR_LOG_FILE
		error "$cmd:\n$VALUE"
	fi  

	return $rv
}

yesno()
{
	# $1 - question
	# $2 - height
	# $3 - width

	$DIALOG --backtitle "AVG Rescue CD" --colors --yesno "$1" $2 $3
}

usb_devices_radiolist()
{
	local path
	local dev
	local model
	local vendor
	local removable
	local devices
	local size
	local enable="on"

	for path in `find /sys/block/ -maxdepth 1 -iname sd* 2>/dev/null`
	do
		removable=`cat $path/removable`
		[ $removable -eq 0 ] && continue
		
		devices=`find $path/sd* -maxdepth 1 -iname sd* 2>/dev/null`

		# no partition
		[ -z "$devices" ] && devices=$path

		for device in $devices
		do
			dev=`basename $device`

			VALUE="`awk '/^.* '$dev'$/{print $3}' /proc/partitions`"
			size="$(($VALUE/1024))"

			if [ $size -lt 510 ]
			then
				continue
			elif [ $size -lt 1024 ]
			then 
				size="${size}MB"
			else
				size="$(($size/1024))GB"
			fi

			model=`cat $path/device/model`
			vendor=`cat $path/device/vendor`

			echo "$dev \"$vendor $model $size\" $enable" >> $USB_DEVICES_FILE

			enable="off"
		done
	done
}

usb_show_devices()
{
	VALUE=`$DIALOG --backtitle "AVG Rescue CD" \
		--title "USB Flash Drive" \
		--cancel-label "Return" \
		--radiolist "Please select drive" 15 70 10 \
		--file $USB_DEVICES_FILE \
		2>&1 1>&3`

	return $?
}

usb_conf_devices()
{
	local rv

	rm -f $USB_DEVICES_FILE
	usb_devices_radiolist
	if [ ! -f "$USB_DEVICES_FILE" ]
	then
		error "No USB Flash Drive found!"
		return 1
	fi

	usb_show_devices
	rv=$?
	rm -f $USB_DEVICES_FILE

	if [ $rv -eq $DIALOG_CANCEL -o $rv -eq $DIALOG_ESC -o $rv -eq $DIALOG_PSCR ]
	then
		return 1
	fi

	if [ $rv -ne $DIALOG_OK ]
	then

		echo "error: usb devices dialog: $rv"
		exit
	fi

	return 0
}

usb_umount()
{
	# $1 - dev name, e.g. sdb

	local mntpoints
	local mntpoint
	local dev="$1"

	mntpoints=`cat /proc/mounts | grep $dev | cut -d' ' -f2`
	for mntpoint in $mntpoints
	do
		cmd "umount $mntpoint" "unmounting $dev" || return 1
	done
}

usb_copy()
{
	# $1 - arl pck files path e.g. /mnt/usb, /mnt/sr0, ...
	# $2 - usb mountpoint

	local arl_files="f1.msg syslinux.cfg memtest menu.c32 vmlinuz CHANGELOG arl-version initrd.lzm arl.conf"
	local f

	for f in $arl_files 
	do
		cmd "cp $1/$f $2" "Copying file: $f" || return 1
	done

	return 0
}

usb_ext2_show_size()
{
	# $1 - available free space

	VALUE=`$DIALOG --backtitle "AVG Rescue CD" \
		--title "Select" \
		--cancel-label "Return" \
		--inputbox \
		"Please enter ext2 file system size. Minimum size is 500MB.
Maximum free space size available on USB drive is $1MB" \
		10 65 \
		"500" \
		2>&1 1>&3`

	return $?
}

usb_ext2_size()
{
	# $1 - dev full path

	local dev_escp="`echo $1 | sed 's%\/%\\\/%g'`"
	local avail=`df $1 | grep -v Filesystem |  awk '/'$dev_escp'/{print $4}'`

	usb_ext2_show_size $avail
	if [ $? -eq $DIALOG_CANCEL -o $? -eq $DIALOG_ESC -o $? -eq $DIALOG_PSCR ]
	then
		return 1
	fi

	if [ $? -ne $DIALOG_OK ]
	then
		quit "error: ext2 file system size dialog: $?"
	fi

	if [ $VALUE -gt $(($avail/1024)) ]
	then
		error "There is not enough free space on the USB drive"
		return 1
	fi

	if [ 500 -gt $VALUE ]
	then
		error "The minimum ext2 file system size is 500MB"
		return 1
	fi

	return 0
}

usb_ext2()
{
	# $1 - dst. usb size
	# $2 - dst. usb mount point
	# $3 - src. arl_rootfs.ext2 mount point

	local rv=0
	local loopdev="`losetup -f`"

	cmd "dd if=/dev/zero of=$2/arl_rootfs.ext2 bs=1M count=$1" "Creating arl_rootfs.ext2 loopback device" &&
	cmd "losetup $loopdev $2/arl_rootfs.ext2" "Setting up loopback device" &&
	cmd "mkfs.ext2 $loopdev" "Creating ext2 file system" &&
	cmd "mkdir -p /mnt/ext2" "Creating /mnt/ext2 directory" &&
	cmd "mount $loopdev /mnt/ext2" "Mounting ext2 file system" || return 1

	cmd "cp -a $3/* /mnt/ext2" "Creating root file system in ext2 fs" || rv=$?

	cmd "umount /mnt/ext2" "Umounting ext2 file system" &&
	cmd "rmdir /mnt/ext2" &&
	cmd "losetup -d $loopdev" "Removing loopback device" || rv=$?

	return $rv
}

usb_create()
{
	# $1 - dev name, e.g. sdb or sdb1,...
	# $2 - arl src files path e.g. /mnt/usb, /mnt/sr0, ...

	local rv=0
	local partition=${1:3:1}
	[ -z "$partition" ] && partition="1"
	local device=${1:0:3}
	local id=`sfdisk --print-id /dev/$device $partition 2>/dev/null`

	cmd "mkdir -p /mnt/usb_origfs" &&
	cmd "mount $ARL_USBDIR/arl_rootfs.ext2 /mnt/usb_origfs" || rv=$?

	if [ $rv -eq 0 ]
	then
		cmd "mkdir -p /mnt/usb_new" "Mounting USB..." &&
		cmd "mount -t vfat /dev/$1 /mnt/usb_new" || rv=$?

		if [ $rv -ne 0 ]
		then
			error "Different file system, FAT16/FAT32, required!"
		else
			usb_ext2_size "/dev/$1" &&
			usb_ext2 "$VALUE" "/mnt/usb_new" "/mnt/usb_origfs" &&
			usb_copy "$2" "/mnt/usb_new" || rv=$?

			cmd "/mnt/usb_origfs/opt/bin/syslinux -i /dev/$1" "Installing syslinux..." &&

			# create mbr if partition specified
			if [ ${#1} -gt 3 ]
			then
				cmd "dd if=/dev/zero of=/dev/$device bs=446 count=1" "Creating MBR..." &&
				cmd "cat /mnt/usb_origfs/usr/lib/syslinux/mbr.bin > /dev/$device" || rv=$?
			fi &&

			# fix patririon ID after is corrupted by installing syslinux
			cmd "sfdisk --change-id /dev/$device $partition $id" || rv=$?
			sleep 1
		fi

		cmd "umount /mnt/usb_new" "Unmounting USB..." &&
		cmd "rmdir /mnt/usb_new"

		cmd "umount /mnt/usb_origfs" &&
		cmd "rmdir /mnt/usb_origfs"
	fi

	return $rv
}

usb()
{
	cmd_INFOBOX_TITLE="USB Flash Drive"

	if [ -z "`which $DIALOG`" ]
	then
		echo "Please install the dialog utility at first."
		echo "For example: \$ sudo apt-get install dialog"
		return
	fi

	if [ ! -f "$ARL_USBDIR/arl-version" ]
	then
		error "arl-version file was not found!"
		return
	fi

	usb_conf_devices 
	if [ $? -eq 1 ]
	then
		return
	fi
	dev=$VALUE

	yesno "All the AVG Rescue CD data on the USB Flash Drive will be rewritten?\n\n\
Do you really want to continue?" 8 70
	if [ $? -eq $DIALOG_CANCEL -o $? -eq $DIALOG_ESC -o $? -eq $DIALOG_PSCR ]
	then
		return
	fi

	usb_umount $dev
	usb_create $dev "$ARL_USBDIR"
	if [ $? -ne 0 ]
	then
		error "Bootable USB Flash Drive was NOT created!"
		return
	fi

	success "Bootable USB Flash Drive successfully created."
}

exec 3>&1

if [ $(id -u) -ne 0 ] 
then
    echo "ERROR: Try to run this as root."
    # LSB specification errorcode
    exit 4
fi

usb
