#!/bin/sh

usage() {
cat <<EOH
Shows the current bank plan, and determines if it is optimal or not.

Usage: $0

Options:
 -q Quiet mode (no messages)

EOH
exit
}

QUIET=N

while getopts :q option; do
  case "${option}" in
    q) QUIET=Y;;
    *) usage;;
  esac
done

GREEN='\033[1;32m'
ORANGE='\033[0;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color
FMT=" -> ${ORANGE}%-16s${NC} : %s "

ACTIVE_BANK=""
BOOTED_BANK=""
OPTIMAL=0

test_optimal() {
  local actual="$1"
  local expected="$2"
  if [ "$actual" = "$expected" ]; then
    OPTIMAL=$(( $OPTIMAL + 1 ))
    [ $QUIET = N ] && echo -e "${GREEN} OK${NC}"
  else
    [ $QUIET = N ] && echo -e "${UNEXPECTED_CLR} !! Should be ${expected} !!${NC}"
  fi
}

handle_content() {
  local config="$1"
  local content="$2"
  if [ $QUIET = N ]; then
    if [ "$config" = "activeversion" -a -n "$content" -a "/$content/" != "Unknown" -a -n "$ACTIVE_BANK" -a -e /overlay/$ACTIVE_BANK/etc/config/versioncusto ]; then
      content="$(echo $content | cut -d- -f1)$(uci -c /overlay/$ACTIVE_BANK/etc/config -q get versioncusto.override.fwversion_suffix)"
    elif [ "$config" = "bootedversion" -a -n "$content" -a "/$content/" != "Unknown" -a -n "$BOOTED_BANK" -a -e /overlay/$BOOTED_BANK/etc/config/versioncusto ]; then
      content="$(echo $content | cut -d- -f1)$(uci -c /overlay/$BOOTED_BANK/etc/config -q get versioncusto.override.fwversion_suffix)"
    elif [ \( "$config" = "passiveversion" -o "$config" = "notbootedversion" \) -a -n "$content" -a "/$content/" != "Unknown" -a -n "$BOOTED_BANK" ]; then
      local passive="bank_1"
      content="$(echo $content | cut -d- -f1)"
      [ "$BOOTED_BANK" = "bank_1" -a "$BOOTED_BANK" != "$ACTIVE_BANK" ] && passive="bank_2"
      [ -e /overlay/$passive/etc/config/versioncusto ] && content="$content$(uci -c /overlay/$passive/etc/config -q get versioncusto.override.fwversion_suffix)"
    fi
    printf "$FMT" "$config" "$content"
  fi
  if [ "$config" = "active" ]; then
    ACTIVE_BANK="$content"
    test_optimal "$content" "$ACTIVE_OPTIMAL"
  elif [ "$config" = "booted" ]; then
    BOOTED_BANK="$content"
    test_optimal "$content" "$BOOTED_OPTIMAL"
  else
    [ $QUIET = N ] && echo
  fi
}

set_optimal_bank_plan() {
  echo
  if [ -e set-optimal-bank-plan ]; then
    echo "Run the following command:"
  else
    echo "Download https://raw.githubusercontent.com/seud0nym/tch-gui-unhide/master/utilities/set-optimal-bank-plan"
    echo "and then run the following command:"
  fi
  echo
  echo -e "${GREEN} sh set-optimal-bank-plan${NC}"
  echo
}

if [ -d /proc/banktable ]; then
  BANKTABLE=Y
  ACTIVE_OPTIMAL="bank_1"
  BOOTED_OPTIMAL="bank_2"
  UNEXPECTED_CLR="$RED"
  for f in $(find /proc/banktable -type f | sort); do
    handle_content "$(basename $f)" "$(cat $f)"
    [ "$(basename $f)" = "booted" ] && handle_content "bootedversion" $(uci -q get env.var.friendly_sw_version_activebank)
  done
else
  BANKTABLE=N
  ACTIVE_OPTIMAL="bank_2"
  BOOTED_OPTIMAL="bank_2"
  UNEXPECTED_CLR="$ORANGE"
  for f in active:partition:active activeversion:tchversion:active booted:partition:booted bootedversion:tchversion:booted inactive:partition:passive notbooted:partition:notbooted notbootedversion:tchversion:notbooted passiveversion:tchversion:passive; do
    set -- $(echo $f | tr ':' ' ')
    CONTENT=$(bootmgr $2 $3 2>/dev/null | grep -v rd_metadata)
    [ "$CONTENT" = "1" -o "$CONTENT" = "2" ] && CONTENT="bank_$CONTENT"
    [ "$1" = activeversion -a -z "$CONTENT" ] && CONTENT=$(uci -q get env.var.friendly_sw_version_activebank)
    [ "$1" = passiveversion -a -z "$CONTENT" ] && CONTENT=$(uci -q get env.var.friendly_sw_version_passivebank)
    [ -z "$CONTENT" ] && CONTENT="Unknown"
    handle_content "$1" "$CONTENT"
  done
fi

[ $QUIET = N ] && echo

if [ $QUIET = N -a -e /usr/lib/lua/processvendorinfo.lua ]; then
  if grep -q acsurl /usr/lib/lua/processvendorinfo.lua; then
    echo -e " == Booted firmware is ${GREEN}VULNERABLE${NC} to tch-exploit${NC} ==\n"
  else
    echo -e " !! ${ORANGE}WARNING${NC}: Booted firmware is ${RED}NOT VULNERABLE${NC} to tch-exploit${NC} !!\n"
  fi
fi

if [ $OPTIMAL -eq 2 ]; then
  [ $QUIET = N ] && echo -e " == Bank Plan is ${GREEN}OPTIMAL${NC} ==\n"
  exit 0
else
  if [ $QUIET = N ]; then
    if [ $BANKTABLE = Y ]; then
      echo -e " !! Current Bank Plan is ${RED}NOT OPTIMAL${NC} !!\n"
      set_optimal_bank_plan
      echo "or manually follow the instructions at:"
      echo "https://hack-technicolor.readthedocs.io/en/stable/Hacking/PostRoot/#bank-planning"
      echo
      echo "NOTE: If you are unable to set an optimal bank plan, you may need to do a complete"
      echo "      reset of the device using the following command:"
      echo
      echo "       mtd -r erase rootfs_data"
      echo
      echo "      The device will reboot, and you WILL lose root, but you will be able to set"
      echo "      the optimal bank plan after re-acquiring root."
      echo
    else
      echo -e " !! Current Bank Plan is ${ORANGE}SUB-OPTIMAL${NC} !!\n"
      echo
      echo "It is strongly recommended that you always boot from bank 2 on this device, and"
      echo "keep a rootable firmware in bank 1 for updating the bank 2 firmware."
    fi
  fi
  exit 1
fi