#!/bin/bash
set -eo pipefail

# Get the root mtd device number (mtdX) from "/dev/ubiblockX_Y on /"
findrootmtd() {
  rootmatch=" on / "
  m="$(mount | grep "${rootmatch}" | grep "ubiblock")"
  m="${m##*ubiblock}"
  m="${m%_*}"
  if [ -z "${m}" ]; then
    # default to bmc mtd (0)
    m=0
  fi
  echo "mtd${m}"
}

findrootubi() {
  rootmatch=" on / "
  m="$(mount | grep "${rootmatch}")"
  m="${m##*ubiblock}"
  m="${m% on*}"
  echo "ubi${m}"
}

# Get the mtd device number (mtdX)
findmtd() {
  m="$(grep -xl "$1" /sys/class/mtd/*/name)"
  m="${m%/name}"
  m="${m##*/}"
  echo "${m}"
}

# Get the mtd device number only (return X of mtdX)
findmtdnum() {
  m="$(findmtd "$1")"
  m="${m##mtd}"
  echo "${m}"
}

# Get the ubi device number (ubiX_Y)
findubi() {
  u="$(grep -xl "$1" /sys/class/ubi/ubi?/subsystem/ubi*/name)"
  u="${u%/name}"
  u="${u##*/}"
  echo "${u}"
}

# Get the ubi device number (ubiX_Y) on a specific mtd
findubi_onmtd() {
  u="$(grep -xl "$1" /sys/class/ubi/ubi"$2"/subsystem/ubi"$2"*/name)"
  u="${u%/name}"
  u="${u##*/}"
  echo "${u}"
}

# Get all ubi device names on a specific mtd that match requested string
findubiname_onmtd() {
  u="$(grep -h "$1" /sys/class/ubi/ubi"$2"/subsystem/ubi"$2"*/name)"
  u="${u%/name}"
  u="${u##*/}"
  echo "${u}"
}

# Get the name from the requested ubiX_Y volume
findname() {
  n="$(cat /sys/class/ubi/$1/name)"
  echo "${n}"
}

# Set the u-boot envs that perform a side switch on failure to boot
set_wdt2bite() {
  if ! fw_printenv wdt2bite 2>/dev/null; then
    fw_setenv wdt2bite "mw.l 0x1e785024 0xa 1; mw.b 0x1e78502c 0xb3 1"
    fw_setenv bootalt "run wdt2bite"
    fw_setenv obmc_bootcmd "ubi part obmc-ubi; run do_rwreset; ubi read \
\${loadaddr} \${kernelname}; bootm \${loadaddr} || run bootalt"
  fi
}

# Make space on flash before creating new volumes. This can be enhanced
# determine current flash usage. For now only keep a "keepmax" number of them
ubi_remove_volumes()
{
  rootubi="$(findrootubi)"
  rootname="$(findname "${rootubi}")"
  rootversion="${rootname##*-}"
  rootkernel="kernel-${rootversion}"

  # Just keep max number of volumes before updating, don't delete the version
  # the BMC is booted from, and when a version is identified to be deleted,
  # delete both the rofs and kernel volumes for that version.
  rmnames="$(findubiname_onmtd "${name%-*}-" "${ro}")"
  rmnames=(${rmnames})
  ubicount="${#rmnames[@]}"
  while [ ${ubicount} -ge ${keepmax} ]; do
    # Loop through existing volumes and skip currently active ones
    for (( index=0; index<${#rmnames[@]}; index++ )); do
      rmname="${rmnames[${index}]}"
      rmversion="${rmname##*-}"
      [ "${rmversion}" == "${version}" ] && continue
      rmubi="$(findubi_onmtd "rofs-${rmversion}" "${ro}")"
      if [[ ( "${rmubi}" != "${rootubi}" ) &&
            ( "${rmname}" != "${rootkernel}" ) ]]; then
        ubi_remove "rofs-${rmversion}" "${ro}"
        ubi_remove "kernel-${rmversion}" "${ro}"
        # Remove priority value
        fw_setenv "${rmversion}"
        break
      fi
    done
    # Decrease count regardless to avoid an infinite loop
    (( ubicount-- ))
  done
}

ubi_rw() {
  rwmtd="$(findmtd "${reqmtd}")"
  rw="${rwmtd#mtd}"
  ubidev="/dev/ubi${rw}"

  # Update rwfs_size, check imgsize was specified, otherwise it'd clear the var
  if [ ! -z "$imgsize" ]; then
    rwsize="$(fw_printenv -n rwfs_size 2>/dev/null)" || true
    if [[ "${imgsize}" != "${rwsize}" ]]; then
      fw_setenv rwfs_size "${imgsize}"
    fi
  fi

  vol="$(findubi "${name}")"
  if [ -z "${vol}" ]; then
    ubimkvol "${ubidev}" -N "${name}" -s "${imgsize}"
  fi
}

ubi_ro() {
  keepmax=2 # Default 2 volumes per mtd
  romtd="$(findmtd "${reqmtd}")"
  romtd2="$(findmtd "${reqmtd2}")"

  if [ ! "${romtd}" == "${romtd2}" ]; then
    # Request to use alternate mtd device, choose the non-root one
    keepmax=1 # 1 volume on each of the requested mtds
    rootmtd="$(findrootmtd)"
    if [ "${rootmtd}" == "${romtd}" ]; then
      romtd="${romtd2}"
    fi
  fi
  ro="${romtd#mtd}"
  ubidev="/dev/ubi${ro}"

  ubi_remove_volumes

  if [ -z "${imgfile}" ]; then
    echo "Unable to create read-only volume. Image file not specified."
    return 1
  fi

  # Create a ubi volume, dynamically sized to fit BMC image if size unspecified
  img="/tmp/images/${version}/${imgfile}"
  imgsize="$(stat -c '%s' ${img})"

  vol="$(findubi "${name}")"
  if [ ! -z "${vol}" ]; then
    # Allow a duplicate kernel volume on the alt mtd
    if [[ "${name}" =~ "kernel" ]]; then
      vol="$(findubi_onmtd "${name}" "${ro}")"
    fi
  fi
  if [ -z "${vol}" ]; then
    ubimkvol "${ubidev}" -N "${name}" -s "${imgsize}" --type=static
    vol="$(findubi "${name}")"
  fi
}

# Squashfs images need a ubi block
ubi_block() {
  vol="$(findubi "${name}")"
  ubidevid="${vol#ubi}"
  block="/dev/ubiblock${ubidevid}"
  if [ ! -e "$block" ]; then
    ubiblock --create "/dev/ubi${ubidevid}"
  fi
}

ubi_updatevol() {
  vol="$(findubi "${name}")"
  ubidevid="${vol#ubi}"
  img="/tmp/images/${version}/${imgfile}"
  ubiupdatevol "/dev/ubi${ubidevid}" "${img}"
}

ubi_remove() {
    rmname="$1"
    rmmtd="$2"
    if [ ! -z "${rmmtd}" ]; then
      vol="$(findubi_onmtd "${rmname}" "${rmmtd}")"
    else
      vol="$(findubi "${rmname}")"
    fi

    if [ ! -z "$vol" ]; then
        vol="${vol%_*}"

        if grep -q "$rmname" /proc/mounts; then
            mountdir=$(grep "$rmname" /proc/mounts | cut -d " " -f 2)
            umount "$mountdir"
            rm -r "$mountdir"
        fi

        ubirmvol "/dev/${vol}" -N "$rmname"
    fi
}

ubi_cleanup() {
    # When ubi_cleanup is run, it expects one or no active version.
    activeVersion=$(busctl --list --no-pager tree \
            xyz.openbmc_project.Software.BMC.Updater | \
            grep /xyz/openbmc_project/software/ | tail -c 9)

    if [[ -z "$activeVersion" ]]; then
        vols=$(ubinfo -a | grep "rofs-" | cut -c 14-)
        vols=(${vols})
    else
        vols=$(ubinfo -a | grep "rofs-" | \
                grep -v "$activeVersion" | cut -c 14-)
        vols=(${vols})
    fi

    for (( index=0; index<${#vols[@]}; index++ )); do
         ubi_remove ${vols[index]}
    done
}

mount_alt_rwfs() {
  altNum="$(findmtdnum "alt-bmc")"
  if [ ! -z "${altNum}" ]; then
    altVarMount="/media/alt/var"
    mkdir -p "${altVarMount}"
    if mount ubi"${altNum}":rwfs "${altVarMount}" -t ubifs -o defaults; then
      mkdir -p "${altVarMount}"/persist/etc
    fi
  fi
}

remount_ubi() {
  bmcmtd="$(findmtd "bmc")"
  altbmcmtd="$(findmtd "alt-bmc")"
  mtds="${bmcmtd: -1}","${altbmcmtd: -1}"

  IFS=',' read -r -a mtds <<< "$mtds"
  mtds=($(echo "${mtds[@]}" | tr ' ' '\n' | sort -u | tr '\n' ' '))
  for mtd in ${mtds[@]}; do
    # Get information on all ubi volumes
    ubinfo=$(ubinfo -d ${mtd})
    presentVolumes=${ubinfo##*:}
    IFS=', ' read -r -a array <<< "$presentVolumes"
    for element in ${array[@]}; do
      elementProperties=$(ubinfo -d $mtd -n $element)
      # Get ubi volume name by getting rid of additional properties
      name=${elementProperties#*Name:}
      name="${name%Character*}"
      name="$(echo -e "${name}" | tr -d '[:space:]')"

      if [[ ${name} == rofs-* ]]; then
        mountdir="/media/${name}"

        if [ ! -d ${mountdir} ]; then
          mkdir -p "${mountdir}"
          # U-Boot will create the ubiblock for the running version, but not
          # for the version on the other chip
          if [ ! -e "/dev/ubiblock${mtd}_${element}" ]; then
            ubiblock --create /dev/ubi${mtd}_${element}
          fi
          mount -t squashfs -o ro "/dev/ubiblock${mtd}_${element}" "${mountdir}"
        fi
      fi
    done
  done

  set_wdt2bite
}

# Read the current env variable and set it on the alternate boot env
copy_env_var_to_alt() {
  varName=$1
  value="$(fw_printenv -n "${varName}")"
  fw_setenv -c /etc/alt_fw_env.config "${varName}" "${value}"
}

# When the alternate bmc chip boots, u-boot thinks its the primary mtdX.
# Therefore need to swap the chip numbers when copying the ubiblock and root to
# alternate bmc u-boot environment.
copy_ubiblock_to_alt() {
  value="$(fw_printenv -n ubiblock)"
  bmcNum="$(findmtdnum "bmc")"
  altNum="$(findmtdnum "alt-bmc")"
  replaceAlt="${value/${altNum},/${bmcNum},}"

  if [[ "${value}" == "${replaceAlt}" ]]; then
    replaceBmc="${value/${bmcNum},/${altNum},}"
    value=${replaceBmc}
  else
    value=${replaceAlt}
  fi

  fw_setenv -c /etc/alt_fw_env.config ubiblock "${value}"
}

copy_root_to_alt() {
  value="$(fw_printenv -n root)"
  bmcNum="$(findmtdnum "bmc")"
  altNum="$(findmtdnum "alt-bmc")"
  replaceAlt="${value/${altNum}_/${bmcNum}_}"

  if [[ "${value}" == "${replaceAlt}" ]]; then
    replaceBmc="${value/${bmcNum}_/${altNum}_}"
    value=${replaceBmc}
  else
    value=${replaceAlt}
  fi

  fw_setenv -c /etc/alt_fw_env.config root "${value}"
}

ubi_setenv() {
    # The U-Boot environment maintains two banks of environment variables.
    # The banks need to be consistent with each other to ensure that these
    # variables can reliably be read from file. In order to guarantee that the
    # banks are both correct, we need to run fw_setenv twice.
    variable=$1
    if [[ "$variable" == *"="* ]]; then
        varName="${variable%=*}"
        value="${variable##*=}"
        # Write only if var is not set already to the requested value
        currentValue="$(fw_printenv -n "${varName}" 2>/dev/null)" || true
        if [[ "${currenValue}" != "${value}" ]]; then
          fw_setenv "$varName" "$value"
          fw_setenv "$varName" "$value"
        fi
    else
        fw_setenv "$variable"
        fw_setenv "$variable"
    fi
}

mtd_write() {
  flashmtd="$(findmtd "${reqmtd}")"
  img="/tmp/images/${version}/${imgfile}"
  flashcp -v ${img} /dev/${flashmtd}
}

backup_env_vars() {
  copy_env_var_to_alt kernelname
  copy_ubiblock_to_alt
  copy_root_to_alt
}

update_env_vars() {
  vol="$(findubi rofs-"${version}")"
  if [ -z "${vol}" ]; then
    return 1
  fi
  ubidevid="${vol#ubi}"
  block="/dev/ubiblock${ubidevid}"
  if [ ! -e "${block}" ]; then
    return 1
  fi
  ubi_setenv "kernelname=kernel-${version}"
  ubi_setenv "ubiblock=$(echo "${ubidevid}" | sed 's/_/,/')"
  ubi_setenv "root=${block}"
}

#TODO: Replace the implementation with systemd-inhibitors lock
#      once systemd/systemd#949 is resolved
rebootguardenable() {
  dir="/run/systemd/system/"
  file="reboot-guard.conf"
  units=("reboot" "poweroff" "halt")

  for unit in "${units[@]}"; do
    mkdir -p ${dir}${unit}.target.d
    echo -e "[Unit]\nRefuseManualStart=yes" >> ${dir}${unit}.target.d/${file}
  done
}

#TODO: Replace the implementation with systemd-inhibitors lock
#      once systemd/systemd#949 is resolved
rebootguarddisable() {
  dir="/run/systemd/system/"
  file="reboot-guard.conf"
  units=("reboot" "poweroff" "halt")

  for unit in "${units[@]}"; do
    rm -rf ${dir}${unit}.target.d/${file}
  done
}

# Create a copy in the alt mtd
create_vol_in_alt() {
  alt="alt-bmc"
  altmtd="$(findmtd "${alt}")"
  if [ ! -z "${altmtd}" ]; then
    reqmtd="${alt}"
    reqmtd2="${alt}"
    ubi_ro
    ubi_updatevol
  fi
}

# Copy contents of one MTD device to another
mtd_copy() {
  in=$1
  out=$2

  # Must erase MTD first to prevent corruption
  flash_eraseall "${out}"
  dd if="${in}" of="${out}"
}

mirroruboot() {
  bmc="$(findmtd "u-boot")"
  bmcdev="/dev/${bmc}"
  alt="$(findmtd "alt-u-boot")"
  altdev="/dev/${alt}"

  checksum_bmc="$(md5sum "${bmcdev}")"
  checksum_bmc="${checksum_bmc% *}"
  checksum_alt="$(md5sum "${altdev}")"
  checksum_alt="${checksum_alt% *}"

  if [[ "${checksum_bmc}" != "${checksum_alt}" ]]; then
    bmcenv="$(findmtd "u-boot-env")"
    bmcenvdev="/dev/${bmcenv}"
    altenv="$(findmtd "alt-u-boot-env")"
    altenvdev="/dev/${altenv}"

    echo "Mirroring U-boot to alt chip"
    mtd_copy "${bmcdev}" "${altdev}"
    mtd_copy "${bmcenvdev}" "${altenvdev}"

    copy_ubiblock_to_alt
    copy_root_to_alt
  fi
}

case "$1" in
  mtduboot)
    reqmtd="$2"
    version="$3"
    imgfile="image-u-boot"
    mtd_write
    ;;
  ubirw)
    reqmtd="$2"
    name="$3"
    imgsize="$4"
    ubi_rw
    ;;
  ubiro)
    reqmtd="$(echo "$2" | cut -d "+" -f 1)"
    reqmtd2="$(echo "$2" | cut -d "+" -f 2)"
    name="$3"
    version="$4"
    imgfile="image-rofs"
    ubi_ro
    ubi_updatevol
    ubi_block
    ;;
  ubikernel)
    reqmtd="$(echo "$2" | cut -d "+" -f 1)"
    reqmtd2="$(echo "$2" | cut -d "+" -f 2)"
    name="$3"
    version="$4"
    imgfile="image-kernel"
    ubi_ro
    ubi_updatevol
    create_vol_in_alt
    ;;
  ubiremove)
    name="$2"
    ubi_remove "${name}"
    ;;
  ubicleanup)
    ubi_cleanup
    ;;
  ubisetenv)
    ubi_setenv "$2"
    ;;
  ubiremount)
    remount_ubi
    mount_alt_rwfs
    ;;
  createenvbackup)
    backup_env_vars
    ;;
  updateubootvars)
    version="$2"
    update_env_vars
    ;;
  rebootguardenable)
    rebootguardenable
    ;;
  rebootguarddisable)
    rebootguarddisable
    ;;
  mirroruboot)
    mirroruboot
    ;;
  *)
    echo "Invalid argument"
    exit 1
    ;;
esac