#!/bin/sh

## defaults:
FIXIT=0			## default to no repair
FORCEDISK=0		## default to no-force
WIPEDISK=0		## default to no-wipe of disks


usage() {
	echo "Usage: $0 [-f] [drive1 [drive2 ...]]"
	echo " Verify boot status for drives (or all if no args)."
	echo "  -f          fixes any found wrong"
	echo "  --wipe      will wipe all disk boot-info (and disklabels)"
	echo "  --force     will blindly use specified drive arg(s) as device file(s)"
	echo " drives can be specified as:"
	echo "  /dev/daN, daN, portN, pN, unitN, uN, lnumN, lN, bayN, bayXY, N[bay], XY[bay]"
	exit
}

if [ "$1" = "-h" -o "$1" = "-H" ]; then
	usage
fi


## arg1 is filename to create
## arg2 is disk/slice/partition designator
## arg3 is extra dd flags
getregion() {
	dd if=$2 of=$1 $3 >/dev/null 2>&1
	if [ "$?" -ne 0 ]; then
		#dd errored... redo with outpu
		dd if=$2 of=$1 bs=512 $3
		exit
	fi
}

## arg1 is the disk to be fixed on error
## arg2 is source filename
## arg3 is template filename
## arg4 is extra cmp flags
## arg5 is description for diff
diffregion() {
	cmp $2 $3 $4
	if [ "$?" -ne 0 ]; then
		#cmp errored...
		echo "  *** $5 is damaged or out of date"

		if [ "$FIXIT" -ne 0 ]; then
			#fix MBR
			fdisk -IB $1
			#fix boot1/boot2
			disklabel -B ${1}s1

### this was needed to fix a dongle:  disklabel -wrB ad0s1 auto
### that's still the WRONG disklabel
### isi_dongle_sync worked though (failed to sync on first run tho)
		fi
	else
		echo "      $5 is OK"
	fi
}

## arg1 is device to check
checkdisk() {
    disk="$1"

    ## so many files... do them in a subdir.
    tmp="/tmp/chkdisk.$$"
    mkdir $tmp

    echo "Checking disk $disk..."

    ## sector zero is mbr and fdisk partition table, at offsets 0x1be-0x1fd)
    getregion   $tmp/mbr    "$disk"     "bs=512 count=1"

    ### check mbr, before fdisk partition table
    dd if=$tmp/mbr  of=$tmp/diskmbr_start bs=0x1be count=1 >/dev/null 2>&1
    dd if=/boot/mbr of=$tmp/bootmbr_start bs=0x1be count=1 >/dev/null 2>&1
    diffregion $disk $tmp/diskmbr_start $tmp/bootmbr_start ""  "mbr start"

    ### check partition table (partition 0)
    ### sys/sys/disklabel.h has the struct_dos_partition for IBMPC
    ### (beware ifdefs), showing 16 bytes.
    ### The first 5 are flags, starting {head,sector,cylinder}, type.
    ###   Fixed at 0x80 0x01 0x01 0x00 0xa5
    ### The next 3 are end {head,sector,cylinder}.  Depend on disk size.
    ### The next 4 are starting LBA sector.
    ###   Fixed at 0x3f 0x00 0x00 0x00  (0x0000003f with endianness)
    ### The last 4 are partition size in sectors.  Depend on disk size.

    ## sh's echo has a "-e" that lets you use c-type escape sequences,
    ## and "-n" to suppress final newline.  But -n doesn't work :<!
    dd if=$tmp/mbr of=$tmp/diskmbr_pt1 iseek=0x1be bs=1 count=5 >/dev/null 2>&1
    echo -e "\0200\0001\0001\0000\0245" | tr -d "\n" > $tmp/bootmbr_pt1
    diffregion $disk $tmp/diskmbr_pt1 $tmp/bootmbr_pt1  ""     "fdisk pt1"

    dd if=$tmp/mbr of=$tmp/diskmbr_pt2 iseek=0x1c6 bs=1 count=4 >/dev/null 2>&1
    echo -e "\0077\0000\0000\0000"      | tr -d "\n" > $tmp/bootmbr_pt2
    diffregion $disk $tmp/diskmbr_pt2 $tmp/bootmbr_pt2  ""     "fdisk pt2"

    ### check mbr, after (and end of) fdisk partition table
    ### we don't use partitions 1-3, so they should remain zero
    diffregion $disk $tmp/mbr   /boot/mbr    "0x1ce 0x1ce"     "mbr end  "

    ### check boot1 (slice s1, block 0)
    getregion   $tmp/boot1 "${disk}s1"  "bs=512 count=1"
    diffregion $disk $tmp/boot1 /boot/boot1  ""                "boot1    "

    ### check disklabel (slice s1, block 1)
    # nothing to check for now

    ### check boot2 (slice s1, blocks 2-15)
    getregion   $tmp/boot2 "${disk}s1"  "bs=512 iseek=2 count=14"        
    diffregion $disk $tmp/boot2 /boot/boot2  "0 512"           "boot2    "

    ### cleanup
    rm -rf $tmp
}

## arg1 is device to clear
wipedisk() {
	disk="$1"
	diskname=`basename "$disk"`
	if gpart show $diskname 2>&1 | grep -qi gpt; then
	    /usr/sbin/isi_wipe_disk $diskname
	else
	    #79 = 63 blocks of fdisk + 16 blocks of boot1/disklabel/boot2
	    dd if=/dev/zero of=${disk} bs=512 count=79
	fi
}

if [ "`/usr/sbin/isi_bootdisk_type --future`" != imdd ]; then
    echo "$0 is for IMDD systems only." >&2
    exit 1
fi

if [ "$#" -ne 0 -a "$1" = "-f" ]; then
	FIXIT=1
	shift
fi
if [ "$#" -ne 0 -a "$1" = "--wipe" ]; then
	WIPEDISK=1
	shift
fi
if [ "$#" -ne 0 -a "$1" = "--force" ]; then
	FORCEDISK=1
	shift
fi

if [ "$#" -eq 0 ]; then
    ## no args -- test all drives
	bay=1
	maxbay=`isi_drivenum -Mb`

	if [ "$WIPEDISK" -ne 0 ]; then
		echo "Clearing bays $bay through $maxbay"
	else
		echo "Checking bays $bay through $maxbay and dongle"
	fi

	while [ $bay -le $maxbay ]; do
		DEV=`isi_drivenum -t -- $bay`
		if [ "$DEV" = "N/A" ]; then
			echo "Bay $bay is empty.  Nothing to check."
		else
			if [ "$WIPEDISK" -ne 0 ]; then
				wipedisk $DEV
			else
				checkdisk $DEV
			fi
		fi

		# and on to the next bay
		bay=$(($bay+1))
	done

	if [ "$WIPEDISK" -eq 0 ]; then
    ## and the dongle
		if [ -e /dev/ad0 -a "`isi_drivenum -b /dev/ad0`" = "N/A" ];then
			checkdisk /dev/ad0
		else
			echo "Dongle not found."
		fi
	fi

else
    ## some args... test those specific drives
	while [ "$#" -ne 0 ]; do
		if [ $FORCEDISK -eq 1 ]; then
			DEV=$1
			# now at least test that that exists, and error/exit
			# below
			test -e $DEV
		else
			DEV=`isi_drivenum -t -- $1`
			# if that failed, we'll error/exit below
		fi

		if [ "$?" -ne 0 ]; then
			# if either the forcedisk version doesn't exist, or
			# the standard drivenum version failed, error
			echo "No such drive \"$1\"."
			echo ""
			usage
		fi

		if [ "$WIPEDISK" -ne 0 ]; then
			wipedisk $DEV
		else
			checkdisk $DEV
		fi

		# and on to the next arg
		shift
	done
fi
