#!/bin/sh

##########################################################################
# Copyright (c) 2015 Fabian Keil <fk@fabiankeil.de>
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
##########################################################################
#
# image-checksum.sh /path/to/memstick.img
#
# Unfortunately the memstick target currently does not create reproducible
# ElectroBSD images due to unreproducible differences in the file system
# layer of the data partition.
#
# To be able to (sort of) compare memstick images anyway, this script
# produces a "partial image checksum" that is based on the partition layout,
# the checksum of the boot code partition and an mtree spec of the data
# partition which includes checksums, sizes and timestamps for all the
# files.
#
# A memstick image whose "partial checksum" matches the one of another
# image can be totally considered to be nearly as trustworthy. Obviously
# that's a somewhat worthless property, it is thus recommended that you
# rebuild the potentionally malicious image using a trusted operating
# system first. After you've done this, potentionally malicious differences
# in the unchecked parts should be gone.
#
# Just kidding, image-checksum.sh is only intended to regression-test
# the ElectroBSD build system.
#
# Also note that this script relies on non-standardized output of other
# tools which might occasionally change. To be able to reproduce partial
# image checksums you thus need a userland that is close enough to the
# one that was used to create the original version.
#
##########################################################################

UFS_PARTITION=s2a
EXPECTED_PARTITIONS=2
MOUNTPOINT=/mnt
VERBOSE=0
MTREE_KEYWORDS=size,time,mode,uid,gid,sha256,device,flags,link,nlink,type
REUSE_EXISTING_CACHE_FILES=false

verbose_log() {
    local message="$*"
    if [ "${VERBOSE}" = 0 ]; then
        return
    fi
    echo "${message}"
}

create_mtree_spec_file() {
    local md_unit spec_file

    md_unit="${1}"
    spec_file="${2}"

    verbose_log "Mounting /dev/md${md_unit}${UFS_PARTITION} at ${MOUNTPOINT}"
    mount -o ro "/dev/md${md_unit}${UFS_PARTITION}" "${MOUNTPOINT}" || return 1

    verbose_log "Running mtree, saving spec in ${spec_file}"
    mtree -c -k "${MTREE_KEYWORDS}" -p "${MOUNTPOINT}" | mtree -C -k all > "${spec_file}" || return 1

    verbose_log "Unmounting ${MOUNTPOINT} ..."
    umount "${MOUNTPOINT}" || return 1
}

create_ls_file() {
    local md_unit ls_file

    md_unit="${1}"
    ls_file="${2}"

    verbose_log "Mounting /dev/md${md_unit}${UFS_PARTITION} at ${MOUNTPOINT}"
    mount -o ro "/dev/md${md_unit}${UFS_PARTITION}" "${MOUNTPOINT}" || return 1

    verbose_log "Running ls to get inodes, saving output in ${ls_file}"
    env -i LC_ALL=C ls -liR "${MOUNTPOINT}" > "${ls_file}" || return 1

    verbose_log "Unmounting ${MOUNTPOINT} ..."
    umount "${MOUNTPOINT}" || return 1
}

partition_count_acceptable() {
    local md_unit="${1}"

    # Verify that there are exactly two partitions present
    partitions=$(gpart show -r -p "md${md_unit}" | grep -c "md${md_unit}"s)
    if [ "${partitions}" != "${EXPECTED_PARTITIONS}" ]; then
        echo "Invalid number of partitions: ${partitions}"
        return 1;
    fi
}

generate_partial_image_checksum() {
    local image_file \
          md_unit spec_file gpart_file ls_file \
          gpart_checksum bootcode_checksum ls_checksum \
          weak_image_checksum mtree_checksum

    image_file="${1}"
    spec_file="${image_file}.mtree"
    if [ -f "${spec_file}" ]; then
        echo "Spec file ${spec_file} already exists"
        ${REUSE_EXISTING_CACHE_FILES} || return 1
    fi
    gpart_file="${image_file}.gpart"
    if [ -f "${spec_file}" ]; then
        echo "gpart file ${gpart_file} already exists"
        ${REUSE_EXISTING_CACHE_FILES} || return 1
    fi
    ls_file="${image_file}.ls"
    if [ -f "${ls_file}" ]; then
        echo "ls file ${ls_file} already exists"
        ${REUSE_EXISTING_CACHE_FILES} || return 1
    fi

    md_unit=$(mdconfig -o readonly -n -f "${image_file}")
    if [ $? != 0 ]; then
        return 1
    fi

    partition_count_acceptable "${md_unit}" || return 1

    if [ ! -f "${spec_file}" ]; then
        create_mtree_spec_file "${md_unit}" "${spec_file}" || return 1
    fi
    if [ ! -f "${gpart_file}" ]; then
        gpart list "md${md_unit}" | sed -E -e "s@(: md)${md_unit}@\1X@" > "${gpart_file}"
    fi
    if [ ! -f "${ls_file}" ]; then
        create_ls_file "${md_unit}" "${ls_file}" || return 1
    fi

    gpart_checksum=$(sha256 -q "${gpart_file}")
    verbose_log "gpart checksum: ${gpart_checksum}"

    mdconfig -d -u "${md_unit}" || return 1

    bootcode_checksum=$(dd if=/dev/md${md_unit}s1 2>/dev/null | sha256)
    verbose_log "Boot code checksum: ${bootcode_checksum}"
    mtree_checksum=$(sha256 -q "${spec_file}")
    verbose_log "mtree checksum: ${mtree_checksum}"
    ls_checksum=$(sha256 -q "${ls_file}")
    verbose_log "ls checksum: ${ls_checksum}"

    weak_image_checksum=$(echo "${gpart_checksum} ${bootcode_checksum} ${mtree_checksum} ${ls_checksum}" | sha256)
    echo "Partial image checksum for ${image_file}: ${weak_image_checksum}"
}

main() {
    local image_file \
          args

    args=$(getopt m:rv $*)
    if [ $? -ne 0 ]; then
        echo 'You are doing it wrong: Invalid flag specified'
        exit 2
    fi
    set -- ${args}
    while true; do
        case "$1" in
            -m)
                shift
                MTREE_KEYWORDS="${1}"
                shift
                ;;
            -r)
                REUSE_EXISTING_CACHE_FILES=true
                shift
                ;;
            -v)
                VERBOSE=1
                shift
                ;;
            --)
                shift; break
                ;;
        esac
    done

    if [ -z "${1}" ]; then
        echo "No image file provided"
        return 1
    fi

    for image_file in "${@}"; do
        generate_partial_image_checksum "${image_file}" || return 1
    done
}

main "${@}"