#!/bin/sh

# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This script can be called during startup to trash the stateful partition
# and possibly reset the other root filesystem

. /usr/sbin/chromeos-common.sh
SCRIPT="$0"

# stateful partition isn't around for logging, so dump to the screen:
set -x

if [ $(id -u) -ne 0 ]; then
  echo 'You must run this as root'
  exit 1
fi

# Run "clobber-state fast" to clobber the state quickly but unsecurely.
FACTORY_WIPE=$(echo "$@" | grep -sq factory && echo "factory")
FAST_WIPE=$(echo "$@" | grep -sq fast && echo "fast")
KEEPIMG=$(echo "$@" | grep -sq keepimg && echo "keepimg")
FACTORY_WIPE_FIRMWARE_PID=""

wipedev() {
  DEV="$1"
  # Wipe the filesystem size if we can determine it. Full partition wipe
  # takes a long time on 16G SSD or rotating media.
  if dumpe2fs ${DEV} ; then
    FS_BLOCKS=$(dumpe2fs -h ${DEV} | grep "Block count" | sed "s/^.*:\W*//")
    FS_BLOCKSIZE=$(dumpe2fs -h ${DEV} | grep "Block size" | sed "s/^.*:\W*//")
  else
    FS_BLOCKS=$(numsectors ${DEV})
    FS_BLOCKSIZE=512
  fi
  BLOCKS_4M=$((4 * 1024 * 1024 / $FS_BLOCKSIZE))  # 4MiB in sectors
  FULL_BLKS=$(($FS_BLOCKS / $BLOCKS_4M))
  REMAINDER_SECTORS=$(($FS_BLOCKS % $BLOCKS_4M))

  if type pv; then
    pv -etp -s $((4 * 1024 * 1024 * ${FULL_BLKS})) </dev/zero 2>/dev/tty1 |
      dd of=${DEV} bs=4M count=${FULL_BLKS} iflag=fullblock oflag=sync
  else
    dd if=/dev/zero of=${DEV} bs=4M count=${FULL_BLKS}
  fi
  dd if=/dev/zero of=${DEV} bs=${FS_BLOCKSIZE} count=${REMAINDER_SECTORS} \
    seek=$(($FULL_BLKS * $BLOCKS_4M))
}

# Wipe (developer) firmware if running in factory wiping mode.
factory_wipe_firmware() {
  if [ -n "$FACTORY_WIPE_FIRMWARE_PID" ]; then
    # wiping already in progress; wait for complete.
    wait "$FACTORY_WIPE_FIRMWARE_PID" || return $?
    return 0
  fi

  local wipe_script='/usr/sbin/firmware-factory-wipe'
  # Factory wipe doesn't occur in developer mode w.r.t firmware wiping.
  if crossystem "devsw_boot?1" || [ ! -x "${wipe_script}" ]; then
    return 0
  fi
  if [ "$FAST_WIPE" = "fast" ]; then
    echo "Starting to wipe firmware in parallel..."
    "${wipe_script}" & FACTORY_WIPE_FIRMWARE_PID="$!"
  else
    echo "Starting to wipe firmware..."
    "${wipe_script}"
  fi
}

# Root devs are sda3, sda5.
# Kernel devs to go along with these are sda2, sda4 respectively.
ROOT_DEV=$(rootdev -s)
STATE_DEV=${ROOT_DEV%[0-9]*}1
OTHER_ROOT_DEV=$(echo $ROOT_DEV | tr '35' '53')
KERNEL_DEV=$(echo $ROOT_DEV | tr '35' '24')
OTHER_KERNEL_DEV=$(echo $OTHER_ROOT_DEV | tr '35' '24')
WIPE_PART_NUM=${OTHER_ROOT_DEV##[/a-z]*[/a-z]}

if [ "$WIPE_PART_NUM" != "3" ] && [ "$WIPE_PART_NUM" != "5" ]
then
  echo "Invalid partition to wipe, $WIPE_PART_NUM (${OTHER_ROOT_DEV})"
  exit 1
fi

# Preserve the log file
clobber-log --preserve "$SCRIPT" "$@"

umount /mnt/stateful_partition

if [ "$FAST_WIPE" = "fast" ]; then
  # factory_wipe_firmware can be executed in parallel.
  if [ "$FACTORY_WIPE" = "factory" ]; then
    factory_wipe_firmware
  fi

  # Just wipe the start of the partition and remake the fs on
  # the stateful partition.
  dd bs=4M count=1 if=/dev/zero of=${STATE_DEV}
  if [ -z "${KEEPIMG}" ]; then
    dd bs=4M count=1 if=/dev/zero of=${OTHER_ROOT_DEV}
    dd bs=4M count=1 if=/dev/zero of=${OTHER_KERNEL_DEV}
  fi
  /sbin/mkfs.ext3 "$STATE_DEV"
else

  if [ -z "${KEEPIMG}" ]; then
    wipedev ${OTHER_ROOT_DEV}
    wipedev ${OTHER_KERNEL_DEV}
  fi

  # Wipe everything on the stateful partition.
  wipedev ${STATE_DEV}

  /sbin/mkfs.ext3 "$STATE_DEV"
  # TODO(wad) tune2fs.
fi

# Mount the fresh image for last minute additions.
mount -n -t ext3 "${STATE_DEV}" /mnt/stateful_partition

# Restore the log file
clobber-log --restore "$SCRIPT" "$@"

# Tag that we're in developer mode otherwise we may get wiped again.
if crossystem devsw_boot?1 recoverysw_boot?0; then
  touch /mnt/stateful_partition/.developer_mode
fi

# Flush linux caches.
sync
echo 3 > /proc/sys/vm/drop_caches

# Do any board specific wiping here.
# board_factory_wipe.sh will be installed by the board overlay if necessary.
if [ "$FACTORY_WIPE" = "factory" ]; then
  factory_wipe_firmware
  BOARD_WIPE=/usr/sbin/board_factory_wipe.sh
  if [ -x "${BOARD_WIPE}" ]; then
    ${BOARD_WIPE}
  fi
fi

/sbin/shutdown -r now
sleep 1d  # Wait for shutdown