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

# `bash -x cros_chroot $ACTION(s)` recommended for debugging.
set -eu

# Some possible improvements to this script include a sudo keepalive forked
# process and using set -E and trap ERR.

#
# SUPPORTING FUNCTIONS
#

# Run specified command in the chroot. ~/.bash_profile sourced manually because
# cov couldn't figure out how else to switch users before sudo is emerged.
function chroot_cmd {
  if [ "${1}" = root ]; then
    shift
    sudo chroot chroot su -l -c ". /home/${USER}/.bash_profile; ${*}"
  elif [ -z "${*}" ]; then
    sudo chroot chroot su -l "${USER}"
  else
    sudo chroot chroot su -l -c ". /home/${USER}/.bash_profile; ${*}" "${USER}"
  fi
}

# Set the network and filesystem paths based on user input and stable defaults.
function define_archives {
  AROPTS="--checkpoint=10000"
  for e in "./dev/*" "./home/*" "./proc/*" "./sys/*" "./tmp/*"; do
    AROPTS="${AROPTS} --exclude ${e}"
  done

  if $(sedflags "--stage[ =]*3" true false); then
    local MIRRORPATH=$(sedflags "--mirrorpath[ =]*(\w+)" "\1" localmirror)
    local STABLE=$(sed -nr "s/stage3\s*(.*)/\1/p" "${CF}/stable_versions")
    local AR=$(sedflags "--stage[ =]*3[ =]+([0-9]{8})" \
      "\1/stage3-amd64-\1.tar.bz2" "${STABLE}")
    ARCHIVE[${#ARCHIVE[@]}]="${MIRRORPATH}/releases/amd64/autobuilds/${AR}"
    UNPACKDIR[${#UNPACKDIR[@]}]="chroot"

    if $(sedflags "--portage[ =]*(none)" false true); then
      STABLE=$(sed -nr "s/portage\s*(.*)/\1/p" "${CF}/stable_versions")
      AR=$(sedflags "--portage[ =]*([0-9]{8})" "portage-\1.tar.bz2" \
        "${STABLE}")
      ARCHIVE[${#ARCHIVE[@]}]="${MIRRORPATH}/snapshots/${AR}"
      UNPACKDIR[${#UNPACKDIR[@]}]="chroot/usr"
    fi
  else
    local MIRRORPATH=$(sedflags "--mirrorpath[ =]*(\w+)" "stage4mirror/\1" \
      stage4mirror/stable)
    local RELEASE=$(sedflags "--stage[ =]*4[ =]+([0-9]+)" \
      "\1/release.txt" "latest.txt")
    local DIR="${ARCHIVEDIR}/$(dirname ${RELEASE})"
    mkdir -p "${DIR}" || die "create ${DIR}"
    if $(sedflags "fetch" true false); then
      wget -NP "${DIR}" "${MIRROR}/${RELEASE}" || die "download version info"
    fi
    local FILE=$(cat "${ARCHIVEDIR}/${MIRRORPATH}/${RELEASE}")
    ARCHIVE[${#ARCHIVE[@]}]="${MIRRORPATH}/${FILE}"
    UNPACKDIR[${#UNPACKDIR[@]}]="chroot"
  fi
}

# Print error and return nonzero status.
function die {
  echo "Could not ${1}."
  umount_list
  return 30
}

# Conditionally mount everything in our array of mounts.
function mount_list {
  for i in $(seq 0 $((${#MNTPATH[@]}-1))); do
    if ! mount | grep -q "$(pwd -P)/chroot/${MNTPATH[$i]}"; then
      sudo mount ${MNTOPT[$i]} "chroot/${MNTPATH[$i]}" ||
      die "mount ${MNTPATH[$i]}"
    fi
  done
}

# Functionalized implementation of cov's sed-based argument processing.
# Takes an extended regular expression to find, an expression optionally
# containing named groups to replace it with, and a default value, as the
# first, second, and third arguments, respectively. Slashes should not be
# escaped in arguments passed to this function.
function sedflags {
  local FIND=".* $(echo ${1} | sed 's/\//\\\//g')[ =].*"
  local REPLACE="$(echo ${2} | sed 's/\//\\\//g')"
  local RETURN=$(echo "${ARGS}" | sed -nr "s/${FIND}/${REPLACE}/p")
  echo -n "${RETURN:=${3}}"
}

# Unmount all mounted paths lazily
function umount_all_lazily {
  local ERROR=0
  # XXX: This assumes there are no pipe characters in the paths
  for m in $(mount | sed -nr "s|.* on ($(pwd -P)/.*) type .*|\1|p"); do
    if mount | grep -q "${m}"; then
      sudo umount -dl "${m}" || ERROR=40
    fi
  done
  return ${ERROR}
}

# Unmount listed mount paths in reverse order.
function umount_list {
  local ERROR=0
  for i in $(seq $((${#MNTPATH[@]}-1)) -1 0); do
    local MY_MNTPATH="$(pwd -P)/chroot/${MNTPATH[$i]}"
    if mount | grep -q "${MY_MNTPATH}"; then
      sudo umount -d "${MY_MNTPATH}" || ERROR=40
    fi
  done
  return ${ERROR}
}

#
# ACTIONS
#

# Make sure there are no mounts and clean the chroot.
# XXX: We should really check the lockfile before cleaning.
function clean {
  echo "Removing ./chroot"
  if [ -d chroot ]; then
    umount_all_lazily
    sudo rm -fr chroot || die "remove old chroot"
  fi
}

# Download bootstrapping essentials.
# XXX: It'd be nice to check hashes.
# Note release info fetching in define_archives.
function fetch {
  for i in $(seq 0 $((${#ARCHIVE[@]}-1))); do
    local DIR="${ARCHIVEDIR}/$(dirname ${ARCHIVE[$i]})"
    local URL="${MIRROR}/${ARCHIVE[$i]}"
    wget -NP "${DIR}" "${URL}" || die "download ${URL}"
  done
}

# Unpack the downloaded archives.
function unpack {
  echo "Unpacking tarballs"
  [ ! -d chroot ] || die "unpack to \`./chroot'. Directory already exists. \
Include the clean action to delete the existing chroot"
  local RELEASE=""
  for i in $(seq 0 $((${#ARCHIVE[@]}-1))); do
    local AR="${ARCHIVEDIR}/${ARCHIVE[$i]}"
    [ -n $(basename "${ARCHIVE[$i]}") ] || die "determine file to unpack.\
 Please try re-running \`${0} fetch\`"
    [ -d "${UNPACKDIR[$i]}" ] || mkdir -p "${UNPACKDIR[$i]}" ||
      die "create ${UNPACKDIR[$i]}"
    sudo tar -x ${AROPTS} -C ${UNPACKDIR[$i]} -f "${AR}" ||
      die "unpack ${AR}"
    # XXX: We could save time if the tarballs came with better permissions
    if [ -z ${AR/*portage-*/} ]; then
      sudo find chroot/usr/portage -type f -exec chmod 744 {} \; &&
      sudo find chroot/usr/portage -type d -exec chmod 755 {} \; ||
      die "fix Portage permissions"
    fi
    [ ${i} -eq 0 ] || RELEASE="${RELEASE}-"
    RELEASE="${RELEASE}"$(echo "${ARCHIVE[$i]}" | sed -nr \
      "s/.*\/(\w).*-([0-9]+).*/\1\2/p")
  done
  echo "${RELEASE}" | sudo tee chroot/etc/debian_chroot >/dev/null ||
    die "record chroot information"
}

# Setup the user, source directory, and Portage configuration and packages.
function configure {
  echo "Configuring chroot"

  # Necessary files from the native system and configuration folder
  local FILES="/etc/hosts /etc/localtime /etc/resolv.conf"
  sudo install -m644 ${FILES} chroot/etc ||
    die "install name resolution and timezone configuration"
  sudo install -m644  "${CF}/.bashrc" "${CF}/.bash_profile" chroot/etc/skel ||
    die "install default configuration files"
  sudo install -groot -m440 -oroot  "${CF}/sudoers" chroot/etc ||
    die "install sudo configuration"

  # Add user, but don't die if she already exists
  chroot_cmd root "useradd -m -G portage,users,wheel -u $(id -u) ${USER}" ||
    [ $? -eq 9 ] || die "add user ${USER}"

  # Necessary Portage directories
  local DIRS="chroot/home/${USER}/trunk chroot/usr/local/portage
    chroot/var/lib/portage/distfiles chroot/var/lib/portage/distfiles-target
     chroot/var/lib/portage/pkgs"
  sudo install -dm775 ${DIRS} || die "install Portage directories"

  # Create configuration file links
  # XXX: The toolchain overlays will very soon be, or are, obselete.
  local OL="usr/local/portage/chromiumos"
  local LINKS=(etc/make.conf
    etc/make.profile
    proc/mounts
    "${OL}")
  local TARGETS=("../${OL}/chromeos/config/make.conf.amd64-host"
    "../${OL}/profiles/default/linux/amd64/10.0"
    ../etc/mtab
    "../../../home/${USER}/trunk/src/third_party/chromiumos-overlay")

  for i in $(seq 0 $((${#LINKS[@]}-1))); do
    local LINKDIR="chroot/$(dirname ${LINKS[$i]})"
    local RELTARGET="$(basename ${LINKS[$i]})"
    pushd "${LINKDIR}" 1>/dev/null &&
      sudo ln -fns "${TARGETS[$i]}" "${RELTARGET}" &&
      popd 1>/dev/null || die "create link for ${TARGETS[$i]}"
  done
}

# Install all chroot packages.
function hostpkg {
  echo "Updating chroot packages"
  mount_list
  # XXX: A package blacklist file is in order if this list expands.
  chroot_cmd root "emerge -C dhcpcd" || die "remove dhcpcd"
  chroot_cmd root "emerge ${EFLAGS} -DNu hard-host-depends world" ||
    die "install and update packages"
  umount_list
}

# Create a stage4 tarball of the chroot.
function archive {
  echo "Archiving chroot"
  umount_all_lazily
  local BUILD=$(sedflags "--buildno[ =]*([0-9]+)" "\1" "")
  [ -n "${BUILD}" ] || die "determine build number. Please supply a\
 numerical argument to --buildno"
  local FILE=$(cat chroot/etc/debian_chroot)
  [ -n "${FILE}" ] || die "determine chroot release information"
  local S4DIR="${ARCHIVEDIR}/stage4mirror/"$(sedflags \
    "--mirrorpath[ =]*(\w+)" "\1" stable)
  local AR="${S4DIR}/${BUILD}/${FILE}.tar.bz2"
  mkdir -p $(dirname "${AR}") || die "create directory for archive"
  [ ! -f "${AR}" ] || die "archive to ${AR}.
Delete or move the existing file and re-run \`${0} archive\` to continue"
  pushd chroot 1>/dev/null || die "enter chroot directory"
  sudo tar -ac ${AROPTS} . -f "${AR}" &&
  popd 1>/dev/null || die "make archive"
  cp chroot/etc/debian_chroot "${S4DIR}/${BUILD}/release.txt" &&
  echo "${BUILD}/${FILE}.tar.bz2" > "${S4DIR}/latest.txt" ||
    die "update latest release information"
}

# Mount /dev and company, but don't mount until all PID's have left.
function enter {
  echo "Entering chroot"
  sudo install -dm1777 "chroot/var/lock" || die "install /var/lock"
  local LOCKFILE="chroot/var/lock/cros_chroot-enter"
  (
    flock 200
    echo $$ >> "${LOCKFILE}"
    mount_list
  ) 200>> "${LOCKFILE}" || die "setup environment"
  chroot_cmd "${CMDS}"
  # Only teardown if we're the last enter_chroot to die
  (
    flock 200

    # check each pid in $LOCKFILE to see if it's died unexpectedly
    local TMP_LOCKFILE="${LOCKFILE}.tmp"

    echo -n > "${TMP_LOCKFILE}"  # Erase/reset temp file
    for pid in $(cat "${LOCKFILE}"); do
      [ "${pid}" != $$ ] && [ -n $(ps --pid "${pid}" -o comm=) ] &&
        echo "${pid}" >> "${TMP_LOCKFILE}"
    done
    # Remove any dups from lock file while installing new one
    sort -n "${TMP_LOCKFILE}" | uniq > "${LOCKFILE}"

    if [ -s "${LOCKFILE}" ]; then
      echo "At least one other pid is running in the chroot, so not"
      echo "tearing down environment."
    else
      umount_list
    fi
  ) 200>> "${LOCKFILE}" || die "teardown environment"
}

#
# PROCESS COMMANDLINE ARGUMENTS AND SET GLOBAL VARIABLES
#

# Split commandline arguments at " -- " and pad them
set +u; STAR="${*}"; set -u
ARGS=" ${STAR% -- *} "
CMDS=${STAR#* -- }
[ "${CMDS}" = "${STAR}" ] && CMDS=

# Print help
if $(sedflags "(-h|--help)?" true false); then
  echo 'cros_chroot - A chroot fetch, creation, and archival script.
Usage: ./cros_choot [options] [actions]

All operations will be performed on a directory "chroot" in the current
directory. To work with multiple chroots, launch the script from different
directories. By default, the latest stable tarball versions will be used.

Actions:
  clean
    Remove the chroot directory
  fetch
    Download a stage tarball and Portage snapshot. Versions and paths can be
    controlled via options. The host can be controlled with an environment
    variable. By default, the latest stable versions are fetched.
  unpack
    Decompress the stage tarball to the chroot directory. If the stage
    tarball did not contain a portage snapshot, decompress the Portage
    snapshot as well.
  configure
    Configure the chroot with DNS, sudo, timezone, user, and other
    settings.
  hostpkg
    Install and update packages necessary to build Chromium OS.
  archive
    Compress and store the chroot to ${CROS_CHROOT_ARCHIVE}.
  enter
    Enter the chroot as your current user, either running commands specified
    after "--" in the commandline arguments, or providing an interactive shell.
Virtual targets:
  make: fetch, unpack, configure, essentialpkg, hostpkg
  all: fetch, unpack, configure, essentialpkg, hostpkg, archive, enter

Options:
  No arguments, -h, --help
    Print this help.
  --archivedir PATH
    Location to store archived chroots in. Defaults to
    "${SCRIPTDIR}/../archive", where ${SCRIPTDIR} is the directory in which
    cros_chroot resides. This path cannot contain spaces.
  --buildno NUMBER
    Number to name the directory.
  --mirrorpath DIRECTORY
    Directory on the mirror to use. For stage3, currently, "localmirror"
    (stable) and "gentoo" (upstream) are available. The stage4 stable branch is
    named "stable".
  --portage VERSION
    Version of the Portage tarball to use, or "none" to not fetch a Portage
    snapshot.
  --stage [3,4] [VERSION]
    Version of the stage3 or stage4 prebuilt chroot to use. If stage3 is not
    specified, stage4 will be assumed.
  --usepkg
    Fetch and use binary packages where possible, rather than building from
    source.

Environment variables:
  ${CROS_CHROOT_MIRROR}
    Host from which to fetch stage and portage tarballs. Defaults to
    "http://build.chromium.org/mirror/chromiumos".

Remote control file:
  cros_chroot sources "${HOME}/.cros_chrootrc", if that file exists,
  exporting the variable ${CROS_CHROOT_ACTION} set to the name of the current
  action being executed. Use this file for site-specific customization like
  disabling GNOME automounting.

Examples:
  ./cros_chroot clean all
  ./cros_chroot make --stage 20100617 --portage 20100617 --mirrorpath gentoo
  ./cros_chroot enter -- echo -n "\"Hello, \"; whoami"'
  exit 0
fi

# Current directory, script directory, and storage directories
SCRIPTDIR="$(dirname ${0})"
[ -L "${0}" ] && SCRIPTDIR="${SCRIPTDIR}/$(dirname $(readlink ${0}))"
SCRIPTDIR="$(cd ${SCRIPTDIR}; pwd -P)"
# XXX: This path can't contain spaces.
ARCHIVEDIR=$(sedflags "--archivedir[ =]*(.*)" "\1" \
  "${SCRIPTDIR/%scripts/archive}")
CF="${SCRIPTDIR/%scripts/config}"
SRCDIR="${SCRIPTDIR/%src\/third_party\/chromiumos-overlay\/chromeos\/scripts/}"
SCRIPTDIR="${SCRIPTDIR/#.\//}"

# Mount path and option lists
MNTPATH=(dev          dev/pts            proc           sys
  "home/${USER}/trunk")
MNTOPT=("--bind /dev" "devpts -t devpts" "proc -t proc" "sysfs -t sysfs"
  "--bind ${SRCDIR}")

# Tarball path information
MIRROR="${CROS_CHROOT_MIRROR:=http://build.chromium.org/mirror/chromiumos}"
declare -a ARCHIVE UNPACKDIR
define_archives

# Set emerge options
EFLAGS=$(sedflags --usepkg -g "")
JOBS=$(($(cat /proc/cpuinfo | sed -nr '/^processor/ s/.*([0-9]+$)/\1/p' \
  | tail -n 1)+2))
EFLAGS="${EFLAGS} --jobs "$(sedflags "--jobs[ =]*([0-9]+)" "\1" "${JOBS}")

# Expand virtual targets
ARGS=$(sedflags make "${ARGS}fetch unpack configure hostpkg archive " "${ARGS}")
ARGS=$(sedflags all "${ARGS}fetch unpack configure hostpkg archive enter " \
  "${ARGS}")

#
# EXECUTE ACTIONS
#

# Inital source of the RC file
touch "${HOME}/.cros_chrootrc" || die "touch ${HOME}/.cros_chrootrc"
CROS_CHROOT_ACTION= && . "${HOME}/.cros_chrootrc"

# Execute actions in proper order, sourcing RC file
for a in clean fetch unpack configure hostpkg archive enter
do
  if $(sedflags ${a} true false); then
    CROS_CHROOT_ACTION="pre-${a}" && . "${HOME}/.cros_chrootrc"
    ${a}
    CROS_CHROOT_ACTION="post-${a}" && . "${HOME}/.cros_chrootrc"
  fi
done

exit 0