#!/bin/bash
set +o posix
shopt -s extglob
#BLURB="lx2160acex7 Firmware flashing tool"
# Note: the BLURB above is consumed by 'pkgtools' menu within the installed OS,
# when you elect to re-run some of the package setup scripts.
#
# Copyright 2023 Stuart Winter, Earth, Milky Way, "".
# All rights reserved.
#
# Redistribution and use of this script, with or without modification, is
# permitted provided that the following conditions are met:
#
# 1. Redistributions of this script must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
#  THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``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 AUTHOR 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.
#
################################################################################
# Script.....: /usr/sbin/bootloader-flash-lx2160acex7
# Called from: /usr/lib/setup/SeTconfig as
#              T_PX/var/lib/pkgtools/setup/setup.12.bootloader-flash-lx2160acex7
#              or from the shell within the installed OS.
# Purpose....: Write firmware to SPI flash of supported Hardware Models.
# Version....: 1.00
# Date.......: 18-Nov-2023
# Authors....: Stuart Winter <mozes@slackware.com>
################################################################################
# References and Credits:
# [1] https://github.com/Wooty-B/LX2K_Guide
#     Contains file naming convention info and other useful information.
################################################################################
# Change log
# 1.00 - 18-Nov-2023
# * Initial version.
################################################################################

# From within the Installer, this script is called from /usr/lib/setup/SeTconfig
#
# Two arguments are provided:
# 1 -- the target prefix (normally /, but ${T_PX} from the bootdisk)
# 2 -- the name of the root device.
# Location of the root file system:
T_PX=${1}

# Temporary space for decompression:
TMPDIR=${T_PX}/var/lib/pkgtools/setup

# Initial sanity checks:
[ $( id -u ) -ne 0 ] && { echo "ERROR: You must be root to use this tool" ; exit 1 ;}

# Sanity check.  This directory should exist within the OS and the Installer.
[ ! -d "${TMPDIR}" ] && { echo "ERROR: pkgtools temporary directory '${TMPDIR}' does not exist." ; exit 1 ;}
# Append our temporary directory to it.  '/var/lib/pkgtools/setup/tmp' is the standard location
# for temporary space for Slackware package setup scripts.
TMPDIR+=/tmp/bootloader-flash-lx2160acex7

# Determine the Hardware Model name:
[ -x ${T_PX}/sbin/slk-hwm-discover ] && { export HWM=$( slk-hwm-discover ) ;} || {
  err_report_dialog 0 \
  "Missing OS component" \
  "\n'slk-hwm-discover' is missing from your system. This is a critical component of the Kernel package.\n" ;}

# Determine archtecture:
ARCH=$( slk-hwm-discover --print-arch )

# We may require slight behaviour changes depending on the context in which
# this script is launched:
[ -f /.installer-version ] && within_installer=Yes

# Location of the firmware:
# $T_PX is where /usr/lib/setup/SeTpartitions stores the mount point for the new OS' root fs:
# On the OS it'll be a null string.
firmwaredir=${T_PX}/usr/share/hwm-bw-lx2160acex7/bootloader-firmware/
firmware_chksums=${firmwaredir}/checksums.xz
[ ! -d "${firmwaredir}" ] && { echo "ERROR: Unable to find firmware directory '${firmwaredir}'" ; exit 1 ;}
[ ! -f "${firmware_chksums}" ] && { echo "ERROR: Unable to find firmware checksum inventory '${firmware_chksums}'" ; exit 1 ;}

#########################################################################################################################
# Functions
#########################################################################################################################

# Clean the temporary directory, if it was used to house a decompressed firmware file:
function cleanup() {
   rm -rf $TMPDIR
}

# Create the temporary holding space and ensure that there's sufficient space
# available to hold the decompressed firmware file:
function preallocate_tmpspace() {
  local tmpdir=$1
  local size=$2
  local exitval
  mkdir -p $tmpdir
  fallocate -l $size $tmpdir/testfile-$size > /dev/null 2>&1
  estat=$?
  # Wipe the temp file but not the directory, as we'll use it.
  rm -f $tmpdir/testfile-$size
  return $estat
}

# Discover the packaged firmware file size:
function fwfilesize_pkg(){
   # Handle LZMA compressed firmware:
   file ${1} | grep -q 'XZ compressed' && {
      # It's compressed:
      xz --robot --list ${1} | grep -E '^file' | awk '{print $5}' ;} || {
      # Uncompressed; report size natively:
      find ${1} -printf "%s" ;}
}

# Decompress firmware:
# This function is used by 'mtd-utils' flashing tool, as it requires an
# uncompressed file.
function decompress_pkg_fw() {
   local fw_basename="${1##*/}" # rk3399_rockpro64-spi-idbloader.img.xz
   local out_file
   # Decompress the firmware to a temporary location:
   # XZ-compressed files:
   [ "${fw_basename##*.}" = "xz" ] && {
      out_file=${TMPDIR}/${fw_basename%.xz} # lop off the file extension
      # Decompress to temporary holding space:
      xz -dc ${1} > ${out_file} ;}
   # Return the path + file name:
   echo ${out_file} ;}

# Generate a checksum of the first 100K of a locally-held (packaged)
# firmware file.  This is used to determine whether the firmware stored in
# SPI == any of those we package:
function chksum_pkgd_fw() {
   local local_fwfile=$1
   local furry estat
   # Handle LZMA compressed firmware:
   file ${local_fwfile} | grep -q 'XZ compressed' && furry="xz"
   local local_fwfile_chksum="$( ${furry}cat ${local_fwfile} | dd status=none bs=4k count=100k iflag=count_bytes 2>/dev/null | md5sum | awk '{print $1}' )"
   estat=$?
   [ ${estat} -gt 0  ] && return ${estat}
   echo "${local_fwfile_chksum}"
   return 0 ;}

#lx2160acex7_2000_700_2400_8_5_2_flexspi_nor_ee5c233.img.xz
# Describe firmware.  This is used in the dialog boxes when selecting
# the Hardware Profile:
function fw_describe() {
   #File format:
   #<SoC speed>_<bus_speed>_<DDR RAM speed>_<SERDES config>_flexspi_nor_
   #SERDES = Fibre port configuration
   local soc_speed bus_speed ram_speed serdes_conf
   IFS="_" read -r soc_speed bus_speed ram_speed serdes_conf <<< "$1"
   echo  ${1} \"SoC speed: ${soc_speed}, Bus speed: ${bus_speed}, RAM speed: ${ram_speed}, SERDES conf: ${serdes_conf}\" \\
}

# Dump out the range of Hardware Profiles in a format consumable by 'dialog':
function fw_dialoglist_populate() {
   local hw_profile
   while read hw_profile; do
      fw_describe ${hw_profile}
   done< <( find ${firmwaredir} -name "${SOC}_*_*_*flexspi_nor*img*" -printf "%f\n" | cut -d_ -f2-7 | sort | uniq )
}

# Dump out the available firmware options for the chosen Hardware Configuration Profile:
function fw_avail_dialoglist_populate() {
   local hw_profile=$1
   local fw_file

   while read fw_file; do
     echo \"${fw_file}\" \"${fw_file}\" off \\
   done< <( find ${firmwaredir} -type f -name "${SOC}_${hw_profile}*img*" -printf "%f\n" )
}

# Determine the number of Hardware Profiles:
function hwp_total() {
   find ${firmwaredir} -type f -name "${SOC}_*_*_*flexspi_nor*img*" -printf "%f\n" | cut -d_ -f2-7 | sort | uniq | wc -l
}

# Flash management:
function write_flash() {
   local flash_dev=${1}
   local flash_bin=${2}
   # 'flashcp' takes care of erase & write directly:
   printf "\nFlash device: ${flash_dev}\n"
   printf "Image file..: ${flash_bin##*/}\n\n"
   printf "This will take approximately two (2) minutes\n\n"
   printf "Please wait...\n\n"
   flashcp -vA ${flash_bin} ${flash_dev}
   estat=$?
   if [ $estat -gt 0 ]; then
      read -p "An error occurred.  Press ENTER to continue"
      return $estat
    else
      return 0
   fi
}

function offerretry() {
   dialog \
      --backtitle "Bootware" \
      --title "INSTALL FIRMWARE INTO ONBOARD FLASH (SPI FLASH)" --yesno \
"\nA failure was encountered during the flashing process\n\n
Without the Boot Loader being successfully installed to the onboard
flash of this Hardware Model, it may fail to boot.\n\n

\n\nDo you want to retry? (recommendation is 'Yes')\n \

\n" 15 69
   # If they selected no, exit now.
   return $?
}

function err_report_dialog() {
  local exitcode=$1
  local msg_title="$2"
  local msg="$3"
  dialog \
      --backtitle "Error" \
      --title "${msg_title}" \
      --ok-button "OK" \
      --msgbox "${3}" 0 0
  clear
  exit ${exitcode} ;}

function offerready() {
   dialog \
      --backtitle "Bootware" \
      --title "INSTALL FIRMWARE INTO ONBOARD FLASH (SPI FLASH)" \
      --yes-label "FLASH" \
      --no-label "ABORT" \
      --yesno \
"\nYour machine must remain powered up during the process.\n
\nPress ABORT to abort without flashing.\n\n" 0 0
   return $?
}

function offersuccess() {
   dialog \
      --backtitle "Bootware" \
      --title "INSTALL FIRMWARE INTO ONBOARD FLASH (SPI FLASH)" --ok-button "OK" \
      --msgbox "\n\nFlashing was successful.\n" 9 60
   clear
   return 0
}

# Ensure that the flash device exists and we can read it.
# An md5 sum of the selected size will be returned.  This is used for Hardware Models
# where we ship a variety of firmware to support physical hardware configuration profiles.
# The MD5sum of the first X bytes is compared to a store of all available firmware files
# in order to determine the hardware profile.
function spi_flash_test () {
   local flash_dev=$1
   local flash_md5_readsize=$2
   local flash_md5 estat
   [ ! -c ${flash_dev} ] && return 3
   flash_md5=$( dd status=none bs=4k count=${flash_md5_readsize} iflag=count_bytes if=${flash_dev} 2>/dev/null | md5sum | awk '{print $1}' )
   estat=$?
   [ ${estat} -gt 0 ] && return ${estat}
   echo "${flash_md5}"
   return 0
}

# Generate a checksum of the first 100K of the SPI flash.
# This checksum will be used to determine the Hardware Configuration Profile:
function spi_chksum() {
   local spi_chksum=$( spi_flash_test ${BOARDFLASHDEV} 100k )
   estat=$?
   [ ${estat} -gt 0 ] && return ${estat}
   echo "${spi_chksum}"
   return 0
}

# Read the first 100k from the first SD card we
# find.  Most (at least  the Hardware Models we support)
# only have one SD card slot.  This code also handles where
# eMMC and SD devices change address between reboots.
function sd_chksum() {
  local device sd_md5 estat
#  local sd_md5_readsize=$1
  local sd_md5_readsize=100k
  grep -qE 'mmcblk[0-9]' /proc/partitions && {
#    lsblk -o name,type -Mripnd $( grep -E 'mmcblk[0-9]' /proc/partitions | awk '$2 == 0 {print "/dev/" $4}' ) | grep -E 'disk$' | awk '{print $1}' | \
    grep -E 'mmcblk[0-9]' /proc/partitions | awk '$2 == 0 {print "/dev/" $4}' | \
     while read device ; do
       [ "$( < /sys/block/${device##*/}/device/type )" = "SD" ] && {
          sd_md5=$( dd status=none bs=4k count=${sd_md5_readsize} iflag=count_bytes if=${device} 2>/dev/null | md5sum | awk '{print $1}' )
          estat=$?
          [ ${estat} -gt 0 ] && return ${estat}
          [ ! -z "$2" ] && echo $device || echo "${sd_md5}"
          return 0 ;}
    done ;} || return 3
}

#########################################################################################################################

# Set the variables on a per-device basis.
platform_detected=0
case $HWM in
   "SolidRun CEX7 Platform"*|"SolidRun LX2160A Honeycomb"*)
     platform_detected=1
     BOARDFLASHDEV=/dev/mtd0
     SOC=lx2160acex7 ;;
esac

# Sanity checks:
if [ $platform_detected -ne 1 ]; then
   # Writing a new boot loader to the SPI flash is not a requirement
   # for Slackware ARM / AArch64.
   # Within the Installer we don't output any messages since they are irrelevant and look like an error.
   # Within the OS, we report that the user isn't running this tool on the correct Hardware Model.
   [ -z "${within_installer}" ] && { echo "This boot loader flashing tool does not support this Hardware Model." ;}
   # Exit with 0 status to avoid any breakages within the Installer:
   exit 0
fi

# Ensure the flash device is present and that we can read it:
spi_flash_test ${BOARDFLASHDEV} 1k > /dev/null || \
   { echo "$0 : ERROR - unaddressable flash device ${BOARDFLASHDEV}" ; exit 1 ;}

# Scan the storage devices that may contain firmware, searching for
# a known firmware.  This informs our Hardware Profile Configuration:
for scan_dev in spi_chksum sd_chksum ; do
  chksum=$( eval ${scan_dev} )
  estat=$?
#  echo "Checksum for ${scan_dev}: ${chksum}"
  if [ ${estat} -eq 0 ]; then
     fw_filename=$( xz -dc ${firmware_chksums} | grep -E ":${chksum}$" | awk -F: '{print $1}' )
     [ ! -z "${fw_filename}" ] && break
  fi
done

# If we found a match, populate variables with the Hardware Configuration Profile, extracted from
# elements of the determined file name:
if [ ! -z "${fw_filename}" ]; then
   # Hardware profile name: e.g. "2000_700_2400_8_5_2"
   hardware_profile=$( echo ${fw_filename} | cut -d'_' -f2-7)
   IFS="_" read -r soc_speed bus_speed ram_speed serdes_conf <<< "${hardware_profile}"
   # If we have found a Hardware Configuration Profile, we need to determine whether the
   # SPI flash contains one of the packaged (within the OS) versions of the firmware:
   #
   # Generate a checksum of the SPI flash:
   chksum_spi=$( spi_chksum )
   while read pkgd_firmware; do
      # Compare the SPI's checksum to all of the packaged firmware for the discovered
      # Hardware Configuration Profile.  If we matched, store the packaged firmware file name:
      [ "$( chksum_pkgd_fw ${pkgd_firmware} )" = "${chksum_spi}" ] && { discovered_firmware_file=${pkgd_firmware##*/} ; break ;}
   done< <( find ${firmwaredir} -name "${SOC}_${hardware_profile}*" -type f )
fi

# debug/test to test how menus are populated.
#unset hardware_profile
#unset discovered_firmware_file

# If the firmware within the SPI flash matches one of the packaged firmwares, we'll default
# to 'no' for the flashing option.  Users would only want to re-flash if they had some issue
# or just fancied it for l0lz!
[ ! -z "${discovered_firmware_file}" ] && {
  default_flash_opt="--defaultno" # set dialog's default response option to 'No'
  fw_sync_msg="\Zb\Zr\Z2Firmware is up to date - flashing \Zunot required\Z1\Zn" ;} || {
  fw_sync_msg="\Zb\Zr\Z1Firmware update available - flashing recommended\Z1\Zn" ;}

# Even if we haven't found a match between the firmware on the SPI flash and any of
# those that we have packaged (available within the OS), if we _have_ discovered the Hardware
# Configuration Profile, we'll surface the file name from the checksum inventory ($fw_filename)
# rather than the locally packaged one in the OS.
[ ! -z "${hardware_profile}" ] && {
   firmware_info="Installed Firmware..: \Zb\Zr\Z8 ${discovered_firmware_file:-${fw_filename}} \Z1\Zn
\nSoC speed...........: \Zb\Zr\Z8 ${soc_speed} \Z1\Zn
\nBus speed...........: \Zb\Zr\Z8 ${bus_speed} \Z1\Zn
\nRAM speed...........: \Zb\Zr\Z8 ${ram_speed} \Z1\Zn
\nSERDES configuration: \Zb\Zr\Z8 ${serdes_conf} \Z1\Zn" ;} || {
   # We couldn't discover a Hardware Configurartion Profile - God knows how they booted the
   # Slackware or the Installer! ;-)  I guess they use their own firmware rather than one
   # from SolidRun?
   # Let's give them the news.
   fw_sync_msg="\Zb\Zr\Z1 Flashing is required \Z1\Zn"
   firmware_info="\n\n\Zb\Zr\Z1 Unable to detect your Hardware Configuration Profile or discover the firmware \Z1\Zn\n" ;}

# Display current firmware information:
dialog \
   ${default_flash_opt} \
   --colors \
   --backtitle "Bootware" \
   --title "INSTALL FIRMWARE INTO ONBOARD FLASH (SPI FLASH)" --yesno \
"
\nHardware Model      : \Zb\Zr\Z8 ${HWM} \Z1\Zn
\nSystem On Chip      : \Zb\Zr\Z8 ${SOC} \Z1\Zn
\n\n${firmware_info}
\n\n${fw_sync_msg}
\n\nDo you want to flash an available version now?\n

\n" 0 0

# If they selected "no", exit now:
[ $? = 1 ] && { clear ; cleanup ; exit 0 ;}

# Workflow to select a Hardware Configuration Profile, then pick a corresponding firmware
# file.
#

# Select the Hardware Configuration Profile:
#
# If we've discovered the existing Hardware Configuration Profile, we'll default the menu
# choice to this one and surface the existing profile for reference:
[ ! -z "${hardware_profile}" ] && {
   menu_default_hwp="--default-item ${hardware_profile}"
   # Remove the installed firmware info.  We don't need that when we're selecting a profile:
   firmware_info="Your existing profile is:\n $(echo "$firmware_info" | sed '/Installed Firmware/d')" ;} ||
   { # Unset the "unable to detect.." message set earlier for the previous dialog:
      unset firmware_info ;}

# 'tail -n1' is required to handle noise from libgpm when
# running under 'screen' within the Installer.
exec 3>&1
hwp_choice=$( cat << EOF | bash 2>&1 1>&3 | tail -n1
dialog \
   --colors \
   --backtitle "Bootware" \
   --title "SELECT HARDWARE CONFIGURATION PROFILE" \
   --no-tags \
   --no-cancel \
   ${menu_default_hwp} \
   --menu \
"\nThese settings relate to the physical hardware specifications and settings on the main board of the device.
\n${firmware_info}
\n\nChoose the Hardware Configuration Profile:" $(( $( hwp_total ) * 2 + 2 + $( echo "$firmware_info" | grep -o '\\n' | wc -l ) )) 0 0 \
$( fw_dialoglist_populate )
EOF
)
retval=$?
exec 3>&-

#read -p "HWP chosen: $hwp_choice"

# Select the appropriate firmware to flash, based on the Hardware Configuration Profile:
#

# Thought: if there's <2 firmware options available, skip the menu entirely since there's
# no actual choice. Not sure. I like the idea of selecting/knowing it.
while true; do
   exec 3>&1
   fw_choice=$( cat << EOF | bash 2>&1 1>&3 | tail -n1
   dialog \
      --colors \
      --backtitle "Bootware" \
      --title "SELECT FIRMWARE" \
      --no-tags \
      --no-cancel \
      --radiolist \
"\nThe list below contains the available firmware for the chosen Hardware Configuration Profile.\n
Most profiles contain only one option, but others may have more.\n
\n\nChoose your firmware:" 0 0 0 \
$( fw_avail_dialoglist_populate ${hwp_choice} )
EOF
)
   retval=$?
   exec 3>&-
    # Ensure a firmware was selected:
    [ -z "${fw_choice}" ] && { dialog --title "Error" --msgbox "\nPlease select at least one option\n\n" 0 0 ;} || break
done

#read -p "firmware chosen: $fw_choice"

# Construct an absolute path to the firmware file and perform a final sanity check:
fw_flash_bin=${firmwaredir}/${fw_choice}
[ ! -s "${fw_flash_bin}" ] && { echo "ERROR: Cannot locate firmware file ${fw_flash_bin}" ; exit 1 ;}
#read -p "flash file: ${fw_flash_bin}"


# If it's compressed, decompress it ready for flashing:
if file ${fw_flash_bin} | grep -q 'XZ compressed'; then
   # Determine the uncompressed file size:
   fw_flash_bin_size=$( fwfilesize_pkg ${fw_flash_bin} )
   # Allocate and ensure sufficient temporary space:
   preallocate_tmpspace $TMPDIR ${fw_flash_bin_size}
   if [ $? -gt 0 ]; then
      dialog \
       	 --backtitle "Bootware" \
       	 --title "Error" --ok-button "OK" \
       	 --msgbox "\n\nInsufficient temporary storage: ${fw_flash_bin_size} bytes required.\n" 9 60
      cleanup > /dev/null
      exit 1
   fi
   # Decompress, storing the uncompressed version's path+file name:
   fw_flash_bin="$( decompress_pkg_fw ${fw_flash_bin} )"
fi

# Final notice and chance to bail out:
offerready || { clear ; cleanup ; exit 0 ;}

# Loop until either flashing was successful or the user abandons it (after being offered to retry):
while true; do
   clear
   write_flash $BOARDFLASHDEV ${fw_flash_bin}
   estat=$?
   if [ $estat -gt 0 ]; then
      offerretry || { clear ; cleanup ; exit 1 ;} # Failed to flash, user didn't want to retry
    else
      # Flashing was successful, and user acknowledged:
      offersuccess ; cleanup
      break;
   fi
done