#!/usr/bin/bash
# Invoked inside a managed Linux guest by Libreswan; no host route changes.
set -eu
cd /etc/ipsec.d
IPTABLES=$(command -v iptables)
IP6TABLES=$(command -v ip6tables)
chain=MYKOS_VPN_OUT
allow_peer() {
  local address=$1
  [[ "$address" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]] || exit 1
  if "$IPTABLES" -w 5 -nL "$chain" >/dev/null 2>&1; then
    "$IPTABLES" -w 5 -C "$chain" -d "$address" -p udp -m multiport --dports 500,4500 -j ACCEPT 2>/dev/null ||
      "$IPTABLES" -w 5 -I "$chain" 1 -d "$address" -p udp -m multiport --dports 500,4500 -j ACCEPT
  fi
}
dns_guard() {
  # Keep local stub resolvers local; redirect their non-loopback upstream
  # queries through the same full-tunnel policy as all other guest traffic.
  local protocol operation=${1:--I}
  for protocol in udp tcp; do
    if [ "$operation" = -D ]; then
      "$IPTABLES" -w 5 -t nat -C OUTPUT ! -d 127.0.0.0/8 -p "$protocol" --dport 53 -m comment --comment mykos-full-vpn-dns -j DNAT --to-destination 8.8.8.8 2>/dev/null &&
        "$IPTABLES" -w 5 -t nat -D OUTPUT ! -d 127.0.0.0/8 -p "$protocol" --dport 53 -m comment --comment mykos-full-vpn-dns -j DNAT --to-destination 8.8.8.8 || true
    else
      "$IPTABLES" -w 5 -t nat -C OUTPUT ! -d 127.0.0.0/8 -p "$protocol" --dport 53 -m comment --comment mykos-full-vpn-dns -j DNAT --to-destination 8.8.8.8 2>/dev/null ||
        "$IPTABLES" -w 5 -t nat -I OUTPUT 1 ! -d 127.0.0.0/8 -p "$protocol" --dport 53 -m comment --comment mykos-full-vpn-dns -j DNAT --to-destination 8.8.8.8
    fi
  done
}
guard() {
  [ "${MYKOS_HOST_GATEWAY:-0}" != 1 ] || return 0
  # These private chains are owned by MYKOS and have no host-wide effect.
  if ! "$IPTABLES" -w 5 -nL "$chain" >/dev/null 2>&1; then
    "$IPTABLES" -w 5 -N "$chain"
    "$IPTABLES" -w 5 -A "$chain" -o lo -j ACCEPT
    "$IPTABLES" -w 5 -A "$chain" -m policy --dir out --pol ipsec -j ACCEPT
    "$IPTABLES" -w 5 -A "$chain" -j REJECT
  fi
  for file in *.vpn-peer; do
    [ -f "$file" ] && [ ! -L "$file" ] || continue
    read -r address < "$file";allow_peer "$address"
  done
  "$IPTABLES" -w 5 -C OUTPUT -m comment --comment mykos-full-vpn -j "$chain" 2>/dev/null ||
    "$IPTABLES" -w 5 -I OUTPUT 1 -m comment --comment mykos-full-vpn -j "$chain"
  # The managed network and tunnels are IPv4. Prevent an IPv6 escape path.
  if ! "$IP6TABLES" -w 5 -nL "$chain" >/dev/null 2>&1; then
    "$IP6TABLES" -w 5 -N "$chain"
    "$IP6TABLES" -w 5 -A "$chain" -o lo -j ACCEPT
    "$IP6TABLES" -w 5 -A "$chain" -j REJECT
  fi
  "$IP6TABLES" -w 5 -C OUTPUT -m comment --comment mykos-full-vpn -j "$chain" 2>/dev/null ||
    "$IP6TABLES" -w 5 -I OUTPUT 1 -m comment --comment mykos-full-vpn -j "$chain"
  dns_guard
}
record_order() {
  [ ! -L .mykos-connected-order ] && [ ! -L .mykos-selection.lock ] || exit 1
  (
    umask 077
    exec 8>.mykos-selection.lock
    flock 8
    grep -Fqx -- "$1" .mykos-connected-order 2>/dev/null || printf '%s\n' "$1" >> .mykos-connected-order
  )
}
forget_order() {
  [ ! -L .mykos-connected-order ] && [ ! -L .mykos-selection.lock ] || exit 1
  (
    umask 077
    exec 8>.mykos-selection.lock
    flock 8
    if [ -f .mykos-connected-order ]; then
      temporary=$(mktemp .mykos-order.XXXXXXXX)
      grep -Fxv -- "$1" .mykos-connected-order > "$temporary" || true
      mv -f -- "$temporary" .mykos-connected-order
    fi
  )
}
schedule_default() {
  # Run after this updown handler returns; never block Pluto on its own IPC.
  setsid /bin/bash -c '
        exec 9>.mykos-default.lock
        flock -n 9 || exit 0
        sleep 1
        wanted=$1
        if ipsec status | grep -F "\"$wanted\"" | grep -Eq "STATE_V2_ESTABLISHED_CHILD_SA|STATE_QUICK_I2|STATE_QUICK_R2|IPsec SA established"; then exit 0; fi
        ipsec auto --add "$wanted" >/dev/null 2>&1 || :
        ipsec auto --asynchronous --up "$wanted"
      ' _ "$1-mykos-default" </dev/null >/dev/null 2>&1 &
}
promote_remaining() {
  # This action is called by the administrator, outside Pluto's updown callback.
  # Never query Pluto while holding a lock that its updown callback may need.
  [ ! -L .mykos-connected-order ] || exit 1
  local removed=$1 candidate temporary status
  status=$(timeout 5 ipsec status) || return 0
  while read -r candidate; do
    [[ "$candidate" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,69}$ ]] || continue
    [ "$candidate" != "$removed" ] && [ -f "$candidate.conn.conf" ] && [ ! -L "$candidate.conn.conf" ] && [ -f "$candidate.updown" ] || continue
    printf '%s\n' "$status" | grep -F "\"$candidate\"" | grep -Eq 'STATE_V2_ESTABLISHED_CHILD_SA|STATE_QUICK_I2|STATE_QUICK_R2|IPsec SA established' || continue
    umask 077
    temporary=$(mktemp .mykos-connected.XXXXXXXX)
    printf '%s\n' "$candidate" > "$temporary"
    ln "$temporary" .mykos-first-tunnel 2>/dev/null || true
    rm -f -- "$temporary"
    [ ! -L .mykos-first-tunnel ] || exit 1
    read -r candidate < .mykos-first-tunnel
    guard
    schedule_default "$candidate"
    return 0
  done < <(if [ -f .mykos-connected-order ]; then cat .mykos-connected-order; fi)
}
if [ "${1:-}" = --remove ]; then
  name=$2
  [[ "$name" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,69}$ ]] || exit 1
  forget_order "$name"
  ipsec auto --delete "$name-mykos-default" >/dev/null 2>&1 || true
  if [ -f .mykos-first-tunnel ] && [ ! -L .mykos-first-tunnel ]; then
    read -r primary < .mykos-first-tunnel
    if [ "$primary" = "$name" ]; then rm -f .mykos-first-tunnel;promote_remaining "$name";fi
  fi
  # Keep the guard while any other connection remains; reconnect selects a new primary.
  remaining=0
  for file in *.conn.conf; do
    [ -f "$file" ] && [ "$file" != "$name.conn.conf" ] && remaining=1
  done
  if [ "$remaining" = 0 ] && [ "${MYKOS_HOST_GATEWAY:-0}" != 1 ]; then
    "$IPTABLES" -w 5 -D OUTPUT -m comment --comment mykos-full-vpn -j "$chain" 2>/dev/null || true
    "$IP6TABLES" -w 5 -D OUTPUT -m comment --comment mykos-full-vpn -j "$chain" 2>/dev/null || true
    dns_guard -D
  fi
  exit
fi
if [ "${1:-}" = --allow-peer ]; then allow_peer "$2";exit;fi
if [ "${1:-}" = --restore ]; then [ ! -f .mykos-first-tunnel ] || guard;exit;fi
name=${PLUTO_CONNECTION:-}
[[ "$name" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,69}$ ]] || exit 1
case "${PLUTO_VERB:-}" in
  up-client|up-host|up-client-v6|up-host-v6)
    record_order "$name"
    # A hard link makes the first successful CHILD_SA the atomic winner.
    umask 077
    temporary=$(mktemp .mykos-connected.XXXXXXXX)
    trap 'rm -f -- "$temporary"' EXIT
    printf '%s\n' "$name" > "$temporary"
    ln "$temporary" .mykos-first-tunnel 2>/dev/null || true
    [ ! -L .mykos-first-tunnel ] || exit 1
    read -r primary < .mykos-first-tunnel
    if [ "$primary" = "$name" ]; then
      guard
      schedule_default "$name"
    fi
    ;;
esac
# Keep the distribution routing/firewall helper behavior for the specific-site SA.
for original in /usr/libexec/ipsec/_updown.netkey /usr/lib/ipsec/_updown.netkey; do
  if [ -x "$original" ]; then exec "$original"; fi
done
