#!/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}"

  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)"
  pnor="${pnormtd#mtd}"
  ubidev="/dev/ubi${pnor}"
  mountdir="/media/${name}"

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

  # Create a static ubi volume of arbitrary size 24MB,
  # the current pnor image is ~19MB
  # TODO Set size based on file size openbmc/openbmc#1840
  vol="$(findubi "${name}")"
  if [ -z "${vol}" ]; then
    ubimkvol "${ubidev}" -N "${name}" -s 24MiB --type=static
    vol="$(findubi "${name}")"
  fi

  # Create a ubi block needed for read-only volumes,
  # and update the volume with the pnor squashfs image
  ubidevid="${vol#ubi}"
  block="/dev/ubiblock${ubidevid}"
  if [ ! -e "$block" ]; then
    img="/tmp/images/${version}/pnor.xz.squashfs"
    ubiblock --create "/dev/ubi${ubidevid}"
    ubiupdatevol "/dev/ubi${ubidevid}" "${img}"
  fi

  if ! is_mounted "${name}"; then
    mount -t squashfs -o ro "${block}" "${mountdir}"
  fi
}

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

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

  if [ ! -d "${mountdir}" ]; then
    mkdir "${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}")"
  if [ -n "${vol}" ]; then
    ubirmvol "${ubidev}" -N "${name}"
  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
  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 not formatted as ubi.
      return
    fi
  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
}

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
    ;;
  *)
    echo "Invalid argument"
    exit 1
    ;;
esac
rc=$?
if [ ${rc} -ne 0 ]; then
  echo "$0: error ${rc}"
  exit ${rc}
fi