#!/bin/sh -u
# Copyright (c) 2012 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.
#
# A script to install from removable media to hard disk.

# If we're not running as root, restart as root.
if [ "${UID:-$(id -u)}" -ne 0 ]; then
  exec sudo "$0" "$@"
fi

# Load functions and constants for chromeos-install.
. /usr/share/misc/chromeos-common.sh || exit 1
. /usr/share/misc/shflags || exit 1

# Source blocksize
SRC_BLKSIZE=512

# Helpful constants.
HARDWARE_DIAGNOSTICS_PATH=/tmp/hardware_diagnostics.log
TMPMNT=/tmp/install-mount-point
# This is defined later once we have mounted the ROOT.
STATEFUL_FORMAT=
# Partition numbers that have assumptions about them. This list should be kept
# to a minimal. Check copy_partition for most special casing.

# TODO(installer): Clean up all these flags. There are way too many flags in
# this script.

DEFINE_string dst "" "Destination device"
DEFINE_boolean skip_src_removable ${FLAGS_FALSE} \
  "Skip check to ensure source is removable"
DEFINE_boolean skip_dst_removable ${FLAGS_FALSE} \
  "Skip check to ensure destination is not removable"
DEFINE_boolean skip_rootfs ${FLAGS_FALSE} \
  "Skip installing the rootfs; Only set up partition table"
DEFINE_boolean yes ${FLAGS_FALSE} \
  "Answer yes to everything"
DEFINE_boolean preserve_stateful ${FLAGS_FALSE} \
  "Don't create a new filesystem for the stateful partition. Be careful \
using this option as this may make the stateful partition not mountable."
DEFINE_string payload_image "" "Path to a Chromium OS image to install onto \
the device's hard drive"
DEFINE_string gpt_layout "" "Path to a script for pre-defined GPT partition \
layout"
DEFINE_string pmbr_code "" "Path to PMBR code to be installed"
DEFINE_string target_bios "" "Bios type to boot with (see postinst --bios)"
DEFINE_boolean mtd_layout ${FLAGS_FALSE} "This system uses MTD for \
partitioning rather than GPT"
DEFINE_boolean debug ${FLAGS_FALSE} "Show debug output"
DEFINE_boolean large_test_partitions ${FLAGS_FALSE} \
  "Make partitions 9 and 10 large (for filesystem testing)"
DEFINE_boolean skip_postinstall ${FLAGS_FALSE} \
  "Skip postinstall for situations where you're building for a \
non-native arch. Note that this will probably break verity."
DEFINE_string lab_preserve_logs "" "Path to a file containing logs to be \
preserved"
DEFINE_boolean storage_diags "${FLAGS_FALSE}" "Print storage diagnostic \
information on failure"
DEFINE_string oobe_pub_key "" "Path to public key for OOBE auto-configuration \
validation"
DEFINE_string oobe_priv_key "" "Path to private key for OOBE \
auto-configuration signing"


# Parse command line
FLAGS "$@" || exit 1
eval set -- "${FLAGS_ARGV}"

die() {
  echo "$*" >&2
  exit 1
}

fast_dd() {
  # Usage: fast_dd <count> <seek> <skip> other dd args
  # Note: <count> and <seek> are in units of SRC_BLKSIZE, while <skip> is in
  # units of DST_BLKSIZE. Supply 0 for <count> to not include any count=
  # argument in dd.
  local user_count="$1"
  local user_seek="$2"
  local user_skip="$3"
  shift 3
  # Provide some simple progress updates to the user.
  set -- "$@" status=progress
  # Find the largest block size that all the parameters are a factor of.
  local block_size=$((2 * 1024 * 1024))
  while [ $(((user_count * SRC_BLKSIZE) % block_size)) -ne 0 ] || \
         [ $(((user_skip * SRC_BLKSIZE) % block_size)) -ne 0 ] || \
         [ $(((user_seek * DST_BLKSIZE) % block_size)) -ne 0 ]; do

    : $((block_size /= 2))
  done

  # Print a humble info line if the block size is not super, and complain more
  # loudly if it's really small (and the partition is big).
  if [ "${block_size}" -ne $((2 * 1024 * 1024)) ]; then
    echo "DD with block size ${block_size}"
    if [ "${block_size}" -lt $((128 * 1024)) ] && \
        [ $((user_count * SRC_BLKSIZE)) -gt $((128 * 1024 * 1024)) ]; then
      echo
      echo "WARNING: DOING A SLOW MISALIGNED dd OPERATION. PLEASE FIX"
      echo "count=${user_count} seek=${user_seek} skip=${user_skip}"
      echo "SRC_BLKSIZE=${SRC_BLKSIZE} DST_BLKSIZE=${DST_BLKSIZE}"
      echo
    fi
  fi

  # Convert the block counts in their respective sizes into the common block
  # size, and blast off.
  local count_common=$((user_count * SRC_BLKSIZE / block_size))
  local seek_common=$((user_seek * DST_BLKSIZE / block_size))
  local skip_common=$((user_skip * SRC_BLKSIZE / block_size))
  local count_arg=""
  if [ "${count_common}" -ne 0 ]; then
    count_arg="count=${count_common}"
  fi

  dd "$@" bs="${block_size}" seek="${seek_common}" skip="${skip_common}" \
      "${count_arg}"
}

# Get the specified env var for the specified partition.
#  $1 the field name such as "PARTITION_SIZE", "FS_FORMAT"
#  $2 the partition such as "1", or "ROOT_A"
_get_field() {
  local field part
  field="$1"
  part="$2"
  eval echo \""\${${field}_${part}}"\"
}

get_format() {
  _get_field FORMAT "$@"
}

get_fs_format() {
  _get_field FS_FORMAT "$@"
}

get_partition_size() {
  _get_field PARTITION_SIZE "$@"
}

get_reserved_ebs() {
  _get_field RESERVED_EBS "$@"
}

# Calculate the maximum number of bad blocks per 1024 blocks for UBI.
#  $1 partition number
calculate_max_beb_per_1024() {
  local part_no mtd_size eb_size nr_blocks
  part_no="$1"
  # The max beb per 1024 is on the total device size, not the partition size.
  mtd_size=$(cat /sys/class/mtd/mtd0/size)
  eb_size=$(cat /sys/class/mtd/mtd0/erasesize)
  nr_blocks=$((mtd_size / eb_size))
  reserved_ebs=$(get_reserved_ebs "${part_no}")
  echo $((reserved_ebs * 1024 / nr_blocks))
}

# Format and make UBI volume if it's not already there.
#  $1 partition number such as "1", "2"
#  $2 volume name
init_ubi_volume() {
  local part_no volume_name phy_ubi log_ubi
  part_no="$1"
  volume_name="$2"
  phy_ubi="/dev/ubi${part_no}"
  log_ubi="${phy_ubi}_0"
  if [ ! -e "${phy_ubi}" ]; then
    ubiformat -y -e 0 "/dev/mtd${part_no}"
    ubiattach -d "${part_no}" -m "${part_no}" \
              --max-beb-per1024 "$(calculate_max_beb_per_1024 "${part_no}")"
  fi
  if [ ! -e "${log_ubi}" ]; then
    local volume_size
    volume_size=$(get_partition_size "${part_no}")
    ubimkvol -s "${volume_size}" -N "${volume_name}" "${phy_ubi}"
  fi
}

# Update a specific partition in the destination device.
write_partition() {
  local user_part="$1"
  local user_count="$2"
  local user_seek="$3"
  local user_skip="$4"
  local src="$5"
  local dst="$6"
  local format fs_format

  if [ "${user_count}" -eq 0 ]; then
    echo "Skipping partition as it does not exist"
    return 0
  fi

  format="$(get_format "${user_part}")"
  case ${format} in
  nand)
    flash_erase "/dev/mtd${user_part}" 0 0
    nandwrite --input-skip $((user_skip * SRC_BLKSIZE)) \
      --input-size $((user_count * SRC_BLKSIZE)) \
      "/dev/mtd${user_part}" "${src}"
    ;;

  ubi)
    local phy_ubi="/dev/ubi${user_part}"
    local log_ubi="${phy_ubi}_0"
    local sysfs_name="/sys/class/mtd/mtd${user_part}/name"

    init_ubi_volume "${user_part}" "$(cat "${sysfs_name}")"

    fs_format="$(get_fs_format "${user_part}")"
    case ${fs_format} in
    ubifs)
      local src_mnt="${TMPMNT}/src" dst_mnt="${TMPMNT}/dst"
      mkdir -p "${src_mnt}" "${dst_mnt}"

      mkfs.ubifs -y -x none -R 0 "${log_ubi}"
      mount "${log_ubi}" "${dst_mnt}"

      # Have to copy the files over by hand as the source partition is a
      # different filesystem type (like ext4).
      loop_offset_setup "${src}" "${user_seek}" "${SRC_BLKSIZE}"
      TMPMNT="${src_mnt}" mount_on_loop_dev
      cp -a "${src_mnt}"/* "${dst_mnt}"/
      TMPMNT="${src_mnt}" umount_from_loop_dev
      loop_offset_cleanup

      umount "${log_ubi}"
      ;;
    *)
      ubiupdatevol --skip $((user_skip * SRC_BLKSIZE)) \
        --size $((user_count * SRC_BLKSIZE)) "${log_ubi}" "${src}"
      ;;
    esac
    ;;

  *)
    fast_dd "${user_count}" "${user_seek}" "${user_skip}" \
      if="${src}" of="${dst}" conv=notrunc
    ;;
  esac
}

# Find root partition of the block device that we are installing from
get_root_device() {
  rootdev -s
}

# Check for optional payload image
check_payload_image() {
  if [ "${FLAGS_skip_rootfs}" -eq "${FLAGS_TRUE}" ] && \
      [ -s "${FLAGS_gpt_layout}" ]; then
    # Usually this is used for partition setup.
    SRC=""
    ROOT=""
  elif [ -z "${FLAGS_payload_image}" ]; then
    # Find root partition of the root block device
    SRC=$(get_block_dev_from_partition_dev "$(get_root_device)")
    ROOT=""
  else
    if [ ! -e "${FLAGS_payload_image}" ]; then
      die "Error: No payload image found at ${FLAGS_payload_image}"
    fi

    # Needed to copy PMBR code off image
    SRC="${FLAGS_payload_image}"
    ROOT="$(mktemp -d)"
  fi
}

# Clean any mounts that might be present to avoid
# aliasing access to block devices.
prepare_disk() {
  if [ -e /etc/init/cros-disks.conf ]; then
    initctl stop cros-disks || true
  fi
  # Often times, nothing is mounted, so swallow the warnings.
  umount -f /media/*/* 2>&1 | \
    grep -v -i -F \
      -e 'no mount point specified' \
      -e 'not mounted' \
      -e 'No such file or directory' \
      -e 'not found' || true
}

# Like mount but keeps track of the current mounts so that they can be cleaned
# up automatically.
tracked_mount() {
  local last_arg
  eval last_arg=\$$#
  MOUNTS="${last_arg}${MOUNTS:+ }${MOUNTS:-}"
  mount "$@"
}

# Unmount with tracking.
tracked_umount() {
  # dash does not support ${//} expansions.
  local new_mounts
  for mount in $MOUNTS; do
    if [ "$mount" != "$1" ]; then
      new_mounts="${new_mounts:-}${new_mounts+ }$mount"
    fi
  done
  MOUNTS=${new_mounts:-}

  umount "$1"
}

# Create a loop device on the given file at a specified (sector) offset.
# Remember the loop device using the global variable LOOP_DEV.
# Invoke as: command
# Args: FILE OFFSET BLKSIZE
loop_offset_setup() {
  local filename=$1
  local offset=$2
  local blocksize=$3

  set --
  if [ "${blocksize}" -ne 512 ]; then
    set -- "$@" -b "${blocksize}"
  fi

  LOOP_DEV=$(losetup -f "$@" --show -o $((offset * blocksize)) "${filename}")
  if [ -z "$LOOP_DEV" ]; then
    die "No free loop device. Free up a loop device or reboot. Exiting."
  fi

  LOOPS="${LOOP_DEV}${LOOPS:+ }${LOOPS:-}"
}

# Delete the current loop device.
loop_offset_cleanup() {
  # dash does not support ${//} expansions.
  local new_loops
  for loop in $LOOPS; do
    if [ "$loop" != "$LOOP_DEV" ]; then
      new_loops="${new_loops:-}${new_loops+ }$loop"
    fi
  done
  LOOPS=${new_loops:-}

  # losetup -a doesn't always show every active device, so we'll always try to
  # delete what we think is the active one without checking first. Report
  # success no matter what.
  losetup -d "${LOOP_DEV}" || /bin/true
}

# Mount the existing loop device at the mountpoint in $TMPMNT.
# Args: optional 'readwrite'. If present, mount read-write, otherwise read-only.
mount_on_loop_dev() {
  local rw_flag=${1-readonly}
  set --
  if [ "${rw_flag}" != "readwrite" ]; then
    set -- "$@" -o ro
  fi
  tracked_mount "$@" "${LOOP_DEV}" "${TMPMNT}"
}

# Unmount loop-mounted device.
umount_from_loop_dev() {
  mount | grep -q " on ${TMPMNT} " && tracked_umount "${TMPMNT}"
}

# Check if all arguments are non-empty values
check_non_empty_values() {
  local value
  for value in "$@"; do
    if [ -z "$value" ]; then
      return ${FLAGS_FALSE}
    fi
  done
  return ${FLAGS_TRUE}
}

# Undo all mounts and loops and runs hw diagnostics on failure.
cleanup_on_failure() {
  set +e

  if [ "${FLAGS_storage_diags}" -eq "${FLAGS_TRUE}" ]; then
    if [ -b "${DST-}" ]; then
      # Generate the diagnostics log that can be used by a caller.
      echo "Running a hw diagnostics test -- this might take a couple minutes."
      badblocks -e 100 -sv "${DST}" 2>&1 | tee "${HARDWARE_DIAGNOSTICS_PATH}"
    fi

    if [ -f /usr/share/misc/storage-info-common.sh ]; then
      . /usr/share/misc/storage-info-common.sh
      # Run a few extra diagnostics with output to stdout. These will
      # be stored as part of the recovery.log for recovery images.
      get_storage_info
    fi
  fi

  cleanup
}

# Undo all mounts and loops.
cleanup() {
  set +e

  local mount_point
  for mount_point in ${MOUNTS:-}; do
    umount "${mount_point}" || /bin/true
  done
  MOUNTS=""

  local loop_dev
  for loop_dev in ${LOOPS:-}; do
    losetup -d "${loop_dev}" || /bin/true
  done
  LOOPS=""

  if [ -n "${ROOT}" ]; then
    rmdir "${ROOT}"
  fi
}

check_removable() {
  if [ "${FLAGS_skip_dst_removable}" -eq "${FLAGS_TRUE}" ]; then
    return
  fi

  local removable

  if ! removable="$(cat "/sys/block/${DST#/dev/}/removable")"; then
    die "Error: Invalid destination device (must be whole device): ${DST}"
  fi

  if [ "${removable}" != "0" ]; then
    die "Error: Attempt to install to a removeable device: ${DST}"
  fi
}

# Wipes and expands the stateful partition.
wipe_stateful() {
  echo "Clearing the stateful partition..."
  local stateful_fs_format
  stateful_fs_format="$(get_fs_format "${PARTITION_NUM_STATE}")"
  # state options are stored in $@.
  set --

  case "${STATEFUL_FORMAT}" in
  ubi)
    local phy_ubi="/dev/ubi${PARTITION_NUM_STATE}"
    local log_ubi="${phy_ubi}_0"
    local sysfs_name="/sys/class/mtd/mtd${PARTITION_NUM_STATE}/name"

    init_ubi_volume "${PARTITION_NUM_STATE}" "$(cat "${sysfs_name}")"
    ;;
  *)
    loop_offset_setup "${DST}" "${START_STATEFUL}" "${DST_BLKSIZE}"
    ;;
  esac

  # Check if the kernel we are going to install support ext4 crypto.
  if ext4_dir_encryption_supported; then
    set -- "$@" -O encrypt
  fi

  local num_4k_sectors
  if [ "${DST_BLKSIZE}" -gt 4096 ]; then
    num_4k_sectors=$(( NUM_STATEFUL_SECTORS * (DST_BLKSIZE / 4096) ))
  else
    num_4k_sectors=$(( NUM_STATEFUL_SECTORS / (4096 / DST_BLKSIZE) ))
  fi

  # We always make any ext* stateful partitions ext4.
  case "${stateful_fs_format}" in
  ext[234])
    mkfs.ext4 -F -b 4096 -L "H-STATE" "$@" "${LOOP_DEV}" \
      ${num_4k_sectors}
    ;;
  ubifs)
    mkfs.ubifs -y -x none -R 0 "/dev/ubi${PARTITION_NUM_STATE}_0"
    ;;
  esac

  case ${STATEFUL_FORMAT} in
  ubi) ;;
  *)
    # Need to synchronize before releasing loop device, otherwise calling
    # loop_offset_cleanup may return "device busy" error.
    sync
    loop_offset_cleanup
    ;;
  esac

  # When the stateful partition is wiped the TPM ownership must be reset. This
  # command will not work on older devices which do not support it. In that case
  # it will be ignored.
  crossystem clear_tpm_owner_request=1 || true
}

# Install the stateful partition content
# Method handles copying data over to the stateful partition. This is done
# differently than other partitions due to the EXPAND option i.e. src partition
# and dst partitions are of different sizes. In addition, there are some special
# tweaks we do for stateful here for various workflows.
install_stateful() {
  # In general, the system isn't allowed to depend on anything
  # being in the stateful partition at startup.  We make some
  # exceptions for dev images (only), as enumerated below:
  #
  # var_overlay
  #   These are included to support gmerge, and must be kept in
  #   sync with those listed in /etc/init/var-overlay.conf:
  #      db/pkg
  #      lib/portage
  #
  # dev_image
  #   This provides tools specifically chosen to be mounted at
  #   /usr/local as development only tools.
  #
  # Every exception added makes the dev image different from
  # the release image, which could mask bugs.  Make sure every
  # item you add here is well justified.
  echo "Installing the stateful partition..."
  case "${STATEFUL_FORMAT}" in
  ubi)
    # We modify the global used here as it affects how we unmount later.
    LOOP_DEV="/dev/ubi${PARTITION_NUM_STATE}_0"
    ;;
  *)
    loop_offset_setup "${DST}" "${START_STATEFUL}" "${DST_BLKSIZE}"
    ;;
  esac
  mount_on_loop_dev readwrite

  # Move log files listed in FLAGS_lab_preserve_logs from stateful_partition to
  # a dedicated location. This flag is used to enable Autotest to collect log
  # files before reimage deleting all prior logs.
  if crossystem 'cros_debug?1' && [ -n "${FLAGS_lab_preserve_logs}" ]; then
    local gatherme="${TMPMNT}/.gatherme"
    touch "${gatherme}"
    local prior_log_dir="${TMPMNT}/unencrypted/prior_logs"
    mkdir -p "${prior_log_dir}"
    local log_path
    for log_path in $(sed -e '/^#/ d' -e '/^$/ d' "${FLAGS_lab_preserve_logs}"); do
      case "${log_path}" in
        /dev/* | /sys/*)
          ;;
        /*)
          echo "${log_path}" >> "${gatherme}"
          continue
          ;;
        *)
          log_path="${TMPMNT}/${log_path}"
          ;;
      esac
      if [ -d "${log_path}" ]; then
        cp -au -r --parents "${log_path}" "${prior_log_dir}" || true
      elif [ -f "${log_path}" ]; then
        cp -au "${log_path}" "${prior_log_dir}" || true
      fi
    done
  fi

  # Whitelist files to copy onto the stateful partition.
  #
  # When adding to the whitelist, consider the need for related changes in
  # src/platform/init/chromeos_startup, and in src/platform/dev/stateful_update.
  #
  local dirlist="
    unencrypted/cros-components/offline-demo-mode-resources
    unencrypted/import_extensions
  "

  if crossystem 'cros_debug?1'; then
    dirlist="
      ${dirlist}
      var_overlay/db/pkg
      var_overlay/lib/portage
      dev_image
    "
  fi

  if crossystem 'devsw_boot?1' ; then
    # This is a base build, and the dev switch was on when we booted;
    # we assume it will be on for the next boot.  We touch
    # ".developer_mode" to avoid a pointless delay after reboot while
    # chromeos_startup wipes an empty stateful partition.
    #
    # See chromeos_startup for the companion code that checks for this
    # file.
    #
    touch ${TMPMNT}/.developer_mode
  fi

  if [ -n "${IS_RECOVERY_INSTALL-}" ] ; then
    # This is a recovery install; write some recovery metrics to the stateful
    # partition to be reported after next boot. See:
    # init/upstart/send-recovery-metrics.conf
    local recovery_histograms="${TMPMNT}/.recovery_histograms"
    metrics_client -W "${recovery_histograms}" -e "Installer.Recovery.Reason" \
      "$(crossystem recovery_reason)" 255
  fi

  local dir
  for dir in ${dirlist}; do
    if [ ! -d "${ROOT}/mnt/stateful_partition/${dir}" ]; then
      continue
    fi
    local parent
    parent=$(dirname "${dir}")
    mkdir -p "${TMPMNT}/${parent}"
    cp -au "${ROOT}/mnt/stateful_partition/${dir}" "${TMPMNT}/${dir}"
  done

  if [ -n "${FLAGS_oobe_pub_key}" ] && [ -n "${FLAGS_oobe_priv_key}" ]; then
    echo "Finalizing OOBE auto-config setup..."

    # Start udevd since it might not be running, and finish_oobe_auto_config
    # needs it for walking /dev/disk/by-id/.
    udevd --daemon
    udevadm trigger
    udevadm settle

    local stateful_device
    stateful_device="$(cgpt find -l STATE "${FLAGS_payload_image}")"

    finish_oobe_auto_config \
      --private_key="${FLAGS_oobe_priv_key}" \
      --public_key="${FLAGS_oobe_pub_key}" \
      --src_stateful_dev="${stateful_device}" \
      --src_stateful="${ROOT}/mnt/stateful_partition" \
      --dst_stateful="${TMPMNT}"
  fi

  umount_from_loop_dev
  case ${STATEFUL_FORMAT} in
  ubi) ;;
  *)
    loop_offset_cleanup
    ;;
  esac
}

# Copy partition from src to dst (figures out partition offsets). Note, this
# has some special casing for rootfs, kernel, and stateful partitions. In
# addition, it only copies partitions that are equally sized over one another.
# $1 - Partition number we are copying to.
# $2 - src image
# $3 - dst image.
copy_partition() {
  local part_num=$1
  local src=$2
  local dst=$3
  local part_size
  local src_offset
  local dst_offset

  part_size="$(partsize "${src}" "${part_num}")"
  src_offset="$(partoffset "${src}" "${part_num}")"
  dst_offset="$(partoffset "${dst}" "${part_num}")"

  echo "Installing partition ${part_num} to ${dst}"

  case "${part_num}" in
  "${PARTITION_NUM_STATE}")
    install_stateful
    ;;
  "${PARTITION_NUM_ROOT_A}"|"${PARTITION_NUM_ROOT_B}")
    # Always copy from ROOT_A for rootfs partitions.
    part_size=$(partsize "${src}" "${PARTITION_NUM_ROOT_A}")
    src_offset=$(partoffset "${src}" "${PARTITION_NUM_ROOT_A}")
    write_partition "${part_num}" "${part_size}" "${dst_offset}" \
      "${src_offset}" "${src}" "${dst}"
    ;;
  "${PARTITION_NUM_KERN_A}"|"${PARTITION_NUM_KERN_B}")
    # Use kernel B from the source into both kernel A and B in the destination.
    part_size="$(partsize "${src}" "${PARTITION_NUM_KERN_B}")"
    src_offset="$(partoffset "${src}" "${PARTITION_NUM_KERN_B}")"
    write_partition "${part_num}" "${part_size}" "${dst_offset}" \
      "${src_offset}" "${src}" "${dst}"
    ;;
  *)
    local src_part_size
    local dst_part_size
    src_part_size="$((part_size * SRC_BLKSIZE))"
    dst_part_size="$(partsize "${dst}" "${part_num}")"
    dst_part_size="$((dst_part_size * DST_BLKSIZE))"
    if [ "${src_part_size}" -ne "${dst_part_size}" ] || \
        [ "${src_part_size}" -le 4096 ]; then
      # We only copy partitions that are equally sized and greater than the
      # min fs block size. This matches the build_image logic.
      return
    fi
    write_partition "${part_num}" "${part_size}" "${dst_offset}" \
      "${src_offset}" "${src}" "${dst}"
    ;;
  esac
}

# Remove partitions 1 to 12 from MTD device. Recreate partitions 1 to 12 with
# information from the current GPT table.
# $1 is the device node, such as "/dev/mtd0"
recreate_nand_partitions() {
  local dst=$1
  local blocksize=$2
  local part_no

  for part_no in /dev/mtd*; do
    part_no="$(echo "${part_no}" | cut -b9- | grep '^[0-9]*$' || :)"
    if [ -z "${part_no}" ] || [ "${part_no}" = "0" ]; then
      continue
    fi
    # Ignore any error in case no UBI volume is attached.
    ubidetach -m "${part_no}" > /dev/null 2>&1 || true
    nand_partition del "${dst}" "${part_no}"
  done

  local gpt_file mtd_size mtd_gpt_file part_size part_offset
  gpt_file=$(mktemp)
  flashrom -r "-iRW_GPT:${gpt_file}"
  mtd_size=$(cat "/sys/class/mtd/$(basename "${dst}")/size")
  mtd_gpt_file="-D ${mtd_size} ${gpt_file}"
  for part_no in $(seq "${PARTITION_NUM_STATE}" \
      "${PARTITION_NUM_EFI_SYSTEM}"); do
    part_size=$(partsize "${mtd_gpt_file}" "${part_no}")
    : $(( part_size *= blocksize ))
    part_offset=$(partoffset "${mtd_gpt_file}" "${part_no}")
    : $(( part_offset *= blocksize ))
    nand_partition add "${dst}" "${part_no}" \
      "${part_offset}" "${part_size}"
  done
  rm -f "${gpt_file}"
}

# Find our destination device.
# If the user hasn't selected a destination,
# we expect that the disk layout declares it for us.
check_dst() {
  if [ -z "${DST}" ]; then
    die "Error: can not determine destination device. Specify --dst yourself."
  fi

  if [ "${DST}" = "/dev/mtd0" ]; then
    FLAGS_mtd_layout=${FLAGS_TRUE}
  fi

  # Check out the dst device.
  if [ "${FLAGS_mtd_layout}" -eq "${FLAGS_TRUE}" ]; then
    FLAGS_skip_dst_removable=${FLAGS_TRUE}
  elif [ ! -b "${DST}" ]; then
    die "Error: Unable to find destination block device: ${DST}"
  fi

  if [ "${DST}" = "${SRC}" ]; then
    die "Error: src and dst are the same: ${SRC} = ${DST}"
  fi
}

# Gets the right PMBR (protective master boot record) code (either from
# FLAGS_pmbr_code, source or destination media) by printing the file path
# containing PMBR code in standard out.
get_pmbr_code() {
  local pmbr_code="/tmp/gptmbr.bin"

  if [ -n "${FLAGS_pmbr_code}" ]; then
    echo "${FLAGS_pmbr_code}"
  elif [ "${FLAGS_mtd_layout}" -eq "${FLAGS_TRUE}" ]; then
    # We don't use PMBR if this is on MTD.
    dd bs="${DST_BLKSIZE}" \
      count=1 if=/dev/zero of="${pmbr_code}" >/dev/null 2>&1
    echo "${pmbr_code}"
  else
    # Steal the PMBR code from the source MBR to put on the dest MBR, for
    # booting on legacy-BIOS devices.
    dd bs="${DST_BLKSIZE}" count=1 if="${SRC}" of="${pmbr_code}" >/dev/null 2>&1
    echo "${pmbr_code}"
  fi
}

# Reload the system partitions after the partition table was modified (so the
# device nodes like /dev/sda1 can be accessed).
reload_partitions() {
  if [ "${FLAGS_mtd_layout}" -eq "${FLAGS_FALSE}" ]; then
    # Reload the partition table on block devices only.
    # On MTD, the ChromeOS kernel loads the partition table at boot time.
    #
    # In some cases, we may be racing with udev for access to the
    # device leading to EBUSY when we reread the partition table.  We
    # avoid the conflict by using `udevadm settle`, so that udev goes
    # first.  cf. crbug.com/343681.
    udevadm settle
    /sbin/blockdev --rereadpt "${DST}"
  else
    # On NAND, we need to recreate the partition table.
    recreate_nand_partitions "${DST}" "${DST_BLKSIZE}"
  fi
}

# Post partition copying work and special casing
do_post_install() {
  set --
  if [ -n "${FLAGS_target_bios}" ]; then
    set -- "$@" --bios "${FLAGS_target_bios}"
  fi
  if [ "${FLAGS_debug}" -eq "${FLAGS_TRUE}" ]; then
    set -- "$@" --debug
  fi
  local dst_rootfs_dev=""

  # Now run the postinstall script on one new rootfs. Note that even though
  # we're passing the new destination partition number as an arg, the postinst
  # script had better not try to access it, for the reasons we just gave.
  # We can't run this if the target arch isn't the same as the host arch
  if [ "${FLAGS_skip_postinstall}" -eq "${FLAGS_FALSE}" ]; then
    if [ ${FLAGS_mtd_layout} -eq ${FLAGS_TRUE} ]; then
      if ! [ -b "/dev/ubiblock${PARTITION_NUM_ROOT_A}_0" ]; then
        ubiblock -c "/dev/ubi${PARTITION_NUM_ROOT_A}_0"
      fi
      LOOP_DEV="/dev/ubiblock${PARTITION_NUM_ROOT_A}_0"
      # We need to pass the __writable__ device to postinst, hence ubiX_0.
      dst_rootfs_dev="/dev/ubi${PARTITION_NUM_ROOT_A}_0"
    else
      loop_offset_setup "${DST}" "${START_ROOTFS_A}" "${DST_BLKSIZE}"
      dst_rootfs_dev="$(make_partition_dev "${DST}" "${PARTITION_NUM_ROOT_A}")"
    fi

    mount_on_loop_dev
    IS_INSTALL="1" "${TMPMNT}/postinst" "${dst_rootfs_dev}" "$@"
    umount_from_loop_dev
    if [ "${FLAGS_mtd_layout}" -eq "${FLAGS_FALSE}" ]; then
      loop_offset_cleanup
    fi
  fi
}

legacy_offset_size_export() {
  # Exports all the variables that install_gpt did previously.
  # This should disappear eventually, but it's here to make existing
  # code work for now.

  START_STATEFUL="$(partoffset "$1" "${PARTITION_NUM_STATE}")"
  START_ROOTFS_A="$(partoffset "$1" "${PARTITION_NUM_ROOT_A}")"

  NUM_STATEFUL_SECTORS="$(partsize "$1" "${PARTITION_NUM_STATE}")"
}

main() {
  # Be aggressive.
  set -eu
  if [ "${FLAGS_debug}" = "${FLAGS_TRUE}" ]; then
    set -x
  fi

  check_payload_image
  mkdir -p "${TMPMNT}"

  # We untrap on success and run cleanup ourselves. Otherwise, on any failure,
  # run our custom trap method to gather any diagnostic data before cleaning up.
  trap cleanup_on_failure EXIT

  # Clean media browser mounts if they've popped up.
  prepare_disk
  locate_gpt

  # Special handling for payload_image. This is passed in for recovery images
  # and USB installs. This is done first so we can read the gpt partition
  # file below.
  if [ -n "${FLAGS_payload_image}" ]; then
    PARTITION_NUM_ROOT_A=$(cgpt find -n -l ROOT-A "${FLAGS_payload_image}")
    PARTITION_NUM_STATE=$(cgpt find -n -l STATE "${FLAGS_payload_image}")

    SRC="${FLAGS_payload_image}"
    # Mount files that are required to be referenced (when not already mounted).
    loop_offset_setup "${SRC}" \
      "$(partoffset "${SRC}" "${PARTITION_NUM_ROOT_A}")" 512
    tracked_mount -o ro "${LOOP_DEV}" "${ROOT}"
    loop_offset_setup "${SRC}" \
      "$(partoffset "${SRC}" "${PARTITION_NUM_STATE}")" 512
    tracked_mount -o ro "${LOOP_DEV}" "${ROOT}"/mnt/stateful_partition
  fi

  # Reload the GPT helper functions and the image settings from target root.
  . "${ROOT}/usr/sbin/write_gpt.sh"
  load_base_vars

  # This was moved out of check_payload_image because DEFAULT_ROOTDEV
  # is not defined until after the GPT helper functions are loaded.
  if [ "${FLAGS_skip_src_removable}" -eq "${FLAGS_FALSE}" ]; then
    if [ "$(cat "/sys/block/${SRC#/dev/}/removable")" != "1" ]; then
      # Removable flag is implemented inconsistently for ARM sdcard reader.
      # Allow all devices except the default fixed drive.
      if [ "${SRC}" = "$(get_fixed_dst_drive)" ]; then
        trap - EXIT
        cleanup
        die "Error: Source can not be the destination device: ${SRC}"
      fi
    fi
  fi

  # Now that we have loaded the partition table we can actually read the format
  # and partition information.
  STATEFUL_FORMAT="$(get_format "${PARTITION_NUM_STATE}")"

  DST=${FLAGS_dst:-$(get_fixed_dst_drive)}
  check_dst
  check_removable

  DST_BLKSIZE="$(blocksize "${DST}")"

  # Ask for confirmation to be sure.
  echo "This will install from '${SRC}' to '${DST}'."
  echo "This will erase all data at this destination: ${DST}"
  local sure
  if [ "${FLAGS_yes}" -eq "${FLAGS_FALSE}" ]; then
    read -r -p "Are you sure (y/N)? " sure
    if [ "${sure}" != "y" ]; then
      # Don't run diagnostics if the user explicitly bailed out.
      trap - EXIT
      cleanup
      die "Ok, better safe than sorry; you answered '${sure}'."
    fi
  fi

  # Write the GPT using the board specific script. The parameters are ignored
  # on MTD devices.
  write_base_table "${DST}" "$(get_pmbr_code)"
  legacy_offset_size_export "${DST}"
  reload_partitions

  if [ "${FLAGS_skip_rootfs}" -eq "${FLAGS_TRUE}" ]; then
    echo "Done installing partitions."
    exit 0
  fi

  if [ "${FLAGS_preserve_stateful}" -eq "${FLAGS_FALSE}" ] && \
      [ -z "${FLAGS_lab_preserve_logs}" ]; then
    wipe_stateful
  fi

  # Install the content. Do so in reverse order to have stateful get installed
  # last. The order shouldn't matter but legacy behavior has us go in reverse
  # order.
  local current_part
  for current_part in $(seq "${PARTITION_NUM_EFI_SYSTEM}" -1 \
      "${PARTITION_NUM_STATE}"); do
    copy_partition "${current_part}" "${SRC}" "${DST}"
  done

  do_post_install

  # Force data to disk before we declare done.
  sync
  cleanup
  trap - EXIT

  echo "------------------------------------------------------------"
  echo ""
  echo "Installation to '${DST}' complete."
  echo "Please shutdown, remove the USB device, cross your fingers, and reboot."
}

main "$@"