#!/bin/bash
#
# 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.
#
# Script which ensures that a given image has an up-to-date
# kernel partition, rootfs integrity hashes, and legacy bootloader configs.

# --- BEGIN COMMON.SH BOILERPLATE ---
# Load common CrOS utilities.  Inside the chroot this file is installed in
# /usr/lib/crosutils.  Outside the chroot we find it relative to the script's
# location.
find_common_sh() {
  local common_paths=("$(dirname "$(readlink -f "$0")")/.." /usr/lib/crosutils)
  local path

  SCRIPT_ROOT="${common_paths[0]}"
  for path in "${common_paths[@]}"; do
    if [ -r "${path}/common.sh" ]; then
      SCRIPT_ROOT="${path}"
      break
    fi
  done
}

find_common_sh
. "${SCRIPT_ROOT}/common.sh" || exit 1
# --- END COMMON.SH BOILERPLATE ---

# Need to be inside the chroot to load chromeos-common.sh
assert_inside_chroot

# Load functions and constants for chromeos-install
. /usr/lib/installer/chromeos-common.sh || exit 1
. "${BUILD_LIBRARY_DIR}/toolchain_util.sh" || exit 1
. "${BUILD_LIBRARY_DIR}/build_image_util.sh" || exit 1

switch_to_strict_mode

if [ $# -lt 2 ]; then
  echo "Usage: ${0} /PATH/TO/IMAGE IMAGE.BIN [shflags overrides]"
  exit 1
fi

IMAGE_DIR="$(readlink -f "${1}")"
BOOT_DESC_FILE="${IMAGE_DIR}/boot.desc"
IMAGE="${IMAGE_DIR}/${2}"
shift
shift
FLAG_OVERRIDES="${@}"

if [ ! -r "${BOOT_DESC_FILE}" ]; then
  warn "${BOOT_DESC_FILE} cannot be read!"
  warn "Falling back to command line parsing"
  BOOT_DESC="${@}"
else
  BOOT_DESC="$(cat ${BOOT_DESC_FILE} | tr -s '\n' ' ')"
  info "Boot-time configuration for $(dirname "${IMAGE}"): "
  cat ${BOOT_DESC_FILE} | while read line; do
    info "  ${line}"
  done
fi

if [ ! -r "${IMAGE}" ]; then
  die "${IMAGE} cannot be read!"
fi


locate_gpt
set +e

# Now parse the build settings from ${OUTPUT_DIR}/boot.desc
DEFINE_string board "${DEFAULT_BOARD}" \
  "Board we're building for."
DEFINE_string image "coreos_base.img" \
  "Full path to the coreos image to make bootable."
DEFINE_string arch "x86" \
  "Architecture to make bootable for: arm, x86, or amd64"
DEFINE_boolean enable_rootfs_verification ${FLAGS_FALSE} \
  "Default all bootloaders to NOT use kernel-based root fs integrity checking."
DEFINE_string au_key "" \
  "Filename of the au_key to install"
DEFINE_string production_track "" \
    "Use production values and a given track for update service."
DEFINE_boolean fsck_rootfs ${FLAGS_FALSE} \
  "Check integrity of the rootfs on the modified image."


# Parse the boot.desc and any overrides
eval set -- "${BOOT_DESC} ${FLAG_OVERRIDES}"
FLAGS "${@}" || exit 1

. "${BUILD_LIBRARY_DIR}/board_options.sh" || exit 1

# Only now can we die on error.  shflags functions leak non-zero error codes,
# so will die prematurely if 'switch_to_strict_mode' is specified before now.
switch_to_strict_mode -u

mount_gpt_cleanup() {
  "${SCRIPTS_DIR}/mount_gpt_image.sh" \
    -u -r "${rootfs_mountpoint}" \
    -s "${statefs_mountpoint}" -e "${espfs_mountpoint}"
}

make_image_bootable() {
  local image="$1"

  # Default to non-verified
  local enable_rootfs_verification_flag=--noenable_rootfs_verification
  if [[ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]]; then
    enable_rootfs_verification_flag=--enable_rootfs_verification
  fi

  trap "mount_gpt_cleanup" EXIT
  "${SCRIPTS_DIR}/mount_gpt_image.sh" --from "$(dirname "${image}")" \
    --image "$(basename ${image})" -r "${rootfs_mountpoint}" \
    -s "${statefulfs_mountpoint}"

  legacy_offset_size_export ${image}

  if [ -n "${FLAGS_production_track}" ]; then
    # Replace /etc/lsb-release on the image.
    "${BUILD_LIBRARY_DIR}/set_lsb_release" \
      --production_track="${FLAGS_production_track}" \
      --root="${rootfs_mountpoint}" \
      --board="${BOARD}"
  fi

  # Install an auto update key on the root before sealing it off
  if [ ! -z "${FLAGS_au_key}" ]; then
    local key_location=${rootfs_mountpoint}"/usr/share/update_engine/"
    sudo mkdir -p "${key_location}"
    sudo cp "${FLAGS_au_key}" "$key_location/update-payload-key.pub.pem"
    sudo chown root:root "$key_location/update-payload-key.pub.pem"
    sudo chmod 644 "$key_location/update-payload-key.pub.pem"
    echo "AU verification key was installed. Do not forget to resign the image!"
  fi

  # The rootfs should never be mounted rw again after this point without
  # re-calling make_image_bootable.
  sudo mount -o remount,ro "${rootfs_mountpoint}"
  # Newer `mount` will decode the filename backing the loop device,
  # so we need to dig deeper and find the answer ourselves.
  root_dev=$(awk -v mnt="${rootfs_mountpoint}" \
             '$2 == mnt { print $1 }' /proc/mounts)

  # Make the filesystem un-mountable as read-write.
  # mount_gpt_image.sh will undo this as needed.
  # TODO(wad) make sure there is parity in the signing scripts.
  if [ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]; then
    # TODO(wad) this would be a good place to reset any other ext2 metadata.
    warn "Disabling r/w mount of the root filesystem"
    disable_rw_mount "$root_dev"
  fi

  # We should update the esp in place in the image.
  local bootloader_to="${image}"
  local esp_offset="$(partoffset ${image} ${NUM_ESP})"
  esp_offset=$((esp_offset * 512))  # sectors to bytes
  local esp_size="$(partsize ${image} ${NUM_ESP})"
  esp_size=$((esp_size * 512))  # sectors to bytes
  local bootloader_to_flags="--to_offset=${esp_offset} --to_size=${esp_size}"

  # Update ESP partition
  # NOTE: Boot kernel is identical to regular kernel for now
  ${SCRIPTS_DIR}/update_bootloaders.sh \
    --arch=${FLAGS_arch} \
    --to="${bootloader_to}" \
    --from="${rootfs_mountpoint}"/boot \
    --vmlinuz_boot_kernel="${rootfs_mountpoint}"/boot/vmlinuz \
    --vmlinuz="${rootfs_mountpoint}"/boot/vmlinuz \
    ${bootloader_to_flags}

  trap - EXIT
  ${SCRIPTS_DIR}/mount_gpt_image.sh -u -r "${rootfs_mountpoint}" \
    -s "${statefulfs_mountpoint}"
}

verify_image_rootfs() {
  local image=$1
  local rootfs_offset="$(partoffset ${image} 3)"

  local rootfs_tmp_file=$(mktemp)
  trap "rm ${rootfs_tmp_file}" EXIT
  sudo dd if="${image}" of="${rootfs_tmp_file}" bs=512 skip="${rootfs_offset}" \
          status=none

  # This flips the read-only compatibility flag, so that
  # e2fsck does not complain about unknown file system capabilities.
  enable_rw_mount "${rootfs_tmp_file}"
  info "Running e2fsck to check root file system for errors"
  sudo e2fsck -fn "${rootfs_tmp_file}" ||
    die "Root file system has errors, please ensure boot.desc and/or \
command line parameters are correct"
}

# Store output and temporary files next to image.
rootfs_mountpoint="${IMAGE_DIR}/rootfs_dir"
statefulfs_mountpoint="${IMAGE_DIR}/stateful_dir"
espfs_mountpoint="${IMAGE_DIR}/esp"

# Create the directories if they don't exist.
mkdir -p ${rootfs_mountpoint}
mkdir -p ${statefulfs_mountpoint}
mkdir -p ${espfs_mountpoint}

make_image_bootable "${IMAGE}"
if [ ${FLAGS_fsck_rootfs} -eq ${FLAGS_TRUE} ]; then
  verify_image_rootfs "${IMAGE}"
fi

rmdir ${rootfs_mountpoint}
rmdir ${statefulfs_mountpoint}
rmdir ${espfs_mountpoint}