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

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

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

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() {
  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
  local flash_root=${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 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

  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

  cd /var/volatile

  # rootfs should always be flashed last, so always keep it last here
  local files=(${bootstrap_file} ${uboot_file} ${config_file} ${oem_file} ${kernel_file} ${rootfs_file} ${install_file})
  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)
    if [[ ! "${hw_version}" =~ ^${firmware_model}- ]]; then
      echo "Wrong firmware for this hardware"
      echo "hw version: ${hw_version}"
      echo "firmware model: ${firmware_model}"
      return
    fi

    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
      # Extract files early if we are going to
      # execute the upgrade script.
      if [[ ${file} == ${install_file} ]] ; then
        ((upgrade_script=1))
      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 "^${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

  # 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
          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

  if [ -f "${upgrade_file}" ]; then
    for (( i = 0; i < ${#files[@]}; i++ )); do
      local file=${files[i]}
      local dev=${devs[i]}
      if tar -t -f ${upgrade_file} | grep -F -q "${file}"; then
        if [ "${file}" == "${bootstrap_file}" ] || [ "${file}" == "${uboot_file}" ]; then
          local file_size=$(tar -xO -f ${upgrade_file} ${file} | wc -c)
          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
        fi
        echo "Flashing ${dev} with ${file}..."

        flash_erase -j ${dev} 0 0
        tar -xO -f ${upgrade_file} ${file} | ${flash_dir}/nandwrite.static -p ${dev}
      fi
    done
  else
    if [ -f ${uboot_solo_file} ]; then
      echo "Flashing ${uboot_mtd} (u-boot) with ${uboot_solo_file}..."

      flash_erase ${uboot_mtd} 0 0
      ${flash_dir}/nandwrite.static -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 ${kernel_mtd} 0 0
      ${flash_dir}/nandwrite.static -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 -j ${rootfs_mtd} 0 0
      ${flash_dir}/nandwrite.static -p ${rootfs_mtd} ${rootfs_solo_file}
    else
      echo "rootfs file ${rootfs_solo_file} not found"
    fi
  fi

  echo "Rebooting..."

  ${flash_dir}/upgrade-reboot

  # 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
if [[ -f ${upgrade_fname} ]] ; then
    owner=$(stat -c%u "${upgrade_fname}")
    if ((${#owner} > 0)) && ((owner == 0)) ; then
      flash_upgrade /var/volatile
      flash_upgrade /media/card
    fi
fi

umount_all

: exit 0