#!/bin/sh

SCRIPT="$(basename $0)"

if [ "o$1" = "o-?" ]; then
cat <<EOH
Sets the optimal bank plan.

For devices with /proc/banktable, it is effectively the same as
https://hack-technicolor.rtfd.io/en/stable/Hacking/PostRoot/#bank-planning.

For devices that use bootmgr and ubi filesystems, the device will be set
to boot the current firmware from bank 2, but the firmware in bank 1 will
be retained.

Usage: $0

EOH
exit
fi

GREEN='\033[1;32m'
ORANGE='\033[0;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

[ -d /proc/banktable ] && BANKTABLE=Y || BANKTABLE=N

echo -e "${ORANGE}Checking active and booted banks...${NC}"
if [ $BANKTABLE = Y ]; then
  ACTIVE="$(cat /proc/banktable/active 2>/dev/null | cut -d_ -f2)"
  BOOTED="$(cat /proc/banktable/booted 2>/dev/null | cut -d_ -f2)"
else
  ACTIVE="$(bootmgr partition active 2>/dev/null | grep -v rd_metadata)"
  BOOTED="$(bootmgr partition booted 2>/dev/null | grep -v rd_metadata)"
fi
if [ -z "$ACTIVE" -o -z "$BOOTED" ]; then
  echo -e "${RED}==> ABORTING: Cannot determine active or booted bank!!${NC}"
  exit 2
elif [ \( \( $BANKTABLE = Y -a "$ACTIVE" = "1" \) -o \( $BANKTABLE = N -a "$ACTIVE" = "2" \) \) -a "$BOOTED" = "2" ]; then
  echo -e "${GREEN}==> OKAY: Active = bank ${ACTIVE} Booted = bank ${BOOTED}${NC}"
  echo -e "${GREEN}==> EXITING: Bank Plan is already OPTIMAL!!${NC}"
  exit
else
  echo -e "${ORANGE}==> SUB-OPTIMAL: ${RED}Active = bank ${ACTIVE} Booted = bank ${BOOTED}${NC}"
fi

echo -e "${ORANGE}Checking firmware versions...${NC}"
if [ $BANKTABLE = Y ]; then
  ACTIVE_VERSION="$(cat /proc/banktable/activeversion 2>/dev/null)"
  PASSIVE_VERSION="$(cat /proc/banktable/passiveversion 2>/dev/null)"
fi
[ -z "$ACTIVE_VERSION" ] && ACTIVE_VERSION=$(uci -q get env.var.friendly_sw_version_activebank)
[ -z "$PASSIVE_VERSION" ] && PASSIVE_VERSION=$(uci -q get env.var.friendly_sw_version_passivebank)
if [ -n "$ACTIVE_VERSION" ]; then
  BOOTED_SUFFIX="$(uci -q get versioncusto.override.fwversion_suffix)"
  [ -n "$BOOTED_SUFFIX" ] && ACTIVE_VERSION="$(echo $ACTIVE_VERSION | cut -d- -f1)$BOOTED_SUFFIX"
fi
if [ -n "$PASSIVE_VERSION" ]; then
  [ "$BOOTED" = "2" ] && NOT_BOOTED="bank_1" || NOT_BOOTED="bank_2"
  if [ -e /overlay/$NOT_BOOTED/etc/config/versioncusto ]; then
    NOT_BOOTED_SUFFIX="$(uci -c /overlay/$NOT_BOOTED/etc/config -q get versioncusto.override.fwversion_suffix)"
    [ -n "$NOT_BOOTED_SUFFIX" ] && PASSIVE_VERSION="$(echo $PASSIVE_VERSION | cut -d- -f1)$NOT_BOOTED_SUFFIX"
  fi
fi

if [ $BANKTABLE = Y -a "$BOOTED" = "2" -a \( -z "$ACTIVE_VERSION" -o "$ACTIVE_VERSION" = "Unknown" \) ]; then
  echo -e "${RED}==> ABORTING: Firmware version in bank 1 is Unknown!!${NC}"
  exit 2
else
  echo -e "${GREEN}==> Active  = ${ACTIVE_VERSION}${NC}"
  echo -e "${GREEN}==> Passive = ${PASSIVE_VERSION}${NC}"
fi

manual_bank_planning() {
  echo -e "${ORANGE}"
  echo -e "########################################################"
  echo -e "# You will need to manually set the optimal bank plan. #"
  echo -e "#                                                      #"
  echo -e "# 1. Ensure you have a rootable firmware in bank 1.    #"
  echo -e "# 2. Ensure you are running from bank 1:               #"
  echo -e "#    a. Use ${NC}./show-bank-plan${ORANGE} to see active/booted bank #"
  echo -e "#    b. If necessary, swap to bank 1 using             #"
  echo -e "#       ${NC}./reset-to-factory-defaults-with-root -s${ORANGE}       #"
  echo -e "# 3. Use ./safe-firmware-upgrade to load required      #"
  echo -e "#    firmware into bank 2                              #"
  echo -e "# 4. Confirm bank layout with ${NC}./show-bank-plan${ORANGE}         #"
  echo -e "########################################################"
  echo -e "${NC}"
}

echo -e "${ORANGE}Ensuring the banks match in size...${NC}"
if [ $BANKTABLE = Y ]; then
  BANK_COUNT=$(grep -c bank_ /proc/mtd)
  if [ $BANK_COUNT -ne 2 ]; then
    echo -e "${RED}==> ABORTING: Unexpected bank count - Expected 2 but found ${BANK_COUNT}!${NC}"
    exit 2
  fi
  BANK1_SIZE="$(grep bank_1 /proc/mtd | cut -d' ' -f2)"
  BANK2_SIZE="$(grep bank_2 /proc/mtd | cut -d' ' -f2)"
else
  BANK1_SIZE="$(grep rootfs1 /proc/mtd | cut -d' ' -f2)"
  BANK2_SIZE="$(grep rootfs2 /proc/mtd | cut -d' ' -f2)"
fi
if [ "$BANK1_SIZE" != "$BANK2_SIZE" ]; then
  echo -e "${RED}==> ABORTING: Bank 1 and Bank 2 are different sizes${NC}"
  echo " - Bank 1 = $BANK1_SIZE"
  echo " - Bank 2 = $BANK2_SIZE"
  if [ $BANKTABLE = N ]; then
    manual_bank_planning
  fi
  exit 2
else
  echo -e "${GREEN}==> OKAY: Bank 1 = ${BANK1_SIZE} Bank 2 = ${BANK2_SIZE}${NC}"
  if [ $BANKTABLE = N ]; then
    BANK1_SIZE="$(grep bootfs1 /proc/mtd | cut -d' ' -f2)"
    BANK2_SIZE="$(grep bootfs2 /proc/mtd | cut -d' ' -f2)"
    if [ "$BANK1_SIZE" != "$BANK2_SIZE" ]; then
      echo -e "${RED}==> ABORTING: Boot 1 and Boot 2 are different sizes${NC}"
      echo " - Boot 1 = $BANK1_SIZE"
      echo " - Boot 2 = $BANK2_SIZE"
      manual_bank_planning
      exit 2
    else
      echo -e "${GREEN}==> OKAY: Boot 1 = ${BANK1_SIZE} Boot 2 = ${BANK2_SIZE}${NC}"
    fi
  fi
fi
if [ $BANKTABLE = N ]; then
  echo -e "${ORANGE}Checking /tmp space...${NC}"
  rm -f /tmp/ubi*
  sync
  TMP_AVAIL_BYTES=$(( $(df -k /tmp | tail -n1 | tr -s ' ' | cut -d' ' -f4) * 1024 ))
  ROOTFS1_BYTES=$(ubinfo -d 0 -N rootfs1 | grep ^Size: | cut -d'(' -f2 | cut -d' ' -f1)
  if [ $TMP_AVAIL_BYTES -lt $ROOTFS1_BYTES ]; then
    echo -e "${RED}==> ABORTING: Insufficient /tmp space to copy root partition${NC}"
    echo " - rootfs1 = $ROOTFS1_BYTES bytes"
    echo " - /tmp    = $TMP_AVAIL_BYTES bytes"
    exit 2
  else
    echo -e "${GREEN}==> OKAY: rootfs1 = ${ROOTFS1_BYTES} /tmp = ${TMP_AVAIL_BYTES}${NC}"
  fi
fi

clone_dev() {
  local source="/dev/$1"
  local tmpimg="/tmp/$1"
  local target="/dev/$2"
  local dd_err=$(mktemp)
  local abort=n

  echo -e "${ORANGE}Imaging bank 1 $3 partition into $tmpimg...${NC}"
  if dd if=$source bs=1M 2>$dd_err >$tmpimg; then
    echo -e "${ORANGE}Copying $tmpimg into bank 2 $3 partition...${NC}"
    if ! ubiupdatevol $target $tmpimg; then
      abort=y
      echo -e "${RED}==> ABORTING: Update of $target failed!${NC}"
    fi
  else
    abort=y
    cat $dd_err
    echo -e "${RED}==> ABORTING: Imaging of $source failed!${NC}"
  fi
  rm -f $tmpimg $dd_err
  if [ $abort = y ]; then
    exit 2
  fi
}

echo -n "Enter y to attempt to set the optimal bank plan, or anything else to exit now: "
read
if [ "$REPLY" = "y" -o "$REPLY" = "Y" ]; then
  if [ "$BOOTED" = "1" ]; then
    REBOOT=Y
    SUCCESS_COUNT=-1
    if [ $BANKTABLE = Y ]; then
      SUCCESS_COUNT=1
      echo -e "${ORANGE}Copying bank 1 firmware into bank 2...${NC}"
      mtd -e bank_2 write /dev/$(grep bank_1 /proc/mtd | cut -d: -f1) bank_2
      echo -e "${ORANGE}Verifying firmware copy...${NC}"
      HASHES=$(sha256sum $(grep bank_ /proc/mtd | cut -d: -f1 | xargs -n1 echo /dev | tr ' ' '/'))
    elif [ $BANKTABLE = N ]; then
      SUCCESS_COUNT=2
      BOOTFS1_DEV="ubi0_$(ubinfo -d 0 -N bootfs1 | grep "Volume ID" | tr -s ' ' | cut -d' ' -f3)"
      BOOTFS2_DEV="ubi0_$(ubinfo -d 0 -N bootfs2 | grep "Volume ID" | tr -s ' ' | cut -d' ' -f3)"
      ROOTFS1_DEV="ubi0_$(ubinfo -d 0 -N rootfs1 | grep "Volume ID" | tr -s ' ' | cut -d' ' -f3)"
      ROOTFS2_DEV="ubi0_$(ubinfo -d 0 -N rootfs2 | grep "Volume ID" | tr -s ' ' | cut -d' ' -f3)"
      clone_dev $BOOTFS1_DEV $BOOTFS2_DEV boot
      clone_dev $ROOTFS1_DEV $ROOTFS2_DEV root
      echo -e "${ORANGE}Verifying firmware copy...${NC}"
      HASHES=$(sha256sum /dev/$BOOTFS1_DEV /dev/$BOOTFS2_DEV /dev/$ROOTFS1_DEV /dev/$ROOTFS2_DEV)
    fi
    if [ $(echo "$HASHES" | cut -d' ' -f1 | sort -u | wc -l) -eq $SUCCESS_COUNT ]; then
      echo -e "${GREEN}==> Copy successful!${NC}"
      echo -e "${ORANGE}Making a temporary copy of the bank 1 overlay for the booted firmware...${NC}"
      tar -czf /tmp/bank_overlay_backup -C /overlay/bank_1 .
      echo -e "${ORANGE}Cleaning up space by removing any existing old overlays...${NC}"
      rm -rf /overlay/*
      echo -e "${ORANGE}Using the previously made temp copy as overlay for bank 2...${NC}"
      mkdir -p /overlay/bank_2
      tar -xzf /tmp/bank_overlay_backup -C /overlay/bank_2 
    else
      echo -e "${RED}==> ABORTING: Verification of bank copy FAILED!${NC}"
      echo "$HASHES"
      exit 2
    fi
  else
    REBOOT=N
    echo -e "${GREEN}==> Already booted from bank 2 - Firmware copy ${ORANGE}NOT REQUIRED${NC}"
    echo -e "${ORANGE}Cleaning up space by removing any existing old bank 1 overlay...${NC}"
    rm -rf /overlay/bank_1
  fi

  if [ $BANKTABLE = Y ]; then
    echo -e "${ORANGE}Activating bank 1...${NC}"
    echo bank_1 > /proc/banktable/active
    echo -e "${ORANGE}Erasing firmware in bank 1...${NC}"
    mtd erase bank_1
  else
    echo -e "${ORANGE}Switching to bank 2...${NC}"
    bootmgr switch_over 2
  fi

  echo -e "${ORANGE}Making sure changes are synced to flash memory...${NC}"
  sync
  
  if [ $REBOOT = Y ]; then
    echo -e "${GREEN}"
    echo "################################################"
    echo "# Operation Complete.                          #"
    echo "#                                              #"
    if [ -e /proc/sysrq-trigger ]; then
      echo "# About to emulate system crash to hard reboot #"
      echo "# for the new bank plan to take effect!        #"
      echo "################################################"
      echo -e "${NC}"
      echo c > /proc/sysrq-trigger;
    else
      echo "# Rebooting to apply new bank plan!            #"
      echo "################################################"
      echo -e "${NC}"
      sync
      reboot
    fi
  else
    echo -e "${GREEN}Done! No reboot required...${NC}"
  fi
fi