#!/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/share/misc/chromeos-common.sh
SCRIPT="$0"

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

# Log file to store the output of this run.
CLOBBER_STATE_LOG="/tmp/clobber-state.log"

# Redirect stdout to our log.
exec > "$CLOBBER_STATE_LOG" 2>&1

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 -sqw factory && echo "factory")
FAST_WIPE=$(echo "$@" | grep -sqw fast && echo "fast")
KEEPIMG=$(echo "$@" | grep -sqw keepimg && echo "keepimg")
SAFE_WIPE=$(echo "$@" | grep -sqw safe && echo "safe")
FACTORY_WIPE_FIRMWARE_PID=""

PRESERVED_TAR="/tmp/preserve.tar"
STATE_PATH="/mnt/stateful_partition"
POWERWASH_COUNT="/mnt/stateful_partition/unencrypted/preserve/powerwash_count"

# List of files to preserve relative to /mnt/stateful_partition/
PRESERVED_FILES=""

# Preserve these files in safe mode.
if [ "$SAFE_WIPE" = "safe" ]; then
  PRESERVED_FILES="
    unencrypted/preserve/attestation.epb
    unencrypted/preserve/powerwash_count
    unencrypted/preserve/enrollment_state.epb
    unencrypted/preserve/update_engine/prefs/rollback-version
    home/.shadow/install_attributes.pb
  "

  # Powerwash count is only preserved for "safe" powerwashes.
  COUNT=1
  if [ -f $POWERWASH_COUNT ]; then
    COUNT_UNSANITIZED=$(head -1 $POWERWASH_COUNT | cut -c1-4)
    if [ $(expr "$COUNT_UNSANITIZED" : "^[0-9][0-9]*$") -ne 0 ]; then
      COUNT=$(( COUNT_UNSANITIZED + 1 ))
    fi
  fi
  echo $COUNT > $POWERWASH_COUNT
fi

# For a factory wipe (and only factory) preserve seed for CRX cache.
if [ "$FACTORY_WIPE" = "factory" ]; then
  IMPORT_FILES="$(cd ${STATE_PATH};
                  echo unencrypted/import_extensions/extensions/*.crx)"
  if [ "$IMPORT_FILES" != \
       "unencrypted/import_extensions/extensions/*.crx" ]; then
    PRESERVED_FILES="${PRESERVED_FILES} ${IMPORT_FILES}"
  fi
fi

PRESERVED_LIST=""
if [ -n "$PRESERVED_FILES" ]; then
  # We want to preserve permissions and recreate the directory structure
  # for all of the files in the PRESERVED_FILES variable. In order to do
  # so we run tar --no-recurison and specify the names of each of the
  # parent directories. For example for home/.shadow/install_attributes.pb
  # we pass to tar home home/.shadow home/.shadow/install_attributes.pb
  for file in $PRESERVED_FILES; do
    if [ ! -e "$STATE_PATH/$file" ]; then
      continue
    fi
    path=$file
    while [ "$path" != '.' ]; do
      PRESERVED_LIST="$path $PRESERVED_LIST"
      path=$(dirname $path)
    done
  done
  tar cf $PRESERVED_TAR -C $STATE_PATH --no-recursion -- $PRESERVED_LIST
fi

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
    # Opening a TTY device with no one else attached will allocate and reset
    # terminal attributes. To prevent user input echo and ctrl-breaks, we need
    # to allocate and configure the terminal before using tty for output.
    # Note &2 is currently used for debug messages dumping so we can't redirect
    # subshell by 2>/dev/tty1.
    ( stty -F /dev/tty1 raw -echo -cread
      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
    ) >/dev/tty1
  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
}

# Make sure the active kernel is still bootable after being wiped.
# The system may be in AU state that active kernel does not have "successful"
# bit set to 1 (only tries).
ensure_bootable_kernel() {
  local kernel_num="$1"
  local dst="$(rootdev -d -s)"
  local active_flag="$(cgpt show -S -i "${kernel_num}" "${dst}")"
  local priority="$(cgpt show -P -i "${kernel_num}" "${dst}")"

  if [ "${active_flag}" -lt 1 ]; then
    cgpt add -i "${kernel_num}" "${dst}" -S 1
  fi
  if [ "${priority}" -lt 1 ]; then
    cgpt prioritize -i "${kernel_num}" "${dst}" -P 3
  fi
  sync
}

# 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')
KERNEL_PART_NUM=${KERNEL_DEV##[/a-z]*[/a-z]}
WIPE_PART_NUM=${OTHER_ROOT_DEV##[/a-z]*[/a-z]}
# Save the option before wipe-out
WIPE_OPTION_FILE="${STATE_PATH}/factory_wipe_option"
if [ -O ${WIPE_OPTION_FILE} ]; then
  WIPE_OPTION=$(cat ${WIPE_OPTION_FILE})
else
  WIPE_OPTION=
fi
# Discover type of device holding the stateful partition; assume SSD.
# Since there doesn't seem to be a good way to get from a partition name
# to the base device name beyond simple heuristics, just find the device
# with the same major number but with minor 0.
ROTATIONAL=0
MAJOR=$(stat -c %t ${STATE_DEV})
for i in $(find /dev -type b); do
  if [ "$(stat -c %t:%T $i)" = "${MAJOR}:0" ]; then
    ROTATIONAL_PATH="/sys/block/$(basename ${i})/queue/rotational"
    if [ -r "${ROTATIONAL_PATH}" ]; then
      ROTATIONAL=$(cat "${ROTATIONAL_PATH}")
      break
    fi
  fi
done

# Sanity check root device partition number.
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" "$@"

# On a non-fast wipe, rotational drives take too long. Override to run them
# through "fast" mode, with a forced delay. Sensitive contents should already
# be encrypted.
if [ "$FAST_WIPE" != "fast" ] && [ "$ROTATIONAL" = "1" ]; then
  # If the stateful filesystem is available, do some best-effort content
  # shredding. Since the filesystem is not mounted with "data=journal", the
  # writes really are overwriting the block contents (unlike on an SSD).
  if grep -q " ${STATE_PATH} " /proc/mounts ; then
    (
      # Directly remove things that are already encrypted (which are also the
      # large things), or are static from images.
      rm -rf "${STATE_PATH}/encrypted.block" \
             "${STATE_PATH}/var_overlay" \
             "${STATE_PATH}/dev_image"
      find "${STATE_PATH}/home/.shadow" -maxdepth 2 -type d \
                                        -name vault -print0 |
        xargs -r0 rm -rf
      # Shred everything else. We care about contents not filenames, so do not
      # use "-u" since metadata updates via fdatasync dominate the shred time.
      # Note that if the count-down is interrupted, the reset file continues
      # to exist, which correctly continues to indicate a needed wipe.
      find "${STATE_PATH}"/. -type f -print0 | xargs -r0 shred -fz
      sync
    ) &
  fi
  # Since the above rm/shred combo can be very fast, force a minimum of a 5
  # minute delay for this mode.
  delay=300
  ( stty -F /dev/tty1 raw -echo -cread
    while [ "${delay}" -ge 0 ]; do
      printf "%2d:%02d\r" $(( delay / 60 )) $(( delay % 60 ))
      sleep 1
      : $(( delay -= 1 ))
    done
    echo
  ) >/dev/tty1
  wait
  FAST_WIPE="fast"
fi
# Make sure the stateful partition has been unmounted.
umount -n "${STATE_PATH}"

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
    ensure_bootable_kernel "${KERNEL_PART_NUM}"
    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
else

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

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

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

# Mount the fresh image for last minute additions.
mount -n -t ext4 "${STATE_DEV}" "${STATE_PATH}"


# If there were preserved files, restore them.
if [ -n "$PRESERVED_LIST" ]; then
  tar xfp $PRESERVED_TAR -C $STATE_PATH
  touch /mnt/stateful_partition/unencrypted/.powerwash_completed
fi

# 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" && ! crossystem "mainfw_act?recovery"; 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} ${WIPE_OPTION}
  fi
fi

# Stop appending to the log and preserve it.
exec > /dev/null 2>&1
mv -f "$CLOBBER_STATE_LOG" "$STATE_PATH/unencrypted/clobber-state.log"

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