#!/bin/sh

# Copyright (c) 2009 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 is called after an AutoUpdate or USB install. The first argument
# is the partition where the new rootfs is installed or empty. If non-empty
# the rootfs should be updated w/ the new bootloader config. If empty, the
# rootfs is mounted-read only and should not be updated.

# Setup INSTALL_ROOT
INSTALL_ROOT=$(dirname "$0")
# Load helper functions
. "${INSTALL_ROOT}"/usr/sbin/chromeos-common.sh
. "${INSTALL_ROOT}"/usr/lib/shflags

# Process flags now, couldn't earlier since we needed INSTALL_ROOT
# to know where to load shflags from
DEFINE_string esp_part_file "" "File containing the ESP partition"
DEFINE_boolean postcommit ${FLAGS_FALSE} "Run postcommit"

# WARNING: do NOT change this default value without consulting adlr/nsanders.
# If you want to temporary enable firmware update for testing, create a
# tag file in ${INSTALL_ROOT}/root/.force_update_firmware
FLAGS_firmware_default=${FLAGS_FALSE}
if [ -f "${INSTALL_ROOT}/root/.force_update_firmware" ]; then
  FLAGS_firmware_default=${FLAGS_TRUE}
fi
DEFINE_boolean update_firmware $FLAGS_firmware_default "Enable firmware update."
DEFINE_boolean debug ${FLAGS_FALSE} "Show debug output. Default: false"

# 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

# Update syslinux configuration file and install syslinux itself.
INSTALL_DEV="$1"

# Find whole disk device.
ROOT_DEV=$(get_block_dev_from_partition_dev ${INSTALL_DEV})
NEW_PART_NUM=${INSTALL_DEV##*/*[a-z]}
NEW_K_DEV=$(make_partition_dev ${ROOT_DEV} $(( $NEW_PART_NUM - 1)) )

case ${NEW_PART_NUM} in
  "3" )
    BOOT_SLOT="A"
    ;;
  "5" )
    BOOT_SLOT="B"
    ;;
  * )
    # Not a valid boot location.
    echo "Not a valid target parition number: ${NEW_PART_NUM}"
    exit 1
    ;;
esac
# For legacy and efi boot, we only have one boot-capable partition.
LEGACY_PART_NUM=12

if [ -z "${IS_FACTORY_INSTALL}" -a \
     -z "${IS_RECOVERY_INSTALL}" -a \
     -z "${IS_INSTALL}" ]; then
  IS_UPDATE=1
fi

# Turn on/off x86-specific stuff
IS_INTEL_ARCH=$(uname -m | grep "^i.86\|^x86_..\$" | wc -l)

# Usage: if $(version_less A.B.C.D W.X.Y.Z); then...
version_less() {
  local a=$(printf '%010d%010d%010d%010d' $(echo "$1" | sed 's/\./ /g'))
  local b=$(printf '%010d%010d%010d%010d' $(echo "$2" | sed 's/\./ /g'))
  [ $a \< $b ]
}

src_version() {
  cat /etc/lsb-release | grep ^CHROMEOS_RELEASE_VERSION | cut -d = -f 2-
}

MAKE_DEV_READONLY=${FLAGS_FALSE}
if [ -n "$IS_UPDATE" ]; then
  if $(version_less $(src_version) 0.10.156.2); then
    echo "Patching new rootfs"
    dd if=/dev/zero of="${INSTALL_DEV}" bs=1 count=2 conv=notrunc seek=1400
    MAKE_DEV_READONLY=${FLAGS_TRUE}
  fi
fi

update_firmware() {
  # See if we need to update firmware. NOTE: we must activate new firmware
  # only after new kernel is actived (installed and made bootable),
  # otherwise new firmware with all old kernels may lead to recovery screen
  # (due to new key).
  FIRMWARE_UPDATE_SCRIPT="${INSTALL_ROOT}/usr/sbin/chromeos-firmwareupdate"
  if [ ${FLAGS_update_firmware} -eq "${FLAGS_TRUE}" -a \
       -x "${FIRMWARE_UPDATE_SCRIPT}" ]; then
    FIRMWARE_UPDATE_MODE=""
    if [ -n "${IS_FACTORY_INSTALL}" ]; then
      # factory-mode
      FIRMWARE_UPDATE_MODE="--mode=factory_install"
    elif [ -n "${IS_RECOVERY_INSTALL}" ]; then
      # recovery-mode
      FIRMWARE_UPDATE_MODE="--mode=recovery"
    elif [ -n "${IS_INSTALL}" ]; then
      # installation from chromeos-install
      FIRMWARE_UPDATE_MODE="--mode=recovery"
    else
      # Default: background update by Update Engine.
      # "--background-update" tells it not to perform any dangerous
      # thing that would require/cause reboot, for example re-writing
      # EC firmware. Those tasks will be done in next boot.
      FIRMWARE_UPDATE_MODE="--mode=autoupdate"
      # FIXME(wfrichar): This is platform-specific. See below.
      RECHECK_AFTER_UPDATE="true"
    fi
    FIRMWARE_UPDATE_ARGS=" $FIRMWARE_UPDATE_MODE "
    echo "Updating firmware ($FIRMWARE_UPDATE_ARGS)"
    env INSTALL_ROOT="${INSTALL_ROOT}" \
      "${FIRMWARE_UPDATE_SCRIPT}" ${FIRMWARE_UPDATE_ARGS}
    # Next step after postinst may take a lot of time (eg, disk wiping)
    # and people may confuse that as 'firmware update takes a long wait',
    # we explicitly prompt here.
    if [ "$?" -eq 0 ]; then
      echo "Firmware update complete."
      if [ "$(mosys platform name)" = "Mario" ] &&
          [ -n "${RECHECK_AFTER_UPDATE}" ]; then
        echo "Check for kernel subkey changes"
        # FIXME(wfrichar): This is a platform-specific workaround for
        # chrome-os-partner:1654. It should probably go into the firmware
        # update script instead of here, but that script doesn't have all the
        # tools available to make the necessary changes.
        SCRATCHDIR=$(mktemp -d /tmp/XXXXXX)
        cd $SCRATCHDIR
        # Note: This code only runs when we're autoupdating, so we know that
        # we're running from the newly updated partition. We need to use the
        # various firmware utilities from that partition too.
        PATH=${DEFAULT_CHROOT_DIR}/usr/sbin:${PATH}
        PATH=${DEFAULT_CHROOT_DIR}/usr/bin:${PATH}
        # Extract the A & B kernel subkeys from the newly written BIOS.
        flashrom -r bios.rom
        dump_fmap -x bios.rom
        gbb_utility -g --rootkey rootkey.vbpubk GBB_Area 
        vbutil_firmware --verify Firmware_A_Key \
          --signpubkey rootkey.vbpubk --fv Firmware_A_Data \
          --kernelkey kernel_subkey_a.vbpubk
        vbutil_firmware --verify Firmware_B_Key \
          --signpubkey rootkey.vbpubk --fv Firmware_B_Data \
          --kernelkey kernel_subkey_b.vbpubk
        # FIXME: I want to use 'cmp', but we don't have it because it pulls
        # in all the locale stuff and is too big. When/if we switch to
        # busybox, we should use cmp here instead.
        hash_a=$(md5sum kernel_subkey_a.vbpubk | cut -f1 -d' ')
        hash_b=$(md5sum kernel_subkey_b.vbpubk | cut -f1 -d' ')
        if [ "$hash_a" != "$hash_b" ]; then
          echo "Kernel subkeys are different. Using original boot order."
          guid=$(sed 's/.*kern_guid=\([0-9a-fA-F-]\+\).*/\1/' /proc/cmdline)
          kern_dev=$($sudo "$GPT" find -1 -u $guid)
          base_dev=$(get_block_dev_from_partition_dev $kern_dev)
          kern_num=$(get_partition_number $kern_dev)
          "$GPT" prioritize $base_dev -i $kern_num
        fi
        cd - >/dev/null 2>&1
        rm -rf $SCRATCHDIR
        echo "Kernel subkey check completed."
      fi
    else
      return 1  # error
    fi
  fi
  return 0  # success
}

if [ "${FLAGS_postcommit}" -eq "${FLAGS_FALSE}" ]; then
  # Pre-commit. Returning an error here will prevent ever booting into the
  # installed system.

  echo "Postinst running"

  # If this FS was mounted read-write, we can't do deltas from it. Mark the
  # FS as such
  (touch "${INSTALL_ROOT}/.nodelta" || true) > /dev/null 2>&1

  echo " Set boot target to ${INSTALL_DEV}:" \
    "Partition ${NEW_PART_NUM}, Slot ${BOOT_SLOT}"

  # TODO(rkc): This is a bit of a hack - we're depending on the fact that our
  # build machine will always be Intel based - if it isn't, VM builds will fail
  # Install syslinux files on EFI partition
  if [ -z "${IS_INSTALL}" -a "${IS_INTEL_ARCH}" -eq "1" ]; then
    EFI_MNT=$(mktemp -d /tmp/EFI.XXXXXX)
    if [ -z "${FLAGS_esp_part_file}" ]; then
      EFI_PART=$(make_partition_dev "${ROOT_DEV}" "${LEGACY_PART_NUM}")
    else
      EFI_PART=$(losetup -f)
      if [ -z "${EFI_PART}" ]; then
        # It's okay to die here; this parameter will never be given on live
        # system - this is only in case of mounted image postinst runs
        die "No free loop device. Free up a loop device or reboot. Exiting."
      fi
      losetup "${EFI_PART}" "${FLAGS_esp_part_file}"
    fi
    # TODO: scary hack to work around disk corruption
    # Note, we mean to run fsck.vfat twice, fsck.vfat returns status 1 for
    # "Recoverable errors have been detected or dosfsck has discovered an
    # internal inconsistency."  If the first fsck fixes errors the first time,
    # the second fsck will succeed.  Otherwise, nuke it from orbit.
    if ! "${INSTALL_ROOT}/sbin/fsck.vfat" -a "${EFI_PART}"; then
      if ! "${INSTALL_ROOT}/sbin/fsck.vfat" -a "${EFI_PART}"; then
        echo "Rebuilding FAT filessytem."
        "${INSTALL_ROOT}/usr/sbin/mkfs.vfat" "${EFI_PART}"
      fi
    fi
    mount "${EFI_PART}" "${EFI_MNT}"
    set +e
    # Copy over the templates, but don't clobber anything.
    # Clobbering is boot_slot specific and handled by setimage below.
    cp -nR "${INSTALL_ROOT}"/boot/syslinux "${EFI_MNT}"
    ERR=${?}
    echo "Updating syslinux"
    # Update all legacy bootloaders.
    if [ "${ERR}" -eq "0" ]; then
      syslinux_cfg="${INSTALL_ROOT}/boot/syslinux/root.${BOOT_SLOT}.cfg"
      grub_cfg="${INSTALL_ROOT}/boot/efi/boot/grub.cfg"
      ROOT_DEV_PARAM=""
      if [ -n "$IS_UPDATE" ]; then
        ROOT_DEV_PARAM="--rootfs_image=$INSTALL_DEV --kernel_image=$NEW_K_DEV"
      fi
      "${INSTALL_ROOT}"/usr/sbin/chromeos-setimage ${BOOT_SLOT} \
          --dst=${ROOT_DEV} --run_as_root \
          --install_root="${INSTALL_ROOT}" \
          --update_syslinux_cfg="${syslinux_cfg}" \
          --update_grub_cfg="${grub_cfg}" \
          --esp_mounted_at="${EFI_MNT}" \
          --update_vmlinuz=${INSTALL_ROOT}/boot/vmlinuz \
          ${SUB_CMD_DEBUG_FLAG} ${ROOT_DEV_PARAM}
      ERR=${?}
    fi
    umount "${EFI_MNT}"
    rmdir "${EFI_MNT}"
    if [ "${ERR}" -ne "0" ]; then
      echo "Failed to copy syslinux files in and make them bootable."
      exit 1
    fi
    set -e
    mount -o bind /dev "${INSTALL_ROOT}"/dev
    set +e
    mount -o bind /proc "${INSTALL_ROOT}"/proc
    ERR=${?}
    if [ "${ERR}" -ne "0" ]; then
      umount "${INSTALL_ROOT}"/dev
      echo "Failed to bind mount /proc"
      exit 1
    fi
    mount -o bind /tmp "${INSTALL_ROOT}"/tmp
    ERR=${?}
    if [ "${ERR}" -ne "0" ]; then
      umount "${INSTALL_ROOT}"/dev
      umount "${INSTALL_ROOT}"/proc
      echo "Failed to bind mount /tmp"
      exit 1
    fi
    chroot "${INSTALL_ROOT}" /usr/bin/syslinux "${EFI_PART}"
    ERR=${?}
    umount "${INSTALL_ROOT}"/tmp
    umount "${INSTALL_ROOT}"/proc
    umount "${INSTALL_ROOT}"/dev
    if (expr match "${EFI_PART}" "/dev/loop*" >/dev/null); then
      losetup -d "${EFI_PART}"
    fi
    set -e
    if [ "${ERR}" -ne "0" ]; then
      echo "Failed to run syslinux on: ${EFI_PART}"
      exit 1
    fi
  fi

  # This export is so that locate_gpt will look within this install
  # partition to find that copy of the cgpt program.
  # Note, this only works when running $INSTALL_ROOT/postinst.
  export DEFAULT_CHROOT_DIR=$(dirname "$0")
  locate_gpt  # sets $GPT to cgpt executable (or fails and exits immediately)

  # If we're working on a real machine, not a mounted image, only then
  # run the gpt tool to manipulate the partitions - if we're touching
  # the ESP from a file, it means we are not on the target hardware, making
  # the GPT and firmware scripts inapplicable nor are we auto-updating
  if [ -z "${FLAGS_esp_part_file}" ]; then
    echo "Updating the root filesystem hashes and legacy bootloaders"
    ROOT_DEV_PARAM=""
    if [ -n "$IS_UPDATE" ]; then
      ROOT_DEV_PARAM="--rootfs_image=$INSTALL_DEV --kernel_image=$NEW_K_DEV"
    fi
    if [ ${IS_INTEL_ARCH} -eq 1 ]; then
      # Ensure we use a template from the INSTALL_ROOT for both grub
      # and syslinux otherwise we can never update them with new rootfs
      # vboot configuration data.
      syslinux_cfg="${INSTALL_ROOT}/boot/syslinux/root.${BOOT_SLOT}.cfg"
      grub_cfg="${INSTALL_ROOT}/boot/efi/boot/grub.cfg"
      "${INSTALL_ROOT}"/usr/sbin/chromeos-setimage ${BOOT_SLOT} \
          --dst=${ROOT_DEV} --run_as_root \
          --install_root="${INSTALL_ROOT}" \
          --update_syslinux_cfg="${syslinux_cfg}" \
          --update_grub_cfg="${grub_cfg}" \
          --update_vmlinuz=${INSTALL_ROOT}/boot/vmlinuz \
          ${SUB_CMD_DEBUG_FLAG} ${ROOT_DEV_PARAM}
    else
      "${INSTALL_ROOT}"/usr/sbin/chromeos-setimage ${BOOT_SLOT} \
          --dst=${ROOT_DEV} --run_as_root \
          --install_root="${INSTALL_ROOT}" \
          ${SUB_CMD_DEBUG_FLAG} ${ROOT_DEV_PARAM}
    fi

    echo "Updating Partition Table Attributes..."
    NEW_KERN_NUM=$(( $NEW_PART_NUM - 1))
    # Make updated partition highest priority for next reboot
    "$GPT" prioritize -i $NEW_KERN_NUM $ROOT_DEV
    # Set kernel as successful=1 except when autoupdating
    new_as_successful=0
    if [ -n "${IS_FACTORY_INSTALL}" -o \
         -n "${IS_RECOVERY_INSTALL}" -o \
         -n "${IS_INSTALL}" ]; then
      new_as_successful=1
    fi
    "$GPT" add -i $NEW_KERN_NUM -S $new_as_successful -T 6 $ROOT_DEV

    # Set up gpt boot selection for legacy devices.
    # flush linux caches; seems to be necessary
    sync
    echo 3 > /proc/sys/vm/drop_caches

    echo "Updating gpt PMBR target for legacy BIOS"
    # Configure the PMBR to boot the new image.
    echo "marking part $LEGACY_PART_NUM bootable (new)"
    "$GPT" boot -i $LEGACY_PART_NUM ${ROOT_DEV} 2>&1

    if [ ${MAKE_DEV_READONLY} -eq ${FLAGS_TRUE} ]; then
      echo "Making dev ${INSTALL_DEV} read-only"
      hdparm -r 1 "${INSTALL_DEV}"
    fi

    # At this point in the script, the new partition has been marked bootable
    # and a reboot will boot into it. Thus, it's important that any future
    # errors in this script do not cause this script to return failure unless
    # in factory mode.
    if [ -z "${IS_FACTORY_INSTALL}" ]; then
      set +e
    fi

    # We have a new image, making the ureadahead pack files
    # out-of-date.  Delete the files so that ureadahead will
    # regenerate them on the next reboot.
    # WARNING: This doesn't work with upgrade from USB, rather than full
    # install/recovery. We don't have support for it as it'll increase the
    # complexity here, and only developers do upgrade from USB.
    rm -f /var/lib/ureadahead/*

    # Create a file indicating that the install is completed. The file
    # will be used in /sbin/chromeos_startup to run tasks on the next boot.
    # See comments above about removing ureadahead files.
    touch /mnt/stateful_partition/.install_completed

    # TODO: support firmware update for ARM
    if [ ${IS_INTEL_ARCH} -eq 1 ] && ! update_firmware; then
      # Note: This will only rollback the ChromeOS verified boot target.
      # The assumption is that systems running firmware autoupdate
      # are running H2C Bios. If the firmware updater crashes or writes
      # corrupt data rather than gracefully failing, we'll probably need
      # to recover with a recovery image.
      echo "Firmware update failed."
      echo "Rolling back update due to failure installing required firmware."
      "$GPT" add -i $NEW_KERN_NUM -P 0 -S 0 -T 0 $ROOT_DEV
      exit 1
    fi
  fi
else
  # Post-commit. At this point an unexpected reboot may boot the installed
  # system, but returning an error here will cause the updater to try to
  # not boot the installed system, instead keeping the existing system.
  echo -n  # this is a noop, required for 'sh'

  # We are deprecating post-commit postinst, so do not put anything in here.

  echo "Updating gpt PMBR target for legacy BIOS (post-commit)"
  # I know we aren't supposed to put anything here, but here is the
  # justification:
  # The updater in 0.7.42.* sets the rootfs bootable after postinst (precommit)
  # is run, and then runs postinst (postcommit). Later versions of the updater
  # neither mark partitions bootable, nor run postinst (postcommit). In those
  # newer versions, this code won't be run, but it won't need to be run.
  # If we didn't care about clients as old as 0.7.42.*, we could leave this
  # area empty. Unfortunately, we can't update old updaters.

  # At some point, when we don't care about clients that old (0.7.42.*),
  # We should remove this whole code path in this file.


  if [ -x "${INSTALL_ROOT}"/usr/bin/cgpt ]; then
    echo "marking part $LEGACY_PART_NUM bootable (new)"
    "${INSTALL_ROOT}"/usr/bin/cgpt boot -i $LEGACY_PART_NUM ${ROOT_DEV} 2>&1
  else
    # The old gpt tool requires a -b arg to specify the PMBR bootcode. We don't
    # want to change the code, so we have to extract it, then put it back.
    dd if=${ROOT_DEV} bs=512 count=1 of=/tmp/oldpmbr.bin
    echo "marking part $LEGACY_PART_NUM bootable (old)"
    gpt -S boot -i $LEGACY_PART_NUM -b /tmp/oldpmbr.bin ${ROOT_DEV} 2>&1
  fi
fi