#!/bin/bash
### BEGIN INIT INFO
# Provides:          umountfs
# Required-Start:
# Required-Stop:     
# Default-Start:
# Default-Stop:      0 6
# Short-Description: Turn off swap and unmount all local file systems.
# Description:
### END INIT INFO

# The firmware upgrade script initializes the flash_upgrade directory
# for doing the upgrade.
#
# If you wish to erase /var/volatile and any U-Boot environment, along
# with user-data, do touch /var/volatile/do_wipe
#
# Note that /var/oem is not touched.  /var/oem could be erased manually
# as follows:
#
# Verify the partition:
#
# grep 'OEM Config' /proc/mtd
#
# Assuming /var/config is /dev/mtd7:
#
# umount -l /var/oem
# flash_erase -j /dev/mtd7 0 0

PATH=/sbin:/bin:/usr/sbin:/usr/bin
upgrade_fname="/var/volatile/do_flash_upgrade"
wipe_fname="/var/volatile/do_wipe"

umount_all() {
  echo "Deactivating swap..."
  swapoff -a

  # We leave /proc mounted.
  echo "Unmounting local filesystems..."
  grep -q /mnt/ram /proc/mounts && mount -o remount,ro /mnt/ram
  umount -f -a -r > /dev/null 2>&1

  mount -o remount,ro /
}

get_file_size() {
    local file_size=0
    local regex='[^[:space:]]+[[:space:]]+[^[:space:]]+[[:space:]]+([^[:space:]]+)'
    local upgrade_file="$1"
    local file="$2"
    if [[ $(tar -tvf ${upgrade_file} ${file}) =~ ${regex} ]] ; then
        file_size=${BASH_REMATCH[1]}
        echo "${file_size}"
    else
        echo "0"
    fi
}

blink_leds() {
  led_dir=/sys/class/leds
  # blink all programmable LEDs except status
  leds=$(ls $led_dir/ | grep -v status)
  nleds=$(echo "$leds" | wc -w)

  # turn LEDs off
  for led in $leds; do
      echo "timer" > $led_dir/$led/trigger
      echo 0  > $led_dir/$led/delay_off
  done

  int=200
  total=$(( int * (nleds+1) ))
  on=$(( total - int ))
  off=$(( total - on ))

  # make a pattern
  for led in $leds; do
      echo "$on"  > $led_dir/$led/delay_on
      echo "$off"  > $led_dir/$led/delay_off
      on=$(( on - int ))
      off=$(( total - on ))
  done
}

err_leds() {
  led_dir=/sys/class/leds
  # blink all programmable LEDs except status
  leds=$(ls $led_dir/ | grep -v status)
  nleds=$(echo "$leds" | wc -w)

  # turn LEDs off
  for led in $leds; do
      echo "timer" > $led_dir/$led/trigger
      echo 0  > $led_dir/$led/delay_off
  done

  int=200
  total=$(( int * (nleds+1) ))
  on=$(( total - int ))
  off=$(( total - on ))

  # make a pattern
  for led in $leds; do
      echo "$on"  > $led_dir/$led/delay_on
      echo "$off"  > $led_dir/$led/delay_off
  done
  sleep 5
}



flash_upgrade() {
  supercap=$(mts-io-sysfs show capability/supercap 2>/dev/null)
  if ((supercap == 1)) ; then
    count=0
    while ((count < 180)) ; do
      powerfail=$(mts-io-sysfs show power-fail)
      if ((powerfail == 1)) ; then
	 logger -s -t 'flash_upgrade' -p daemon.err 'Power has failed. Skipping flash upgrade. Try again later'
	 return
      fi
      full=$(mts-io-sysfs show supercap-full)
      if ((full == 1)); then
	break  # Continue with flash
      fi
      logger -s -t 'flash_upgrade' -p daemon.err 'Supercap is not yet full. Wait for flash upgrade.'
      ((count++))
      sleep 1
    done # Looping on 3 minute timeout
  fi
  if [ $# -ne 1 ]; then
    echo "need to specify flash-root"
    return
  fi
  # flash_root must be a mountpoint that is not the rootfs and be mounted rw
  # External mount may be on symlink.
  local flash_root=$(readlink -f ${1})

  local flash_dir=${flash_root}/flash-upgrade
  local upgrade_file=${flash_dir}/upgrade.bin
  local kernel_solo_file=${flash_dir}/uImage.bin
  local rootfs_solo_file=${flash_dir}/rootfs.jffs2
  local uboot_solo_file=${flash_dir}/u-boot.bin

  local reboot_cmd=/usr/sbin/upgrade-reboot
  local nandwrite_cmd=/usr/bin/nandwrite.static
  local mode=
  local bootstrap_mtd=
  local config_mtd=
  local oem_mtd=
  local rootfs_mtd=
  local kernel_mtd=
  local uboot_mtd=
  local ubootenv1_mtd=
  local ubootenv2_mtd=
  local userd_mtd=
  local bootstrap_file=bstrap.bin
  local config_file=config.jffs2
  local oem_file=oem.jffs2
  local rootfs_file=rootfs.jffs2
  local kernel_file=uImage.bin
  local uboot_file=uboot.bin
  local install_file=install.sh

  if [ ! -d "${flash_dir}" ]; then
    echo "${flash_dir} not present, skipping"
    return
  fi

  if ! mountpoint -q "${flash_root}"; then
    echo "${flash_root} is not a mountpoint"
    return
  fi

  mode=$(grep "${flash_root}" /proc/mounts | cut -d ' ' -f 4 | cut -d ',' -f 1)
  if [ "${mode}" != "rw" ]; then
    echo "${flash_root} is not mounted rw"
    return
  fi

  if [ ! -x "${reboot_cmd}" ]; then
    echo "${reboot_cmd} is not installed"
    return
  fi

  if [ ! -x "${nandwrite_cmd}" ]; then
    echo "${nandwrite_cmd} is not installed"
    return
  fi

  bootstrap_mtd="/dev/$(cat /proc/mtd | grep '\"at91bootstrap\"' | cut -d : -f 1)"
  if [ ! -c "${bootstrap_mtd}" ]; then
    echo "No valid MTD partition is labeled at91bootstrap"
    return
  fi

  config_mtd="/dev/$(cat /proc/mtd | grep '\"Config\"' | cut -d : -f 1)"
  if [ ! -c "${config_mtd}" ]; then
    echo "No valid MTD partition is labeled Config"
    return
  fi

  oem_mtd="/dev/$(cat /proc/mtd | grep '\"OEM Config\"' | cut -d : -f 1)"
  if [ ! -c "${oem_mtd}" ]; then
    echo "No valid MTD partition is labeled OEM Config"
    return
  fi

  kernel_mtd="/dev/$(cat /proc/mtd | grep '\"uImage\"' | cut -d : -f 1)"
  if [ ! -c "${kernel_mtd}" ]; then
    echo "No valid MTD partition is labeled uImage"
    return
  fi

  uboot_mtd="/dev/$(cat /proc/mtd | grep '\"u-Boot\"' | cut -d : -f 1)"
  if [ ! -c "${uboot_mtd}" ]; then
    echo "No valid MTD partition is labeled u-Boot"
    return
  fi

  ubootenv1_mtd="/dev/$(cat /proc/mtd | grep '\"u-Boot Config\"' | cut -d : -f 1)"
  if [ ! -c "${uboot_mtd}" ]; then
    echo "No valid MTD partition is labeled u-Boot"
    return
  fi

  ubootenv2_mtd="/dev/$(cat /proc/mtd | grep '\"u-Boot Redundant Config\"' | cut -d : -f 1)"
  if [ ! -c "${uboot_mtd}" ]; then
    echo "No valid MTD partition is labeled u-Boot"
    return
  fi
  
  rootfs_mtd="/dev/$(cat /proc/mtd | grep '\"Rootfs\"' | cut -d : -f 1)"
  if [ ! -c "${rootfs_mtd}" ]; then
    echo "No valid MTD partition is labeled Rootfs"
    return
  fi

  userd_mtd="/dev/$(cat /proc/mtd | grep '\"User data\"' | cut -d : -f 1)"
  if [ ! -c "${userd_mtd}" ]; then
    echo "No valid MTD partition is labeled User data"
    return
  fi

  cd /var/volatile

  # rootfs should always be flashed last, so always keep it last here
  if ((factory_up == 0)) ; then
    local files=(${bootstrap_file} ${uboot_file} ${config_file} ${oem_file} ${kernel_file} ${rootfs_file} ${install_file})
  else
    local files=(${bootstrap_file} ${uboot_file} ${config_file} ${oem_file} ${kernel_file} ${rootfs_file})
  fi
  local devs=(${bootstrap_mtd} ${uboot_mtd} ${config_mtd} ${oem_mtd} ${kernel_mtd} ${rootfs_mtd})

  if [ -f "${upgrade_file}" ]; then
    echo "Found ${upgrade_file}"
    # make sure firmware is for this device
    shopt -s nocasematch
    local hw_version=$(cat /sys/devices/platform/mts-io/hw-version)
    local firmware_model=$(tar -xO -f ${upgrade_file} model)
    # Do case insensity character match.
    shopt -s nocasematch
    if [[ ! "${hw_version}" =~ ^(${firmware_model})- ]]; then
      echo "Wrong firmware for this hardware"
      echo "hw version: ${hw_version}"
      echo "firmware model: ${firmware_model}"
      return
    fi
    shopt -u nocasematch

    echo "Checking MD5s"
    ((upgrade_script=0))
    # check md5sum
    for (( i = 0; i < ${#files[@]}; i++ )); do
      local file=${files[i]}
      if tar -t -f ${upgrade_file} | grep -F -q "${file}"; then
        if ! tar -x -f ${upgrade_file} ${file}.md5; then
          echo "MD5 not found for ${file}"
          return
        fi
        if ! tar -xO -f ${upgrade_file} ${file} | md5sum -c ${file}.md5; then
          echo "MD5 check failed for ${file}"
          return
        fi
        if [[ ${file} == ${install_file} ]] ; then
          ((upgrade_script=1))
        fi
      fi
    done
  fi

  echo ""
  if [ -f "${upgrade_file}" ]; then
    echo "Starting flash upgrade from ${upgrade_file}..."
  elif [ -f "${kernel_solo_file}" ] || [ -f "${rootfs_solo_file}" ] || [ -f "${uboot_solo_file}" ] ; then
    echo "Starting flash upgrade from ${flash_dir}..."
  else
    return
  fi

  blink_leds
  
  if [[ ${flash_root} == /var/volatile ]] ; then
    oldIFS="${IFS}"                                                                                                             
    IFS=$'\n' rmlist=($(find ${flash_root} -xdev -maxdepth 1 -print | egrep -v "^.*.md5$|^${flash_root}$|^${flash_dir}$|^${flash_dir}/"))
    IFS="${oldIFS}"                                            
    ((i=0))                        
    while((i < ${#rmlist[@]})) ; do
        rm -rf "${rmlist[$i]}"
        ((i++))
    done      
  fi

  if ! cp ${reboot_cmd} ${flash_dir}/upgrade-reboot ; then
      echo "Aborting upgrade.  Failed \"cp ${reboot_cmd} ${flash_dir}/upgrade-reboot\""
      err_leds
      return 1
  fi
  if ! cp ${nandwrite_cmd} ${flash_dir}/nandwrite.static ; then
     echo "Aborting upgrade.  Failed \"cp ${nandwrite_cmd} ${flash_dir}/nandwrite.static\""
     err_leds
     return 1
  fi

  # Keep flash_root mounted
  cd ${flash_root}
  # Run the upgrade script, if it exists, before we change
  # flash-root to read-only
  if ((upgrade_script == 1)) ; then
      if (cd ${flash_dir};tar -xf ${upgrade_file}) ; then
          if ! [[ -x ${flash_dir}/${install_file} ]] ; then
              /bin/mount -o remount,exec ${flash_root}
          fi
          exec ${flash_dir}/${install_file} ${flash_root}
      else
	  echo "Extraction of ${upgrade_file} failed"
	  echo "Upgrade aborted"
	  umount_all
      fi
  fi
  sync
  sleep 2
  mount -o remount,ro ${flash_root}

  # flash_root is not going to be umounted
  sed -i -e "\\|${flash_root}| d" /etc/mtab

  umount_all
  mount -o remount,rw /var/volatile
  grep volatile /proc/mounts

  if [ -f "${upgrade_file}" ]; then
    for (( i = 0; i < ${#files[@]}; i++ )); do
      local file=${files[i]}
      local dev=${devs[i]}
      echo "Processing file $file"
      if tar -t -f ${upgrade_file} | grep -F -q "${file}"; then
        if [ "${file}" == "${bootstrap_file}" ] || [ "${file}" == "${uboot_file}" ]; then
          file_size=$(get_file_size "${upgrade_file}" "${file}")
          if dd if=${dev} bs=${file_size} count=1 | md5sum -c ${file}.md5; then
            echo "Found ${file} in upgrade but it is the same as current. Continuing..."
            continue
          fi
        elif ((factory_up == 1)) ; then # Not AT91Bootstrap or U-Boot.
            file_size=$(get_file_size "${upgrade_file}" "${file}")
            if [ "${file}" == "${kernel_file}" ]; then
                if dd if=${dev} bs=${file_size} count=1 | md5sum -c ${file}.md5; then
                    echo "Found ${file} in upgrade but it is the same as current. Continuing..."
                    continue
                fi  # Matching md5sum
            fi # kernel_file case
        fi # factory upgrade
        echo "Flashing ${dev} with ${file}..."

        flash_erase ${extraopt} -j ${dev} 0 0
        tar -xO -f ${upgrade_file} ${file} | ${flash_dir}/nandwrite.static ${extraopt} -p ${dev}
        if ((factory_up == 1)) ; then
            if [[ ${file} == rootfs.jffs2 ]] ; then
                pwd
                rm -f ${upgrade_file}
                echo "Block size is $rootbs, count=$rootcount"
                dd if=${dev} bs=${rootbs} count=${rootcount} | md5sum -c ${file}.md5
                result=$?
            else
                dd if=${dev} bs=${file_size} count=1 | md5sum -c ${file}.md5
                result=$?
            fi
            if ((result == 0)); then
                echo "SUCCESS: File ${file} was written successfully"
            else
                if [[ ${file} == rootfs.jffs2 ]] ; then
                    echo "dd if=${dev} bs=${rootbs} count=${rootcount} | md5sum"
                    dd if=${dev} bs=${rootbs} count=${rootcount} | md5sum
                else
                    echo "dd if=${dev} bs=${file_size} count=1 | md5sum -c ${file}.md5"
                    dd if=${dev} bs=${file_size} count=1 | md5sum -c ${file}.md5
                fi
                echo "ERROR: File ${file} is bad"
            fi
        fi # factory upgrade
      fi # Found a file in the upgrade package
    done # Loop through the files in an upgrade package
  else
    if [ -f ${uboot_solo_file} ]; then
      echo "Flashing ${uboot_mtd} (u-boot) with ${uboot_solo_file}..."

      flash_erase ${extraopt} ${uboot_mtd} 0 0
      ${flash_dir}/nandwrite.static ${extraopt} -p ${uboot_mtd} ${uboot_solo_file}
    else
      echo "u-boot file ${uboot_solo_file} not found"
    fi

    if [ -f ${kernel_solo_file} ]; then
      echo "Flashing ${kernel_mtd} (uImage) with ${kernel_solo_file}..."

      flash_erase ${extraopt} ${kernel_mtd} 0 0
      ${flash_dir}/nandwrite.static ${extraopt} -p ${kernel_mtd} ${kernel_solo_file}
    else
      echo "uImage file ${kernel_solo_file} not found"
    fi

    if [ -f ${rootfs_solo_file} ]; then
      echo "Flashing ${rootfs_mtd} (rootfs) with ${rootfs_solo_file}..."

      flash_erase ${extraopt} -j ${rootfs_mtd} 0 0
      ${flash_dir}/nandwrite.static ${extraopt} -p ${rootfs_mtd} ${rootfs_solo_file}
    else
      echo "rootfs file ${rootfs_solo_file} not found"
    fi
  fi

  if ((wipe == 1)) ; then
      echo 'Wiping out all partitions except /var/oem'
      flash_erase ${extraopt} -j ${config_mtd} 0 0
      flash_erase ${extraopt} -j ${userd_mtd} 0 0
      flash_erase ${extraopt} ${ubootenv1_mtd} 0 0
      flash_erase ${extraopt} ${ubootenv2_mtd} 0 0
  fi

  if ((factory_up == 1)) ; then
    echo "Finished upgrade.  Halting."
    while : ; do : ; done  # Wait forever
  else
    echo "Rebooting..."
    ${flash_dir}/upgrade-reboot
  fi

  # Should not get here normally
  echo "upgrade-reboot failed"
  exit 1
}

# do flash on reboot if do_flash_upgrade exists
# and is owned by root
wipe=0
if [[ -f ${wipe_fname} ]] ; then
    owner=$(stat -c%u "${wipe_fname}")
    if ((${#owner} > 0)) && ((owner == 0)) ; then
        wipe=1
        extraopt="-q"
    fi
fi
if [[ -f ${upgrade_fname} ]] ; then
    read rootbs rootcount < ${upgrade_fname}
    echo "Root Block Size is $rootbs, rootcount is $rootcount"
    if [[ -n $rootbs ]] ; then
      factory_up=1
      extraopt="-q"
    else
      factory_up=0
    fi
    owner=$(stat -c%u "${upgrade_fname}")
    if ((${#owner} > 0)) && ((owner == 0)) ; then
      flash_upgrade /var/volatile
      flash_upgrade /media/card
      # Look for upgrade on external media besides
      # SD card.
      cd /run/media
      # See if there is a USB driver
      last=$(dirname $(ls -d */flash-upgrade))
      if ((${#last} > 0)) ; then
        flash_upgrade "/run/media/${last}"
      fi
    fi
fi

umount_all

: exit 0