ATV Bootloader Patchstick, Backup, and Restore Creation Script

#!/bin/sh
#
# Apple TV Patchstick & Backup/Restore Stick
#
# 2008/10/10 - bubbaATbubba.org
# - Added -n option so that full Apple TV isn't erased (only recovery partition updated)
#   This is helpful if you want to update recovery partition to one version (say 2.1), then re-update
#   it with an older version (say 1.0) (in order to get around having to upgrade to the most recent 
#   AppleTV version. 
# 
# 2008/09/24 - bubbaATbubba.org
# - Added ability to update a stick rather than repartition/create each time
# - Added ability to create/update stick to a backup stick (for backing up your ATV)
# - Added ability to create/update stick to a restore stick (for restoring your ATV)
#
# 2008/09/12 - bubbaATbubba.org 
# - Initial Version
# 
# This script assumes that you've downloaded and installed the atv-bootloader
# copies of Parted (http://code.google.com/p/atv-bootloader/wiki/InstallParted) 
# and HFS Tools (http://code.google.com/p/atv-bootloader/wiki/InstallHFSTools)
#
# It also assumes that you've downloaded the recovery tools and extracted them
# to the directory where atvps.sh exists:
# (http://code.google.com/p/atv-bootloader/downloads/list)
# # wget http://atv-bootloader.googlecode.com/files/recovery-0.6.tar.gz
# # sudo tar -xzf recovery-0.6.tar.gz
# 
 
# leave these empty unless you always want to use same device/size
USB=
SIZE=
EFI=
RECOVERY=
 
usage () {
cat<<EOF
Usage: $0 options
 
Options:
  -d   Raw USB device
 
  -c   Create a new patchstick
    or 
  -u   Update pre-configured/pre-partitioned stick
 
  -t   Type of stick: p = patchstick (default)
	              b = backup script
		      r = restore script
  -h   Help
 
 Valid Options for Create:
  -s   Size of Patchstick partition
  -e   Location of EFI Bootloader
  -r   Location of extracted recovery.tar.gz
 
 Valid Options for Restore type:
  -b   Location of ATV Backup files
  -n   No Erase (do not erase disk, just update recovery files)
 
Examples:
 Create a new patchstick on 2GB flash drive. 
 WARNING: create mode will ERASE the disk you specify!
  $0 -d /dev/sdd -s 1800M -e boot.efi -r recovery -t p -c
 
 Update already created patchstick for backup mode
  $0 -d /dev/sdd -t b -u 
 
 Update already created patchstick for restore mode where
 backup copy is not already on the patchstick.
  $0 -d /dev/sdd -t r -u -b my_atv_backupdir 
 
 Update already created patchstick for restore mode where
 backup copy is already on the patchstick.
  $0 -d /dev/sdd -t r -u 
 
Notes:
 For Patchstick Mode:  You will be given a sample patchstick.sh file on
 the Patchstick partition.  You will need to change this script so that it
 processes and installs all necessary patches.  Examples can be found at
 http://bubba.org/ATV
 
 For Backup Mode:  Files will be stored on the Patchstick partition in the 
 backup directory. It is assumed that you will copy these files off after
 a successful backup.
 
 For Restore Mode:  You will need to have already run a backup using the
 above backup method and SAVED the backup files somewhere else or leave
 them on your patchstick.  Keep in mind, if you keep your backup on your
 patchstick you should only use the "update" mode because the "create" mode
 REFORMATS your USB drive.  Keep in mind running the restore mode on your
 Apple TV WILL FORMAT YOUR APPLE TV DRIVE unless you use the -n option!
 
EOF
}
 
restore_script () {
cat<<EOF
#!/bin/bash
#
 
exec 2>/dev/console
exec 1>/dev/console
 
echo
echo "        --- AppleTV Restore Stick by Bubba ---"
GTP="/mnt/rootfs/backup/org_gtp.txt"
if [ -d /mnt/rootfs/backup ]; then
        if [ -e /mnt/rootfs/backup/OS.dmg ] && [ -e \${GTP} ]; then
	     if [ ! -e /mnt/rootfs/noerase ]; then
                P1S=\`grep " 1  " \${GTP} | awk '{print \$2}'\`
                P1E=\`grep " 1  " \${GTP} | awk '{print \$3}'\`
                P2S=\`grep " 2  " \${GTP} | awk '{print \$2}'\`
                P2E=\`grep " 2  " \${GTP} | awk '{print \$3}'\`
                P3S=\`grep " 3  " \${GTP} | awk '{print \$2}'\`
                P3E=\`grep " 3  " \${GTP} | awk '{print \$3}'\`
                P4S=\`grep " 4  " \${GTP} | awk '{print \$2}'\`
                P4E=\`grep " 4  " \${GTP} | awk '{print \$3}'\`
 
                echo "        * Wiping paritition table on /dev/sda ..."
                dd if=/dev/zero of=/dev/sda bs=4096 count=1M
                # create the GPT format
                parted -s /dev/sda mklabel gpt
                echo "        * Creating EFI partition /dev/sda1 - start \${P1S} end \${P1E}"
                parted -s /dev/sda mkpart primary fat32 \${P1S} \${P1E}
                parted -s /dev/sda set 1 boot on
 
                echo "        * Creating Recovery partition /dev/sda2 - start \${P2S} end \${P2E}"
                parted -s /dev/sda mkpart primary HFS \${P2S} \${P2E}
                parted -s /dev/sda set 2 atvrecv on
 
                echo "        * Creating OSBoot partition /dev/sda3 - start \${P3S} end \${P3E}"
                parted -s /dev/sda mkpart primary HFS \${P3S} \${P3E}
 
                echo "        * Creating Media partition /dev/sda4 - start \${P4S} end \${P4E}"
                parted -s /dev/sda mkpart primary HFS \${P4S} \${P4E}
 
                echo "        * Creating file systems on all 4 partitions..."
                # make file systems for the four partitions
                mkfs.msdos -F 32 -n EFI /dev/sda1
                mkfs.hfsplus -v Recovery /dev/sda2
                mkfs.hfsplus -J -v OSBoot /dev/sda3
                mkfs.hfsplus -J -v Media /dev/sda4
 
                partprobe /dev/sda
            fi
            mkdir /dst
            fsck.hfsplus -f /dev/sda2
            mount -t hfsplus -o rw,force /dev/sda2 /dst
            cp -farp /mnt/rootfs/backup/* /dst/
	    rm -f /dst/org_gtp.txt
            sync &>/dev/null
            umount /dst
            echo "        * Restore complete!  Now reboot & perform a factory restore!"
        else 
            echo "        * Backup files or org_gtp.txt not found in backup location."
        fi
fi
sleep 100000
EOF
}
 
backup_script () {
cat<<EOF
#!/bin/bash
#
#
exec 2>/dev/console
exec 1>/dev/console
 
echo
echo "        --- AppleTV Backup Stick by Bubba ---"
 
BKUPDIR="/mnt/rootfs/backup"
 
if [ -d \${BKUPDIR} ]; then
        echo "        * Backup location found!"
        if [ -e \${BKUPDIR}/OS.dmg ] && [ ! -e /mnt/rootfs/forcebu ]; then
                echo "        * ERROR: It appears a backup has already run."
                echo "        You either need to empty out Patchstick/backup"
                echo "        directory or touch Patchstick/forcebu."
        else 
                echo "        * Mounting Recovery partition on /src"
                mkdir /src
                fsck.hfsplus -f /dev/sda2
                mount -t hfsplus -o ro,force /dev/sda2 /src
                echo "        * Backing up recovery partition..."
                cp -arpf /src/* \${BKUPDIR}
                parted -s /dev/sda unit s print > \${BKUPDIR}/org_gtp.txt
                sync &>/dev/null
                umount /src
                echo "        * Backup complete!"
        fi
else 
        echo "        * Backup location not found.  Please create a backup directory on the Patchstick partition."
fi
sleep 100000
EOF
}
 
patchstick_script () {
cat<<EOF
#!/bin/bash
 
exec 2>/dev/console
exec 1>/dev/console
 
echo
echo "        --- AppleTV Patchstick by Bubba ---"
 
echo "        * mounting OSBoot partition"
mkdir /OSBoot
fsck.hfsplus -f /dev/sda3 
mount -t hfsplus -o rw,force /dev/sda3 /OSBoot 
 
# since all of our installer scripts reference /stuff, we need
# to set this up the bomb.
echo "       * symlinking /mnt/rootfs -> /stuff"
ln -s /mnt/rootfs /stuff
 
echo "        * keeping the OSBoot partition r/w for plugins"
touch /OSBoot/.readwrite
 
# patch OS-related things.  
if [ -d /OSBoot/dev ] && [ -d /stuff/installer.d ]; then
        if [ -e /OSBoot/.ospatch ] && [ ! -e /stuff/forceos ]; then
                echo "        * Skipping OS patches since /.ospatch exists.  Use forceos file on Patchstick to override."
        else 
                for i in /stuff/installer.d/*; do
                        echo "        * Installing \${i} ..."
                        /bin/bash "\${i}" 1>> /OSBoot/.upgrade.log 2>&1
                done
                touch /OSBoot/.ospatch
                echo "        * OS Patching Successful!"
        fi
else 
        echo "        * ERROR: Unable to mount /dev/sda3 to apply OS patches."
fi
 
sync &>/dev/null
umount /userspace
umount /OSBoot
echo "        * Please ssh into Apple TV and see /.upgrade.log and check for errors."
echo "        * To ssh: ssh frontrow@appletv.local  Password: frontrow"
echo "        * Please unplug your Apple TV to reboot/reset the device."
sleep 100000
EOF
}
 
cleanup () {
 	sync	
	umount /mnt/Recovery
	umount /mnt/Patchstick
	rmdir /mnt/Recovery
	rmdir /mnt/Patchstick
}
 
TYPE="p"
CREATE=0
UPDATE=0
NOERASE=0
 
while getopts "chund:s:e:r:t:b:" OPTION
do
  case $OPTION in
     h)
	usage
	exit 1
	;;
     d) 
	USB=$OPTARG
	;;
     s)
	SIZE=$OPTARG
	;;
     e)
	EFI=$OPTARG
	;;
     r)
	RECOVERY=$OPTARG
	;;
     t)
	TYPE=$OPTARG
	;;
     c)
	CREATE=1
	;;
     u)
	UPDATE=1
	;;
     b)
	BACKUP=$OPTARG
	;;
     n)
	NOERASE=1
	;;
     ?)
	usage
	exit
	;;
  esac
done 
 
# ( [ -e blah ] || [ -e blah ] ) 
if [ ${UPDATE} -eq 1 ]; then
	if [ ! -e "${USB}" ] || [ -z ${TYPE} ] || [ -n "${BACKUP}" -a ! -e "${BACKUP}/org_gtp.txt" ]; then
		echo "Ooops.  Something went wrong.  Please check your options and verify locations."
		usage
		exit 1
	fi
elif [ ${CREATE} -eq 1 ]; then
	if [ ! -e "${USB}" ] || [ -z "${SIZE}" ] || [ ! -e "${EFI}" ] || 
	[ ! -e "${RECOVERY}/mach_kernel" ] ||  [ -z ${TYPE} ] || 
	[ -n "${BACKUP}" -a ! -e "${BACKUP}/org_gtp.txt" ]; then
		echo "Ooops.  Something went wrong.  Please check your options and verify locations."
		usage
		exit 1
	fi
else 
	usage
	exit 1
fi
 
mount | grep ${USB}
if [ $? -eq 0 ]; then
	echo "You're trying to alter a mounted partition - ${USB} (not good).  Exiting."
	exit 1
fi	
 
if [ "$(id -u)" != "0" ]; then
	echo "Sorry, you must be root to run this."
	exit 1
fi
 
if [ ${CREATE} -eq 1 ]; then
	echo "Zeroing out device: ${USB}..."
	dd if=/dev/zero of=${USB} bs=4096 count=1M
	sleep 5
	partprobe ${USB}
 
	echo "Creating the gpt partition on ${USB}..."
	parted -s ${USB} mklabel gpt
 
	echo "Creating the recovery partition..."
	# create a 25MB "Recovery" partition  (starting at sector 40 is important)
	parted -s ${USB} mkpart primary HFS 40s 25M
	parted -s ${USB} set 1 atvrecv on
 
	echo "Creating the Patchstick partition with size ${SIZE}..."
	parted -s ${USB} mkpart primary HFS 25M ${SIZE}
	partprobe ${USB}
 
	echo "Printing updated partition info on ${USB}..."
	parted ${USB} print
	sync
 
	if [ -e "${USB}1" ] && [ -e "${USB}2" ]; then
		echo "Formatting hfsplus partitions on ${USB}1 and ${USB}2..."
		mkfs.hfsplus -v Recovery "${USB}1"
		mkfs.hfsplus -v Patchstick "${USB}2"
		sync
	fi
fi
 
if [ ${CREATE} -eq 1 ] || [ ${UPDATE} -eq 1 ]; then
 
	if [ ! -d /mnt/Recovery ]; then
		mkdir -p /mnt/Recovery
	fi
	if [ ! -d /mnt/Patchstick ]; then
		mkdir -p /mnt/Patchstick
	fi
 
	if [ ${UPDATE} -eq 1 ]; then
		echo "Forced checking of hfsplus partitions: ${USB}1 and ${USB}2..."
		fsck.hfsplus -f ${USB}1
		fsck.hfsplus -f ${USB}2
	fi
 
	echo "Mounting hfsplus partitions on ${USB}1 and ${USB}2..."	
	mount -t hfsplus "${USB}1" /mnt/Recovery
	mount -t hfsplus "${USB}2" /mnt/Patchstick
 
	if [ ${CREATE} -eq 1 ]; then
		echo "Copying recovery code to /mnt/Recovery..."
		cp -arp ${RECOVERY}/* /mnt/Recovery/
 
		echo "Copying boot.efi to /mnt/Recovery..."
		cp -ap ${EFI} /mnt/Recovery
	fi
 
	if [ ! -e "/mnt/Recovery/com.apple.Boot.plist" ]; then
		echo "ERROR: Unable to find valid com.apple.Boot.plist in /mnt/Recovery."
		echo "Verify that you have created or updated the correct partition." 
		cleanup
		exit 1
	fi	
 
	if [ ${TYPE} = "r" ];  then
		if [ -n "${BACKUP}" -a ! -e "${BACKUP}/org_gtp.txt" ] &&
			[ ! -e "/mnt/Patchstick/backup/org_gtp.txt" ]; then
			echo "ERROR: No valid backup location was found. "
			echo -n "Tried looking in /mnt/Patchstick/backup "
			if [ -n "${BACKUP}" ]; then
				echo "and in ${BACKUP}."
			else 
				echo
			fi
			cleanup
			exit 1
		elif [ -n "${BACKUP}" -a -e "${BACKUP}/org_gtp.txt" ]; then
			echo "Found valid recovery data, copying to Patchstick..."
			if [ ! -d /mnt/Patchstick/backup ]; then
				mkdir -p /mnt/Patchstick/backup
			fi
			cp -arpfv ${BACKUP}/* /mnt/Patchstick/backup
		elif [ -z "${BACKUP}" -a -e "/mnt/Patchstick/backup/org_gtp.txt" ]; then
			echo "Found existing recovery data.  Leaving..."
		else 
			echo "Error in restore mode. Please check options. Did you mean to specify a source directory (-b)"
			cleanup
			exit 1
		fi
	        if [ -e /mnt/Patchstick/noerase ] && [ ${NOERASE} -eq 0 ]; then
			rm -f /mnt/Patchstick/noerase
		fi
		if [ ! -e /mnt/Patchstick/noerase ] && [ ${NOERASE} -eq 1 ]; then
			touch /mnt/Patchstick/noerase
		fi
		if [ -e /mnt/Recovery/boot_linux.sh ]; then
			rm -f /mnt/Recovery/boot_linux.sh
		fi
		echo "Creating restore script /mnt/Patchstick/boot_linux.sh..."
		restore_script > /mnt/Patchstick/boot_linux.sh
		chmod 755 /mnt/Patchstick/boot_linux.sh	
 
		echo "Updating com.apple.Boot.plist to default to restore mode..."
		cat /mnt/Recovery/com.apple.Boot.plist | sed s/atv-boot=[a-z]*/atv-boot=manual/ > /mnt/Recovery/plist.new
		mv -f /mnt/Recovery/plist.new /mnt/Recovery/com.apple.Boot.plist 
		if [ ${NOERASE} -eq 0 ]; then
			echo "WARNING:  BOOTING YOUR APPLE TV FROM THIS DRIVE WILL COMPLETELY ERASE YOUR APPLE TV!"
		else
			echo "WARNING:  BOOTING YOUR APPLE TV FROM THIS DRIVE WILL OVERWRITE YOUR RECOVERY PARTITION!"
		fi
	elif [ ${TYPE} = "p" ]; then
		echo "Creating patchstick script /mnt/Patchstick/patchstick.sh..."
		if [ -e /mnt/Recovery/patchstick.sh ]; then
			rm -f /mnt/Recovery/patchstick.sh
		fi
		patchstick_script > /mnt/Patchstick/patchstick.sh
		chmod 755 /mnt/Patchstick/patchstick.sh
 
		echo "Updating com.apple.Boot.plist to default to patchstick mode..."
		cat /mnt/Recovery/com.apple.Boot.plist | sed s/atv-boot=[a-z]*/atv-boot=patchstick/ > /mnt/Recovery/plist.new
		mv -f /mnt/Recovery/plist.new /mnt/Recovery/com.apple.Boot.plist 
	elif [ ${TYPE} = "b" ]; then
		echo "Creating Backup script /mnt/Patchstick/boot_linux.sh..."
		if [ -e /mnt/Recovery/boot_linux.sh ]; then
			rm -f /mnt/Recovery/boot_linux.sh
		fi
		backup_script > /mnt/Patchstick/boot_linux.sh
		chmod 755 /mnt/Patchstick/boot_linux.sh	
 
		echo "Updating com.apple.Boot.plist to default to backup mode..."
		cat /mnt/Recovery/com.apple.Boot.plist | sed s/atv-boot=[a-z]*/atv-boot=manual/ > /mnt/Recovery/plist.new
		mv -f /mnt/Recovery/plist.new /mnt/Recovery/com.apple.Boot.plist 
 
		echo "Creating backup directory on Patchstick partition..."
		if [ -e "/mnt/Patchstick/backup/org_gtp.txt" ]; then
			echo "WARNING: Valid backup already found on Patchstick partition.  Please remove these files "
			echo "or create the file \"forcebu\" on the Patchstick partition to force overwriting."
		fi	
		if [ ! -d /mnt/Patchstick/backup ]; then
			mkdir -p /mnt/Patchstick/backup
		fi
	else 
		echo "Invalid option."
	fi
	cleanup
fi