#!/bin/bash

# Copyright (C) 2021 CUSTOM
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

CUSTOM_NAME=metel
CUSTOM_PREFIX=metel-
CUSTOM_FW_TYPE=2

VERSION=0.6.0
CUSTOM_EEPROM=${CUSTOM_PREFIX}eeprom
CUSTOM_PARTITION=${CUSTOM_PREFIX}partition
CUSTOM_IMAGE=${CUSTOM_PREFIX}image
ALL_PARAM="$@"
PROGRESS_FILE=/dev/null

########
# HELP #
########

_help ()
{
  #cat $0 | sed -n 's/^#H# //p'
  cat << EOF
CUSTOM custom-firmware, version $VERSION
Usage: custom-firmware [OPTION] filename.tar
Upgrade firmware in this device.

Example:
  custom-firmware firmware.tar
                            # upgrade firmware using file firmware.tar

Options:
  -f, --file=fw.tar         use filename fw.tar
  --skip-tar-script         skip execution of tar script
  -p, --progress=prg.file   echo progress to prg.file

Miscellaneous:
  -V, --version             display version information and exit
  -h, --help                display this help text and exit

Exit status:
  0  if OK,
  1  otherwise.
EOF
  exit 0
}

###########
# VERSION #
###########

_version ()
{
  cat << EOF
${CUSTOM_PREFIX}firmware $VERSION
EOF
  exit 0
}

#############
# FUNCTIONS #
#############

_image_check ()
{
  echo "FIRMWARE IMAGE"
  echo "=============="

  case $FIRMWARE_FILE in
    /*)
      ;;
    *)
      FIRMWARE_FILE=$PWD/$FIRMWARE_FILE
      ;;
  esac

  if [ ! -f "$FIRMWARE_FILE" ]; then
    echo "firmware file not found!"
    exit 1
  fi
  echo "filename:        $FIRMWARE_FILE"

  FIRMWARE_LIST=`tar -tf $FIRMWARE_FILE`
  if [ $? -ne 0 ]; then
    echo "firmware file is not a tar file!"
    exit 1
  fi

  FIRMWARE_FORMAT=`tar -xOf $FIRMWARE_FILE fw_format`
  if [ $? -ne 0 ]; then
    FIRMWARE_FORMAT=1
    echo "firmware file has no fw_format"
  else
    echo "type:            $FIRMWARE_FORMAT"
  fi

  case $FIRMWARE_FORMAT in
    2 | 3)
      if [ "${CUSTOM_FW_TYPE}" != "${FIRMWARE_FORMAT}" ]; then
        echo "invalid firmware format!"
        exit 1
      fi
      ;;

    *)
      echo "unsupported firmware format!"
      exit 1
      ;;
  esac
}

_image_exe ()
{
  EXE_FILE=`echo "$FIRMWARE_LIST" | grep ${CUSTOM_PREFIX}firmware`
  FIRMWARE_EXECUTED=0

  if [ -z "$EXE_FILE" ]; then
    echo "execute:         none"
  else
    if [ "${OPT_SKIP_TAR_SCRIPT}" == 1 ]; then
      echo "execute:         skip"
    else
      FIRMWARE_EXECUTED=1
      echo "execute:         $EXE_FILE"
      RES=`cd /tmp && tar -xf $FIRMWARE_FILE $EXE_FILE`
      chmod 777 /tmp/$EXE_FILE
      echo "executing:   /tmp/$EXE_FILE --skip-tar-script ${ALL_PARAM}"
      echo ""
      /tmp/$EXE_FILE --skip-tar-script ${ALL_PARAM}
      exit $?
    fi
  fi
  echo ""
}

_eeprom ()
{

  echo "EEPROM"
  echo "======"

  # read EEPROM
  EEPROM_LIST=`$CUSTOM_EEPROM`
  if [ $? -ne 0 ]; then
    echo "could not read eeprom!"
    exit 1
  fi

  # parse EEPROM data
  PRODUCT_NUM=`echo "$EEPROM_LIST" | sed -ne 's/product code\:.*0x\(.*\)/\1/p'`
  ASSEMBLY=`echo "$EEPROM_LIST" | sed -ne 's/assembly\: *\(.*_.*\)/\1/p'`
  ASSEMBLY_PCB=`echo "$EEPROM_LIST" | sed -ne 's/assembly-pcb\: *\(.*_.*\)/\1/p'`
  ASSEMBLY_DTB=`echo "$EEPROM_LIST" | sed -ne 's/assembly-dtb\: *\(.*_.*\)/\1/p'`
  if [[ ! -z $ASSEMBLY ]]; then
    ASSEMBLY_PCB=$ASSEMBLY
    ASSEMBLY_DTB=$ASSEMBLY
  fi
  echo "product number:  $PRODUCT_NUM"
  echo "assembly number: $ASSEMBLY_PCB"
  echo "dtb number:      $ASSEMBLY_DTB"
}

_product ()
{
  PRODUCT_NAME="unknown"

  RES=$(cd /tmp && tar -xf $FIRMWARE_FILE product-asm-main.files)
  if [ $? -ne 0 ]; then
    echo "$RES"
    echo "product-asm files could not be extracted!"
    exit 1
  fi

  FILE_MAIN=/tmp/product-asm-main.files

  if [ ! -f $FILE_MAIN ]; then
    echo "could not file $FILE_MAIN"
    exit 1
  fi

  OLDIFS=$IFS
  IFS=";"
  while read iname ipro iasm iother
  do
    # skip comments
    if [[ "$iname" = \#* ]]; then
      continue
    fi

    # skip empty lines
    if [[ "$iname" = "" ]]; then
      continue
    fi

    #echo "$iname, $ipro, $iasm"
    if [[ "$ipro" = "0x$PRODUCT_NUM" ]]; then
      if [[ "$iasm" = "$ASSEMBLY_DTB" ]]; then
        PRODUCT_NAME=$iname
        break
      fi
    fi
  done < $FILE_MAIN
  IFS=$OLDIFS

  PRODUCT_DTB="${PRODUCT_NAME}_${ASSEMBLY_DTB}"
  echo "product dtb:     $PRODUCT_DTB"
  echo ""
}

_unio ()
{

  RES=`cd /tmp && tar -xf $FIRMWARE_FILE product-asm-io.files product-asm-if.files`
  if [ $? -ne 0 ]; then
    echo "$RES"
    echo "product-asm files could not be extracted!"
    exit 1
  fi

  FILE_IO=/tmp/product-asm-io.files
  FILE_IF=/tmp/product-asm-if.files

  # check if a file exists
  FIRMWARE_LIST=`tar -tf $FIRMWARE_FILE`
  if [ $? -ne 0 ]; then
    echo "firmware file is not a tar file!"
    exit 1
  fi

  if [ ! -f $FILE_IO ]; then
    echo "could not file $FILE_IO"
    exit 1
  fi

  # check if a file exists
  if [ ! -f $FILE_IF ]; then
    echo "could not file $FILE_IF"
    exit 1
  fi

  EE=`$CUSTOM_EEPROM --mode=hex -d/sys/kernel/${CUSTOM_PREFIX}unio/if-module/data_string`
  IF_PRO=`echo "$EE" | sed -ne 's/product code\:.*\(0x.*\)/\1/p'`
  IF_ASM=`echo "$EE" | sed -ne 's/assembly-pcb\: *\(.*_.*\)/\1/p'`

  EE=`$CUSTOM_EEPROM --mode=hex -d/sys/kernel/${CUSTOM_PREFIX}unio/io-module/data_string`
  IO_PRO=`echo "$EE" | sed -ne 's/product code\:.*\(0x.*\)/\1/p'`
  IO_ASM=`echo "$EE" | sed -ne 's/assembly-pcb\: *\(.*_.*\)/\1/p'`

  IO_FNAME="io-none"
  IF_FNAME="if-none"
  PRODUCT_DTB1="${PRODUCT_DTB}_${IO_FNAME}_${IF_FNAME}"

  OLDIFS=$IFS
  IFS=";"
  while read iname ipro iasm iother
  do
    # skip comments
    if [[ "$iname" = \#* ]]; then
      continue
    fi

    # skip empty lines
    if [[ "$iname" = "" ]]; then
      continue
    fi

    #echo "$iname, $ipro, $iasm"
    if [[ "$ipro" = "$IO_PRO" ]]; then
      if [[ "$iasm" = "$IO_ASM" ]]; then
        IO_FNAME=$iname
        break
      fi
    fi
  done < $FILE_IO

  while read iname ipro iasm iother
  do
    # skip comments
    if [[ "$iname" = \#* ]]; then
      continue
    fi

    # skip empty lines
    if [[ "$iname" = "" ]]; then
      continue
    fi

    #echo "$iname, $ipro, $iasm"
    if [[ "$ipro" = "$IF_PRO" ]]; then
      if [[ "$iasm" = "$IF_ASM" ]]; then
        IF_FNAME=$iname
        break
      fi
    fi
  done < $FILE_IF
  IFS=$OLDIFS

  echo "IO-MODULE"
  echo "========="
  echo "product number:  $IO_PRO"
  echo "assembly number: $IO_ASM"
  echo "dtb number:      $IO_FNAME"
  echo ""

  echo "IF-MODULE"
  echo "========="
  echo "product number:  $IF_PRO"
  echo "assembly number: $IF_ASM"
  echo "dtb number:      $IF_FNAME"
  echo ""

}

_image_files ()
{
  OPTEXT=
  if [ -f "/etc/signed_kernel_dtb" ]; then
    OPTEXT=.sig
  fi

  echo "IMAGE FILES"
  echo "==========="

  VERIFY_FILE=`echo "$FIRMWARE_LIST" | grep "^verify.txt$"`
  if [ "${FIRMWARE_FORMAT}" = "3" ]; then
    KERNEL_FILE=`echo "$FIRMWARE_LIST" | grep "^kernel${OPTEXT}$"`
  else
    KERNEL_FILE=`echo "$FIRMWARE_LIST" | grep "^zImage${OPTEXT}$"`
  fi
  ROOTFS_FILE=`echo "$FIRMWARE_LIST" | grep "^rootfs.ubi$"`

  if [ "$(cat /dev/null | sha256sum)" != "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855  -" ]; then
    echo "sha256sum verification failed!"
    exit 1
  fi

  VERIFY_SRC=z
  if [ -z "$VERIFY_FILE" ]; then
    if [ "${FIRMWARE_FORMAT}" = "3" ]; then
      echo "image verification file is required and it could not be found!"
      exit 1
    fi
  else
    VERIFY_SRC=`tar -xOf $FIRMWARE_FILE verify.txt | sha256sum`
  fi

  VERIFY_DST=z
  if [ ! -f "/etc/verify.txt" ]; then
    if [ "${FIRMWARE_FORMAT}" = "3" ]; then
      V=$(cat /etc/metel/fw_version)
      if [ "$V" == "12.0.13195" ] && [ -f "/etc/signed_kernel_dtb" ]; then
        VERIFY_DST=${VERIFY_SRC}
      else
        echo "target verification file is required and it could not be found!"
        exit 1
      fi
    fi
  else
    VERIFY_DST=`cat /etc/verify.txt | sha256sum`
  fi

  if [ "${VERIFY_SRC}" != "${VERIFY_DST}" ]; then
    echo "image/target verification failed!"
    exit 1
  fi

  if [ -z "$KERNEL_FILE" ]; then
    echo "kernel file could not be found!"
    exit 1
  fi
  echo "kernel file:     $KERNEL_FILE"

  if [ -z "$ROOTFS_FILE" ]; then
    echo "rootfs file could not be found!"
    exit 1
  fi
  echo "rootfs file:     $ROOTFS_FILE"

  echo "dtb filename1:   ${PRODUCT_DTB1}"
  DTB_FILE1=`echo "$FIRMWARE_LIST" | grep ${PRODUCT_DTB1}.dtb${OPTEXT}`
  if [ -z "$DTB_FILE1" ]; then
    echo "dtb file could not be found!"
    exit 1
  fi
  echo "dtb file1:       $DTB_FILE1"

  PRODUCT_DTB2="${PRODUCT_DTB}_${IO_FNAME}_${IF_FNAME}"
  echo "dtb filename2:   ${PRODUCT_DTB2}"
  DTB_FILE2=`echo "$FIRMWARE_LIST" | grep ${PRODUCT_DTB2}.dtb${OPTEXT}`
  if [ -z "$DTB_FILE2" ]; then
    echo "dtb file could not be found!"
    exit 1
  fi
  echo "dtb file2:       $DTB_FILE2"
  echo ""
}

_partition ()
{
  echo "PARTITION"
  echo "========="

  IMAGE_LIST=`$CUSTOM_PARTITION | grep image`
  if [ $? -ne 0 ]; then
    echo "could not read partitions!"
    echo "please execute: $CUSTOM_PARTITION -c init"
    exit 1
  fi

  IMAGE_COUNT=`echo "$IMAGE_LIST" | wc -l`
  if [ $IMAGE_COUNT -eq 0 ]; then
    echo "no image partition found!"
    echo "please execute: $CUSTOM_PARTITION -c init"
    exit 1
  fi
  echo "total:           $IMAGE_COUNT"

  NUM_IMAGE=`cat /proc/cmdline | sed -e 's/.*'$CUSTOM_NAME'.partition=\([0-9]\).*/\1/'`
  echo "current:         $NUM_IMAGE"


  PART_UPGRADE=0
  PART_VALID=0

  IMAGE_NUMS=`echo "$IMAGE_LIST" | sed -e 's/[^0-9]*\([0-9]\).*/\1/'`
  for p in $IMAGE_NUMS; do

    RES=`$CUSTOM_IMAGE -p $p -c is_valid_boot`
    if [ $? -eq 0 ]; then
      PART_VALID=$p
      echo "partition $p:     valid"
    else
      PART_UPGRADE=$p
      echo "partition $p:     invalid"
    fi

  done

  if [ $PART_UPGRADE -eq 0 ]; then
    echo "image for upgrade is unavailable!"
    exit 1
  fi
  echo "upgrade:         $PART_UPGRADE"
  echo ""
}

_upgrade ()
{
  echo "UPGRADE"
  echo "======="

  echo "0%" >> $PROGRESS_FILE

  RES=`$CUSTOM_IMAGE -p $PART_UPGRADE -c init`
  if [ $? -ne 0 ]; then
    echo "$RES"
    echo "partition $PART_UPGRADE could not be initialized!"
    exit 1
  fi
  echo "partition $PART_UPGRADE has been initialized"

  if [ -f "/etc/signed_kernel_dtb" ]; then
    RES=$(cd /tmp && tar -xf $FIRMWARE_FILE $KERNEL_FILE $ROOTFS_FILE $DTB_FILE1 $DTB_FILE2 rootfs.iv rootfs.sv)
  else
    RES=$(cd /tmp && tar -xf $FIRMWARE_FILE $KERNEL_FILE $ROOTFS_FILE $DTB_FILE1 $DTB_FILE2)
  fi
  if [ $? -ne 0 ]; then
    echo "$RES"
    echo "firmware could not be extracted!"
    exit 1
  fi
  echo "firmware has been extracted"

  echo "5%" >> $PROGRESS_FILE

  # specific dtb
  FILE=/tmp/$DTB_FILE2
  echo "writing dtb $FILE"
  RES=`$CUSTOM_IMAGE -p $PART_UPGRADE -c write_file_image -r 5 -f $FILE`
  if [ $? -ne 0 ]; then
    echo "could not write $FILE!"
    exit 1
  fi


  # none dtb
  FILE=/tmp/$DTB_FILE1
  echo "writing dtb $FILE"
  RES=`$CUSTOM_IMAGE -p $PART_UPGRADE -c write_file_image -r 6 -f $FILE`
  if [ $? -ne 0 ]; then
    echo "could not write $FILE!"
    exit 1
  fi


  FILE=/tmp/$KERNEL_FILE
  echo "writing kernel $FILE"
  $CUSTOM_IMAGE -p $PART_UPGRADE -c write_file_image -r 9 -f $FILE 1>/dev/null 2>&1 &
  PID=$!

  echo "10%" >> $PROGRESS_FILE
  sleep 0.6
  echo "15%" >> $PROGRESS_FILE
  sleep 0.6
  echo "20%" >> $PROGRESS_FILE
  sleep 0.6
  echo "25%" >> $PROGRESS_FILE

  wait $PID
  if [ $? -ne 0 ]; then
    echo "could not write $FILE!"
    exit 1
  fi


  FILE=/tmp/$ROOTFS_FILE
  echo "writing rootfs $FILE"

  $CUSTOM_IMAGE -p $PART_UPGRADE -c write_file_image -r 10 -f $FILE 1>/dev/null 2>&1 &
  PID=$!

  echo "30%" >> $PROGRESS_FILE
  sleep 0.6
  echo "35%" >> $PROGRESS_FILE
  sleep 0.6
  echo "40%" >> $PROGRESS_FILE
  sleep 0.6
  echo "45%" >> $PROGRESS_FILE
  sleep 0.6
  echo "50%" >> $PROGRESS_FILE
  sleep 0.6
  echo "55%" >> $PROGRESS_FILE
  sleep 0.6
  echo "60%" >> $PROGRESS_FILE
  sleep 0.6
  echo "65%" >> $PROGRESS_FILE
  sleep 0.6
  echo "70%" >> $PROGRESS_FILE
  sleep 0.6
  echo "75%" >> $PROGRESS_FILE
  sleep 0.6
  echo "80%" >> $PROGRESS_FILE
  sleep 0.6
  echo "85%" >> $PROGRESS_FILE
  sleep 0.6
  echo "90%" >> $PROGRESS_FILE
  sleep 0.6
  echo "95%" >> $PROGRESS_FILE
  sleep 0.6

  wait $PID
  if [ $? -ne 0 ]; then
    echo "could not write $FILE!"
    exit 1
  fi

  if [ -f "/etc/signed_kernel_dtb" ]; then
    FILE=/tmp/$ROOTFS_FILE
    FS=$(stat -c %s $FILE)
    RES=$($CUSTOM_IMAGE -p $PART_UPGRADE -c add_valid --rootfs-len $FS --rootfs-iv /tmp/rootfs.iv --rootfs-sig /tmp/rootfs.sv)
  else
    RES=$($CUSTOM_IMAGE -p $PART_UPGRADE -c add_valid)
  fi
  if [ $? -ne 0 ]; then
    echo "could not add valid boot header!"
    exit 1
  fi

  echo "valid boot header to partition $PART_UPGRADE has been added"

  if [ $PART_VALID -eq 0 ]; then
    echo "previous image unavailable!"
    # 20221116 - this is not an error, but a valid state when recording the first image during production
    #exit 1
  elif [ $PART_VALID -eq 2 ]; then
    echo "recovery is not invalidated"
  else
    RES=`$CUSTOM_IMAGE -p $PART_VALID -c add_invalid`
    if [ $? -ne 0 ]; then
      echo "could not add invalid boot header!"
      exit 1
    fi
    echo "invalid boot header to partition $PART_VALID has been added"
  fi

  echo "=== UPGRADE DONE ==="
  echo "100%" >> $PROGRESS_FILE

}

#################
# PARSE OPTIONS #
#################

while getopts ":p:hV-:" opt; do
  case "$opt" in
    -)
      case "$OPTARG" in
        help)
          _help
          ;;
        version)
          _version
          ;;
        skip-tar-script)
          OPT_SKIP_TAR_SCRIPT=1
          ;;
        *)
          echo "invalid option: --$OPTARG"
          exit 1
          ;;
      esac
      ;;
    V)
      _version
      ;;
    h)
      _help
      ;;
    p)
      PROGRESS_FILE=$OPTARG
      ;;
    \?)
      echo "invalid option: -$OPTARG"
      exit 1
      ;;
  esac
done

touch $PROGRESS_FILE

#############
# MAIN CODE #
#############

shift $((OPTIND-1))
CMD="$1"
case "$CMD" in
  *.tar)
    FIRMWARE_FILE="$CMD"
    _image_check
    _image_exe
    _eeprom
    _product
    _unio
    _image_files
    _partition
    _upgrade
    ;;

  *)
    echo "please specify a filename.tar"
    exit 1
    ;;
esac

###############
# END OF FILE #
###############
