#!/bin/bash
#
# Script Version: 1.0
# Author Contact: robin@twilightcoders.net
#
# Copyright (C) 2009-2010 Twilight Coders
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
# 3. The name of the author may not be used to endorse or promote
#    products derived from this software without specific prior
#    written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR `AS IS'' AND ANY EXPRESS
# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
# GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


# Get the version number. This is needed for certain programs and
# compatibility.
osxversion_full=`sw_vers -productVersion`;
osxversion_major=`sw_vers -productVersion | grep -oE "[0-9]+\.[0-9]+"`;

#sudo kextcache -v 1 -a i386 -a x86_64 -m /System/Library/Caches/com.apple.kext.caches/Startup/Extensions.mkext /System/Library/Extensions

if [ -z "$1" ]; then
	echo "Usage: efi [[init boot0 file... boot1h file... boot file...] [[mount] [open]]|unmount] [rebuild|update] [install file...]]";
	exit
elif [ "$USER" != "root" -a "$USER" != "" ]; then
   echo "Error: must run as root. try 'sudo -s'"
   exit
fi

mounted_efi=false;
efi_dir="/Volumes/EFI";

for option in $* ; do
	if [ "${option:0:1}" = "-" ] ; then
		# Remove Dash
		options="${option:1:${#option}}";
		while [ "${#options}" -gt "0" ] ; do
			case ${options:0:1} in
				d) disk="${options:1:${#options}}"; break;;
				*) usage;;
			esac
			# Shift Array
			options="${options:1:${#options}}";
		done
	fi
done

if [ "$disk" == "" ]; then
	disk=`diskutil info / | grep "Part Of Whole" | grep -o "disk[0-9]"`
	
	read -p "Warning: No disk specified. Using startup disk ($disk). Continue? (y[n]) ";

	if [ "$REPLY" != "y" ]; then
		echo "Aborting.";
		exit;
	fi
fi

disk=${disk%/};
disk=${disk##*/};
disk_main=`diskutil list | grep Apple_HFS | grep -o "${disk}[a-z][0-9]"`
disk_efi=`diskutil list | grep EFI | grep -o "${disk}[a-z][0-9]"`
X=`echo "$disk_efi" | grep -o "[0-9]" | head -n1`

function mount_efi
{
	if [ ! -e "$efi_dir" ]; then
		mkdir $efi_dir;
	fi

	# Check to see if it's already mounted
	if [ "`mount | grep $efi_dir | cut -d" " -f3`" != "$efi_dir" ]; then
		mount_hfs /dev/$disk_efi $efi_dir;
		mounted_efi=true;
	fi
}

function unmount_efi
{
	if [ "`mount | grep $efi_dir | cut -d" " -f3`" == "$efi_dir" ]; then
		rm -rf $efi_dir/.DS_Store $efi_dir/.Trashes
		umount -f /dev/$disk_efi;
	fi
	if [ -e "$efi_dir" -a "`ls $efi_dir`" == "" ]; then
		rm -rf $efi_dir;
	fi	
	mounted_efi=false;
}

function rebuild
{
	mount_efi;
	
	echo -n "Updating EFI boot cache ... "

	mv $efi_dir/Extra/Extensions.mkext $efi_dir/Extra/Extensions.mkext.previous
	chmod -R 755 $efi_dir/Extra/Extensions
	chown -R root:wheel $efi_dir/Extra/Extensions
	kextcache -a i386 -a x86_64 -m $efi_dir/Extra/Extensions.mkext $efi_dir/Extra/Extensions > $efi_dir/Extra/update.log 2>&1
	chmod 755 $efi_dir/Extra/Extensions.mkext
	chown root:wheel $efi_dir/Extra/Extensions.mkext

	# No news is good news
	echo "";
}

function install
{
	# Strip trailing slash!
	f=${1%/};
	name=${f##*/};

	if [ ! -e "$f" ]; then
		echo "Aborting: kext $f not found.";
		exit;
	fi

	mount_efi
	cp -R "$f" $efi_dir/Extra/Extensions/

	read -p "Installed $name to $efi_dir/Extensions. Would you like to rebuild the cache (y/n)? "
	[ "$REPLY" == "y" ] && rebuild;
}

if [ "$1" == "init" ]; then
	if [ -z "$4" ]; then
		echo "Please specify boot0, boot1h and boot files.";
		exit;
	fi
	echo "Initialization will erase all of $disk_efi. This is unrecoverable."
	read -p "Are you sure you want to proceed? (y/[n])? "
	if [ "$REPLY" == "y" ]; then

		# Why check before it's even initialized?
		#fsck_hfs /dev/$disk_efi


		if [ `echo "$osxversion_major >= 10.6" | bc` == 1 ]; then
			# >= 10.6
			newfs_hfs -v EFI /dev/$disk_efi
		else
			# < 10.6
			diskutil eraseVolume "HFS+" "EFI" /dev/$disk_efi
		fi
		
		# Install boot0 to the MBR
		fdisk -f "$2" -u -y /dev/rdisk${X}

		# Install boot1h to the bootsector of the EFI partiton
		dd if="$3" of=/dev/rdisk${X}s1

		# Check EFI partition for consistency
		fsck_hfs /dev/$disk_efi

		# If good (TODO) mount!
		mount_efi

		# Copy boot file to root of boot root of EFI partition
		cp "$4" /Volumes/EFI/boot

		# Prevent File System Events Daemon (fseventsd) from logging onto the
		# EFI partition (could cause to make unmountable)
		mkdir /Volumes/EFI/.fseventsd/
		touch /Volumes/EFI/.fseventsd/no_log

		# TODO: look up fdisk -e /dev/$disk_efi (f 1, w, q) inline
		mkdir -p /Volumes/EFI/Extra/Extensions
	else
		echo "Aborting..."
	fi
elif [ "$1" == "install" ]; then
	install "$2";
elif [ "$1" == "mount" ]; then
	mount_efi;
	mounted_efi=false;
	if [ "$2" == "open" ]; then
		# Note: Have to do this because for some reason OS X doesn't open it the first time.
		open -g $efi_dir;
		open $efi_dir;
	fi
	exit;
elif [ "$1" == "unmount" ]; then
	unmount_efi;
	exit;
elif [ "$1" == "rebuild" -o "$1" == "update" ]; then
	rebuild;
fi

if [ $mounted_efi == true ]; then
	read -p "Would you like to unmount and remove $efi_dir? (y/[n])? ";
	if [ "$REPLY" == "y" ]; then
		unmount_efi;
	fi
fi