#!/bin/ksh
#
# Script to verify that the correct EFI & GPT boot code is installed on 
# the servers, and potentielly repair broken ones
#
# Version: 1.0 (2018-09-14)
# Author:  Peter Eriksson <peter.x.eriksson@liu.se>

TMPFILE="/tmp/check-bootcode.tmp.$$"

EFIBOOTBLK=/boot/boot1.efifat
GPTBOOTBLK=/boot/gptzfsboot

checksum() {
    dd status=none if="$1" bs=1024 count="$2" | sum | awk '($2 > 0) {print $1}'
}

label2dev() {
    gpart list | awk -v L="$1" '($2 == "Name:") { N=$3 } ($1 == "label:" && $2 == L) { print  "/dev/" N; exit 0 }' | sed -e 's/p[0-9]*$//'
}

dir2zpool() {
    df "$1" | awk -v D="$1" '($6 == D) { print $1 }' | sed -e 's:/.*$::'
}

zpool2devs() {
    zpool status "$1" | awk '($2 == "ONLINE" && $3 >= 0 && $4 >= $0 && $5 >= 0) { print $1 }' | while read DEV; do
	if [ -c "/dev/${DEV}" ]; then
	    case "$DEV" in
		gpt/*)
		    echo "$DEV" | sed -e 's:^gpt/::'
		    ;;
		*)
		    echo "/dev/${DEV}" | sed -e 's/p[0-9]*$//'
		    ;;
	    esac
	fi
    done

}

fileksize() {
    ls -l "$1" | awk '{printf "%u\n", $5/1024}'
}

EFILEN="`fileksize \"$EFIBOOTBLK\"`"
GPTLEN="`fileksize \"$GPTBOOTBLK\"`"

EFISUM="`checksum \"$EFIBOOTBLK\" $EFILEN`"
GPTSUM="`checksum \"$GPTBOOTBLK\" $GPTLEN`"

RC=0

VERBOSE=no
REPAIR=no

MORE=yes
while [ $MORE = yes ]; do
    case "$1" in
	-v|--verbose)
	    VERBOSE=yes
	    ;;
	-r|--repair)
	    REPAIR=yes
	    ;;
	-h|--help)
	    echo "Usage: $0 [-v|--verbose] [-r|--repair] [-h|--help]"
	    exit 0
	    ;;
	-*)
	    echo "$0: Error: $1: Invalid switch" >&2
	    exit 1
	    ;;
	*)
	    MORE=no
	    ;;
    esac
    if [ $MORE = yes ]; then
	shift
    fi
done

if [ "$*" = "" ]; then
    ARGS="/"
else
    ARGS="$*"
fi

DEVS=""
for D in $ARGS; do
    case "$D" in
	/dev/*)
	    DEVS="$DEVS $D"
	    ;;
	/*)
	    ZPOOL="`dir2zpool \"$D\"`"
	    DEVS="$DEVS `zpool2devs $ZPOOL`"
	    ;;
	*[0-9])
	    DEVS="$DEVS `label2dev \"$D\"`"
	    ;;
	*)
	    DEVS="$DEVS `zpool2devs $D`"
	    ;;
    esac
done

XRC=0
for DEV in $DEVS; do
    case "$DEV" in
	/*)
	    ;;
	*[0-9])
	    DEV="`label2dev \"$DEV\"`"
	    ;;
    esac

    RC=0
    gpart show "$DEV" | awk '($3 > 0 && ($4 == "efi" || $4 == "freebsd-boot")) { print $3, $4 }' | while read P T; do
        DEVSUM="-"
	case "$T" in
	    "freebsd-boot")
	        DEVSUM="`checksum \"${DEV}p${P}\" $GPTLEN`"
		if [ "$DEVSUM" != "$GPTSUM" ]; then
		    if [ "$REPAIR" = yes ]; then
			gpart bootcode -b /boot/pmbr -p /boot/gptzfsboot -i $P "$DEV"
                        RC=$?
                        if [ $VERBOSE = yes ]; then
			    echo "${DEV}p${P}: Corrected GPT boot code" >&2
                        fi
		    else
			echo "$0: Error: ${DEV}p${P}: Invalid GPT boot code" >&2
			RC=1
		    fi
		fi
		;;
	    "efi")
	        DEVSUM="`checksum \"${DEV}p${P}\" $EFILEN`"
		if [ "$DEVSUM" != "$EFISUM" ]; then
		    if [ "$REPAIR" = yes ]; then
			gpart bootcode -p /boot/boot1.efifat -i $P "$DEV"
                        RC=$?
                        if [ $VERBOSE = yes ]; then
			    echo "${DEV}p${P}: Corrected EFI boot code" >&2
                        fi
		    else
			echo "$0: Error: ${DEV}p${P}: Invalid EFI boot code" >&2
			RC=1
		    fi
		fi
		;;
	    *)
		echo "$0: Error: $T: Invalid partition type" >&2
		RC=1
                ;;
	esac
	if [ $RC = 0 -a "$VERBOSE" = yes ]; then
	    echo "${DEV}p${P}: OK (type=$T, checksum=$DEVSUM)"
	fi
    done
    if [ $RC != 0 ]; then
	XRC="$RC"
    fi
done

exit $XRC