#!/bin/bash

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

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

# Get the mount information
function is_mounted() {
    grep -q "$1" /proc/mounts
    return $?
}

# Attach the pnor mtd device to ubi.
function 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
}

function mount_squashfs() {
    pnormtd="$(findmtd pnor)"
    ubidev="/dev/ubi${pnormtd#mtd}"
    mountdir="/media/${name}"
    vol="$(findubi "${name}")"
    img="/tmp/images/${version}/pnor.xz.squashfs"
    # shellcheck disable=SC2012 # ls provides the size in human-readable form
    filesize="$(ls -sh "$img" | awk -F " " '{print $1}')"

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

    if [ -n "${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
    if ! vol="$(findubi "${name}")"; then
        echo "Unable to create RO volume!"
        return 1
    fi

    ubidevid="${vol#ubi}"
    if ! ubiupdatevol "/dev/ubi${ubidevid}" "${img}"; then
        echo "Unable to update RO volume!"
        return 1
    fi

    if ! ubiblock --create "/dev/ubi${ubidevid}"; then
        echo "Unable to create UBI block for RO volume!"
        return 1
    fi

    if ! mount -t squashfs -o ro "/dev/ubiblock${ubidevid}" "${mountdir}"; then
        echo "Unable to mount RO volume!"
        return 1
    fi
}

function 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
}

function 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
}

function 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
}

function 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-)
        mapfile -t array <<< "${vols}"
    else
        vols=$(ubinfo -a | grep -e "pnor-ro-" -e "pnor-rw-" | \
            grep -v "$activeVersion" | cut -c 14-)
        mapfile -t array <<< "${vols}"
    fi

    for (( index=0; index<${#array[@]}; index++ )); do
        name=${array[index]}
        umount_ubi
    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
        ;;
    ubicleanup)
        ubi_cleanup
        ;;
    *)
        echo "Invalid argument"
        exit 1
        ;;
esac
rc=$?
if [ ${rc} -ne 0 ]; then
    echo "$0: error ${rc}"
    exit ${rc}
fi