loaders-update

#!/bin/sh

#### Functions & data ####

readonly WMN='shoot-me'

TestRoot()
{
	if [ "$(id -u)" -ne 0 ]; then 
		echo 'Consider to run this utility as root.'
	fi
}

# PossiblyRun CmdToExecute VarUpdatable VarUpdated
PossiblyRun()
{
	if ($SM); then
		unset rep
		echo "About to execute: $1"
		read -p 'Are you sure (y/N)? ' rep
		case "$rep" in
			[yY]*) 	;;
			*) echo 'Nothing has been writted.'
			   eval "$2=\$(($2+1))"
	   	  	   return;;
		esac
		eval "$1"
		if [ $? -eq 0 ]; then
			eval "$3=\$(($3+1))"
		else
			eval "$2=\$(($2+1))"
			Err=$((Err+1))
		fi
	else
		echo "Would run: $1"
		eval "$2=\$(($2+1))"
	fi 
}

# CheckSourceFile SourceFile
CheckSourceFile()
{
	# SL var is global
	SL="$1" 

	if ! [ -e "$SL" ]; then
		echo "Cannot find $SL on your system! Exiting."
		exit 1
	fi
}

# CopyEfiLoader SourceFile WorkingFile
CopyEfiLoader()
{
	local s wf md
	s="$1"
	wf="$2"

	# Checks if it's an executable efi file
	if [ -z "$(file -b "$wf" | grep PE32+)" ]; then
		return
	fi
	# Checks if the file is a FreeBSD loader
	if [ "$(grep -c FreeBSD "$wf")" -gt 0 ]; then			
		md="$(md5 -q "$wf")"
		if  [ "$md" = "$Efimd5" ]; then
			echo "EFI loader $wf is up to date."
		else
			PossiblyRun "cp $s $wf" UpEFI mEFI
		fi
	fi
}

# EfiUpdate SourceFile TargetPartition MountedDest
EfiUpdate()
{
	local s p m mt cmd e
	s="$1"
	p="$2"
	m="$3"
	mt=false

	# If m is empty, the target partition isn't mounted, we must mount it
	if [ -z "$m" ]; then
		m='/mnt'
		cmd="mount -t msdosfs /dev/$p $m"
		echo "$cmd"
		e="$(eval "$cmd" 2>&1)"
		if [ -n "$e" ]; then
			echo $e
			echo "Cannot mount $p, so cannot looking for its loader(s)."
			case "$e" in
				*Invalid*) echo 'Is this partition formatted?'
					   Ferr=$((Ferr+1));;
				*) Err=$((Err+1));;
			esac
			return
		fi
		mt=true
	fi
	lf="$(find "$m" -type f)"
	for f in $lf; do
		CopyEfiLoader "$s" "$f"
	done
	if ($mt); then
		cmd="umount $m"
		echo "$cmd"
		eval "$cmd"
	fi	
}

#### Main ####

echo 'This utility aims to show the automatic update of FreeBSD loaders is possible.'
SM=false
if [ -n "$1" ]; then
	if [ "$1" = "$WMN" ]; then
		SM=true
		if [ "$(id -u)" -ne 0 ]; then 
			echo "You need to be root to operate in $WMN mode."
			return 1
		fi
	else
		echo 'Required mode is unknown.'
		return 1
	fi
fi
if ! ($SM); then
	echo 'Default mode: it tells what it would do but changes nothing.'
else
	echo "$WMN mode selected!"
fi

echo

a="$(sysctl -n hw.machine)"
if ! [ "$a" = "amd64" ]; then
	echo 'This utility works only with amd64 architecture. Exiting.'
	echo
	return 1
fi

# Get the list of disks
LD="$(sysctl -n kern.disks)"
if [ -z "$LD" ]; then
	echo 'No disk has been detected. Exiting.'
	echo
	return 1
fi

# Search for efi & freebsd-boot partitions
GPT=false
for d in $LD; do
	# To avoid an error if the disk is amovible and absent (cdrom)
	gpart show "$d" > /dev/null 2>&1

	if [ $? -eq 0 ]; then
		gp="$(gpart show "$d" | head -n1 | awk '{ print $5 }')"
		if [ "$gp" = "GPT" ]; then
			GPT=true

			# Looking for a efi type partition
			p="$(gpart show -p "$d" | grep efi | awk '{ print $3 }')"
			if [ -n "$p" ]; then
				EFIP="$EFIP $p"
			fi
			unset p

			# Looking for a freebsd-boot type partition
			pi="$(gpart show "$d" | grep freebsd-boot | awk '{ print $3 }')"
			if [ -n "$pi" ]; then
				# BIOSD lists the disks and BIOSI, the corresponding indexes
				BIOSD="$BIOSD $d"
				BIOSI="$BIOSI $pi"
			fi
			unset pi		
		fi
	fi
done

if ! ($GPT); then
	echo 'This machine has no disk with GPT scheme. Exiting.'
	echo
	return 1
fi

Err=0
Ferr=0
UpEFI=0
UpBIOS=0
mEFI=0
mBIOS=0
if [ -n "$EFIP" ]; then
	echo 'One or more efi partition(s) have been found.'
	CheckSourceFile /boot/loader.efi
	Efimd5="$(md5 -q "$SL")"
	echo
	
	for p in $EFIP; do
		# Try to see if the efi partition is already mounted
		# (typically in /boot/efi)

		# Search first by the geom name (e.g. ada0p1)
		m="$(mount -p | grep "$p" | awk '{ print $2 }')"

		# If not found, try by the label name (e.g. efiboot0)
		if [ -z "$m" ]; then
			lp="$(gpart show -pl | grep "$p" | awk '{ print $4 }')"
			m="$(mount -p | grep "$lp" | awk '{ print $2 }')"
		fi
		if [ -n "$m" ]; then
			echo "Efi partition $p is already mounted in $m."
		fi	
		EfiUpdate "$SL" "$p" "$m"
		unset m	
	done
	echo
fi

if [ -n "$BIOSD" ]; then

	echo 'One or more freebsd-boot partition(s) have been found.'

	# Check the root file system
	rfs="$(mount -p | grep "	/	" | awk '{ print $3 }')"
	if [ -n "$rfs" ]; then
		echo "The root file system is $rfs."
		CheckSourceFile /boot/pmbr
		Spmbr="$SL"
		Pmbrmd5="$(head -c 446 "$Spmbr" | md5 -q)"
		if [ "$rfs" = "zfs" ]; then
			CheckSourceFile /boot/gptzfsboot
		else	
			CheckSourceFile /boot/gptboot
		fi
		# SL contains the loader file name suited for the system
		Gptmd5="$(md5 -q "$SL")"
		Lgpt="$(stat -f "%z" "$SL")"
		echo

		i=1
		for d in $BIOSD; do
			echo "Examining $d..."
		
			# Retrieving the corresponding partition index in BIOSI
			index=$(eval "echo \$BIOSI | awk '{ print \$$i }'")
			i=$((i+1))

			# Try to determine whether the partition content is gptboot or gptzfsboot
			p="/dev/${d}p$index"
			r1="$(grep -c ZFS "$p")"
			if [ $? -gt 1 ]; then
				echo "Error during access to $p. Won't update these loaders."
				Err=$((Err+1))
				continue
			fi
			r2="$(grep -c zfs "$p")"
			if [ "$r1" -gt 0 ] && [ "$r2" -gt 0 ]; then
				nfs="zfs"
			else
				nfs="ufs"
			fi
			if ! [ "$rfs" = "$nfs" ]; then
				echo 'There is a mismatch between the root fs and the current loader.'
				echo "Root fs: $rfs / Partition has: $nfs."
				echo "--> No loader update on $p."
				Ferr=$((Ferr+1))
				continue
			fi
			cpmbr="$(head -c 446 /dev/"$d" | md5 -q)"
			bcode="-b $Spmbr "
			if [ "$cpmbr" = "$Pmbrmd5" ]; then
				echo 'The pmbr on this disk is up to date.'
				unset bcode
			fi
			cfbb="$(head -c "$Lgpt" /dev/"${d}"p"$index" | md5 -q)"
			pcode="-p $SL -i $index "
			if [ "$cfbb" = "$Gptmd5" ]; then
				echo "The freebsd-boot partition ${d}p$index is up to date."
				unset pcode
			fi
			if [ -n "$bcode" ] || [ -n "$pcode" ]; then		
				PossiblyRun "gpart bootcode $bcode$pcode$d" UpBIOS mBIOS
			fi
		done
	else
		echo 'Cannot determine the root file system.'
		echo "Won't work on the freebsd-boot partition(s)."
		Err=$((Err+1))
	fi
echo
fi


echo '-------------------------------'
BM=$(sysctl -n machdep.bootmethod)
echo "Your current boot method is $BM."

# Try to figure out partition & efi file on which the system is currently booting (EFI only)
if [ "$BM" = "UEFI" ] && [ "$(id -u)" -eq 0 ]; then
	cb="$(efibootmgr -v | grep +)"
	echo -n 'Boot device: '

	# It may be impossible for efibootmgr to convert to a unix path
	efibootmgr -E > /dev/null 2>&1
	if [ $? -gt 0 ]; then
		# Just print the line reported by efibootmgr
		echo $cb | cut -f3- -d ' '
	else
		# There, we can identify the boot partition
		cbp="$(efibootmgr -E)"
		# Case where the boot partition is a label, convert it to a geom partition
		if [ -n "$(echo $cbp | grep gpt/)" ]; then
			lp="$(echo $cbp | cut -f2 -d /)"
			if [ -n "$lp" ]; then	
				cbp="$(gpart show -lp | grep "$lp" | awk '{ print $3 }')" 
			fi
		fi
		echo "$cbp $(echo $cb |cut -f2 -d /)"
	fi
fi

if [ -z "$EFIP" ] && [ -z "$BIOSD" ]; then
	echo 'Found no efi partition and no freebsd-boot partition.'
	echo 'Nothing seems updatable.'
else
	if [ "$mEFI" -eq 0 ] && [ "$mBIOS" -eq 0 ] && [ "$UpEFI" -eq 0 ] && [ "$UpBIOS" -eq 0 ]; then
		echo 'One or more target partition(s) have been found...'
		if [ "$Err" -gt 0 ]; then
			echo 'But no loader seems to be updatable.'
			echo "$Err error(s) occured during the scan."
			TestRoot
		else
			if [ "$Ferr" -eq 0 ]; then
				echo 'All loaders are up to date.'
			fi
		fi
	else
		if [ "$UpEFI" -gt 0 ]; then
			echo "Updatable EFI loader: $UpEFI"
		fi
		if [ "$UpBIOS" -gt 0 ]; then
			echo "Updatable BIOS loader: $UpBIOS"
		fi
		if [ "$mEFI" -gt 0 ]; then
			echo "Updated EFI loader: $mEFI"
		fi
		if [ "$mBIOS" -gt 0 ]; then
			echo "Updated BIOS loader: $mBIOS"
		fi
		if  [ "$Err" -gt 0 ]; then
			echo "One or more loaders may be updatable, but encountered $Err error(s)."
			TestRoot
		fi
	fi
	if [ "$Ferr" -gt 0 ]; then
		echo "$Ferr serious error(s) occured. See texts above."
	fi
fi

echo '-------------------------------'