#!/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.
. "$(dirname "$0")/chromeos-common.sh" || exit 1
. /usr/share/misc/shflags || exit 1

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 skip_vblock ${FLAGS_FALSE} \
  "Skip copying the HD vblock to install destination"
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 arch "" \
  "Architecture for this image, must be one of 'ARM' or 'INTEL'.  If \
unset, auto-detect."
DEFINE_string payload_image "" "Path to a Chromium OS image to install onto \
the device's hard drive"
DEFINE_boolean use_payload_kern_b ${FLAGS_FALSE} \
  "Copy KERN-B instead of KERN-A from payload_image."
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 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."


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

set -e

SUB_CMD_DEBUG_FLAG=""
if [ "$FLAGS_debug" -eq "${FLAGS_TRUE}" ]; then
  set -x
  SUB_CMD_DEBUG_FLAG="--debug"
fi

# Determine our architecture
if [ -z "$FLAGS_arch" ]; then
  case $(uname -m) in
  i?86|x86_64)
    ARCH="INTEL"
    ;;
  arm*)
    ARCH="ARM"
    ;;
  *)
    echo "Error: Failed to auto detect architecture" >&2
    exit 1
  esac
else
  case ${FLAGS_arch} in
  INTEL|ARM)
    ;;
  *)
    echo "Error: Unknown architecture '$FLAGS_arch'." >& 2
    exit 1
  esac
  ARCH="$FLAGS_arch"
fi

fast_dd() {
  # Usage: fast_dd <block size> <count> <seek> <skip> other dd args
  local user_block_size="$1"
  shift
  local user_count="$1"
  shift
  local user_seek="$1"
  shift
  local user_skip="$1"
  shift
  local ideal_block_size=$((2 * 1024 * 1024))  # 2 MiB
  if [ $(($ideal_block_size % $user_block_size)) -eq 0 ]; then
    local factor=$(($ideal_block_size / $user_block_size))
    if [ $(($user_count % $factor)) -eq 0 -a \
         $(($user_seek % $factor)) -eq 0 -a \
         $(($user_skip % $factor)) -eq 0 ]; then
      local count_arg=""
      if [ "$user_count" -ne 0 ]; then
        count_arg="count=$(($user_count / $factor))"
      fi
      sudo dd $* bs="$ideal_block_size" seek=$(($user_seek / $factor)) \
          skip=$(($user_skip / $factor)) $count_arg
      return
    fi
  fi
  # Give up and do the user's slow dd
  echo
  echo WARNING: DOING A SLOW dd OPERATION. PLEASE FIX
  echo
  local count_arg=""
  if [ "$user_count" -ne 0 ]; then
    count_arg="count=$user_count"
  fi
  sudo dd $* bs="$user_block_size" seek="$user_seek" skip="$user_skip" \
    $count_arg
}

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

# Check for optional payload image
if [ "$FLAGS_skip_rootfs" -eq "$FLAGS_TRUE" -a -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=""

  if [ "$FLAGS_skip_src_removable" -eq "${FLAGS_FALSE}" ]; then
    if [ "$(cat /sys/block/${SRC#/dev/}/removable)" != "1" ]; then
      # Work around ARM kernel bug http://crosbug.com/14871
      # Removable flag is implemented inconsistantly for ARM sdcard reader.
      if [ "${SRC}" != "/dev/mmcblk1" ]; then
        echo "Error: Source does not look like a removable device: $SRC"
        exit 1
      fi
    fi
  fi
else
  if [ ! -e "$FLAGS_payload_image" ]; then
    echo "Error: No payload image found at $FLAGS_payload_image"
    exit 1
  fi

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

# Find our destination device
if [ -z "$FLAGS_dst" ]; then
  if [ "$ARCH" = "INTEL" ]; then
    # This finds the first ATA device listed by parted.
    SD_INFO=$(sudo parted -l -s -m 2> /dev/null)
    DST=$(expr "${SD_INFO}" : ".*\(/dev/sd[^:]*\)[^;]*:ATA\ ")
    if [ -z "$DST" ]; then
      # As a backup, install to /dev/sda
      DST=/dev/sda
    fi
  else
    DST=/dev/mmcblk0
  fi
else
  DST="$FLAGS_dst"
fi

# Check out the dst device.
if [ ! -b "$DST" ]; then
  echo "Error: Unable to find destination block device: $DST"
  exit 1
fi

DST_REMOVABLE=$(cat /sys/block/${DST#/dev/}/removable)
if [ $? -ne 0 ]; then
  echo "Error: Invalid destination device (must be whole device): $DST"
  exit 1
fi

if [ "$FLAGS_skip_dst_removable" -eq "${FLAGS_FALSE}" ]; then
  if [ "$DST_REMOVABLE" != "0" ]; then
    echo "Error: Attempt to install to a removeable device: $DST"
    exit 1
  fi
fi

if [ "$DST" = "$SRC" ]; then
  echo "Error: src and dst the same: $SRC = $DST"
  exit 1
fi

# Ask for root password to be sure.
echo "This will install from '$SRC' to '$DST'. If you are sure this is"
echo "what you want then feel free to enter the root password to proceed."
sudo -K

echo "This will erase all data at this destination: $DST"
if [ "${FLAGS_yes}" -eq "$FLAGS_FALSE" ]; then
  read -p "Are you sure (y/N)? " SURE
  if [ "$SURE" != "y" ]; then
    echo "Ok, better safe than sorry; you answered '$SURE'."
    exit 1
  fi
fi

##############################################################################
# Helpful constants and functions.

PMBRCODE=/tmp/gptmbr.bin
TMPFILE=/tmp/install-temp-file
TMPMNT=/tmp/install-mount-point
mkdir -p ${TMPMNT}

# Clean any mounts that might be present to avoid
# aliasing access to block devices.
prepare_disk() {
  if [ -e /etc/init/cros-disks.conf ]; then
    sudo initctl stop cros-disks || true
  fi
  # Often times, nothing is mounted, so swallow the warnings.
  sudo umount -f /media/*/* 2>&1 | \
    grep -v -i \
      -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:-}"
  sudo mount "$@"
}

# Unmount with tracking.
tracked_umount() {
  # dash doesnt 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:-}

  sudo 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
loop_offset_setup() {
  local filename=$1
  local offset=$2

  LOOP_DEV=$(sudo losetup -f)
  if [ -z "$LOOP_DEV" ]; then
    echo "No free loop device. Free up a loop device or reboot. Exiting."
    exit 1
  fi
  sudo losetup -o $(($offset * 512)) ${LOOP_DEV} ${filename}

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

# Delete the current loop device.
loop_offset_cleanup() {
  # dash doesnt 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.
  sudo 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}
  local mount_flags=""
  if [ "${rw_flag}" != "readwrite" ]; then
    mount_flags="-o ro"
  fi
  tracked_mount ${mount_flags} ${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.
cleanup() {
  set +e

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

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

  if [ ! -z "$ROOT" ]; then
    rmdir "$ROOT"
  fi
}
trap cleanup EXIT

##############################################################################

KERNEL_IMG_OFFSET=0
ROOTFS_IMG_OFFSET=0
OEM_IMG_OFFSET=0
ESP_IMG_OFFSET=0

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

# What do we expect & require to have on the source device?
if [ -z "$FLAGS_payload_image" ]; then
  STATEFUL_IMG=$(make_partition_dev ${SRC} 1)
  KERNEL_IMG=$(make_partition_dev ${SRC} 2)
  ROOTFS_IMG=$(make_partition_dev ${SRC} 3)
  OEM_IMG=$(make_partition_dev ${SRC} 8)
  ESP_IMG=$(make_partition_dev ${SRC} 12)
else
  KERNEL_IMG=${FLAGS_payload_image}
  ROOTFS_IMG=${FLAGS_payload_image}
  OEM_IMG=${FLAGS_payload_image}
  ESP_IMG=${FLAGS_payload_image}

  locate_gpt

  STATEFUL_IMG_OFFSET=$(partoffset "${FLAGS_payload_image}" 1)
  KERNEL_IMG_OFFSET=$(partoffset "${FLAGS_payload_image}" 2)
  ROOTFS_IMG_OFFSET=$(partoffset "${FLAGS_payload_image}" 3)
  OEM_IMG_OFFSET=$(partoffset "${FLAGS_payload_image}" 8)
  ESP_IMG_OFFSET=$(partoffset "${FLAGS_payload_image}" 12)

  if [ ${FLAGS_use_payload_kern_b} -eq ${FLAGS_TRUE} ]; then
    KERNEL_IMG_SECTORS=$(partsize "${FLAGS_payload_image}" 4)
    KERNEL_IMG_OFFSET=$(partoffset "${FLAGS_payload_image}" 4)
  fi

  # Mount the src image
  loop_offset_setup "${FLAGS_payload_image}" $STATEFUL_IMG_OFFSET
  STATEFUL_IMG_LOOP=$LOOP_DEV
  loop_offset_setup "${FLAGS_payload_image}" $ROOTFS_IMG_OFFSET
  ROOTFS_IMG_LOOP=$LOOP_DEV

  tracked_mount -o ro "$ROOTFS_IMG_LOOP" "$ROOT"
  tracked_mount -o ro "$STATEFUL_IMG_LOOP" "$ROOT"/mnt/stateful_partition
fi

if [ -n "${FLAGS_pmbr_code}" ]; then
  PMBRCODE="${FLAGS_pmbr_code}"
elif [ "$ARCH" = "ARM" ]; then
  # Store the existing PMBR from the destination in the temp pmbr file so we can
  # restore it when we're installing the GPT. This ensures that we preserve
  # settings set by crossystem (such as dev_boot_usb).
  sudo dd bs=512 count=1 if=$DST of=$PMBRCODE
else
  # Steal the PMBR code from the source MBR to put on the dest MBR, for booting
  # on legacy-BIOS devices.
  sudo dd bs=512 count=1 if=$SRC of=$PMBRCODE
fi

# Write the GPT using the board specific script.
. "/usr/sbin/write_gpt.sh"
write_base_table ${DST} ${PMBRCODE}
sudo /sbin/blockdev --rereadpt ${DST}
legacy_offset_size_export "${DST}"

# TODO(tgao): add support for arm recovery

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

# Install the content.
echo "Copying kernel..."
fast_dd 512 ${KERNEL_IMG_SECTORS} ${START_KERN_A} ${KERNEL_IMG_OFFSET} \
  if=${KERNEL_IMG} of=${DST} conv=notrunc,fsync

# Copy kernel verification block to target HD
copy_kernel_vblock() {
  local ret=0

  if [ -s $VERIFY_BLOB ]; then
    fast_dd 512 0 ${START_KERN_A} 0 if=${VERIFY_BLOB} of=${DST} conv=notrunc
    fast_dd 512 0 ${START_KERN_B} 0 if=${VERIFY_BLOB} of=${DST} conv=notrunc
    ret=1
  fi
  echo ${ret}
}

# The recovery image is signed with different keys from the hard-disk image. We
# need to copy the hard disk verification block from the stateful partition so
# that reboot from HD can be verified.
VERIFY_BLOB="${ROOT}/mnt/stateful_partition/vmlinuz_hd.vblock"
if [ "$FLAGS_skip_vblock" -eq "$FLAGS_FALSE" ]; then
  VBLOCK_RET=$(copy_kernel_vblock)

  if [ "${VBLOCK_RET}" -eq "1" ]; then
    echo "Copied kernel verification blob from ${VERIFY_BLOB}"
  else
    echo "Error: kernel verification blob not found in stateful partition"
    exit 1
  fi
fi

echo "Copying rootfs..."
# We can no longer update the label on the rootfs because that changes bits
# that will break both the delta updater and verified boot.  We must do a
# straight copy now.  The GPT labels and UUIDs are the only mutable naming
# areas we have after a build.
fast_dd 512 ${ROOTFS_IMG_SECTORS} ${START_ROOTFS_A} ${ROOTFS_IMG_OFFSET} \
  if=${ROOTFS_IMG} of=${DST} conv=notrunc

echo "Copying OEM customization..."
fast_dd 512 ${OEM_IMG_SECTORS} ${START_OEM} ${OEM_IMG_OFFSET} \
  if=${OEM_IMG} of=${DST} conv=notrunc
echo "Copying ESP..."
fast_dd 512 ${ESP_IMG_SECTORS} ${START_ESP} ${ESP_IMG_OFFSET} \
  if=${ESP_IMG} of=${DST} conv=notrunc

# If postinst fails, we should still clear stateful.
if [ "${FLAGS_preserve_stateful}" -eq "${FLAGS_FALSE}" ]; then
  echo "Clearing the stateful partition..."
  loop_offset_setup $DST $START_STATEFUL
  sudo mkfs.ext4 -F -b 4096 -L "H-STATE" ${LOOP_DEV} \
    $(($NUM_STATEFUL_SECTORS / 8))
  # Need to synchronize before releasing loop device, otherwise calling
  # loop_offset_cleanup may return "device busy" error.
  sync
  loop_offset_cleanup
fi

POSTINST_BIOS_ARGS=""
if [ -n "${FLAGS_target_bios}" ]; then
  POSTINST_BIOS_ARGS="--bios ${FLAGS_target_bios}"
fi

# 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
  loop_offset_setup ${DST} ${START_ROOTFS_A}
  mount_on_loop_dev
  sudo IS_INSTALL="1" ${TMPMNT}/postinst $(make_partition_dev ${DST} 3) \
    ${SUB_CMD_DEBUG_FLAG} "${POSTINST_BIOS_ARGS}"
  umount_from_loop_dev
  loop_offset_cleanup
fi

#
# Install the stateful partition content
#
# 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..."
loop_offset_setup $DST $START_STATEFUL
mount_on_loop_dev readwrite
if [ -f "${ROOT}"/root/.dev_mode ]; then
  DIRLIST="
      var_overlay/db/pkg
      var_overlay/lib/portage
      dev_image
  "
  for DIR in ${DIRLIST}; do
    if [ ! -d "${ROOT}/mnt/stateful_partition/${DIR}" ]; then
      continue
    fi
    OPT_U=""
    if cp -u /dev/null /dev/zero 2>/dev/null; then
      OPT_U="u"
    fi
    PARENT=$(dirname ${DIR})
    sudo mkdir -p ${TMPMNT}/${PARENT}
    sudo cp -a${OPT_U} "${ROOT}/mnt/stateful_partition/${DIR}" ${TMPMNT}/${DIR}
  done
elif crossystem 'devsw_boot?1' ; then
  # 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
umount_from_loop_dev
loop_offset_cleanup

# Force data to disk before we declare done.
sync

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