#! /bin/bash
#
#: Title       : updateGrub2
#: Date Created: Sat Jan 14 22:34:57 PST 2012
#: Last Edit   : Mon Nov 19 23:48:09 PST 2012
#: Author      : Agnelo de la Crotche (please_try_again)
#: Version     : 2.3.1
#: Description : install grub2, grub2-efi and refresh Grub2 menu
#: Requires    : os-prober
#: Usage       : updategrub2
#
# -----------------------------------------------

# variables you may modify. 

grub_timeout=60

# check the VESA resolutions available to your graphics card with hwinfo --framebuffer
# before changing this value. If you don't specify a video mode here, 'auto' will be used
# - it is often the the best choice. 

grub_gfx=800x600
#grub_gfx=1024x768x16
#grub_gfx=128x1024x16

# menu colors

COLOR_NORMAL="black"
COLOR_HIGHLIGHT="white"

# To use your own background image, put its path in the variable grub_bg
# and uncomment the line below 
# grub_bg=/usr/local/share/images/stellarcart.png

# You shouldn't need to change these values
unifont_url=http://unifoundry.com
unifont=unifont-5.1.20080820.bdf.gz

# -----------------------------------------------

# Do not modify the variables below!!!

# current version
version="2.3"
prg=$(basename $0)
declare -l	NOCONFIRM INSTALL BSD DISPLAYONLY BLID USEBG noconfirm autoimport gdev distcode grub_mkconfig_lib
GSIGN="020"
GRID="grubx64.efi"
W7ID="bootmgfw.efi"
USEBG=yes
G=$(tput setaf 2)
B=$(tput bold)
N=$(tput sgr0)
grub_mkconfig_lib=/usr/lib/grub/grub-mkconfig_lib
ALTERNATE=0

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Grub Bootsector ID at offset 0x80,81
BSaa75="Legacy Grub"
BS5272="Legacy Grub"
BS48b4="Grub 1.96"
BS488="Grub2 s core.img"
BS7c3c="Grub2 1.98"
BS020="Grub2 1.99"
BS8053="Lilo"

# Grub signature offset
declare G377 G383 G385 G384 G392
BS392="Ubuntu"
BS384="Debian"
BS377="Fedora"
BS383="Mandriva/ArchLinux/Debian"
BS384="openSUSE"
BS385="openSUSE"

# Grub2 core offset in MBR (don't know how reliable)
CO3784="1.99" # Ubuntu, openSUSE 12.1
CO3356="2.00" # openSUSE 12.2

# Grub2 Core size
CS61520="Fedora"
CS53352="openSUSE"
CS44324="openSUSE"
CS41452="Ubuntu/Mint"
CS42508="Ubuntu"
CS44192="ArchLinux"

# Grub2 lzma core sign.
# The distro hint is not 100% reliable)
# You might want to remove if appears to be wrong.
Ld1e9dffeffff0000="Grub2 1.99"
Ld1e9dffeffff8db6="openSUSE Grub2 2.00"
Ld1e9dffeffff6690="Fedora Grub2 2.00"
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# Exit if we are not root
[ $(id -u) -gt 0 ] && exec echo "You need to run this script as root"

# We need lsb_release
which lsb_release &> /dev/null || exec cat << EOFLSB
Please install Linux Standard Base Release Tools
using this command: 

- under openSUSE:
  zypper in lsb_release
- under Fedora:
  yum install redhat-lsb  
- under ArchLinux:
  pacman -S lsb-release

EOFLSB
eval $(lsb_release -ircs | awk '{ sub(/SUSE LINUX/,"openSUSE", $0) ; printf "dist=%s;distver=%s;distcode=%s;", $1, $2, $3 }')
BLID=$dist

case $BLID in
opensuse|fedora|archlinux)
	[ "$distver" == "rolling" ] && distnum=0 || distnum=$((${distver/./}*1))
	case $distnum in
		121) bgsuse=upwind ;;
		122) bgsuse=lightrays ; grub_mkconfig_lib=/usr/share/grub2/grub-mkconfig_lib ;;
	esac
	
	# Default distro background image
	opensuse_bg=/usr/share/backgrounds/$bgsuse/morning-1280x1024.jpg
	fedora_bg=/usr/share/backgrounds/$distcode/default/standard/${distcode}.png
	archlinux_bg=/usr/share/archlinux/wallpaper/archlinux-arrival.jpg

	dist_bg=${BLID}_bg ; dist_bg=${!dist_bg} 
	[ "$grub_bg" ] && [ -f "$grub_bg" ] && dist_bg=$grub_bg
	[ -f $dist_bg ] && grub_bg=$dist_bg || unset USEBG
	
	opensuse_noconfirm="--non-interactive"
	fedora_noconfirm="--assumeyes"
	archlinux_noconfirm="--noconfirm"
	dist_noconfirm=${BLID}_noconfirm ; dist_noconfirm=${!dist_noconfirm}
;;
esac

rootdev=$(cat /etc/mtab | awk '/^\/dev/ { if ( $2 == "/" ) print $1 }')
eval $(udevadm info --query=all --name=$rootdev | awk '/UDISKS_PARTITION_SCHEME/ { print $2 }')


function install_osprober {
	case $BLID in
	opensuse) 
		PTA_repo="http://download.opensuse.org/repositories/home:/please_try_again/${dist}_${distver}/"
		zypper lr -u | grep -q ${PTA_repo} || zypper ar $PTA_repo PTA
		zypper -n $autoimport refresh -r PTA
		zypper $noconfirm in -r PTA os-prober
	;;
	fedora)
		yum $noconfirm install os-prober	
	;;
	archlinux)
		pacman -S os-prober
	;;
	esac
}

function install_pkg {
	case $BLID in
	opensuse) 
		zypper $noconfirm install $grub
	;;
	fedora)
		rpm -qa | grep -q "grub-efi" && yum $noconfirm remove grub-efi
		yum $noconfirm install $grub $grub2
	;;
	archlinux)
		pacman -S $grubpkg
	;;
	esac
}

function install_grub {
	if (  which $ginstall &>/dev/null ); then 
		if [ "$EFI" ] ; then
			[ -f /boot/$ESP/$GRID ] || $ginstall --bootloader-id=$BLID --no-floppy
		else
			eval $(udevadm info --query=property --name=$gdev | awk '/DEVTYPE/ { print }')
			case $DEVTYPE in
			disk)
				declare -l YESNO
				read -n 1 -s -p "Are you sure that you want to install Grub2 in MBR of $gdev? [yn] " YESNO
				[ "$YESNO" == "y" ] && $ginstall $gdev
			;;
			partition)
				if (  which $gprobe &>/dev/null ); then 
					[ "$($gprobe -d $gdev)" == "ext2" ] || exec echo "Destination is not an ext4 or ext3 partition" 
					$ginstall --force $gdev
				else
					exec printf "%s not found.\n" $gprobe
				fi
			;;
			esac
		fi
	else
		exec printf "%s not found.\n" $ginstall
	fi
}

function grub_font {
	case $BLID in
	opensuse)
		if [ ${distver/./} -lt 122 ] ; then
			zypper $noconfirm in gnu-unifont
			grub2-mkfont -o $gdir/unicode.pf2 /usr/share/fonts/uni/unifont.pcf.gz
		fi
	;;
	fedora)
		cd /tmp
		[ -f $unifont ] || wget ${unifont_url}/${unifont}
		[ -f $unifont ] || return
		gunzip $unifont
		grub2-mkfont -o $gdir/unicode.pf2 $(basename $unifont .gz)
	;;
	esac	
}

function gfx_menu {
[ -f $gdir/unicode.pf2 ] || grub_font

if [ "$USEBG" == "yes" ] ; then
	splashsrc=$grub_bg ; splash=splash.${splashsrc##*.}
	[ -f $gdir/$splash ] || cp $grub_bg $gdir/$splash
fi

if [ ! -f /etc/default/grub.orig -a -f /etc/default/grub  ] ; then
	echo "- patching /etc/default/grub ..."
	grub_timeout=${grub_timeout:-10}
	cp /etc/default/grub{,.orig}
	grub_gfx=${grub_gfx:-auto}
	grep -q "GRUB_GFXMODE" /etc/default/grub || echo "#GRUB_GFXMODE=640x480" >> /etc/default/grub
	cat << EOFGRUBDEFAULT | sed -i -f - /etc/default/grub
s|GRUB_TIMEOUT=.*|GRUB_TIMEOUT=$grub_timeout|
s|#GRUB_GFXMODE=640x480|GRUB_GFXMODE=$grub_gfx|
/GRUB_GFXMODE/ a\
GRUB_GFXPAYLOAD_LINUX=keep
EOFGRUBDEFAULT
if [ -f $gdir/$splash ] ; then
	cat << EOFGRUBDEFAULT2 | sed -i -f - /etc/default/grub
/GRUB_GFXMODE/ a\
GRUB_BACKGROUND=$gdir/$splash
EOFGRUBDEFAULT2
fi
fi

if [ ! -f /etc/grub.d/05_menu_color ] ; then
	cat > /etc/grub.d/05_menu_color << EOF05MENUCOLOR
#!/bin/sh -e
set -e

prefix=/usr
exec_prefix=/usr

. $grub_mkconfig_lib

COLOR_NORMAL="$COLOR_NORMAL/black"
COLOR_HIGHLIGHT="$COLOR_HIGHLIGHT/black"

if [ "\${GRUB_TERMINAL_OUTPUT}" = "gfxterm" ] ; then
        cat <<EOF
set color_normal=\${COLOR_NORMAL}
set color_highlight=\${COLOR_HIGHLIGHT}
EOF
else
  cat << EOF
set menu_color_normal=white/black
set menu_color_highlight=black/light-gray
EOF
fi
EOF05MENUCOLOR
chmod 0755 /etc/grub.d/05_menu_color
fi
}

function grub2ver {
	tmpdir=${TMPDIR:-/tmp}
	gdef="grub2 1.99/2.00"
	if [ -d $tmpdir ]; then
		core=$tmpdir/core.img
		part=$1
		dsk=${part:0:8}
		if [ -b  $part ]; then
			cpos=$(hexdump -v -s 92 -n 4 -e '"%u"' $part)
			if [ "$cpos" ]; then
				dd if=$dsk of=$core skip=$cpos count=1024 2>/dev/null
        		eval $(hexdump -v -n 10000 -e '1/1 "%02x"' $core | awk '{ lzp=match($0, "d1e9dffeffff" ); if (lzp == "0") { print "lzpos=0" } else { print "lzpos=" ((lzp - 1 ) / 2 ) + 8 "; lzsign=" substr($0,lzp,16) ";" } }')
				grub2sign=L${lzsign} 
				if [ "${!grub2sign}" ]; then
					gver="${!grub2sign}"
					if [ "$lzpos" ] ; then
						if [ $lzpos -gt 1 ]; then
							eval $(hexdump -v -s 520 -n 13 -e '1/4 "m=%u; " 1/4 "k=%u; " 1/4 "csize=%u; " 1 "prt=%d;"' $core)
				 			[ $(hexdump -v -s ${lzpos} -n 8 -e '1/1 "%02x"' $core) = '0000000000000000' ] && lzpos=$(($lzpos + 8))
			      			mod_info=$(($k - $lzpos + 512)) ; csig="CS$(($m + $mod_info))"
						fi
					fi
					[ "${!csig}" ] && gver="${!csig} ${gver}"
				else
					gver="${gdef}"
				fi
			fi
		fi
	fi
	gver=${gver:-$gdef}
	echo $gver
}

function invalidHD {
	str="[37;1m$devmap contains an orphan link: [31;1m"$1"[37;1m. Maybe you removed a hard disk since you installed Grub2. Either delete $devmap or edit the devices in this file to achieve a correct drive mapping. updateGrub2 can not work with an incorrect device.map.[37;0m"
	printf "\n%s\n\n" "$str" | fmt
	exit
}

function efi_chainload {
# look for Linux bootloaders in the ESP on UEFI systems
[ $ALTERNATE -gt 0 ] && return
for GREFI in $(find /boot/efi/EFI -name "$GRID" -o -name "$W7ID" 2>/dev/null | sed "/$BLID/d;s|Microsoft/Boot|windows|") ; do
	UUID=$($gprobe --target=fs_uuid $GREFI)
	t=$(echo ${GREFI%/*} | sed 's|.*/\(.\)\(.*\)|\U\1\E\2|')
	[ "$t" == "Windows" ] && T="$t x86_64 UEFI-GPT" || T="$t Grub"
	grep -q "$T" /etc/grub.d/40_custom && continue
	echo "chainloading $t ..."
	if [ "$DISPLAYONLY" ] ; then
		printf "\nmenuentry \"%s\" {\n    insmod part_gpt\n    insmod fat\n    insmod search_fs_uuid\n    insmod chain\n    search --fs-uuid --no-floppy --set=root %s\n    chainloader (\${root})%s\n}\n" "$T" "$UUID" "${GREFI/\/boot\/efi/}"
	else
		printf "\nmenuentry \"%s\" {\n    insmod part_gpt\n    insmod fat\n    insmod search_fs_uuid\n    insmod chain\n    search --fs-uuid --no-floppy --set=root %s\n    chainloader (\${root})%s\n}\n" "$T" "$UUID" "${GREFI/\/boot\/efi/}" >> $custom
	fi
done
}

function mbr_chainload {
# DEVICE/BIOS MAPPING
HDDEV=($(cat /proc/partitions | awk '!/major|dm|sr/{ gsub(/[0-9]/,"",$4) ; print $4 }' | sort -u))

i=0
while [ $i -lt ${#HDDEV[*]} ] ; do eval ${HDDEV[$i]}=hd$i ; let i++ ; done

# map hard disk devices to BIOS drives according to /boot/grub2/device.map (if this file exists)
if [ -f $devmap ] ; then
	devlinks=($(sed -n '/(hd/s|(\(.*\))[ 	][ 	]*\(.*\)|\2=\1|p' $devmap))
	if [ ${#devlinks[*]} -ge 1 ] ; then
		for dev in ${devlinks[*]} ; do
			sdx=${dev%=*} ; hdn=${dev#*=}
			if [ -b $sdx ] ; then
				[ -h $sdx ] && sdx=$(readlink $sdx)
			else
				invalidHD $sdx
			fi
			eval $(basename $sdx)=$hdn
		done
	fi
fi
devs=$(/sbin/fdisk -l 2>/dev/null | awk '/83/ { print $1 }' | sed -e "$SKIPDEV")
DEVMAP=$(for d in ${HDDEV[*]} ; do printf "s|\\$\\$%s|%s|;" $d ${!d} ; done)

if ( which $gprobe &>/dev/null ) ; then
	ROOTDEV=$($gprobe /boot -t device)
else
	ROOTDEV=$(df -hl / /boot | awk '/^\/dev/ { print $1 }')
fi

for dev in $ROOTDEV ; do
	[ "$(hexdump -v -s 128 -n  2 -e '/1 "%x"' $dev)" == "$GSIGN" ] && SKIPDEV="$dev $SKIPDEV"
done

SKIPDEV=$(echo $SKIPDEV | tr " " "\n" | sort -u | sed 's|/dev\(.*\)|\1\$/d;|')
devs=$(/sbin/fdisk -l 2>/dev/null | awk '/83/ { print $1 }' | sed -e "$SKIPDEV")

for dev in ${HDDEV[*]} $devs ; do
	if [ ${#dev} -eq 3 ] ; then
		CMT="$dev MBR"
		dev="/dev/$dev"
	else
		CMT="${dev##*/}"
	fi
	dd if=$dev bs=512 count=1 2>/dev/null | grep -q GRUB || continue
	D=${dev:0:8}
	if [ "$dev" == "$D" ]; then
		D=${D##*/}; N="${!D}"
	else
		D=${D##*/}; N="${!D},msdos${dev:8}" 
	fi
	BS="BS$(hexdump -v -s 128 -n 2 -e '/1 "%x"' $dev)" 
	GRUB=${!BS}; 
	if [ "$GRUB" == "Legacy Grub" ]; then
		BS=$(dd if=$dev bs=512 count=1 2>/dev/null | tr -c "[:alnum:]" "0")
		POS=${BS%GRUB*}; POS=${#POS}; SIGN=BS$POS ; SIGN=${!SIGN}
		[ "$SIGN" ] && GRUB="(${SIGN}) $GRUB" 
	else
		if [ "$BS" == "BS020" ]; then
			GRUB="$(grub2ver $dev)"
		else
			BS=${!BS} ; GRUB=${BS:-Grub}
		fi
	fi
	Label="${GRUB// /_}_in_${CMT%% *}"
	grep -q "${Label}" $custom && continue
	if [ "$DISPLAYONLY" ] ; then
		printf "\n###Don't change this comment - UpdateGrub2 identifier: Original name: %s###\nmenuentry \"%s in %s\" {\n    set root=(%s)\n    chainloader +1\n}\n" "$Label" "$GRUB" "$CMT" "$N"
	else
		printf "\n###Don't change this comment - UpdateGrub2 identifier: Original name: %s###\nmenuentry \"%s in %s\" {\n    set root=(%s)\n    chainloader +1\n}\n" "$Label" "$GRUB" "$CMT" "$N" >> $custom
	fi
done
}

function alt_linux {
	if [ -f /etc/grub.d/30_os-prober_alt ]; then
		case $1 in
		on)
			if [ -x /etc/grub.d/10_linux ]; then
				cp /etc/grub.d/10_linux{,.org}
				chmod 644 /etc/grub.d/10_linux.org
				sed -i '/gettext_/s/, with Linux/ - kernel/;s/OS="${GRUB_DISTRIBUTOR} Linux"/OS="${GRUB_DISTRIBUTOR}"/' /etc/grub.d/10_linux
			fi
			if [ -x /etc/grub.d/30_os-prober -a -f /etc/grub.d/30_os-prober_alt ] ; then
				chmod 644 /etc/grub.d/30_os-prober
				chmod 755 /etc/grub.d/30_os-prober_alt
			fi
		;;
		off)
			if [ -f /etc/grub.d/10_linux.org ]; then
				mv /etc/grub.d/10_linux{.org,}
				chmod 755 /etc/grub.d/10_linux
			fi
			if [ -x /etc/grub.d/30_os-prober_alt -a -f /etc/grub.d/30_os-prober ] ; then
		 		if [ $ALTERNATE -eq 1 ] ; then
					chmod 755 /etc/grub.d/30_os-prober
					chmod 644 /etc/grub.d/30_os-prober_alt
				fi
			fi
		;;
		esac
	else
		printf "WARNING: /etc/grub.d/30_os-prober_alt not found. Alternate script not used.\n"
	fi
}

function update_menu {
	[ "$mkconfig" ] || exec echo "grub2 not found"
	if [ "$DISPLAYONLY" ] ; then
		echo "Scanning..."
		$mkconfig
		[ "$ESP" ] && efi_chainload || mbr_chainload 
	else
		if [ -d $gdir ] ; then
			[ "$ESP" ] && efi_chainload || mbr_chainload 
			$mkconfig -o $gmenu
			[ "$ESP" ] && cp $gmenu /boot/${ESP}/ 
		else
	  	    echo "Boot directory $gdir not found"
		fi
	fi
}


function restore_elilo {
[ $(efibootmgr | grep -c -i $BLID) -gt 1 ] && return

eliloPath=$(find /boot/efi -name "elilo.efi")
dev=$($gprobe -t device $eliloPath)
disk=${dev:0:8} ; part=${dev:8} 
eliloBldr=$(echo $eliloPath | sed 's|/boot/efi||;s|/|\\\\|g') 
efibootmgr --create --gpt --disk $disk --part $part --write-signature --label  "$DIST (elilo)" --loader "$eliloBldr"
}


function syntax {
cat << EOFHELP

$prg $version  - Grub2 install and boot menu update.

${G}Syntax:${N}  
  ${B}updateGrub2 [ options ]${N}

${G}Options:${N} 
    ${B}-n --noconfirm${N}           Install packages and import keys without asking for confirmation.
                             Requires a version of os-prober which mounts ufs1 and ufs2 correctly
                             and includes os probes for BSDs (like 1.53a in my repo!).
    ${B}-i --install [device]${N}    Install Grub2 on UEFI systems.
                             On BIOS systems you have to specify a device. This option is ignored if
                             the root partition already contains Grub2 signature. Otherwise you will
                             be able to install Grub2 bootloader in any partition, including the MBR. 
                             Use this option with care! 
    ${B}-d --display${N}             Refresh and print the boot menu to standard output.
    ${B}-a --alternate${N}           Use alternate os-prober script.
    ${B}-A --Alternate${N}           Permanently use alternate os-prober script.
    ${B}-h --help${N}                Displays this help.

                             Without options, updateGrub2 scans for other OSes (provided os-prober is  
                             installed) and refreshes the boot entries in grub.cfg. The first time it
                             is run, it also modifies default settings, adds a script to /etc/grub.d, 
                             builds a missing font which prevents graphic mode from working, lets you
                             set a background image or uses the system default one.
EOFHELP
exit
}

# -------------------------------------------------

args=`getopt -q -u -o haAnbdi:: -l help,alternate,Alternate,noconfirm,bsd,display,install:: -- "$@"`
set -- $args

for i; do
	case "$i" in
	-n|--noconfirm) noconfirm="${dist_noconfirm}" ; autoimport="--gpg-auto-import-keys" ; shift ;;
	-i|--install) INSTALL=yes ; shift ;;
	-d|--display) DISPLAYONLY=yes ; shift ;;
	-a|--alternate) ALTERNATE=1 ; shift ;;
	-A|--Alternate) ALTERNATE=2 ; shift ;;
	-h|--help) syntax ; shift ;;
	--) shift ; break ;;
	esac
done

gdev=$1 ; shift

case $UDISKS_PARTITION_SCHEME in
mbr)
EFI=""
ESP=""
grub2=""
[ "$(hexdump -v -s 128 -n  2 -e '/1 "%x"' $rootdev)" == "$GSIGN" ] && unset INSTALL
if [ "x$INSTALL" == "xyes" ] ; then
	[ "$gdev" ] || exec echo "Option --install requires an argument"
	[ -b  $gdev ] || exec echo "Destination is not a block device"
fi
;;
gpt)
EFI="-efi"
ESP="/efi/EFI/$BLID"
DIST="$dist $distver"
grub2="grub2"
;;
esac

case $BLID in 
	opensuse|fedora) grub="grub2$EFI" ;;
	archlinux)
		grub="grub$EFI"
		mac=$(uname -m)
		[ "$mac" == "x86_64" ] || mac=i386
		EFI=${EFI:+${mac}${EFI}}
		BIOS=${EFI:--bios}
		grubpkg="grub$BIOS"
		;;
	*) grub="grub$EFI" ;;
esac
	
mkconfig="${grub}-mkconfig"
ginstall="${grub}-install"
gprobe="${grub}-probe"
gdir="/boot/${grub}"
gmenu="$gdir/grub.cfg"
devmap="$gdir/device.map"
custom=/etc/grub.d/40_custom

if [ "x$INSTALL" == "xyes" ] ; then
	if [ $distnum -ge 122 ] ; then
		printf "updateGrub2 should not be used to install Grub2 on openSUSE >= 12.2\n"
	else
		which os-prober &>/dev/null || install_osprober
		case $BLID in
			opensuse|fedora) rpm -qa | grep -q $grub || install_pkg ;;
			archlinux) pacman -Q | grep -q $grubpkg || install_pkg ;;
		esac
		[ "$INSTALL" ] && install_grub
	fi
fi

[ "$BLID" == "opensuse" -o "$BLID" == "fedora" ] && gfx_menu

if [ $ALTERNATE -gt 0 ] ; then
	alt_linux on
	update_menu
	alt_linux off
else
	update_menu
fi	

