#!/bin/bash

###############################################################################
#
# Copyright 2025 NVIDIA Corporation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
# the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
###############################################################################

usage() {
    cat << EOF
    Usage: $(basename $0) <pldm-image-file>
EOF
}
#cleanup
cleanup() {
    # Remove the temp directory.
    rm -rf $TMP_DIR
}

if [ $# -lt 1 ]; then
    echo "Error: $(basename $0) requires 1 argument."
    exit 1
fi

pldm_image_file=$1

if [ ! -e "$pldm_image_file" ]; then
    echo "Error: $pldm_image_file file not found."
    exit 1
fi

TMP_DIR=$(mktemp -d)

if [ ! -d "${TMP_DIR}" ]; then
    echo "Error: TMP_DIR not found."
    exit 1
fi

trap cleanup EXIT INT TERM

# NIC FW & BSP
fwpkg_unpack.py --show_all_metadata --verbose $pldm_image_file | grep "ComponentVersionString\"" | awk 'BEGIN { product[1] = "BlueField NICFW version: "; product[2] = "BlueField BSP   version: ";} {print product[NR] $2}' | tr -d '"' 
if [ $? -ne 0 ]; then
    echo "Error: failed to get NICFW version."
fi

# Unpack
fwpkg_unpack.py --unpack --outdir ${TMP_DIR} $pldm_image_file 
if [ $? -ne 0 ]; then
    echo "Error: failed to unpack $pldm_image_file."
    exit 1
fi

# UEFI
print_capsule_file_vers () {
    # Print versions stored in files.

    # UEFI image header
    pattern="4266021321003005"
    uefi_image_file="$TMP_DIR/temp_uefi_image.bin"

    atf_version=$(strings "$IMAGE_PATH" | grep -m 1 "(\(release\|debug\))")
    if [ -n "$atf_version" ]; then
        echo "BlueField ATF   version: $atf_version"
    fi

    # Search for the UEFI header pattern
    xxd -p "$IMAGE_PATH" | tr -d "\n" | grep -b -o "$pattern" | cut -d: -f1 | while read -r header_offset2; do
        header_offset=$(($header_offset2/2))
        image_offset=$((header_offset + 24))

        header_string=$(xxd -p -s $header_offset -l 24 "$IMAGE_PATH")

        if echo "$header_string" | grep -q "^$pattern"; then
            # Extract the image length (next 4 bytes after the header)
            image_len_hex=$(echo "$header_string" | cut -c 17-24)
            image_len=$(echo "$image_len_hex" | sed 's/\(..\)\(..\)\(..\)\(..\)/\4\3\2\1/')
            image_len_dec=$(printf "%d" "0x${image_len#0*}")

            # Check the image offset and length against the input file length
            input_file_length=$(stat -c%s "$IMAGE_PATH")
            end_position=$((image_offset + image_len_dec))
            if [ "$end_position" -gt "$input_file_length" ]; then
                echo "Error: Image offset + length is greater than input file length"
                exit 1
            fi

            # Extract the UEFI image
            tail -c +$((image_offset + 1)) "$IMAGE_PATH" | head -c $image_len_dec > "$uefi_image_file"

            # Fetch the version from the UEFI image
            gzipped=$(file $uefi_image_file | grep gzip)
            if [ -n "$gzipped" ]; then
                mv $uefi_image_file $uefi_image_file.gz
                gunzip $uefi_image_file.gz
                uefi_version="$(strings -el $uefi_image_file | grep "BlueField" | cut -d':' -f 2)"
            else
                echo "Warning: UEFI image not compressed and no version info"
            fi

            echo "BlueField UEFI  version: $uefi_version"
            break
        fi
    done
}

# ATF & UEFI
pushd . > /dev/null
cd ${TMP_DIR} > /dev/null
for file in ${TMP_DIR}/*; do
    IMAGE_PATH="$file"
    print_capsule_file_vers
    # dump images if it contains them
    mlx-mkbfb -x $file &> /dev/null
done

# BMC & CEC
for file in ${TMP_DIR}/dump*; do
    fwpkg_unpack.py --show_all_metadata --verbose $file  | grep "ComponentVersionString\"" | awk -v substrBMC="ApFw" -v substrCEC="Ecfw" '{ if (index($2, substrBMC)) {print "BlueField BMC   version: " $2} else if (index($2, substrCEC)) { print "BlueField CEC   version: " $2}}' | tr -d '"'

    if [ $? -ne 0 ]; then
        echo "Error: failed to read $file."
        exit 1
    fi
done

popd > /dev/null

exit 0
