#!/bin/sh
# Copyright 2011 CyberLeo, All Rights Reserved
# http://wiki.cyberleo.net/wiki/CyberLeo/COPYRIGHT

: ${ORIG_USER=$(id -un)}

# Need root beyond here
[ "$(id -u)" -eq 0 ] || exec sudo env "J_ARCH=${J_ARCH}" "J_BASE=${J_BASE}" "J_NAME=${J_NAME}" "J_USER=${J_USER:-${USER}}" "ORIG_USER=${ORIG_USER}" "DISPLAY=${DISPLAY}" "${0}" "${@}"

meh() { printf " \033[1;32m*\033[0m %s%s\n" "${jname:+${jname}: }" "${*}"; }
omg() { printf " \033[1;33m*\033[0m %s%s\n" "${jname:+${jname}: }" "${*}"; }
wtf() { printf " \033[1;31m*\033[0m %s%s\n" "${jname:+${jname}: }" "${*}"; exit 1; }
pebkac() {
  [ "${*}" ] && printf "%s\n\n" "${*}"
  cat <<EOF
Usage:
  <command> <name> [arguments]
  list
  ls      list available chroots

  status  show jail status

  start   prepare an existing chroot for use

  stop    undo what 'start' did

  kill    Forcefully kill all processes; only applicable with cgroups

  enter
  shell   spawn a shell or command within the chroot

  eval    evaluate a shell command line within the chroot

  mount   mount chroot filesystems according to chroot fstab

  umount  umount chroot filesystems that were mounted with 'mount'

  init <dir> <seed> [seed-specific options]
          initialize a chroot using the specified seed in seeds/

EOF
  exit 1
}

cmd="$(basename "${0}")"
j="$(readlink -f "$(dirname "${0}")")"
jarch="${J_ARCH:-$(uname -m)}"
jbase="${J_BASE:-$(readlink -f "${j}/../")}"
jseed="${j}/seeds"
jname="${J_NAME:-$(basename "${1}")}" #"
juser="${J_USER}"

# Find chroot binary, since it can be in many different places on different distros
for chroot in /usr/sbin/chroot /usr/bin/chroot /sbin/chroot /bin/chroot
do
  [ -f "${chroot}" ] && break
  unset chroot
done
[ "${chroot}" ] || wtf "Cannot find chroot in /usr/sbin:/usr/bin:/sbin:/bin"

# cgroup to use; if not found, cgroup support will not be used
jcgroup="freezer"
# cgroup parent to use (${cgroup}/${parent}/${jname})
jcgparent="j"

# Remove chroot name from argument stack, if passed in
[ "${J_NAME}" ] || shift && true

# Propagate certain environment variables; sterilize the rest of the environment
jenv="
  LANG=${LANG}
  TERM=${TERM}
  USER=${USER}
"

# Create a new chroot, using a seed script in seeds/
j_init() {
  # Make sure this does NOT exist
  jname="${1}"; shift
  [ "${jname}" ] || pebkac "j_init: no name provided"
  seed="$(echo "${1}" | tr 'A-Z' 'a-z')"; shift
  [ -f "${jbase}/j/seeds/${seed}.sh" ] || pebkac "Unknown seed '${seed}'"
  jdir="${jbase}/${jname}"
  [ -d "${jdir}" ] && pebkac "j_init: ${jname} already exists"
  jroot="${jdir}/root"
  set -e
  . "${jbase}/j/seeds/${seed}.sh" "${@}"
}

# Figure out and set chroot parameters; needed for all functions that follow
j_params() {
  ( jname="${1:-jname}"

    # Make sure jname is not empty
    if [ -z "${jname}" ]
    then
      printf "jerror='%s'\n" "jname empty"
      return 1
    fi

    # Given a chroot name, find and set up the chroot dir
    jdir="${jbase}/${jname}"
    jroot="${jdir}/root"
    if [ ! -d "${jdir}" -o ! -d "${jroot}" ]
    then
      printf "jerror='%s'\n" "not a directory"
      return 1
    fi

    # Does this have a valid fstab file?
    [ -f "${jdir}/fstab" ] && jfstab="${jdir}/fstab"

    # Where is the shell?
    jshell=""
    for shell in /bin/bash /usr/bin/bash /usr/local/bin/bash /bin/sh
    do
      if [ -f "${jroot}/${shell}" ]
      then
        jshell=${shell}
        break
      fi
    done
    if [ -z "${jshell}" ]
    then
      printf "jerror='%s'\n" "unable to locate usable shell; is this a jail?"
      return 1
    fi

    printf "jerror='' jname='%s' jdir='%s' jroot='%s' jfstab='%s' jshell='%s'\n" "${jname}" "${jdir}" "${jroot}" "${jfstab}" "${jshell}"
  )
}

# Find the 'freezer' cgroup, to determine whether the chroot can be further isolated
j_locate_cgroup_mount() {
  cgroup="${1:-freezer}"
  cat /proc/mounts | while read src dir type opts fsck dump extra
  do
    [ "${type}" = "cgroup" ] || continue
    oldifs="${IFS}"
    IFS=','
    set -- ${opts}
    IFS=${oldifs}
    while [ "${1}" ]
    do
      [ "${1}" = "${cgroup}" ] && echo "${dir}" && break 2
      shift
    done
  done
}

# Cache and expose cgroup mountpoint for given chroot; make sure it exists
j_cgroup() {
  [ "${jcg}" ] && { echo "${jcg}"; return; }
  jcg="$(j_locate_cgroup_mount "${cgroup}")/${jcgparent}/${jname}"
  mkdir -p "${jcg}"
  echo "${jcg}"
}

j_cg_trap() {
  [ "$(j_cgroup)" ] || return
  while [ "${1}" ]
  do
    echo "${1}" > "$(j_cgroup)/tasks"
    shift
  done
}

# Run a command within a cgroup, if available
j_cg_eval() {
  j_cg="$(j_cgroup)"
  if [ "$(j_cgroup)" ]
  then
    sh -ex <<EOF
echo $$ > "$(j_cgroup)/tasks"
eval "${*}"
EOF
  else
    eval "${*}"
  fi
}

# Copy ipcc into the jail and add helpful symlinks, so it can communicate
# simple commands to the outside
j_ipc_setup() {
  ipcc="$(dirname "${0}")/ipcc"
  cp -f "${ipcc}" "${jroot}/bin"
  ( cd "${jroot}/bin"
    [ -e ee ] || ln -s ipcc ee
    [ -e ff ] || ln -s ipcc ff
    [ -e gitk ] || ln -s ipcc gitk
    [ -e gitka ] || ln -s ipcc gitka
  )
}

# Launch ipcd on jail
j_ipc_start() {
  ipcd_pid="${jdir}/ipcd.pid"
  # Kill any running IPCd; remove stale pid file and socket
  if [ -f "${ipcd_pid}" ]
  then
    pid="$(cat "${ipcd_pid}")"
    if [ "${pid}" ]
    then
      if ps ouser=,command= "${pid}" | grep "^${ORIG_USER}" | grep -q "ipcd"
      then
        omg "IPCd: already running! Killing"
        kill -TERM "${pid}"
      fi
      omg "IPCd: removing stale pid file"
      rm -f "${ipcd_pid}" "${jroot}/tmp/jipc"
    fi
  else
    [ -e "${jroot}/tmp/jipc" ] && {
      omg "IPCd: not running, but pipe exists; removing"
      rm -f "${jroot}/tmp/jipc"
    }
  fi
  su "${ORIG_USER}" "${jbase}/j/ipcd" "${jdir}" &
  pid="${!}"
  echo "${pid}" > "${ipcd_pid}"
  j_cg_trap "${pid}"
}

# Stop ipcd on jail
j_ipc_stop() {
  ipcd_pid="${jdir}/ipcd.pid"
  [ -f "${ipcd_pid}" ] && pid="$(cat "${ipcd_pid}")"
  [ "${pid}" ] && kill -TERM ${pid}
  rm -f "${ipcd_pid}"
}

# Is this a chroot?
j_is() {
  eval $(j_params "${1}")
  [ "${jerror}" ] && return 1 || return 0
}

# List available chroots
j_ls() {
  ( cd "${jbase}"; ls -1 ) | while read jname
  do
    j_is "${jname}" && echo "${jname}"
  done
}

# Chroot is 'up' if /dev/pts and /proc are mounted
j_up() {
  jname="${1:-${jname}}"
  eval "$(j_params "${jname}")"
  [ "${jerror}" ] && wtf "${jerror}"
  grep -q "^[^ ]* ${jroot}/dev/pts devpts" /proc/mounts || return 1
  grep -q "^[^ ]* ${jroot}/proc proc" /proc/mounts || return 1
  return 0
}

# Poll chroot status (j_up)
j_status() {
  [ -z "${1}" ] && set - $(j_ls)
  while [ "${1}" ]
  do
    j_up "${1}" && meh "$(printf '\033[1;32mup\033[0m')" || meh "$(printf '\033[1;31mdown\033[0m')"
    shift
  done
}

# Copy a mount from the host (for things like /dev, /proc, /sys, et alia)
j_fstab_copy_mount() {
  grep -Eq "${1}[[:space:]]" "${tmp_fstab}" || grep -E \
    "[[:space:]]${1}[[:space:]]" /proc/mounts | head -n 1 >> "${tmp_fstab}"
}

# Prepare fstab for use:
# * Strip comments
# * Make sure certain required mounts exist (/dev/pts, /proc)
# * Sort by mountpoint, to ensure mounts are performed in the correct order
j_fstab_synthesize() {
  local tmp_fstab="$(mktemp "${jdir}/.fstab.XXXXXXXX")"
  [ -f "${jfstab}" ] && sed -e 's/#.*$//; /^\W*$/d' "${jfstab}" > "${tmp_fstab}"
  j_fstab_copy_mount "/dev/pts"
  j_fstab_copy_mount "/proc"
  sort -k2 "${tmp_fstab}"
  rm -f "${tmp_fstab}"
}

# Process fstab file and mount the filesystems therein; if no fstab file (or if
# excludes certain critical filesystems) then augment it. This will also maintain
# an mtab file so they can be unmounted later (and so df works inside chroot).
#
# All mountpoints listed in this file will be prepended with ${jroot} prior to
# their being mounted or umounted! Make sure the fstab mountpoints are relative
# to the chroot root!
j_fstab_mount() {
  echo "rootfs / rootfs rw 0 0" > "${jdir}/mtab"
  j_fstab_synthesize | while read src dir type opts fsck dump
  do
    mkdir -p "${dir}" || break
    mount -t "${type}" -o "${opts}" "${src}" "${jroot}/${dir}" || break
    printf "%s %s %s %s 0 0\n" "${src}" "/${dir##/}" "${type}" "${opts}" >> "${jdir}/mtab"
  done
}

# Process fstab and mtab files and unmount the filesystems therein (opposite j_fstab_mount)
j_fstab_umount() {
  [ -s "${jdir}/mtab" ] || j_fstab_synthesize > "${jdir}/mtab"
  tac "${jdir}/mtab" | while read src dir unused
  do
    # Mounted?
    grep -Eq "[[:space:]]${jroot}${dir%%/}[[:space:]]" /proc/mounts || continue
    umount -f "${jroot}/${dir}" || umount -l "${jroot}/${dir}" || break
  done
  :>"${jdir}/mtab"
}

# Start up chroot
j_start() {
  jname="${1:-${jname}}"
  j_up "${jname}" && return 0
  eval "$(j_params "${jname}")"
  meh "starting ${jname} ..."

  # Mount filesystems and copy in /etc/mtab
  j_fstab_mount
  cp "${jdir}/mtab" "${jroot}/etc/mtab"

  # Copy in ipcc and launch ipcd
  j_ipc_setup
  j_ipc_start

  # Start all services in /etc/rcJ.d
  [ -d "${jroot}/etc/rcJ.d" ] || return 0
  j_root_eval "${jname}" '( ls -1 /etc/rcJ.d/* 2>&- | grep /S | sort | sed -e "s/$/ start/" | sh )'
}

# Execute command in chroot as root
j_root_eval() {
  jname="${1:-${jname}}"; shift
  j_up "${jname}" || wtf "chroot not running"
  eval "$(j_params "${jname}")"
  j_cg_trap $$
  env -i ${jenv} "${chroot}" "${jroot}" /bin/sh -c "${*}"
}

# Execute command in chroot
j_eval() {
  jname="${1:-${jname}}"; shift
  j_up "${jname}" || wtf "chroot not running"
  eval "$(j_params "${jname}")"
  j_cg_trap $$
  env -i ${jenv} "${chroot}" "${jroot}" /bin/su "${juser:-${USER}}" -c "${*}"
}

# Invoke a shell within the chroot
j_shell() {
  jname="${1:-${jname}}"
  eval "$(j_params "${jname}")"
  j_eval "${jname}" "cd; exec ${jshell} -l"
}

# Shut down chroot
j_stop() {
  jname="${1:-${jname}}"
  eval "$(j_params "${jname}")"
  j_up "${jname}" || return 0
  meh "stopping ${jname} ..."

  # Stop all services in /etc/rcJ.d
  j_root_eval "${jname}" '[ -d /etc/rcJ.d ] && ( ls -1 /etc/rcJ.d/* 2>&- | grep /S | sort -r | sed -e "s/$/ stop/" | sh )'

  # Terminate ipcd
  j_ipc_stop

  # Unmount filesystems
  j_fstab_umount
}

# Send a signal to all processes inside a chroot, if possible
j_signal() {
  sig="${1:-HUP}"
  jname="${2:-${jname}}"

  [ "$(j_cgroup)" ] || { echo "No cgroup support; cannot signal"; return 1; }
  pids="$(cat "$(j_cgroup)/tasks")"
  if [ "${pids}" ]
  then
    echo Sending ${sig} to ${pids} ...
    echo "FROZEN" > "$(j_cgroup)/freezer.state"
    kill -${sig} ${pids}
    echo "THAWED" > "$(j_cgroup)/freezer.state"
  fi
}

# Kill all processes inside chroot (TERM, then KILL)
j_kill() {
  jname="${1:-${jname}}"
  eval "$(j_params "${jname}")"
  meh "killing ${jname} ..."

  j_signal "TERM" "${jname}"
  sleep 2
  j_signal "KILL" "${jname}"
}

# Hook to allow mounting chroot mounts from command line
j_mount() {
  jname="${1:-${jname}}"
  eval "$(j_params "${jname}")"
  j_fstab_mount
}

# Hook to allow unmounting chroot mounts from command line
j_umount() {
  jname="${1:-${jname}}"
  eval "$(j_params "${jname}")"
  j_fstab_umount
}

case "${cmd}" in
init|create) j_init "${jname}" "${@}" ;;
ls|list) j_ls ;;
status) j_status "${jname}" "${@}" ;;
start) j_start "${jname}" ;;
shell|enter) j_shell "${jname}" ;;
eval) j_eval "${jname}" "${*}" ;;
stop) j_stop "${jname}" ;;
kill) j_kill "${jname}" ;;
mount) j_mount "${jname}" ;;
umount) j_umount "${jname}" ;;
*) pebkac ;;
esac