loaders-update

#!/bin/sh

# Copyright (c) 2024, Emrion

# 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.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "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 COPYRIGHT HOLDER OR CONTRIBUTORS 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.


#### Functions & data ####

readonly Version="loaders-update v1.2.0"
# Available modes
readonly WMN="shoot-me"
readonly SMN="show-me"
# Loaders file names
readonly EFILOADER="loader.efi"
readonly PMBR="pmbr"
readonly GPTBOOT="gptboot"
readonly GPTZFSBOOT="gptzfsboot"
# Mount point for efi partitions
EMD="/mnt"
# Loaders source dir
SLD="/boot"
# If the root file system must be checked or not (BIOS loaders update)
USE_RFS=true
# If the content of the freebsd-boot partition(s) must be checked or not
USE_FBP=true
# We prefer to use sysctl kern.disk to list the disks
USE_KERNDISK=true


Usage()
{
	local exn
	exn=$(basename "$0")
	echo "Usage: $exn mode [-fgr] [-m efi_mount_dir] [-s loaders_source_dir]"
	echo "mode can be one of:"
	echo "  $SMN: just show the commands to type, change nothing."
	echo "  $WMN: may update the loader(s), but ask for confirmation before each one."
	echo "Options:"
	echo "  -f: won't check the freebsd-boot content for BIOS loaders update."
	echo "  -g: force to use 'gpart show' for disk detection."
	echo "  -r: won't check the root file system for BIOS loaders update."
	exit 1
}

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

# TestCmd cmd
# Warning: the command is executed, not only tested
TestCmd()
{
	local r
	r="$(eval "$1" 2>&1)"
	if [ $? -ne 0 ]; then
		echo "$r"
	fi
}

# PossiblyRun CmdToExecute VarUpdatable VarUpdated
PossiblyRun()
{
	local cmd rep
	local VarUpdatable	# Counts what it can be updated
	local VarUpdated	# Counts what it has been updated

	cmd="$1"
	VarUpdatable="$2"
	VarUpdated="$3"

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

# CheckSourceFile SourceFile
CheckSourceFile()
{
	if ! [ -e "$1" ]; then
		echo "Cannot find $1 on your system! Exiting."
		exit 1
	fi
}

# CopyEfiLoader SourceFile TargetFile
CopyEfiLoader()
{
	local SourceFile TargetFile TFmd5
	# global FBSDL: is a loader found in TargetPartition?
	SourceFile="$1"
	TargetFile="$2"

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

# EfiUpdate SourceFile TargetPartition MountedDest
EfiUpdate()
{
	local SourceFile TargetPartition MountedDest
	local mt cmd e lf f
	# global FBSDL: is a loader found in TargetPartition?

	SourceFile="$1"
	TargetPartition="$2"
	MountedDest="$3"
	mt=false

	# If $MountedDest is empty, the target partition isn't mounted, we must mount it
	if [ -z "$MountedDest" ]; then
		MountedDest="$EMD"
		cmd="mount -t msdosfs /dev/$TargetPartition $MountedDest"
		echo "$cmd"
		e="$(TestCmd "$cmd")"
		if [ -n "$e" ]; then
			echo "$e"
			echo "Cannot mount $TargetPartition, 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	# Don't forget to umount before returning
	fi

	FBSDL=false
	lf="$(find "$MountedDest" -type f)"
	for f in $lf; do
		CopyEfiLoader "$SourceFile" "$f"
	done
	if ! ("$FBSDL"); then
		echo "There is no FreeBSD loader in $TargetPartition"
	fi
	if ($mt); then
		cmd="umount $MountedDest"
		echo "$cmd"
		eval "$cmd"
	fi	
}

# SelectBiosLoader (ufs or zfs)
SelectBiosLoader()
{
	# global BiosLoader BCmd5 LBC	# Path/name, MD5 and size of the selected source loader
	
	if [ "$1" = "zfs" ]; then
		BiosLoader="$Szfs"
		BCmd5="$ZfsMD5"
		LBC="$Lzfs"
	else
		BiosLoader="$Sgpt"
		BCmd5="$GptMD5"
		LBC="$Lgpt"
	fi
}

BiosUpdate()
{
	local Spmbr PmbrMD5		# Path/name & md5 of source pmbr
	# global Szfs ZfsMD5 Lzfs	# Path/name, md5 and size of source gptzfsboot
	# global Sgpt GptMD5 Lgpt	# Path/name, md5 and size of source gptboot
	# global BiosLoader BCmd5 LBC	# Path/name, MD5 and size of the selected source loader
	local d index p			# Disk, index and partition currenlty examined
	local rBTX rZFS rzfs		# Results / offset of strings search
	local eof 			# End of file offset in the freebsd-boot partition
	local pfs			# File system managed by the loader inside the partition
	# global rfs			# Root file system
	local CpmbrMD5 CfbbMD5		# MD5 of the current disk pmbr & partition content
	local bcode pcode i

	# Check the source pmbr 
	Spmbr="$SLD/$PMBR" # Path/name of the source loader
	CheckSourceFile "$Spmbr"
	PmbrMD5="$(head -c 446 "$Spmbr" | md5 -q)" # MD5 computation
	# Check the sources BIOS loaders gptzfsboot & gptboot
	Szfs="$SLD/$GPTZFSBOOT"
	CheckSourceFile "$Szfs"
	ZfsMD5="$(md5 -q "$Szfs")"
	Lzfs="$(stat -f "%z" "$Szfs")" # Size in bytes of the file
	Sgpt="$SLD/$GPTBOOT"
	CheckSourceFile "$Sgpt"
	GptMD5="$(md5 -q "$Sgpt")"
	Lgpt="$(stat -f "%z" "$Sgpt")"
	
	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))
		p="/dev/${d}p$index"

		# Is there a BTX loader in this partition?
		rBTX="$(grep -c 'BTX halted' "$p")"
		# Check first if the partition is readable
		if [ $? -gt 1 ]; then
			echo "Error during access to $p. Won't update these loaders."
			Err=$((Err+1))
			continue
		fi

		if ("$USE_FBP"); then
			if [ "$rBTX" -eq 0 ]; then
				echo "There is currently no loader in $p."
				echo "Cannot update without option -f."
				Ferr=$((Ferr+1))
				continue
			else
				# Try to determine whether the partition content is gptboot or gptzfsboot
				# $rZFS is set at the offset where the first string 'ZFS' has been detected
				rZFS="$(grep -abo ZFS "$p" | head -n1 | cut -f1 -d ':')"
				rzfs="$(grep -c zfs "$p")" # Detection of the string 'zfs'
				# $eof is near the end of the last file copied into this partition (sequence of 24 NUL bytes)
				eof=$(cat "$p" | hexdump -v -e '/1 "%02x"' | grep -bo -E '0{48}' | head -n1 | cut -f1 -d ':')
				if [ -z "$eof" ]; then
					echo "Cannot find eof in $p. It won't be analyzed."
					echo "Use option -f if you want to update it."
					Ferr=$((Ferr+1))
					continue
				fi
				eof=$((eof/2)) # Each byte is 2 characters in hexadecimal

				# The comparison of $rZFS versus $eof allows to detect a special case
				# where a partition has been filled with gptzfsboot, then filled with gptboot
				# in a second time. As the size of gptboot is << to the one of gptzfsboot,
				# the strings 'ZFS' and 'zfs' remain in the partition and would be falsely detected
				# as gptzfsboot otherwise.
				if [ -n "$rZFS" ] && [ "$rzfs" -gt 0 ] && [ "$rZFS" -lt "$eof" ]; then
					pfs="zfs" # gptzfsboot	
				else	
					pfs="ufs" # gptboot	
				fi
				SelectBiosLoader "$pfs"
			fi
		fi
		if ($USE_RFS) && ($USE_FBP) && ! [ "$rfs" = "$pfs" ]; then
			echo "There is a mismatch between the root fs and the current loader."
			echo "Root fs: $rfs / Partition has: $pfs."
			echo "--> No loader update on $p."
			Ferr=$((Ferr+1))
			continue
		fi
		if ! ($USE_FBP) && ($USE_RFS); then
			SelectBiosLoader "$rfs"
		fi	
		CpmbrMD5="$(head -c 446 /dev/"$d" | md5 -q)" # MD5 of current disk pmbr
		bcode="-b $Spmbr "
		if [ "$CpmbrMD5" = "$PmbrMD5" ]; then
			echo "The pmbr on this disk is up-to-date."
			unset bcode
		fi
		# MD5 of current freebsd-boot partition content
		CfbbMD5="$(head -c "$LBC" /dev/"${d}"p"$index" | md5 -q)" 
		pcode="-p $BiosLoader -i $index "
		if ! ("$USE_RFS"); then
			echo "Loader in ${d}p$index is $(basename "$BiosLoader")."
		fi
		if [ "$CfbbMD5" = "$BCmd5" ]; 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
}

#### Main ####

# To avoid some errors related to grep
unset GREP_OPTIONS

echo "$Version"

case "$1" in
	"$WMN") SM=true
		if [ "$(id -u)" -ne 0 ]; then 
			echo "You need to be root to operate in $WMN mode."
			return 1
		fi;;
	"$SMN")  SM=false;;
	*) Usage;;
esac 

shift
while getopts "m:s:rgf" opt; do
	case ${opt} in
		m) EMD="${OPTARG}";;
		s) SLD="${OPTARG}";;
		r) USE_RFS=false;;
		g) USE_KERNDISK=false;;
		f) USE_FBP=false;;
		*) Usage;;
	esac
done

if ! ($USE_RFS) && ! ($USE_FBP); then
	echo "Options -r and -f are mutually exclusive. Exiting."
	return 1
fi
if ! [ -d "$EMD" ]; then
	echo "EFI mount point doesn't exist: $EMD"
	return 1
fi
if ! [ -d "$SLD" ]; then
	echo "Loaders source dir doesn't exist: $SLD"
	return 1
fi

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

# Check some cmd to detect a problem with hardened systems
em=$(TestCmd "mount -p")
egs=$(TestCmd "gpart show")

# Get the list of disks
if ($USE_KERNDISK); then 
	LD="$(sysctl -n kern.disks)"
	if [ $? -ne 0 ]; then
		echo "Cannot get the list of disks."
		echo "'sysctl kern.disks' doesn't work. Exiting."
		echo
		return 1
	fi
else
	if [ -z "$egs" ]; then
		LD="$(gpart show | grep GPT | cut -wf4)"
	else
		echo "Cannot get the list of disks."
		echo "'gpart show' doesn't work. Exiting."
		echo "$egs"
		echo
		return 1
	fi
fi
if [ -z "$LD" ]; then
	echo "No disk has been detected. Exiting."
	echo
	return 1
fi

echo

# 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
			# 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	
		fi
	fi
done

if ! ($GPT); then
	echo "This machine seems to have no disk with GPT scheme. Exiting."
	echo
	return 1
fi

Err=0      # Errors that may be resolved by running this script as root
Ferr=0     # Other errors
UpEFI=0    # Updatable EFI loaders
UpBIOS=0   # Updatable BIOS loaders
mEFI=0     # Updated EFI loaders
mBIOS=0    # Updated BIOS loaders

if [ -n "$EFIP" ]; then
	echo "One or more efi partition(s) have been found."
	CheckSourceFile "$SLD/$EFILOADER"
	Efimd5="$(md5 -q "$SLD/$EFILOADER")"
	echo
	if [ -n "$em" ]; then
			Err=$((Err+1))
			echo "Cannot work on EFI loaders because 'mount -p' doesn't work."
			echo "-> $em"
	else
		for p in $EFIP; do
			echo "Examining $p..."

			# Try to see if the efi partition is already mounted
			# 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
				if [ -n "$egs" ]; then
					echo "$egs"
					Err=$((Err+1))
					continue
				fi 
 				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"
		done
	fi
echo
fi

if [ -n "$BIOSD" ]; then
	echo "One or more freebsd-boot partition(s) have been found."
	if ! ($USE_FBP); then
		echo "The content of the freebsd-boot partition(s) won't be analyzed."
	fi
	# Check the root file system
	if ($USE_RFS); then
		if [ -z "$em" ]; then
			rfs="$(mount -p | grep "\s/\s" | awk '{ print $3 }')"
			if [ -n "$rfs" ]; then
				echo "The root file system is $rfs."
				BiosUpdate
			else
				echo "Cannot determine the root file system."
				echo "Won't work on the freebsd-boot partition(s)."
				Err=$((Err+1))
			fi
		else
			echo "The root fs must be checked, but 'mount -p' doesn't work."
			echo "-> $em"
			Err=$((Err+1))
		fi
	else
		echo "The root file system won't be checked."
		BiosUpdate
	fi
echo
fi

echo "-------------------------------"
BM=$(sysctl -n machdep.bootmethod)
if [ $? -ne 0 ]; then
	echo "Cannot get the boot method."	
else
	echo "Your current boot method is $BM."

	# Try to figure out partition & efi file on which the system is currently booting
	if [ "$BM" = "UEFI" ] && [ "$(id -u)" -eq 0 ]; then
		eb="$(TestCmd efibootmgr)"
		if [ -z "$eb" ]; 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" ] && [ -z "$egs" ]; then	
						cbp="$(gpart show -lp | grep "$lp" | awk '{ print $3 }')" 
					fi
				fi
				echo "$cbp $(echo "$cb" |cut -f2 -d /)"
			fi
		else
			echo "Cannot use efibootmgr."
			echo "-> $eb"
		fi
	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 "-------------------------------"