#!/bin/bash
#
# Packet capture helper, reached from the service account through three fully pinned
# sudoers entries. Nothing the caller supplies crosses the privilege boundary: the
# storage, the output path, the interface, the host filter and the duration cap are
# all decided here.
#
# Shipped via deb-common, so the whole file is compiled by handlebars at build time.
# A literal double brace anywhere in it would break the build.

set -euo pipefail

readonly SERVICE_USER='unifi-talk'
readonly DB_NAME='unifi-talk'
readonly CAPTURE_DIR_NAME='unifi-talk-pcap'
readonly CAPTURE_FILE_NAME='talk_lan.pcap'
readonly STATE_FILE='/run/unifi-talk-pcap.state'

# /srv is the stable handle for an app volume whose own mount point is named after a
# filesystem uuid. Fixed paths come first because the glob expands in uuid order, so it
# only decides where none of them exist. The boot medium is deliberately absent: a capture
# belongs on external storage, so a console without a disk cannot start one at all
readonly STORAGE_CANDIDATES=(/srv /ssd1 /volume/*)
readonly MAX_SECONDS=3900
readonly MIN_FREE_MB=2048
readonly MAX_HOSTS=128

# exit codes the caller maps to messages
readonly E_NO_STORAGE=2
readonly E_UNTRUSTED=3
readonly E_NO_SPACE=4
readonly E_NO_DEVICES=5

fail() {
  echo "$2" >&2
  exit "$1"
}

# a directory we are willing to create root-owned state in: only root may add, remove
# or rename entries, so nothing can be swapped underneath us after this returns.
# stat does not dereference, so a symlink fails here: callers pass a resolved path
require_trusted_dir() {
  local target="$1" uid mode kind

  IFS=' ' read -r uid mode kind < <(stat -c '%u %a %F' -- "$target" 2>/dev/null) || return 1

  [ "$kind" = 'directory' ] || return 1
  [ "$uid" = '0' ] || return 1
  [ "$(( 0$mode & 0022 ))" -eq 0 ] || return 1
}

# first candidate that resolves onto a mount other than the root filesystem, so a
# directory inside an app volume counts while anything on the boot medium does not
select_storage() {
  local candidate resolved holder untrusted=0

  for candidate in "${STORAGE_CANDIDATES[@]}"; do
    resolved=$(realpath -e -- "$candidate" 2>/dev/null) || continue

    [ -d "$resolved" ] || continue

    holder=$(stat -c '%m' -- "$resolved" 2>/dev/null) || continue
    [ "$holder" != '/' ] || continue

    # an untrusted candidate must not end the search, or one bad mount hides the rest
    if ! require_trusted_dir "$resolved"; then
      untrusted=1
      continue
    fi

    printf '%s' "$resolved"
    return 0
  done

  [ "$untrusted" -eq 0 ] || return "$E_UNTRUSTED"

  return "$E_NO_STORAGE"
}

free_mb() {
  local _fs _blocks _used avail='' _rest

  read -r _fs _blocks _used avail _rest < <(df -Pm -- "$1" 2>/dev/null | tail -n 1) || true

  case "$avail" in
    '' | *[!0-9]*) avail=0 ;;
  esac

  printf '%s' "$avail"
}

capture_iface() {
  if ip link show br0 > /dev/null 2>&1; then
    printf 'br0'
  else
    printf 'eth0'
  fi
}

# the state file is the only thing telling us what to kill, so it lives in a root-owned
# directory: a forged pid here would be a root kill of an arbitrary process
read_state() {
  STATE_PID=''
  STATE_OUT=''

  require_trusted_dir '/run' || fail "$E_UNTRUSTED" '/run is not root owned'

  [ -f "$STATE_FILE" ] || return 0

  {
    read -r STATE_PID || true
    read -r STATE_OUT || true
  } < "$STATE_FILE"
}

is_our_capture() {
  local pid="$1" out="$2" cmdline

  case "$pid" in
    '' | *[!0-9]*) return 1 ;;
  esac

  [ -n "$out" ] || return 1
  [ -r "/proc/$pid/cmdline" ] || return 1

  cmdline=$(tr '\0' ' ' < "/proc/$pid/cmdline" 2>/dev/null) || return 1

  case "$cmdline" in
    *tcpdump*"$out"*) return 0 ;;
  esac

  return 1
}

# builds "host A or host B ..." as separate argv tokens, so the addresses never reach
# a shell and each one is validated before it is passed on
build_filter() {
  local rows ip count=0 o1 o2 o3 o4

  rows=$(psql -X -U postgres -d "$DB_NAME" -tAc 'SELECT ip FROM device_discovery WHERE ip IS NOT NULL' 2>/dev/null) || rows=''

  FILTER=()

  while IFS= read -r ip; do

    if [ "$count" -ge "$MAX_HOSTS" ]; then
      break
    fi

    if [[ ! "$ip" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
      continue
    fi

    IFS='.' read -r o1 o2 o3 o4 <<< "$ip"

    if [ "$o1" -gt 255 ] || [ "$o2" -gt 255 ] || [ "$o3" -gt 255 ] || [ "$o4" -gt 255 ]; then
      continue
    fi

    if [ "$count" -gt 0 ]; then
      FILTER+=(or)
    fi

    FILTER+=(host "$ip")
    count=$(( count + 1 ))

  done <<< "$rows"

  [ "$count" -gt 0 ]
}

# timeout runs the capture in a process group of its own, so signalling the group is
# what reaches tcpdump; the pgid is read back from /proc first, because without that
# check a plain background job would leave us signalling our own group
signal_capture() {
  local pid="$1" sig="$2" stat rest pgid

  stat=$(cat "/proc/$pid/stat" 2>/dev/null) || return 0

  rest=${stat#*") "}
  pgid=$(printf '%s' "$rest" | cut -d ' ' -f 3)

  if [ "$pgid" = "$pid" ]; then
    kill -s "$sig" -- "-$pid" 2>/dev/null || true
  else
    kill -s "$sig" -- "$pid" 2>/dev/null || true
  fi
}

do_stop() {
  local waited=0

  read_state

  if is_our_capture "$STATE_PID" "$STATE_OUT"; then

    signal_capture "$STATE_PID" TERM

    # let tcpdump flush and close the savefile before anyone reads it
    while [ "$waited" -lt 50 ] && [ -d "/proc/$STATE_PID" ]; do
      sleep 0.1
      waited=$(( waited + 1 ))
    done

    signal_capture "$STATE_PID" KILL
  fi

  rm -f -- "$STATE_FILE"

  # reject a symlink rather than relabelling it: chmod has no --no-dereference,
  # so a link here would have its target's mode changed
  if [ -n "$STATE_OUT" ] && [ ! -L "$STATE_OUT" ] && [ -f "$STATE_OUT" ]; then
    chown -h "$SERVICE_USER:$SERVICE_USER" -- "$STATE_OUT" 2>/dev/null || true
    chmod 0640 -- "$STATE_OUT" 2>/dev/null || true
  fi
}

# every path a capture has been written to, so the space comes back whichever layout left it
# there: builds before the capture directory wrote the savefile, and the gzip a download left
# beside it, to the storage root. select_storage returns only a root nothing but root may write
reclaim_captures() {
  local storage="$1" dir="$storage/$CAPTURE_DIR_NAME"

  # nothing but root can have created this, so it cannot have been swapped for a link since.
  # A directory we will not touch is not a reason to leave the rest of the volume held
  if require_trusted_dir "$dir"; then
    rm -f -- "$dir/$CAPTURE_FILE_NAME"
  fi

  rm -f -- "$storage/$CAPTURE_FILE_NAME" "$storage/$CAPTURE_FILE_NAME.gz"
}

# the capture sits in a directory only root may write, so the caller cannot unlink it: a discard is
# how the volume gets its space back, once a download has been served whole and once at every start
do_discard() {
  local storage

  do_stop

  # a volume that is gone or untrusted holds nothing we can reclaim
  storage=$(select_storage) || return 0

  reclaim_captures "$storage"
}

# the volume a capture goes to, in one place so the probe verb cannot name one other than what a
# start would pick. Free space is not asked here: only a start reclaims, so only a start can judge it
select_capture_storage() {
  local storage

  storage=$(select_storage) || case "$?" in
    "$E_UNTRUSTED") fail "$E_UNTRUSTED" 'external storage is not root owned' ;;
    *) fail "$E_NO_STORAGE" 'no external storage' ;;
  esac

  printf '%s' "$storage"
}

# the answer a start would reach, without reaching it: nothing is signalled, created or removed,
# so this stays safe to ask while a capture is running. Liveness comes first, ahead of the gate
# that can exit: storage going out from under a running capture is what kills it, and that is the
# case the caller most needs an answer for
do_probe() {
  local storage

  # a state file in a directory anyone can write says nothing either way, and unlike a stop - where
  # the pid on it becomes a root kill - a probe can leave that unanswered rather than refuse, so the
  # storage answer stays the one a start would give, code for code
  if require_trusted_dir '/run'; then
    read_state

    if is_our_capture "$STATE_PID" "$STATE_OUT"; then
      printf 'capturing\n'
    else
      printf 'idle\n'
    fi
  fi

  storage=$(select_capture_storage)

  printf '%s\n' "$storage"
}

do_start() {
  local storage dir out iface capture_pid

  storage=$(select_capture_storage)

  dir="$storage/$CAPTURE_DIR_NAME"
  out="$dir/$CAPTURE_FILE_NAME"

  # root owned and only group readable: the service account can read the capture but
  # cannot plant a symlink in the directory for the privileged tcpdump to follow
  mkdir -p -m 0750 -- "$dir"
  require_trusted_dir "$dir" || fail "$E_UNTRUSTED" 'capture directory is not trusted'
  chown "root:$SERVICE_USER" -- "$dir"
  chmod 0750 -- "$dir"

  # asking for a capture is asking to replace the one on disk, so it goes before the space is
  # measured: gating on space this reclaims leaves the last capture refusing every start after it
  do_stop

  reclaim_captures "$storage"

  if [ "$(free_mb "$storage")" -lt "$MIN_FREE_MB" ]; then
    fail "$E_NO_SPACE" 'not enough free space'
  fi

  build_filter || fail "$E_NO_DEVICES" 'no devices to capture'
  iface=$(capture_iface)

  # noclobber creates with O_EXCL, so this fails rather than following anything
  ( set -o noclobber; : > "$out" ) || fail 1 'could not create capture file'
  chown "$SERVICE_USER:$SERVICE_USER" -- "$out"
  chmod 0640 -- "$out"

  # max dump size = 675k × ~1516B ≈ 1 GB
  nohup timeout "$MAX_SECONDS" tcpdump -Z "$SERVICE_USER" -i "$iface" -w "$out" -s 1500 -c 675000 "${FILTER[@]}" > /dev/null 2>&1 &
  capture_pid=$!

  printf '%s\n%s\n' "$capture_pid" "$out" > "$STATE_FILE"
  chmod 0644 -- "$STATE_FILE"

  # a tcpdump that dies on the spot (missing interface, unusable filter) would otherwise
  # look like a running capture that yields an empty file an hour later
  sleep 0.5

  if ! is_our_capture "$capture_pid" "$out"; then
    rm -f -- "$STATE_FILE"
    fail 1 'capture did not start'
  fi

  printf '%s\n' "$out"
}

case "${1:-}" in
  start) do_start ;;
  stop) do_stop ;;
  probe) do_probe ;;
  discard) do_discard ;;
  *) fail 1 'usage: pcap-capture start|stop|probe|discard' ;;
esac
