#!/usr/bin/bash

# Copyright (c) 2020, Mellanox Technologies
# All rights reserved.
#
# 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 OWNER 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.
#
# The views and conclusions contained in the software and documentation are those
# of the authors and should not be interpreted as representing official policies,
# either expressed or implied, of the FreeBSD Project.

set -e

# To debug set run=echo
run=

PROGNAME=$(basename "$0")

usage()
{
    cat <<EOF
Usage: $PROGNAME    [--help]
                    [--bootctl [FILE]]
                    [--capsule [FILE]]

Description:
  --help                : print help.
  --bootctl [FILE]      : update the boot partition via the
                          kernel path. If FILE isn't specified,
                          /lib/firmware/mellanox/boot/default.bfb
                          will be used by default.
  --capsule [FILE]      : update the boot partition via the
                          capsule path. If FILE isn't specified,
                          /lib/firmware/mellanox/boot/capsule/MmcBootCap
                          will be used by default.
  --policy POLICY       : determines the update policy. may be:
                            single - updates the secondary partition and
                                     swaps to it
                            dual   - updates both boot partitions, does not
                                     swap

                          if this flag is not specified, 'single' policy is
                          assumed.
EOF
    exit 0
}

bootctl_mode=
capsule_mode=

PARSED_OPTIONS=$(getopt -n "$PROGNAME" -o h -l help,bootctl,capsule,policy: -- "$@")
eval set -- "$PARSED_OPTIONS"

while true
do
  case $1 in
      -h | --help)
          usage
          ;;
      --bootctl)
          bootctl_mode=1
          shift
          ;;
      --capsule)
          capsule_mode=1
          shift
          ;;
      --policy)
          policy=$2
          shift 2
          ;;
      --)
          shift
          break
          ;;
  esac
done

# Mount efivars sysfs if not already
efivars=/sys/firmware/efi/efivars
test "$(ls -A $efivars)" || mount -t efivarfs efivarfs $efivars

# Check secure boot
secure_boot=
if [ -e "$efivars/PK-8be4df61-93ca-11d2-aa0d-00e098032b8c" ]; then
    secure_boot=1
fi

# Only allow capsule when secure boot is enabled.
if [ -n "$bootctl_mode" ]; then
    if [ -n "$secure_boot" ]; then
        echo "ERROR: need to use capsule when UEFI secure boot PK key is enrolled" >&2
        exit 1
    fi

    if [ -n "$capsule_mode" ]; then
        echo "ERROR: Use either options '--bootctl' or '--capsule'" >&2
        exit 1
    fi
elif [ -z "$capsule_mode" ]; then
    if [ -n "$secure_boot" ]; then
      capsule_mode=1
    else
      bootctl_mode=1
    fi
fi

#
# Fall back to bootctl mode if capsule is not supported, which is
# needed to be compatible with old boot-image versions.
#
os_cap_var=OsIndicationsSupported-8be4df61-93ca-11d2-aa0d-00e098032b8c
os_var=OsIndications-8be4df61-93ca-11d2-aa0d-00e098032b8c
if [ ! -e "${efivars}/${os_cap_var}" ]; then
    capsule_mode=
    bootctl_mode=1
else
    # Check the EFI_OS_INDICATIONS_FILE_CAPSULE_DELIVERY_SUPPORTED bit.
    cp -f ${efivars}/${os_cap_var} /tmp/
    os_cap=$(hexdump -s 4 -n 1  -e '/1 "%d" "\n"' /tmp/${os_cap_var})
    os_cap=$((os_cap & 4))
    if [ $os_cap -eq 0 ]; then
        capsule_mode=
        bootctl_mode=1
    fi
fi

bfb_location=/lib/firmware/mellanox/boot
capsule_location=/lib/firmware/mellanox/boot/capsule

parse_and_set_capsule_versions()
{
    local capsule_file="$1"

    if [ -z "$capsule_file" ]; then
        echo "ERROR: No capsule file specified"
        exit 2
    fi

    if [ ! -f "$capsule_file" ]; then
        echo "ERROR: Capsule file '$capsule_file' not found"
        exit 2
    fi

    # Check if bfver command is available
    if ! command -v bfver >/dev/null 2>&1; then
        echo "Boot image parsing: bfver command not found, skipping"
        return 0
    fi

    # Use bfver to parse the ATF and UEFI version from the capsule file.
    BFVER_OUTPUT=$(bfver -f "$capsule_file")

    local atf_version=""
    local uefi_version=""
    while IFS= read -r line; do
        if [[ $line == *"BlueField ATF version:"* ]]; then
            atf_version=$(echo "$line" | cut -d':' -f3 | xargs)
        elif [[ $line == *"BlueField UEFI version:"* ]]; then
            uefi_version=$(echo "$line" | cut -d':' -f2 | xargs)
        fi
    done <<< "$BFVER_OUTPUT"

    # Check if both versions were successfully parsed
    if [ -n "$atf_version" ] && [ -n "$uefi_version" ]; then
        # Check if bfver command is available
        if ! command -v bfcfg >/dev/null 2>&1; then
            echo "Boot image parsing: bfcfg command not found, skipping"
            return 0
        fi

        # Transfer the ATF and UEFI pending version to UEFI.
        bfcfg --capatfver "$atf_version"
        bfcfg --capuefiver "$uefi_version"

        echo "Boot image detected: ATF pending version: $atf_version, UEFI pending version: $uefi_version"
    fi

    return 0
}

uefi_capsule_update()
{
    # Set capsule update file variable name
    capsule_efi=MmcBootCap
    if [ -f "$capsule_file" ]; then
        capsule_efi=$(basename $capsule_file)
    fi

    # UEFI will read the capsule image from the EFI System Partition.
    # Check whether an EFI System partition is present. If not, then
    # the capsule update is not supported;
    if ! fdisk -l | grep -q "EFI System" >/dev/null; then
        echo "cannot find EFI System Partition" >&2
        exit 2
    fi

    if [ ! -f "$capsule_file" ]; then
        capsule_file=$capsule_location/$capsule_efi
    fi

    if [ ! -f "$capsule_file" ]; then
        echo "Capsule file not found" >&2
        exit 2
    fi

    has_debug_bfb $capsule_file

    # It is possible that multiple EFI System partitions are present;
    # for example, an eMMC partition and/or an NVMe partition.
    # Thus, check whether a given EFI System Partition is mounted.
    # If not mounted, then create the mountpoint and mount the disk
    # partition to the default location in order to copy the update
    # files. It is harmless to copy the files to multiple EFI System
    # partitions.
    device_efi=$(blkid | grep -i -e 'system-boot' -e 'efi' | cut -d':' -f1 | tr '\n' ' ')
    for curr_device_efi in $device_efi
    do
        # Set EFI System partition mountpoint
        mount_efi=/mnt/efi_system_partition

        esp_dir=$(mount | grep $curr_device_efi | cut -f 3 -d' ')
        if [ -z "$esp_dir" ]; then
            if [ ! -d $mount_efi ]; then
                $run mkdir $mount_efi
            fi
            echo "mount $curr_device_efi at $mount_efi"
            $run mount $curr_device_efi $mount_efi
        else
            mount_efi=$esp_dir
        fi

        # Copy the capsule file into the EFI System Partition.
        # "EFI/UpdateCapsule" is the standard directory defined by UEFI spec.
        # Files in this directory will processed as capsules in alphabetical
        # order once the capsule flag is set in the 'OsIndications' variable
        # which is done by the printf command below.
        $run mkdir -p "$mount_efi/EFI/UpdateCapsule" 2>/dev/null
        $run cp $capsule_file "$mount_efi/EFI/UpdateCapsule/${capsule_efi}1"
        if [ "$policy" = "dual" ]; then
            $run cp $capsule_file "$mount_efi/EFI/UpdateCapsule/${capsule_efi}2"
        fi
        $run sync

        # Doing some cleanup, if needed
        if [ -z "$esp_dir"  ]; then
            echo "umount $curr_device_efi"
            $run umount "$mount_efi"
            $run rmdir "$mount_efi"
            $run sync
            # Wait until eMMC internal cache is fully flushed
            $run sleep 3
        fi
    done

    $run printf "\x07\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00" > \
        "${efivars}/${os_var}"
    $run sync

    # Parse and set capsule firmware versions
    parse_and_set_capsule_versions $capsule_file || true

    cat <<EOF

    ***********************************************************************
    ***                                                                 ***
    ***    Reboot the system to process the platform firmware updates   ***
    ***                                                                 ***
    ***********************************************************************

EOF

    # Enable card-reset
    $run /sbin/mlxbf-bootctl -e

    return 0
}

verify_bfb_signature()
{
    in_bfb=$1
    in_ver=$2

    if [ ! "$(which python3 2>/dev/null)" ] ||
            [ ! "$(which openssl 2>/dev/null)" ]; then
        echo "skip BFB signature verification"
        return 0
    fi

    # Read Arm lifecycle, BFB is validated in lifecycle state secure only.
    lc=$(mlxbf-bootctl | grep lifecycle | cut -f2 -d':' | tr -d ' ')
    if [ "$lc" = "Secured(development)" ] || [ "$lc" = "GASecured" ]; then
        # Determine the secure boot key type.
        if [ "$lc" = "Secured(development)" ]; then
            sbkey="development"
        elif [ "$lc" = "GASecured" ]; then
            sbkey="production"
        fi

        # BFB file verification.
        output=$(bfsbverify --bfb $in_bfb --version $in_ver)
        if [ "$(echo $output | grep 'Unsigned BFB file')" ]; then
            echo "unsigned BFB file!"
            exit 2
        elif [ ! "$(echo $output | grep $sbkey)" ]; then
            echo "bad BFB file signature!"
            exit 2
        fi
    else
        echo "Arm life cycle state \"$lc\""
        echo "skip BFB signature verification"
    fi
}

kernel_bootctl_update()
{
    bfb_device=/dev/mmcblk0

    if ! command -v mlxbf-bootctl >/dev/null; then
        echo "mlxbf-bootctl program is not supported" >&2
        exit 2
    fi

    if [ ! -b "$bfb_device" ]; then
        echo "bad block device $bfb_device" >&2
        exit 2
    fi

    if [ -z "$bfb_file" ]; then
        bfb_file=$bfb_location/default.bfb
    fi

    if [ ! -f "$bfb_file" ]; then
        echo "cannot find file $bfb_file" >&2
        exit 2
    fi

    has_debug_bfb $bfb_file

    # Figure out te platform by searching for an ACPI table
    # specific to the BF version (ACPI trick to learn about
    # the platform version) and verify BFB signature,
    # if needed. BFB signature verification is ignored in
    # BlueField-1 devices.
    if [ -e /sys/firmware/acpi/tables ]; then
        if strings /sys/firmware/acpi/tables/SSDT* | grep -q MLNXBF22; then
            # Verify BFB for BlueField-2 device.
            verify_bfb_signature $bfb_file 1
        elif strings /sys/firmware/acpi/tables/SSDT* | grep -q MLNXBF33; then
            # Verify BFB for BlueField-3 device.
            verify_bfb_signature $bfb_file 2
        fi
    else
        echo "warning: cannot find acpi tables!"
        echo "skip BFB signature verification"
    fi

    $run /sbin/mlxbf-bootctl -s -d $bfb_device -b $bfb_file
    if [ "$policy" = "dual" ]; then
        # in dual case, run command again to update both partitions
        $run /sbin/mlxbf-bootctl -s -d $bfb_device -b $bfb_file
    fi
    $run sync

    cat <<EOF

    ***********************************************************************
    ***                                                                 ***
    ***                 Platform firmware updates complete.             ***
    ***                                                                 ***
    ***********************************************************************

EOF

    return 0
}

has_debug_bfb()
{
    # A debug image consists of a regular BFB file which
    # includes an authorization certificate (a.k.a. AuthCert).
    # The debug image is expected to be installed in Secure
    # LifeCycle state.
    #
    # Implement a dumb check to verify whether the
    # authorization certificate is present within the
    # BFB. It is a basic indication that the BFB is
    # intended for debug.

    in_file=$1

    # Read Arm lifecycle, debug image is ignored in lifecycle state non-secure.
    lc=$(mlxbf-bootctl | grep lifecycle | cut -f2 -d':' | tr -d ' ')
    if [ "$lc" != "Secured(development)" ] && [ "$lc" != "GASecured" ]; then
        return 0
    fi

    authcert_cn="Trusted Debug BF Root Authorization Certificate"
    # Add dummy 'tr' command, because 'grep' might fail and cause the script
    # to exit as '-e' is set. 
    authcert="$(strings "$in_file" | grep -m 1 "$authcert_cn" | tr -d ' ')"
    if [ -z "$authcert" ]; then
        return 0
    fi

    # If we reach here, it means that a debug image is
    # going to be installed, thus request user confirmation.
    cat <<EOF

    ***********************************************************************
    ***                                                                 ***
    ***   You are about to install a debug image on Secured platform.   ***
    ***                                                                 ***
    ***********************************************************************
EOF
    read -p "Continue (Y|n)[n]: " -n 1 -r
    echo # (optional) move to a new line
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        echo "Aborted!"
        exit 0
    fi
}

#
# Parse command arguments
#

if [ -n "$bootctl_mode" ] || [ -n "$capsule_mode" ]; then
    bfb_file=
    capsule_file=
    if [ $# -eq 1 ]; then
        if [ ! -f "$1" ]; then
	    echo "cannot find file $1" >&2
	    exit 1
	fi
        bfb_file=$1
	capsule_file=$1
    fi
fi

if [ -n "$policy" ]; then
    if [ "$policy" != "single" ] && [ "$policy" != "dual" ]; then
        echo "ERROR: policy must be either 'single' or 'dual'" >&2
        exit 1
    fi
else # default policy
    policy="single"
fi

if [ -n "$bootctl_mode" ]; then
    if [ $# -ge 2 ]; then
        echo "too many arguments" >&2
        exit 1
    fi

    # Update the boot partitions.
    kernel_bootctl_update; exit $?

elif [ -n "$capsule_mode" ]; then
    if [ $# -ge 2 ]; then
        echo "too many arguments with option '--capsule'" >&2
        exit 1
    fi

    # Initiate UEFI capsule update.
    uefi_capsule_update; exit $?
fi
