#!/bin/sh

GUEST_ROOT=/usr/local/share/dockerbox
CONF_DIR=/usr/local/etc/dockerbox

: "${IFCONFIG:=/sbin/ifconfig}"
: "${GREP:=/usr/bin/grep}"
: "${GRUB_BHYVE:=/usr/local/sbin/grub-bhyve}"
: "${BHYVE:=/usr/sbin/bhyve}"
: "${SED:=/usr/bin/sed}"
: "${ECHO:=/bin/echo}"
: "${DAEMON:=/usr/sbin/daemon}"
: "${ROUTE:=/sbin/route}"
: "${PFCTL:=/sbin/pfctl}"
: "${SYSCTL:=/sbin/sysctl}"
: "${KLDSTAT:=/sbin/kldstat}"
: "${TAIL:=/usr/bin/tail}"
: "${HEAD:=/usr/bin/head}"
: "${KLDLOAD:=/sbin/kldload}"
: "${KLDUNLOAD:=/sbin/kldunload}"
: "${KILL:=/bin/kill}"
: "${PGREP:=/bin/pgrep}"
: "${TOUCH:=/usr/bin/touch}"
: "${SHA256:=/sbin/sha256}"
: "${TAR:=/usr/bin/tar}"
: "${RM:=/bin/rm}"
: "${TRUNCATE:=/usr/bin/truncate}"
: "${MDCONFIG:=/sbin/mdconfig}"
: "${GPART:=/sbin/gpart}"
: "${RESIZE2FS:=/usr/local/sbin/resize2fs}"
: "${E2FSCK:=/usr/local/sbin/e2fsck}"
: "${FILE:=/usr/bin/file}"
: "${AWK:=/usr/bin/awk}"
: "${FETCH:=/usr/bin/fetch}"
: "${INSTALL:=/usr/bin/install}"
: "${LS:=/bin/ls}"
: "${CU:=/usr/bin/cu}"
: "${MKDIR:=/bin/mkdir}"
: "${MKFS:=/usr/local/sbin/mkfs.ext4}"
: "${SEQ:=/usr/bin/seq}"
: "${SLEEP:=/bin/sleep}"
: "${BHYVECTL:=/usr/sbin/bhyvectl}"
: "${SERVICE:=/usr/sbin/service}"
: "${MOUNT_NULLFS:=/sbin/mount_nullfs}"
: "${UMOUNT:=/sbin/umount}"
: "${MOUNT:=/sbin/mount}"
: "${CAT:=/bin/cat}"

: "${VMM_KO:=/boot/kernel/vmm.ko}"
: "${NMDM_KO:=/boot/kernel/nmdm.ko}"
: "${PF_KO:=/boot/kernel/pf.ko}"
: "${PF_CONF:=/etc/pf.conf}"

DOCKERBOX_IF=dockerbox0
DOCKERBOX_VM=dockerbox
DOCKERBOX_LOG=/var/log/dockerbox.log
DISK_IMAGE="${GUEST_ROOT}/disk.img"
DOCKER_IMAGE="${GUEST_ROOT}/docker.img"
DOCKER_DISK_SIZE=20G
DEVICE_MAP="${GUEST_ROOT}/device.map"
DISK_VERSION_FILE="${GUEST_ROOT}/disk-version"
VM_MANAGER_DAEMON_ID="dockerbox-vm-manager"
CONFIG="${CONF_DIR}/dockerbox.conf"
NAT_IP="10.0.0.254"
DOCKERBOX_IP="10.0.0.1"
DOCKERBOX_PORT=2375
DAEMON_PID_FILE=/var/run/dockerbox-daemon.pid
DOCKERBOX_PID_FILE=/var/run/dockerbox.pid

NMDM_DEVICE=/dev/nmdm-dockerbox
NMDM_A="${NMDM_DEVICE}.1A"
NMDM_B="${NMDM_DEVICE}.1B"

DISK_IMG_SUM=413c9eb72b828eb329657e17b87b42acfd24562cadab8e117621e1456e7c125a
DISKVERSION=0.4.5

help_msg="Usage: dockerbox command

command:
  start             start dockerbox
  stop              stop dockerbox
  console           connect to dockerbox console
  status            show dockerbox status
  fetch             fetch dockerbox root disk image
  resize SIZE       extend docker disk image with size of SIZE
  mount SRC TARGET  mount SRC path to TARGET path under the shared folder
  umount TARGET     umount TARGET path
"

log() {
  local _level="$1"
  local _message="$2"
  local _timestamp

  case ${_level} in
    error) ${ECHO} "ERROR: ${_message}";;
    warn) ${ECHO} "WARNING: ${_message}";;
  esac

  _level="$(${ECHO} $_level | tr "[:lower:]" "[:upper:]")"
  _timestamp="$(date +'%FT%H:%M:%S%z')"
  ${ECHO} "${_timestamp} ${_level} ${_message}" >> ${DOCKERBOX_LOG}
}

output() {
  local _message="$1"

  log info "Output: ${_message}"
  ${ECHO} "${_message}."
}

capture_output() {
  local _type="$1"
  local _id="$2"

  while read -r message; do
    log "${_type}" "[${_id}] ${message}"
  done
}

sysctl_value() {
  ${SYSCTL} -nq "$1"
}

check_configuration() {
  local _file="$1"

  if [ ! -f "${_file}" ]; then
    log error "${_file} is missing, please create it from the sample"
    exit 3
  fi
}

has_bridge_interface() {
  ${IFCONFIG} | ${GREP} -Fq "${DOCKERBOX_IF}: "
}

get_tap_interface() {
  if has_bridge_interface; then
    ${IFCONFIG} "${DOCKERBOX_IF}" | ${GREP} -F member | ${SED} -E 's/^.*member:.*(tap[^ ]*).*$/\1/'
  else
    ${ECHO} ""
  fi
}

check_ext_if() {
  ${IFCONFIG} | ${GREP} ${ext_if} -c
}

assert_kmod_loaded() {
  local _kmod="$1"
  local _kmod_file
  local _kldload_out
  local _rc

  if [ -n "$2" ]; then
    _kmod_file="$2"
  else
    _kmod_file="${_kmod}"
  fi

  log debug "assert loaded: kmod=${_kmod}, kmod_file=${_kmod_file}"

  if ! ${KLDSTAT} -q -m "${_kmod}"; then
    log warn "Kernel module ${_kmod} is not loaded. Loading kernel module ${_kmod}"

    _kldload_out=$(${KLDLOAD} "${_kmod_file}" 2>&1)
    _rc=$?
    ${ECHO} "${_kldload_out}" | capture_output debug kldload
    if [ ${_rc} -ne 0 ]; then
      log error "${_kmod_file} kernel module could not be loaded"
      exit 127
    fi

    log warn "Kernel module ${_kmod} was loaded successfully"
  fi
}

nat_count() {
  local _nat_keyword="$1"
  local _ext_if="$2"
  log debug "assert NAT loaded: nat_if=${_nat_keyword} ext_if=${_ext_if}"
  ${PFCTL} -s nat | ${GREP} nat | ${GREP} ${_nat_keyword} | ${GREP} ${_ext_if} -c
}

detect_ext_if() {
  ${ROUTE} -n get -inet default | ${GREP} 'interface:' | ${AWK} '{ print $2 }'
}

check_pf() {
  assert_kmod_loaded "pf" "${PF_KO}"

  if [ "${ext_if}" = "auto" ]; then
    log info "ext_if=auto configured. Detecting external interface"
    ext_if=$(detect_ext_if)
    if [ -z "${ext_if}" ]; then
      log warn "No external interface available"
    else
      log info "External interface is ${ext_if}"
    fi
  elif [ $(check_ext_if) -lt 1 ]; then
    log warn "External interface specified in settings not available"
  fi

  if [ ! $(${PFCTL} -s Running) = "Enabled" ]; then
    log info "Enabling PF"
    ${PFCTL} -e
  fi

  if [ ! -f "${PF_CONF}" ]; then
    log warn "No pf configuration found"
    log info "Creating empty pf configuration ${PF_CONF}"
    ${TOUCH} ${PF_CONF}
  fi

  nat_keyword="$(${ECHO} ${nat_ip} | cut -d'.' -f1-3).0/${nat_mask}"

  if [ $(nat_count ${nat_keyword} ${ext_if}) -lt 1 ]; then
    log warn "No pf NAT set up. Writing NAT rule into ${PF_CONF}"
    ${ECHO} "nat on ${ext_if} inet from ${nat_keyword} to any -> (${ext_if})" >> ${PF_CONF}
    log info "Loading pf conf ${PF_CONF}"
    ${PFCTL} -f ${PF_CONF}
    if [ $(nat_count ${nat_keyword} ${ext_if}) -lt 1 ]; then
      log warn "Failed to load pf NAT with ${PF_CONF}"
    fi
  fi
}

create_bridge() {
  if ! has_bridge_interface; then
    log info "Creating bridge interface: ${DOCKERBOX_IF}"
    ${IFCONFIG} bridge create name ${DOCKERBOX_IF} up 2>&1 | capture_output debug ifconfig
  else
    log warn "Bridge interface already exists: ${DOCKERBOX_IF}, skipping creation"
  fi

  ${IFCONFIG} ${DOCKERBOX_IF} inet ${nat_ip}/${nat_mask} 2>&1 | capture_output debug ifconfig
}

destroy_bridge() {
  if has_bridge_interface; then
    log info "Destroying bridge interface: ${DOCKERBOX_IF}"
    ${IFCONFIG} ${DOCKERBOX_IF} destroy 2>&1 | capture_output debug ifconfig
  fi
}

create_tap() {
  local _tap="$(get_tap_interface)"

  if [ -z "${_tap}" ]; then
    _tap="$(${IFCONFIG} tap create up)"
    log info "Linking tap interface to ${DOCKERBOX_IF}: ${_tap}"
    ${IFCONFIG} ${DOCKERBOX_IF} addm "${_tap}" 2>&1 | capture_output debug ifconfig
  else
    log warn "Linked tap interface already exists: ${_tap}, skipping creation"
  fi
}

destroy_tap() {
  local _tap="$(get_tap_interface)"

  if [ -n "${_tap}" ]; then
    log info "Destroying linked tap interface: ${_tap}"
    ${IFCONFIG} ${_tap} destroy 2>&1 | capture_output debug ifconfig
  fi
}

check_nat() {
  local _tap

  if [ $(sysctl_value net.inet.ip.forwarding) -eq 0 ]; then
    ${SYSCTL} net.inet.ip.forwarding=1 | capture_output debug sysctl
  fi
  create_bridge
  create_tap
  check_pf
}

get_kmod_path() {
  local _kmod="$1"
  ${KLDSTAT} -v -n "${_kmod}" | ${TAIL} +2 | ${HEAD} -1 | ${SED} -e 's![^(]*(\([^)]*\))!\1!'
}

check_vmm() {
  local _kmod_path

  if ${KLDSTAT} -q -m vmm; then
    _kmod_path="$(get_kmod_path vmm)"
    log info "vmm.ko is expected at path: ${VMM_KO}"
    log info "vmm.ko is found at path: ${_kmod_path}"

    if [ "${_kmod_path}" != "${VMM_KO}" ] && ! ${KLDUNLOAD} vmm; then
      log error "vmm.ko is loaded from a different location, but cannot be replaced"
      exit 127
    fi
  fi

  assert_kmod_loaded "vmm" "${VMM_KO}"
}

check_nmdm() {
  assert_kmod_loaded "nmdm" "${NMDM_KO}"
}

check_nfs() {
  nfs_folder="$1"
  if [ -d "${nfs_folder}" ]; then
    log warn "Configuring NFS. ${nfs_folder} will be exported. If the export fails, dockerbox will not mount this export."
    if ! ${GREP} -Eq "${nfs_folder}[[:blank:]]+${dockerbox_ip}" /etc/exports; then
      ${ECHO} "${nfs_folder} ${dockerbox_ip}" >> /etc/exports
    fi
    ${SERVICE} nfsd onerestart >/dev/null 2>&1
    ${SERVICE} mountd onereload
  fi
}

unexport_nfs() {
  nfs_folder="$1"
  [ -z "${nfs_folder}" ] && return 0

  if ${GREP} -Eq "^${nfs_folder}[[:blank:]]+${dockerbox_ip}[[:blank:]]*$" /etc/exports; then
    log warn "NFS unexporting ${nfs_folder}"
    ${SED} -i '' -E "\#^${nfs_folder}[[:blank:]]+${dockerbox_ip}[[:blank:]]*\$#d" /etc/exports
    ${SERVICE} mountd onereload
  fi
}

get_daemon_pid() {
  [ -f ${DAEMON_PID_FILE} ] && cat ${DAEMON_PID_FILE}
}

get_dockerbox_pid() {
  [ -f ${DOCKERBOX_PID_FILE} ] && cat ${DOCKERBOX_PID_FILE}
}

get_dockerbox_status() {
  if [ -n "$(get_dockerbox_pid)" ]; then
    if ${FETCH} -qo /dev/null --timeout=1 "http://${dockerbox_ip}:${DOCKERBOX_PORT}/_ping" > /dev/null 2>&1; then
      return 0
    else
      return 1
    fi
  else
    return 2
  fi
}

gen_docker_image() {
  local _docker_mddev

  ${TRUNCATE} -s ${DOCKER_DISK_SIZE} ${DOCKER_IMAGE}
  log debug "Truncated new docker image at ${DOCKER_IMAGE}"
  _docker_mddev=$(${MDCONFIG} ${DOCKER_IMAGE})
  ${GPART} create -s mbr ${_docker_mddev}
  ${GPART} add -t linux-data ${_docker_mddev}
  log debug "Added partition to ${DOCKER_IMAGE}"
  ${MKFS} -q /dev/${_docker_mddev}s1
  log debug "Made ext4 file system on ${DOCKER_IMAGE}'s partition"
  ${MDCONFIG} -du ${_docker_mddev}
}

check_device_map() {
  if [ ! -f ${DEVICE_MAP} ]; then
    log warn "Device map not found. Generating device map at ${DEVICE_MAP}"
    ${ECHO} "(hd0) ${DISK_IMAGE}" > ${DEVICE_MAP}
  fi
}

check_docker_image() {
  if [ ! -f ${DOCKER_IMAGE} ]; then
    log warn "Docker image not found. Generating docker image at ${DOCKER_IMAGE}. This may exceed startup time limit check"
    gen_docker_image
    log info "Empty docker image generated"
  fi
}

check_docker_image_fs() {
  _md_dev=$(${MDCONFIG} ${DOCKER_IMAGE})
  if [ ! -e /dev/${_md_dev}s1 ]; then
    log warn "No partition found in ${DOCKER_IMAGE}. Creating one"
    ${GPART} add -t linux-data ${_md_dev}
    log warn "Added partition to ${DOCKER_IMAGE}"
  fi
  if [ -z "$(file -s /dev/${_md_dev}s1 | grep 'Linux rev 1.0 ext4')" ]; then
    log warn "No ext4 file system found in ${DOCKER_IMAGE}. Making one"
    ${MKFS} -q /dev/${_md_dev}s1
    log warn "Made ext4 file system in ${DOCKER_IMAGE}"
  fi
  ${MDCONFIG} -du ${_md_dev}
}

vm_manager() {
  local _tap
  local _grub_args
  local _bhyve_args

  log info "Starting dockerbox"

  check_vmm
  check_nat
  check_nfs "${nfs_folder}"
  check_device_map
  check_docker_image
  check_docker_image_fs

  if [ ! -d "${shared_folder}" ]; then
    log warn "Shared folder ${shared_folder} not found. Creating ${shared_folder}."
    ${MKDIR} -p ${shared_folder}
  fi

  _tap="$(get_tap_interface)"
  _grub_args="-m ${DEVICE_MAP} -r hd0,msdos1 -M ${memory}"
  _bhyve_args="-A -H -P -s 0:0,hostbridge -s 1:0,lpc -s 2:0,virtio-net,${_tap} -s 3:0,virtio-blk,${DISK_IMAGE} -s 4:0,virtio-blk,${DOCKER_IMAGE} -s 5:0,virtio-9p,shared_folder=${shared_folder} -s 6:0,virtio-9p,dockerbox_config=${CONF_DIR} -c ${cpu} -m ${memory}"

  if [ "${console}" = "yes" ]; then
    check_nmdm
    _bhyve_args="${_bhyve_args} -l com1,${NMDM_A}"
  fi

  log debug "Grub arguments: ${_grub_args}"
  grub_output=$(env TERM=xterm-256color ${GRUB_BHYVE} ${_grub_args} ${DOCKERBOX_VM} 2>&1)
  grub_rc=$?
  if [ ${grub_rc} -eq 0 ]; then
    log debug "Bhyve arguments: ${_bhyve_args}"
    bhyve_output=$(${BHYVE} ${_bhyve_args} ${DOCKERBOX_VM} 2>&1)
    bhyve_rc=$?
    if [ ${bhyve_rc} -ne 0 ]; then
      ${ECHO} ${bhyve_output} | capture_output error bhyve
    fi
  else
    ${ECHO} ${grub_output} | capture_output error grub-bhyve
  fi

  destroy_tap
  destroy_bridge
  unexport_nfs "${nfs_folder}"
}

dockerbox_start() {
  local _pid
  local _start_wait_max=30

  if get_dockerbox_status; then
    _pid="$(get_dockerbox_pid)"
    log warn "${DOCKERBOX_VM} is already running as PID ${_pid}, left intact"
    return 1
  fi

  if [ ! -f ${DISK_IMAGE} ]; then
    log error "Disk image not installed. Please run 'dockerbox fetch'"
    exit 1
  fi

  installed_disk_version=$(${CAT} ${DISK_VERSION_FILE})
  if [ ! "${installed_disk_version}" = "${DISKVERSION}" ]; then
    log warn "Disk image version (${installed_disk_version}) mismatch script expected version (${DISKVERSION}). Consider upgrade."
  fi

  log info "Start ${DOCKERBOX_VM} daemon"
  ${DAEMON} -p ${DOCKERBOX_PID_FILE} -P ${DAEMON_PID_FILE} -t "${VM_MANAGER_DAEMON_ID}" "${0}" _manage_vm

  _start_wait_max=${start_wait_timeout}

  for i in $(${SEQ} 1 ${_start_wait_max}); do
    ${SLEEP} 1

    if get_dockerbox_status; then
      _pid="$(get_dockerbox_pid)"
      log info "${DOCKERBOX_VM} has started as PID ${_pid}"
      ${ECHO} "${DOCKERBOX_VM} started"
      exit
    elif [ -z "$(get_daemon_pid)" ]; then
      exit 1
    fi
  done

  log warn "${DOCKERBOX_VM} service is not up after ${_start_wait_max} seconds. $(dockerbox_status)"
}

dockerbox_poststart() {
  if [ -n "$(${SERVICE} -e | ${GREP} dockerbox-broker)" ]; then
    ${SERVICE} dockerbox-broker start
  fi
}

dockerbox_console() {
  local _pid

  log info "Begin: dockerbox console"
  _pid="$(get_daemon_pid)"
  log debug "Dockerbox pid: ${_pid}"

  if [ -z "${_pid}" ]; then
    log error "There is no guest to attach to"
    exit 127
  fi

  if ! ${LS} "${NMDM_DEVICE}"* > /dev/null 2>&1; then
    log error "No null-modem device is configured"
    exit 127
  fi

  ${ECHO} 'Connecting, type "~." to leave the session...'
  log info "Attaching to the guest"
  ${CU} -s 115200 -l ${NMDM_B}
  log info "Detached from the guest"
  ${ECHO} "Finished."
  log info "End: dockerbox console"
}

dockerbox_prestop() {
  if [ -n "$(${SERVICE} -e | ${GREP} dockerbox-broker)" ]; then
    ${SERVICE} dockerbox-broker stop
  fi
}

clean_mount() {
  ${MOUNT} -t nullfs -p | ${AWK} '$2 ~ "^'"${shared_folder}"'/" {print $2}' | \
  while read -r target; do
    log info "Umounting ${target}"
    if ! ${UMOUNT} -t nullfs "${target}"; then
      log warn "Failed to unmount ${target}"
    fi
  done
}

dockerbox_stop() {
  local _pid

  _pid="$(get_daemon_pid)"
  if [ -z "${_pid}" ]; then
    log warn "Dockerbox is not running"
    exit 1
  fi

  log info "Stopping guest ${DOCKERBOX_VM}, managed by PID ${_pid}"
  ${BHYVECTL} --force-poweroff --vm=${DOCKERBOX_VM}
  ${BHYVECTL} --destroy --vm=${DOCKERBOX_VM}

  for i in $(${SEQ} 1 ${start_wait_timeout}); do
    ${SLEEP} 1

    if ! get_dockerbox_status; then
      break
    fi
  done
  log info "Dockerbox stopped: killed pid=${_pid}"

  clean_mount
}

dockerbox_status() {
  get_dockerbox_status
  case $? in
    0) ${ECHO} "dockerbox is running and dockerd is reachable." ;;
    1) ${ECHO} "dockerbox is running, but dockerd is not reachable." ;;
    2) ${ECHO} "dockerbox is not running." ;;
  esac
}

dockerbox_resize() {
  local partition
  local part_index=1

  if [ ! -f ${DOCKER_IMAGE} ]; then
    log error "Docker image not found"
    exit 1
  fi

  size=$(${ECHO} "$1" | ${GREP} -E "[1-9][0-9]*[0-9]*[GgMm]")
  if [ "$?" -ne 0 ]; then
    log error "Invalid size argument"
    exit 1
  fi

  log info "Docker image extending ${size}..."

  ${TRUNCATE} -s "+${size}" ${DOCKER_IMAGE}

  log warn "Disk image extended, and will take effect on next start."
}

dockerbox_fetch() {
  DISKRELEASE=dockerbox-img-${DISKVERSION}
  TMPDIR=/tmp

  if [ -f ${DISK_IMAGE} ]; then
    read -p "${DISK_IMAGE} exists. Overwrite? [y/n]: " yn
    if [ $yn != "y" ]; then
      ${ECHO} "Abort"
      exit
    fi
  fi

  log debug "Fetch dockerbox disk image v${DISKVERSION}"
  ${FETCH} -o ${TMPDIR} https://github.com/leafoliage/freebsd-dockerbox-debian/releases/download/v${DISKVERSION}/dockerbox-img-${DISKVERSION}.tar.zstd
  if [ $? -ne 0 ]; then exit 1; fi

  tarball=${TMPDIR}/${DISKRELEASE}.tar.zstd
  if [ "$(${SHA256} -q ${tarball})" != "${DISK_IMG_SUM}" ]; then
    ${RM} ${tarball}
    log error "Dockerbox image checksum unmatch"
    exit 1
  fi
  ${TAR} -C ${TMPDIR} -xf ${tarball}
  ${MKDIR} -p ${GUEST_ROOT}
  ${INSTALL} -m 0644 ${TMPDIR}/${DISKRELEASE}/* ${GUEST_ROOT}
  ${ECHO} ${DISKVERSION} > ${DISK_VERSION_FILE}
  log info "Dockerbox disk image v${DISKVERSION} installed to ${DISK_IMAGE}"
}

dockerbox_mount() {
  if [ -z "$2" ]; then
    echo "Usage: dockerbox mount SRC_PATH TARGET_PATH"
    echo "  Note that TARGET_PATH would be automatically prepended with shared folder path."
    exit 1
  fi

  src_path="$1"
  target_path="${shared_folder}$2"
  echo "Mounting ${src_path} to ${target_path}."
  ${MKDIR} -p ${target_path}
  ${MOUNT_NULLFS} ${src_path} ${target_path}
}

dockerbox_umount() {
  if [ -z "$1" ]; then
    echo "Usage: dockerbox umount TARGET_PATH"
  fi

  ${UMOUNT} "${shared_folder}$1"
}

load_config() {
  cpu=1
  memory=1024M
  ext_if=auto
  dockerbox_ip=$DOCKERBOX_IP
  nat_ip=$NAT_IP
  nat_mask=24
  console="yes"
  start_wait_timeout=30
  shared_folder="/dockerbox"
  nfs_folder=""

  check_configuration ${CONFIG}
  . ${CONFIG}
}


command="$1"
shift

load_config

case ${command} in
  start) dockerbox_start;;
  stop) dockerbox_stop;;
  console) dockerbox_console;;
  status) dockerbox_status;;
  fetch) dockerbox_fetch;;
  resize) dockerbox_resize $@;;
  poststart) dockerbox_poststart;;
  prestop) dockerbox_prestop;;
  mount) dockerbox_mount $@;;
  umount) dockerbox_umount $@;;
  _manage_vm) vm_manager;;
  *) ${ECHO} "${help_msg}";;
esac
