#!/bin/bash
# ref: http://vincent.bernat.im/en/blog/2011-uml-network-lab.html
# Adopt to Archlinux by dlin.tw at gmail
# user-mode-linux's config should modified
#   CONFIG_UML_NET_VDE=y
#   CONFIG_HOSTFS=y
# arch packages: vde2 user-mode-linux screen start-stop-daemon rsyslog

LABNAME="2vm"
# network layout: sw - C1-2

DEPS="screen ip vde_switch vmlinux bash rsyslogd"
PROGNAME="$0"
PROGARGS="$@"

# Check for dependencies needed by this tool
check_dependencies() {
  for dep in $DEPS; do
    which $dep 2> /dev/null > /dev/null || {
      echo "[!] Missing dependency: $dep"
      exit 1
    }
  done
}

# Run our lab in screen
setup_screen() {
  [ x"$TERM" = x"screen" ] || \
    exec screen -ln -S $LABNAME -c /dev/null -t main "$PROGNAME" "$PROGARGS"
  sleep 1
  screen -X zombie cr
  screen -X caption always "%{= wk}%-w%{= BW}%n %t%{-}%+w %-="
}

# Setup a VDE switch
setup_switch() {
  echo "[+] Setup switch $1"
  screen -t "sw-$1" \
    start-stop-daemon --make-pidfile --pidfile "$TMP/switch-$1.pid" \
    --start --startas $(which vde_switch) -- \
    --sock "$TMP/switch-$1.sock"
  screen -X select 0
}

# Start a VM
start_vm() {
  echo "[+] Start VM $1"
  name="$1"
  shift
  screen -t $name \
  start-stop-daemon --make-pidfile --pidfile "$TMP/vm-$name.pid" \
  --start --startas $(which vmlinux) -- \
  uts=$name mem=64M \
  root=/dev/root rootfstype=hostfs init=$(readlink -f "$PROGNAME") \
  "$@"
  screen -X select 0
}

display_help() {
    cat <<EOF

Some screen commands :
 ^A ?     - display 'screen' help
 ^A ^A    - type original ^A
 ^A d     - Detach the screen (resume with screen -r $LABNAME)
 ^A a "   - Select a window
 ^A space - Next window
 ^A ^A    - Toggle with Previous window
 ^A [     - Copy mode, use PgUp/PgDn to scroll screen
EOF
  echo "Press enter to exit the lab - [$LABNAME]"
  read a
}

cleanup() {
  for pid in $TMP/*.pid; do
    kill $(cat $pid)
  done
  screen -X quit
}

case $$ in
  1)
    # Inside UML. Three states:
    #   1. Provide Job Control (equal to getty)
    #   2. Setup AUFS
    #   3. Remaining setup
    STATE=${STATE:-1}

    case $STATE in
      1)
        echo "[+] Set hostname"
        hostname ${uts}
        echo "[+] Set path"
        export TERM=xterm
        export PATH=/usr/local/bin:/usr/bin:/bin:/sbin:/usr/local/sbin:/usr/sbin

        # Provide Job Control (^C)
        export STATE=3 # FIXME:skip AUFS
        exec setsid python -c '
import os, sys
os.close(0)
os.open("/dev/tty0", os.O_RDWR | os.O_NOCTTY | os.O_NONBLOCK, 0)
os.dup2(0, 1)
os.dup2(0, 2)
# os.tcsetpgrp(0, 1)
os.execv(sys.argv[1], [sys.argv[1]])' "$PROGNAME"
      ;;
      2)
        echo "[+] Setup AUFS"
        mount -n -t proc proc /proc
        mount -n -t sysfs sysfs /sys
        mount -o bind /usr/lib/uml/modules /lib/modules
        mount -n -t tmpfs tmpfs /tmp -o rw,nosuid,nodev
        mkdir /tmp/ro
        mkdir /tmp/rw
        mkdir /tmp/aufs
        mount -n -t hostfs hostfs /tmp/ro -o /,ro
        mount -n -t aufs aufs /tmp/aufs -o noatime,dirs=/tmp/rw:/tmp/ro=ro

        # Chroot inside our new root
        export STATE=3
        exec chroot /tmp/aufs "$PROGNAME"
      ;;
    esac

    echo "[+] Set filesystems"
    # FIXME for aufs: rm /etc/mtab
    mount -t proc proc /proc
    mount -t sysfs sysfs /sys
    mount -t tmpfs tmpfs /dev -o rw && {
      cd /dev
      if [ -f $(dirname "$PROGNAME")/dev.tar ]; then
        tar xf $(dirname "$PROGNAME")/dev.tar
      else
        MAKEDEV null consoleonly
      fi
    }
    mount -o bind /usr/lib/uml/modules /lib/modules
    for fs in /var/run /var/tmp /var/log /tmp; do
      mount -t tmpfs tmpfs $fs -o rw,nosuid,nodev
    done
    mount -t hostfs hostfs $(dirname "$PROGNAME") -o $(dirname "$PROGNAME")

    # Interfaces
    echo "[+] Set interfaces"
    for intf in /sys/class/net/*; do
      intf=$(basename $intf)
      ip a l dev $intf 2> /dev/null >/dev/null && ip link set up dev $intf
    done

    echo "[+] Start syslog"
    rsyslogd

    cd $(dirname "$PROGNAME")
    [ -f dev.tar ] || {
      tar -C /dev -cf dev.tar.$uts . && mv dev.tar.$uts dev.tar
    }

    # Configure each UML
    echo "[+] Setup UML"
    sysctl -w net.ipv4.ip_forward=1
    case ${uts} in
      C*)
        n=${uts#C}  # strip prefix
        ip="192.168.0.$n"
        echo "IP=$ip"
        ip addr add $ip/24 dev eth0
      ;;
    esac

    echo "[+] Drop to a shell"
    exec /bin/bash

  ;;
  *)
    TMP=$(mktemp -d /tmp/net.XXX)
    trap "rm -rf $TMP" EXIT
    check_dependencies
    setup_screen
    echo "[+] clean up tmp:$TMP"

    # Setup switches
    setup_switch sw1

    # Start VM
    start_vm C1 eth0=vde,$TMP/switch-sw1.sock
    start_vm C2 eth0=vde,$TMP/switch-sw1.sock
    display_help
    cleanup
  ;;
esac

# vi:set et sw=2 ts=2 sta: