#!/bin/bash

################################################################################
#      This file is part of OpenELEC - http://www.openelec.tv
#      Copyright (C) 2009-2013 Stephan Raue (stephan@openelec.tv)
#
#  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, 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 OpenELEC.tv; see the file COPYING.  If not, write to
#  the Free Software Foundation, 51 Franklin Street, Suite 500, Boston, MA 02110, USA.
#  http://www.gnu.org/copyleft/gpl.html
################################################################################

check_root(){
	if [ $UID != 0 ];then
	    echo ""
		echo "FAIL: This script *must* be run as root!"
		echo ""		
		exit 1
	fi
}

check_runlocation(){
	SCRIPT=$(pwd)
}

check_localtargetfolder(){
	KEYWORDMATCH=$(echo ${INSTALLKEYWORDS[*]} | grep ${KEYWORD})
	if [ "${KEYWORDMATCH}" != "" ]; then
		if [[ -f target/SYSTEM && -f target/KERNEL ]]; then
			IMAGEFILES="local"
		else
			IMAGEFILES="download"
		fi	
	fi
}

check_distro(){
	DISTRO=$(uname -s)
	if [ $DISTRO != "Darwin" ]; then
		echo ""
		echo "FAIL: This script is only for Mac OS X systems, aborting!"
		echo ""
		exit 1
	fi
	OSXVER=$(sw_vers -productVersion)
	if [ $OSXVER -lt "10.5" ]; then
		echo "FAIL: This script requires Max OS X 10.5 or higher to run, aborting!"
	fi
}

read_input(){
	local MESSAGE="${1}"
	if [ -n "${MESSAGE}" ];then
		echo -n "${MESSAGE}"
	fi
	read INPUT
}

choose_target(){
	echo ""
	echo "INFO: SwissArmyKnife keyword is '${KEYWORD}'"
	echo ""
	echo "INFO: The following drives are available to create the SwissArmyKnife USB on:"
	echo ""
	diskutil list | grep 0: | grep -v disk0 | awk '{print "/dev/"$5 " "$3 $4}'
	echo ""
	while [ -z ${INPUT} ];do
    		read_input "INFO: Please enter the target device (e.g. /dev/disk1): "
    		echo ""
	done
	TARGET=${INPUT}
	STARGET=$(echo ${TARGET} | sed 's:/dev/::g')
	RTARGET=$(echo /dev/r${STARGET})
}

warning(){
	local OUTPUT
	OUTPUT=$(fdisk ${TARGET} 2>/dev/null)
	if [ -z "${OUTPUT}" ];then
		echo "FAIL: Installation aborted! Device ${TARGET} is invalid"
		echo ""
		exit 1
	else
		diskutil list ${TARGET}
		echo ""
	fi

	CONTINUE=""
	echo -n "WARN: Continuing will erase all data from ${TARGET}, continue? (Y/n): "
	read CONTINUE
	if [[ $CONTINUE = "N" || $CONTINUE = "n" || $CONTINUE = "" ]]; then
		echo ""	
		echo "ERROR: Installation aborted!"
		echo ""
		exit 1
	elif [[ $CONTINUE = "Y" || $CONTINUE = "y" ]]; then 
		echo ""
	fi
}	

disksize(){
	SECTORS=$(fdisk ${TARGET} | head -n 1 | awk '{print $5}' | sed 's/\[//g')
	if [ "${SECTORS}" -gt "4194304" ]; then
        	DISKSIZE="BIG"
		echo "INFO: ${TARGET} is >2GB"
    	else
        	DISKSIZE="NORMAL"
        	echo "INFO: ${TARGET} is <2GB"
    	fi
}

check_debug(){
	if [ "$2" = "debug" ]; then
		DEBUG="yes"
	else
		DEBUG="no"
	fi
}

check_verbose(){
    	if [ "${VERBOSE}" = "TRUE" ]; then
        	exec 3>&1
        	exec 4>&2
    	else
        	exec 3> /dev/null
        	exec 4> /dev/null
	fi
}

unmount_target(){
	for i in $(jot 30 1 30); do
	        MOUNTED=$(mount | grep ${TARGET})
	        if [ -z "${MOUNTED}" ]; then
			break
	        else
			diskutil unmountDisk ${TARGET} 1>&3 2>&4
			sleep 1
	        fi
	done
}

mount_target(){
	for i in $(jot 30 1 30); do
	        MOUNTED=$(mount | grep ${TARGET})
		if [ -z "${MOUNTED}" ]; then
	                diskutil mountDisk ${TARGET} 1>&3 2>&4
    			sleep 1
	        else
	                break
	        fi
	done
}

prepare_target(){
	echo "INFO: Erasing existing partition schemes on ${TARGET}"
	dd if=/dev/zero of=${TARGET} bs=512 count=100 1>&3 2>&4
	gpt destroy ${TARGET} 1>&3 2>&4
	gpt create ${TARGET} 1>&3 2>&4
}

create_partitions(){
	echo "INFO: Creating RECOVERY partition"
	gpt add -b 40 -i 1 -s 409600 -t "48465300-0000-11AA-AA11-00306543ECAC" ${TARGET} 1>&3 2>&4
	echo "INFO: Creating OPENELEC partition"
	gpt add -b 411648 -i 2 -t "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7" ${TARGET} 1>&3 2>&4
	diskutil eraseVolume HFSX Recovery ${TARGET}s1 1>&3 2>&4
	diskutil eraseVolume "MS-DOS FAT32" OPENELEC ${TARGET}s2 1>&3 2>&4
}

create_partitions_big(){
	echo "INFO: Creating OPENELEC partition"
	diskutil partitionDisk ${TARGET} 1 GPT "MS-DOS FAT32" OPENELEC R 1>&3 2>&4
	unmount_target
	echo "INFO: Removing EFI partition"
	gpt remove -i 1 ${TARGET} 1>&3 2>&4
	echo "INFO: Creating RECOVERY partition"
	unmount_target
	gpt add -b 40 -i 1 -s 409600 -t "48465300-0000-11AA-AA11-00306543ECAC" ${TARGET} 1>&3 2>&4
	diskutil eraseVolume HFSX Recovery ${TARGET}s1 1>&3 2>&4
}

prepare_recovery(){
	if [ "${TEST}" = "TRUE" ]; then
		echo "INFO: Downloading test/atv_recovery.tar.gz and test/atv_recovery.tar.gz.md5"
	   	curl http://chewitt.openelec.tv/test/atv_recovery.tar.gz -o ${SCRIPT}/atv_recovery.tar.gz 1>&3 2>&4 || error_recoveryfiles
   		curl http://chewitt.openelec.tv/test/atv_recovery.tar.gz.md5 -o ${SCRIPT}/atv_recovery.tar.gz.md5 1>&3 2>&4 || error_recoveryfilesmd5
	else
		echo "INFO: Downloading atv_recovery.tar.gz and atv_recovery.tar.gz.md5"
	   	curl http://chewitt.openelec.tv/atv_recovery.tar.gz -o ${SCRIPT}/atv_recovery.tar.gz 1>&3 2>&4 || error_recoveryfiles
   		curl http://chewitt.openelec.tv/atv_recovery.tar.gz.md5 -o ${SCRIPT}/atv_recovery.tar.gz.md5 1>&3 2>&4 || error_recoveryfilesmd5
	fi
	if [ -a ${SCRIPT}/atv_recovery.tar.gz ]; then
		MD5_LOCAL=$(md5 -q ${SCRIPT}/atv_recovery.tar.gz)
	else
		echo "FAIL: atv_recovery.tar.gz does not exist, aborting!"
		fail
	fi
	if [ -a ${SCRIPT}/atv_recovery.tar.gz.md5 ]; then
		MD5_REMOTE=$(cat ${SCRIPT}/atv_recovery.tar.gz.md5 | awk '{print $1}')
	else
		echo "FAIL: atv_recovery.tar.gz.md5 does not exist, aborting!"
		fail
	fi
	if [ "${MD5_LOCAL}" != "${MD5_REMOTE}" ]; then
		echo "FAIL: Checksum validation error for atv_recovery.tar.gz, aborting!"
		fail
	fi
	mount_target
	if [ -d /Volumes/Recovery ]; then
		echo "INFO: Extracting atv_recovery.tar.gz to RECOVERY"
		tar -xzf "${SCRIPT}/atv_recovery.tar.gz" -C /Volumes/Recovery
		echo fat32 > /Volumes/Recovery/filesystem
		if [ "${KEYWORD}" = "emergency" ]; then
			echo emergency > /Volumes/Recovery/keyword
		else
			echo usbimage > /Volumes/Recovery/keyword
		fi
		chmod 7777 /Volumes/Recovery/patchstick.sh
	else
		echo "FAIL: The RECOVERY partition is not available, aborting!"
		fail
	fi
}

error_recoveryfiles(){
	echo "FAIL: Download of atv_recovery.tar.gz failed, aborting!"
	echo ""
	fail
}

error_recoveryfilesmd5(){
	echo "FAIL: Download of atv_recovery.tar.gz.md5 failed, aborting!"
	echo ""
	fail
}

prepare_openelec(){
	mount_target
	if [ "${IMAGEFILES}" = "local" ]; then
		echo "INFO: Copying SYSTEM and KERNEL files to OPENELEC"
		cp target/SYSTEM /Volumes/OPENELEC/
		cp target/KERNEL /Volumes/OPENELEC/
		cp target/SYSTEM.md5 /Volumes/OPENELEC/
		cp target/KERNEL.md5 /Volumes/OPENELEC/
	fi
	if [ "${DEBUG}" = "yes" ]; then
		touch /Volumes/OPENELEC/debug
	fi
	echo "INFO: Creating KEYWORD and README files on OPENELEC"
	echo $KEYWORD > /Volumes/OPENELEC/KEYWORD.txt
	cp /Volumes/Recovery/README.txt	/Volumes/OPENELEC/
}

change_guid(){
	unmount_target
	echo "INFO: Changing RECOVERY partition GUID"
	gpt remove -i 1 ${TARGET} 1>&3 2>&4
	unmount_target
	gpt add -b 40 -i 1 -s 409600 -t "5265636F-7665-11AA-AA11-00306543ECAC" ${RTARGET} 1>&3 2>&4
}

set_labels(){
	echo "INFO: Creating partition labels"
	unmount_target
	gpt label -i 1 -l "Recovery" ${RTARGET} 1>&3 2>&4
	unmount_target
	gpt label -i 2 -l "OPENELEC" ${RTARGET} 1>&3 2>&4
}

banner(){
	clear
	echo ""
	echo "********************************************************************************"
	echo "*                                                                              *"
	echo "*          Welcome to the OpenELEC AppleTV USB SwissArmyKnife Script!          *"
	echo "*                    created by chewitt@openelec.tv, $VERSION                     *"
	echo "*                                                                              *"
	echo "********************************************************************************"
}

end(){
	echo ""	
	echo "INFO: SwissArmyKnife USB creation is complete!"
	echo ""
}

fail(){
	echo ""	
	echo "FAIL: SwissArmyKnife USB creation failed!"
	echo ""
	exit 1
}

check_keywords(){
INSTALLKEYWORDS=('<install>' '<install-usb>' '<factoryrestore>' '<update>' '<bootloader>' \
				 '<backup>' '<restore>' '<chewitt>' '<chewitt-usb>' '<development>' \
				 '<development-usb>' '<migrate-crystalbuntu>' '<migrate-dynaflash>' \
				 '<singleboot>' 'emergency');
}


main(){
	check_root
	check_keywords
	check_debug
	check_runlocation
	check_localtargetfolder
	choose_target
	warning
	check_verbose
	disksize
	unmount_target
	prepare_target
	if [ "${DISKSIZE}" = "BIG" ]; then
    		create_partitions_big
	else
	    	create_partitions
	fi
	prepare_recovery
	prepare_openelec
	change_guid
   	set_labels
   	sleep 5
	unmount_target
}

if [ "$1" != "" ]; then
	KEYWORD=$1
else
	KEYWORD="install"
fi

VERSION="v2.23"
VERBOSE="FALSE" # set to TRUE to put this script in debug mode
TEST="FALSE" # set to TRUE to download test files
banner
main
end
exit 0