#!/bin/sh

# Get the mtd device number (mtdX)
findmtd() {
  m="$(grep -xl "$1" /sys/class/mtd/*/name)"
  m="${m%/name}"
  m="${m##*/}"
  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 mount information
is_mounted() {
  grep -q "$1" /proc/mounts
  return $?
}

# Attach the pnor mtd device to ubi.
attach_ubi() {
  pnormtd="$(findmtd pnor)"
  pnor="${pnormtd#mtd}"
  pnordev="/dev/mtd${pnor}"

  if [ -d "/sys/class/ubi/ubi${pnor}" ]; then
    # Already attached
    return 0
  fi

  ubiattach /dev/ubi_ctrl -m "${pnor}" -d "${pnor}"
  rc=$?
  if [ ${rc} -ne 0 ]; then
    # Check the pnor mtd device is formatted as ubi by reading the first 3 byes,
    # which should be the ascii chars 'UBI'
    magic="$(hexdump -C -n 3 ${pnordev})"
    if [[ "${magic}" =~ "UBI" ]]; then
      # Device already formatted as ubi, ubiattach failed for some other reason
      return ${rc}
    else
      # Format device as ubi
      echo "Starting ubiformat ${pnordev}"
      ubiformat "${pnordev}" -y -q
      # Retry the ubiattach
      ubiattach /dev/ubi_ctrl -m "${pnor}" -d "${pnor}"
    fi
  fi
}

mount_squashfs() {
  pnormtd="$(findmtd pnor)"
  ubidev="/dev/ubi${pnormtd#mtd}"
  mountdir="/media/${name}"
  vol="$(findubi "${name}")"
  img="/tmp/images/${version}/pnor.xz.squashfs"
  filesize="$(ls -sh $img | awk -F " " {'print $1'})"

  if is_mounted "${name}"; then
    echo "${name} is already mounted."
    return 0
  fi

  if [ ! -z "${vol}" ]; then
    ubirmvol "${ubidev}" -N "${name}"
  fi

  if [ ! -d "${mountdir}" ]; then
    mkdir "${mountdir}"
  fi

  # Set size of read-only partition equal to pnor.xz.squashfs
  ubimkvol "${ubidev}" -N "${name}" -s "${filesize}"KiB --type=static
  vol="$(findubi "${name}")"

  if [ $? != 0 ]; then
    echo "Unable to create RO volume!"
    return 1
  fi

  ubidevid="${vol#ubi}"
  ubiupdatevol "/dev/ubi${ubidevid}" "${img}"

  if [ $? != 0 ]; then
    echo "Unable to update RO volume!"
    return 1
  fi

  ubiblock --create "/dev/ubi${ubidevid}"

  if [ $? != 0 ]; then
    echo "Unable to create UBI block for RO volume!"
    return 1
  fi

  mount -t squashfs -o ro "/dev/ubiblock${ubidevid}" "${mountdir}"

  if [ $? != 0 ]; then
    echo "Unable to mount RO volume!"
    return 1
  fi
}

mount_ubi() {
  pnormtd="$(findmtd pnor)"
  pnor="${pnormtd#mtd}"
  ubidev="/dev/ubi${pnor}"
  pnordev="/dev/mtd${pnor}"

  if [[ "${name}" == "pnor-patch" ]]; then
    if [[ "$(fw_printenv fieldmode 2>/dev/null)" == "fieldmode=true" ]]; then
      return 0
    fi
    if [[ ! "$(hexdump -C -n 3 ${pnordev})" =~ "UBI" ]]; then
      return 0
    fi
    mountdir="/usr/local/share/pnor"
  else
    mountdir="/media/${name}"
  fi

  if [[ "${name}" == "pnor-prsv" ]]; then
    size="2MiB"
  else
    size="16MiB"
  fi

  if [ ! -d "${mountdir}" ]; then
    mkdir -p "${mountdir}"
  fi

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

  if ! is_mounted "${name}"; then
    mountdev="ubi${pnor}:${name}"
    mount -t ubifs "${mountdev}" "${mountdir}"
  fi
}

umount_ubi() {
  pnormtd="$(findmtd pnor)"
  pnor="${pnormtd#mtd}"
  ubidev="/dev/ubi${pnor}"
  mountdir="/media/${name}"

  if is_mounted "${name}"; then
    umount "${mountdir}"
  fi

  vol="$(findubi "${name}")"
  id="${vol##*_}"
  if [ -n "${id}" ]; then
    ubirmvol "${ubidev}" -n "${id}"
  fi

  if [ -d "${mountdir}" ]; then
    rm -r "${mountdir}"
  fi
}

remount_ubi() {
  pnormtd="$(findmtd pnor)"
  pnor="${pnormtd#mtd}"
  pnordev="/dev/mtd${pnor}"

  # Re-Attach the pnor mtd device to ubi
  if [[ $(hexdump -C -n 3 ${pnordev}) =~ "UBI" ]]; then
    ubiattach /dev/ubi_ctrl -m "${pnor}" -d "${pnor}"
  else
    # Device not formatted as ubi.
    return 0
  fi

  # Get information on all ubi volumes
  ubinfo=$(ubinfo -d ${pnor})
  presentVolumes=${ubinfo##*:}
  IFS=', ' read -r -a array <<< "$presentVolumes"
  for element in ${array[@]};
  do
    elementProperties=$(ubinfo -d $pnor -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} == pnor-prsv ]] || [[ ${name} == pnor-rw* ]] || [[ ${name} == pnor-ro* ]]; then
      mountdir="/media/${name}"
      if [ ! -d "${mountdir}" ]; then
        mkdir -p "${mountdir}"
      fi

      if [[ ${name} == pnor-ro* ]]
      then
        ubiblock --create /dev/ubi${pnor}_${element}
        mount -t squashfs -o ro "/dev/ubiblock${pnor}_${element}" "${mountdir}"
      else
        mount -t ubifs "ubi${pnor}:${name}" "${mountdir}"
      fi
    fi
  done
}

update_symlinks() {
  PNOR_ACTIVE_PATH="/var/lib/phosphor-software-manager/pnor/"
  PNOR_RO_ACTIVE_PATH="/var/lib/phosphor-software-manager/pnor/ro"
  PNOR_RO_PREFIX="/media/pnor-ro-"
  PNOR_RW_ACTIVE_PATH="/var/lib/phosphor-software-manager/pnor/rw"
  PNOR_RW_PREFIX="/media/pnor-rw-"
  PNOR_PRSV_ACTIVE_PATH="/var/lib/phosphor-software-manager/pnor/prsv"
  PNOR_PRSV="/media/pnor-prsv"
  PERSISTENCE_PATH="/var/lib/obmc/openpower-pnor-code-mgmt/"
  PNOR_PATCH_LOCATION="/usr/local/share/pnor/"

  # Get a list of all active PNOR versions
  data="$(ls -d ${PNOR_RO_PREFIX}*)"
  IFS=$'\n'  array=(${data})

  currentVersion=""
  lowestPriority=255
  for element in ${array[@]}; do
    #Remove the PNOR_RO_PREFIX from the path to get version ID.
    versionId="${element#${PNOR_RO_PREFIX}}"

    # Get the priority of active versions from persistence files.
    if [[ -f "${PERSISTENCE_PATH}${versionId}" ]]; then
      data="$(grep -r "priority" ${PERSISTENCE_PATH}${versionId})"
      priority="${data: -1}"
      if [[ priority -le lowestPriority  ]]; then
        lowestPriority=${priority}
        currentVersion=${versionId}
      fi
    fi
  done

  # Return if no active version found
  if [ -z $currentVersion ]; then
    return 0;
  fi

  if [ ! -d "${PNOR_ACTIVE_PATH}" ]; then
        mkdir -p "${PNOR_ACTIVE_PATH}"
  fi

  # If the RW or RO active links doesn't point to the version with
  # lowest priority, then remove the symlink and create new ones.
  if [[ $(readlink -f "${PNOR_RO_ACTIVE_PATH}") != ${PNOR_RO_PREFIX}${currentVersion}  ]]; then
    rm -f ${PNOR_RO_ACTIVE_PATH}
    rm -rf ${PNOR_PATCH_LOCATION}*
    ln -sfv ${PNOR_RO_PREFIX}${currentVersion} ${PNOR_RO_ACTIVE_PATH}
  fi

  if [[ $(readlink -f "${PNOR_RW_ACTIVE_PATH}") != ${PNOR_RW_PREFIX}${currentVersion}  ]]; then
    rm -f ${PNOR_RW_ACTIVE_PATH}
    ln -sfv ${PNOR_RW_PREFIX}${currentVersion} ${PNOR_RW_ACTIVE_PATH}
  fi

  if [[ ! -h ${PNOR_PRSV_ACTIVE_PATH}  ]]; then
    ln -sfv ${PNOR_PRSV} ${PNOR_PRSV_ACTIVE_PATH}
  fi
}

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

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

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

clear_volatile() {
  service=$(mapper get-service /org/open_power/control/volatile)
  clearVolatileEnabled=$(busctl get-property $service /org/open_power/control/volatile xyz.openbmc_project.Object.Enable Enabled)
  if [[ "$clearVolatileEnabled" != "b true" ]]; then
    return 0
  fi

  PNOR_TOC_FILE="pnor.toc"
  PNOR_RO_ACTIVE_PATH="/var/lib/phosphor-software-manager/pnor/ro/"
  PNOR_RW_ACTIVE_PATH="/var/lib/phosphor-software-manager/pnor/rw/"
  PNOR_PRSV_ACTIVE_PATH="/var/lib/phosphor-software-manager/pnor/prsv/"

  # toc partition string format:
  # partition27=HB_VOLATILE,0x02ba9000,0x02bae000,00,ECC,VOLATILE,READWRITE
  tocFilePath="${PNOR_RO_ACTIVE_PATH}${PNOR_TOC_FILE}"
  volatiles=($(grep VOLATILE "${tocFilePath}" | grep -Eo '^partition([0-9]+)=([A-Za-z0-9_]+)'))
  for (( index=0; index<${#volatiles[@]}; index++ )); do
    volatileName="$(echo ${volatiles[${index}]} | awk -F '=' '{print $2}')"

    rwVolatile="${PNOR_RW_ACTIVE_PATH}${volatileName}"
    if [ -f "${rwVolatile}" ]; then
      echo "Clear $rwVolatile"
      rm "${rwVolatile}"
    fi
    prsvVolatile="${PNOR_PRSV_ACTIVE_PATH}${volatileName}"
    if [ -f "${prsvVolatile}" ]; then
      echo "Clear $prsvVolatile"
      rm "${prsvVolatile}"
    fi
  done
  # Always reset the sensor after clearing
  busctl set-property $service /org/open_power/control/volatile xyz.openbmc_project.Object.Enable Enabled b false
}

case "$1" in
  ubiattach)
    attach_ubi
    ;;
  squashfsmount)
    name="$2"
    version="$3"
    mount_squashfs
    ;;
  ubimount)
    name="$2"
    mount_ubi
    ;;
  ubiumount)
    name="$2"
    umount_ubi
    ;;
  ubiremount)
    remount_ubi
    ;;
  updatesymlinks)
    update_symlinks
    ;;
  ubicleanup)
    ubi_cleanup
    ;;
  clearvolatile)
    clear_volatile
    ;;
  *)
    echo "Invalid argument"
    exit 1
    ;;
esac
rc=$?
if [ ${rc} -ne 0 ]; then
  echo "$0: error ${rc}"
  exit ${rc}
fi