#!/bin/bash

# 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.

# Script to build a bootable keyfob-based chromeos system image from within
# a chromiumos setup. This assumes that all needed packages have been built into
# the given target's root with binary packages turned on. This script will
# build the Chrome OS image using only pre-built binary packages.

# Load common constants.  This should be the first executable line.
# The path to common.sh should be relative to your script's location.
. "$(dirname "$0")/common.sh"

# Script must be run inside the chroot.
assert_inside_chroot

get_default_board

# Flags.
DEFINE_string board "$DEFAULT_BOARD" \
  "The board to build an image for."
DEFINE_string build_root "/build" \
  "The root location for board sysroots."
DEFINE_integer build_attempt 1 \
  "The build attempt for this image build."
DEFINE_string output_root "${DEFAULT_BUILD_ROOT}/images" \
  "Directory in which to place image result directories (named by version)"
DEFINE_boolean replace $FLAGS_FALSE \
  "Overwrite existing output, if any."
DEFINE_boolean withdev $FLAGS_TRUE \
  "Include useful developer friendly utilities in the image."
DEFINE_boolean installmask $FLAGS_TRUE \
  "Use INSTALL_MASK to shrink the resulting image."
DEFINE_integer jobs -1 \
  "How many packages to build in parallel at maximum."
DEFINE_boolean statefuldev $FLAGS_FALSE \
  "Install development packages on stateful partition rather than the rootfs"
DEFINE_string to "" \
  "The target image file or device"
DEFINE_boolean withtest $FLAGS_FALSE \
  "Include packages required for testing and prepare image for testing"
DEFINE_string factory_server "" \
  "Build a factory install image pointing to given server."

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

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

if [ -z "$FLAGS_board" ] ; then
  error "--board is required."
  exit 1
fi

# Determine build version.
. "${SCRIPTS_DIR}/chromeos_version.sh"

# Use canonical path since some tools (e.g. mount) do not like symlinks.
# Append build attempt to output directory.
IMAGE_SUBDIR="${CHROMEOS_VERSION_STRING}-a${FLAGS_build_attempt}"
OUTPUT_DIR="${FLAGS_output_root}/${FLAGS_board}/${IMAGE_SUBDIR}"
ROOT_FS_DIR="${OUTPUT_DIR}/rootfs"
ROOT_FS_IMG="${OUTPUT_DIR}/rootfs.image"
OUTPUT_IMG=${FLAGS_to:-${OUTPUT_DIR}/chromiumos_image.bin}

BOARD="${FLAGS_board}"
BOARD_ROOT="${FLAGS_build_root}/${BOARD}"

LOOP_DEV=
STATEFUL_LOOP_DEV=
ESP_LOOP_DEV=

# What cross-build are we targeting?
. "${BOARD_ROOT}/etc/make.conf.board_setup"
LIBC_VERSION=${LIBC_VERSION:-"2.10.1-r1"}

# Figure out ARCH from the given toolchain.
# TODO: Move to common.sh as a function after scripts are switched over.
TC_ARCH=$(echo "$CHOST" | awk -F'-' '{ print $1 }')
case "$TC_ARCH" in
  arm*)
    ARCH="arm"
    ;;
  *86)
    ARCH="x86"
    ;;
  *)
    error "Unable to determine ARCH from toolchain: $CHOST"
    exit 1
esac

# Hack to fix bug where x86_64 CHOST line gets incorrectly added.
# ToDo(msb): remove this hack.
PACKAGES_FILE="${BOARD_ROOT}/packages/Packages"
sudo sed -e "s/CHOST: x86_64-pc-linux-gnu//" -i "${PACKAGES_FILE}"

# Handle existing directory.
if [[ -e "$OUTPUT_DIR" ]]; then
  if [[ $FLAGS_replace -eq $FLAGS_TRUE ]]; then
    sudo rm -rf "$OUTPUT_DIR"
  else
    echo "Directory $OUTPUT_DIR already exists."
    echo "Use --build_attempt option to specify an unused attempt."
    echo "Or use --replace if you want to overwrite this directory."
    exit 1
  fi
fi

# Create the output directory.
mkdir -p "$OUTPUT_DIR"

cleanup_rootfs_loop() {
  sudo umount "$LOOP_DEV"
  sleep 1  # in case $LOOP_DEV is in use.
  sudo losetup -d "$LOOP_DEV"
}

cleanup_stateful_fs_loop() {
  sudo umount "${ROOT_FS_DIR}/usr/local"
  sudo umount "${ROOT_FS_DIR}/var"
  sudo umount "${STATEFUL_DIR}"
  sleep 1  # follows from cleanup_root_fs_loop.
  sudo losetup -d "$STATEFUL_LOOP_DEV"
}

cleanup_esp_loop() {
  sudo umount "$ESP_DIR"
  sleep 1
  sudo losetup -d "$ESP_LOOP_DEV"
}

cleanup() {
  # Disable die on error.
  set +e

  if [[ -n "$STATEFUL_LOOP_DEV" ]]; then
    cleanup_stateful_fs_loop
  fi

  if [[ -n "$LOOP_DEV" ]]; then
    cleanup_rootfs_loop
  fi

  if [[ -n "$ESP_LOOP_DEV" ]]; then
    cleanup_esp_loop
  fi

  # Turn die on error back on.
  set -e
}

# ${DEV_IMAGE_ROOT} specifies the location of where developer packages will
# be installed on the stateful dir.  On a Chromium OS system, this will
# translate to /usr/local
DEV_IMAGE_ROOT=

# Sets up symlinks for the stateful partition based on the root specified by
# ${1} and var directory specified by ${2}.
setup_symlinks_on_root() {
  echo "Setting up symlinks on the stateful partition rooted at ${1} with"\
    "var directory located at ${2}"

  for path in usr local; do
    if [ -h "${DEV_IMAGE_ROOT}/${path}" ] ; then
      sudo unlink "${DEV_IMAGE_ROOT}/${path}"
    elif [ -e "${DEV_IMAGE_ROOT}/${path}" ] ; then
      echo "*** ERROR: ${DEV_IMAGE_ROOT}/${path} should be a symlink if exists"
      return 1
    fi
    sudo ln -s ${1} "${DEV_IMAGE_ROOT}/${path}"
  done

  # Setup var.  Var is on the stateful partition at /var for both non-developer
  # builds and developer builds.
  if [ -h "${DEV_IMAGE_ROOT}/var" ] ; then
    sudo unlink "${DEV_IMAGE_ROOT}/var"
  elif [ -e "${DEV_IMAGE_ROOT}/var" ] ; then
    echo "*** ERROR: ${DEV_IMAGE_ROOT}/var should be a symlink if it exists"
    return 1
  fi

  sudo ln -s "${2}" "${DEV_IMAGE_ROOT}/var"
}

trap cleanup EXIT

mkdir -p "$ROOT_FS_DIR"

# Create and format the root file system.

# Check for loop device before creating image.
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

# Create root file system disk image to fit on a 1GB memory stick.
# 1 GB in hard-drive-manufacturer-speak is 10^9, not 2^30.  950MB < 10^9 bytes.
ROOT_SIZE_BYTES=$((1024 * 1024 * 720))
dd if=/dev/zero of="$ROOT_FS_IMG" bs=1 count=1 seek=$((ROOT_SIZE_BYTES - 1))
sudo losetup "$LOOP_DEV" "$ROOT_FS_IMG"
sudo mkfs.ext3 "$LOOP_DEV"

# Tune and mount rootfs.
UUID=$(uuidgen)
DISK_LABEL="C-KEYFOB"
sudo tune2fs -L "$DISK_LABEL" -U "$UUID" -c 0 -i 0 "$LOOP_DEV"
sudo mount "$LOOP_DEV" "$ROOT_FS_DIR"

# Create stateful partition of the same size as the rootfs.
STATEFUL_IMG="$OUTPUT_DIR/stateful_partition.image"
STATEFUL_DIR="$OUTPUT_DIR/stateful_partition"
STATEFUL_LOOP_DEV=$(sudo losetup -f)
if [ -z "$STATEFUL_LOOP_DEV" ] ; then
  echo "No free loop device.  Free up a loop device or reboot.  exiting. "
  exit 1
fi
dd if=/dev/zero of="$STATEFUL_IMG" bs=1 count=1 seek=$((ROOT_SIZE_BYTES - 1))
sudo losetup "$STATEFUL_LOOP_DEV" "$STATEFUL_IMG"
sudo mkfs.ext3 "$STATEFUL_LOOP_DEV"
sudo tune2fs -L "C-STATE" -U "$UUID" -c 0 -i 0 \
    "$STATEFUL_LOOP_DEV"

# Mount the stateful partition.
mkdir -p "$STATEFUL_DIR"
sudo mount "$STATEFUL_LOOP_DEV" "$STATEFUL_DIR"

# Set dev image root now that we have mounted the stateful partition we created
DEV_IMAGE_ROOT="$STATEFUL_DIR/dev_image"

# Turn root file system into bootable image.
if [[ "$ARCH" = "x86" ]]; then
  # Setup extlinux configuration.
  # TODO: For some reason the /dev/disk/by-uuid is not being generated by udev
  # in the initramfs. When we figure that out, switch to root=UUID=$UUID.
  sudo mkdir -p "$ROOT_FS_DIR"/boot
  # TODO(adlr): use initramfs for booting.
  cat <<EOF | sudo dd of="$ROOT_FS_DIR"/boot/extlinux.conf
DEFAULT chromeos-usb
PROMPT 0
TIMEOUT 0

label chromeos-usb
  menu label chromeos-usb
  kernel vmlinuz
  append quiet console=tty2 init=/sbin/init boot=local rootwait root=/dev/sdb3 ro noresume noswap i915.modeset=1 loglevel=1

label chromeos-hd
  menu label chromeos-hd
  kernel vmlinuz
  append quiet console=tty2 init=/sbin/init boot=local rootwait root=HDROOT ro noresume noswap i915.modeset=1 loglevel=1
EOF

  # Make partition bootable and label it.
  sudo extlinux -z --install "${ROOT_FS_DIR}/boot"
fi

# -- Install packages into the root file system --

# We need to install libc manually from the cross toolchain.
# TODO: Improve this? We only want libc and not the whole toolchain.
PKGDIR="/var/lib/portage/pkgs/cross/"
sudo tar jxvpf \
  "${PKGDIR}/${CHOST}/cross-${CHOST}"/glibc-${LIBC_VERSION}.tbz2 \
  -C "$ROOT_FS_DIR" --strip-components=3 \
  --exclude=usr/include --exclude=sys-include --exclude=*.a --exclude=*.o

# We need to install libstdc++ manually from the cross toolchain.
# TODO: Figure out a better way of doing this?
sudo cp -a "${BOARD_ROOT}"/lib/libgcc_s.so* "${ROOT_FS_DIR}/lib"
sudo cp -a "${BOARD_ROOT}"/usr/lib/libstdc++.so* "${ROOT_FS_DIR}/usr/lib"

INSTALL_MASK=""
if [[ $FLAGS_installmask -eq ${FLAGS_TRUE} ]] ; then
  INSTALL_MASK="$DEFAULT_INSTALL_MASK"
fi

if [[ $FLAGS_jobs -ne -1 ]]; then
  EMERGE_JOBS="--jobs=$FLAGS_jobs"
fi

# Prepare stateful partition with some pre-created directories
sudo mkdir -p "${DEV_IMAGE_ROOT}"
sudo mkdir -p "${STATEFUL_DIR}/var"

# Create symlinks so that /usr/local/usr based directories are symlinked to
# /usr/local/ directories e.g. /usr/local/usr/bin -> /usr/local/bin, etc.
setup_symlinks_on_root "${DEV_IMAGE_ROOT}" "${STATEFUL_DIR}/var"

# Perform binding rather than symlinking because directories must exist
# on rootfs so that we can bind at run-time since rootfs is read-only
echo "Binding directories from stateful partition onto the rootfs"
sudo mkdir -p "${ROOT_FS_DIR}/usr/local"
sudo mount --bind "${DEV_IMAGE_ROOT}" "${ROOT_FS_DIR}/usr/local"
sudo mkdir -p "${ROOT_FS_DIR}/var"
sudo mount --bind "${STATEFUL_DIR}/var" "${ROOT_FS_DIR}/var"

# We "emerge --root=$ROOT_FS_DIR --root-deps=rdeps --usepkgonly" all of the
# runtime packages for chrome os. This builds up a chrome os image from binary
# packages with runtime dependencies only.  We use INSTALL_MASK to trim the
# image size as much as possible.
sudo INSTALL_MASK="$INSTALL_MASK" emerge-${BOARD} \
  --root="$ROOT_FS_DIR" --root-deps=rdeps \
  --usepkgonly chromeos $EMERGE_JOBS

# Determine the root dir for development packages.
ROOT_DEV_DIR="$ROOT_FS_DIR"
[ $FLAGS_statefuldev -eq $FLAGS_TRUE ] && ROOT_DEV_DIR="$ROOT_FS_DIR/usr/local"

# Install development packages.
if [[ $FLAGS_withdev -eq $FLAGS_TRUE ]] ; then
  sudo INSTALL_MASK="$INSTALL_MASK" emerge-${BOARD} \
      --root="$ROOT_DEV_DIR" --root-deps=rdeps \
      --usepkgonly chromeos-dev $EMERGE_JOBS

  # TODO(sosa@chromium.org) - Re-hide under statefuldev after switch
  # Flag will mount /usr/local on target device
  sudo mkdir -p "$ROOT_FS_DIR/root"
  sudo touch "$ROOT_FS_DIR/root/.dev_mode"

  # The ldd tool is a useful shell script but lives in glibc; just copy it.
  sudo cp -a "$(which ldd)" "${ROOT_DEV_DIR}/usr/bin"
fi

if [ -n "$FLAGS_factory_server" ]; then
  sudo INSTALL_MASK="$INSTALL_MASK" emerge-${BOARD} \
      --root="$ROOT_DEV_DIR" --root-deps=rdeps \
      --usepkgonly chromeos-factoryinstall $EMERGE_JOBS
fi

# Install packages required for testing.
if [[ $FLAGS_withtest -eq $FLAGS_TRUE ]] ; then
  sudo INSTALL_MASK="$INSTALL_MASK" emerge-${BOARD} \
      --root="$ROOT_DEV_DIR" --root-deps=rdeps \
      --usepkgonly chromeos-test $EMERGE_JOBS
fi

# Perform any customizations on the root file system that are needed.
EXTRA_CUSTOMIZE_ROOTFS_FLAGS=""
if [ $FLAGS_withdev -eq $FLAGS_TRUE ]; then
  EXTRA_CUSTOMIZE_ROOTFS_FLAGS="--withdev"
fi
if [ -n "$FLAGS_factory_server" ]; then
# Indentation off b/c of long line
EXTRA_CUSTOMIZE_ROOTFS_FLAGS="$EXTRA_CUSTOMIZE_ROOTFS_FLAGS"\
" --factory_server=$FLAGS_factory_server"
fi

# Extract the kernel from the root filesystem for use by the GPT image. Legacy
# BIOS will use the kernel in the rootfs (via syslinux), Chrome OS BIOS will
# use the kernel partition.
sudo cp -f "${ROOT_FS_DIR}/boot/vmlinuz" "${OUTPUT_DIR}/vmlinuz.image"

# Create EFI System Partition to boot stock EFI BIOS (but not ChromeOS EFI
# BIOS). We only need this for x86, but it's simpler and safer to keep the disk
# images the same for both x86 and ARM.
ESP_IMG=${OUTPUT_DIR}/esp.image
# NOTE: The size argument for mkfs.vfat is in 1024-byte blocks. We'll hard-code
# it to 16M for now.
ESP_BLOCKS=16384
/usr/sbin/mkfs.vfat -C ${OUTPUT_DIR}/esp.image ${ESP_BLOCKS}
ESP_DIR=${OUTPUT_DIR}/esp
ESP_LOOP_DEV=$(sudo losetup -f)
if [ -z "$ESP_LOOP_DEV" ] ; then
  echo "No free loop device.  Free up a loop device or reboot.  exiting. "
  exit 1
fi
mkdir -p "${ESP_DIR}"
sudo losetup "${ESP_LOOP_DEV}" "${ESP_IMG}"
sudo mount "${ESP_LOOP_DEV}" "${ESP_DIR}"
sudo mkdir -p "${ESP_DIR}/efi/boot"
sudo grub-mkimage -p /efi/boot -o "${ESP_DIR}/efi/boot/bootx64.efi" \
  part_gpt fat ext2 normal boot sh chain configfile linux
sudo cp "${ROOT_FS_DIR}/boot/vmlinuz" "${ESP_DIR}/efi/boot/vmlinuz"
cat <<EOF | sudo dd of="${ESP_DIR}/efi/boot/grub.cfg"
set timeout=2
set default=0

menuentry "normal" {
  linux /efi/boot/vmlinuz quiet console=tty2 init=/sbin/init boot=local rootwait root=/dev/sda3 ro noresume noswap i915.modeset=1 loglevel=1
}

menuentry "serial debug" {
  linux /efi/boot/vmlinuz earlyprintk=serial,ttyS0,115200 console=ttyS0,115200 init=/sbin/init boot=local rootwait root=/dev/sda3 ro noresume noswap i915.modeset=1 loglevel=7
}

EOF

# Run ldconfig for rootfs's ld.so.cache
if [ $FLAGS_statefuldev -eq $FLAGS_TRUE ] ; then
  # Re-run ldconfig to fix /etc/ldconfig.so.cache
  sudo /sbin/ldconfig -r "$ROOT_FS_DIR"

  #TODO(sosa@chromium.org) - /usr/bin/xterm symlink not created in stateful.
  sudo ln -sf "/usr/local/bin/aterm" "/usr/bin/xterm"
fi

"${SCRIPTS_DIR}/customize_rootfs" \
  --root="$ROOT_FS_DIR" \
  --target="$ARCH" \
  --board="$BOARD" \
  $EXTRA_CUSTOMIZE_ROOTFS_FLAGS

# Check that the image has been correctly created.
"${SCRIPTS_DIR}/test_image" \
  --root="$ROOT_FS_DIR" \
  --target="$ARCH"

# Clean up symlinks so they work on a running target rooted at "/".
# Here development packages are rooted at /usr/local.  However, do not
# create /usr/local or /var on host (already exist on target).
setup_symlinks_on_root "/usr/local" "/var"

# Cleanup loop devices.
cleanup

# Create the GPT-formatted image
${SCRIPTS_DIR}/build_gpt.sh \
   --arch=${ARCH} --board=${FLAGS_board} "${OUTPUT_DIR}" "${OUTPUT_IMG}"

# Clean up temporary files.
rm -f "${ROOT_FS_IMG}" "${STATEFUL_IMG}" "${OUTPUT_DIR}/vmlinuz.image" \
  "${ESP_IMG}"
rmdir "${ROOT_FS_DIR}" "${STATEFUL_DIR}" "${ESP_DIR}"

OUTSIDE_OUTPUT_DIR="../build/images/${FLAGS_board}/${IMAGE_SUBDIR}"
echo "Done.  Image created in ${OUTPUT_DIR}"
echo "To copy to USB keyfob, OUTSIDE the chroot, do something like:"
echo "  ./image_to_usb.sh --from=${OUTSIDE_OUTPUT_DIR} --to=/dev/sdb"
echo "To convert to VMWare image, OUTSIDE the chroot, do something like:"
echo "  ./image_to_vmware.sh --from=${OUTSIDE_OUTPUT_DIR}"
echo "from the scripts directory where you entered the chroot."

trap - EXIT