#!/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.0.1'
# Available modes
readonly WMN='shoot-me'
readonly SMN='show-me'
# Mount point for efi partitions
EMD="/mnt"
# Loaders source dir
SLD="/boot"

Usage()
{
	local exn
	exn=$(basename "$0")
	echo "Usage: $exn mode [-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."
	exit 1
}

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

# TestCmd cmd
TestCmd()
{
	local r
	r="$(eval "$1" 2>&1)"
	if [ $? -ne 0 ]; then
		echo "$r"
	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="$EMD"
		cmd="mount -t msdosfs /dev/$p $m"
		echo "$cmd"
		e="$(TestCmd "$cmd")"
		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 ####

# To avoid some errors related to grep
unset GREP_OPTIONS

echo "$Version"

SM=false
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") ;;
	*) Usage;;
esac 

shift
while getopts "m:s:" opt; do
	case ${opt} in
		m) EMD="${OPTARG}";;
		s) SLD="${OPTARG}";;
		*) Usage;;
	esac
done

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

if ! ($SM); then
	echo "$SMN mode: it tells what it would do but changes nothing."
else
	echo "$WMN mode selected!"
fi

echo

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

# Get the list of disks
LD="$(sysctl -n kern.disks)"
if [ $? -ne 0 ]; then
	echo "Cannot get the list of disks. Exiting."
	echo
	return 1
fi
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 seems to have no disk with GPT scheme. Exiting.'
	echo
	return 1
fi

Err=0
Ferr=0
UpEFI=0
UpBIOS=0
mEFI=0
mBIOS=0

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

if [ -n "$EFIP" ]; then
	echo 'One or more efi partition(s) have been found.'
	CheckSourceFile "$SLD/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)
		if [ -n "$em" ]; then
			Err=$((Err+1))
			echo "$em"
			break
		fi
 		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
				Err=$((Err+1))
				echo "$egs"
				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"
		unset m	
	done
	echo
fi

if [ -n "$BIOSD" ]; then

	echo 'One or more freebsd-boot partition(s) have been found.'
	
	if [ -n "$em" ]; then
		Err=$((Err+1))
		echo "$em"
	else
		# Check the root file system
		rfs="$(mount -p | grep "\s/\s" | awk '{ print $3 }')"
		if [ -n "$rfs" ]; then
			echo "The root file system is $rfs."
			CheckSourceFile "$SLD/pmbr"
			Spmbr="$SL"
			Pmbrmd5="$(head -c 446 "$Spmbr" | md5 -q)"
			if [ "$rfs" = "zfs" ]; then
				CheckSourceFile "$SLD/gptzfsboot"
			else	
				CheckSourceFile "$SLD//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
	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 (EFI only)
	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 '-------------------------------'

