#!/bin/sh
# orbital-wg - privileged WireGuard bring-up/tear-down + kill switch for Orbital.
#
# Installed root-owned in /usr/local/bin and invoked by the app through pkexec,
# so the OS shows a graphical auth prompt and no standing privilege is granted.
# It raises/drops one tunnel from a config path and, when asked, installs a
# firewall lockdown (the kill switch) so nothing leaves the machine outside the
# tunnel. The lockdown survives an unexpected tunnel drop, so traffic is blocked
# (not leaked) until Orbital reconnects or the user disconnects.
#
#   orbital-wg up   <config-path> [killswitch]
#   orbital-wg down <config-path>

set -eu

KS_TABLE=orbital_ks

# --- guards -----------------------------------------------------------------
#
# The shipped polkit policy lets an active local session run this helper without
# re-authenticating, which is what makes connecting prompt-free. That only stays
# safe if the helper cannot be talked into doing anything except raising the
# caller's OWN tunnel, so validate hard before touching wg-quick:
#
#   1. wg-quick executes PreUp/PostUp/PreDown/PostDown as root. A config
#      carrying those is a root shell, so refuse outright.
#   2. Only ever act on the calling user's own Orbital config, resolved from
#      PKEXEC_UID, and only when that file is really theirs and not a symlink
#      pointed somewhere else.
assert_safe_config() {
  CONF="$1"

  [ -f "$CONF" ] || { echo "orbital-wg: config not found: $CONF" >&2; exit 2; }
  [ -L "$CONF" ] && { echo "orbital-wg: config must not be a symlink" >&2; exit 3; }

  if grep -qiE '^[[:space:]]*(PreUp|PostUp|PreDown|PostDown)[[:space:]]*=' "$CONF"; then
    echo "orbital-wg: config contains script hooks, refusing" >&2
    exit 3
  fi

  # Under pkexec we know who asked; pin the path and the owner to them.
  if [ -n "${PKEXEC_UID:-}" ]; then
    HOME_DIR="$(getent passwd "$PKEXEC_UID" | cut -d: -f6)"
    [ -n "$HOME_DIR" ] || { echo "orbital-wg: unknown caller" >&2; exit 3; }
    if [ "$CONF" != "$HOME_DIR/.config/Orbital/orbital.conf" ]; then
      echo "orbital-wg: refusing a config outside the caller's Orbital profile" >&2
      exit 3
    fi
    OWNER="$(stat -c %u "$CONF" 2>/dev/null || echo -1)"
    if [ "$OWNER" != "$PKEXEC_UID" ]; then
      echo "orbital-wg: config is not owned by the caller" >&2
      exit 3
    fi
  fi
}

remove_killswitch() {
  command -v nft >/dev/null 2>&1 && nft delete table inet "$KS_TABLE" 2>/dev/null || true
}

# Block all egress except: loopback, the tunnel interface, the encrypted path to
# the relay endpoint, and the local network. Everything else is rejected, so a
# dropped tunnel cannot leak plaintext.
install_killswitch() {
  CONF="$1"
  IFACE="$(basename "$CONF" .conf)"
  # Endpoint host is everything before the final ':port'.
  EP="$(grep -iE '^[[:space:]]*Endpoint[[:space:]]*=' "$CONF" | head -1 | sed 's/.*=[[:space:]]*//' | rev | cut -d: -f2- | rev | tr -d '[]')"
  # Resolve to IPv4 (relays are v4). Keep every A record it resolves to, so a
  # reconnect that lands on any of them stays permitted; fall back to the literal
  # when it is already an address or resolution fails.
  EPIP="$(getent ahostsv4 "$EP" 2>/dev/null | awk '{print $1}' | sort -u | paste -sd, -)"
  [ -n "${EPIP:-}" ] || EPIP="$EP"

  if ! command -v nft >/dev/null 2>&1; then
    echo "orbital-wg: nft not found, kill switch unavailable" >&2
    return 1
  fi
  remove_killswitch
  # Allow, on top of the tunnel and the relay endpoint: loopback, the local
  # network (incl. the DHCP broadcast address), IPv6 link-local/ULA, and DHCP
  # client requests - so the box can still renew its lease and reach the LAN
  # while everything else is dropped. Nothing else escapes if the tunnel drops.
  nft -f - <<EOF
table inet $KS_TABLE {
  chain out {
    type filter hook output priority -100; policy drop;
    oifname "lo" accept
    oifname "$IFACE" accept
    ip daddr { $EPIP } accept
    ip daddr { 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 169.254.0.0/16, 255.255.255.255 } accept
    ip6 daddr { fe80::/10, fc00::/7 } accept
    udp dport 67 accept
    reject with icmpx type admin-prohibited
  }
}
EOF
  echo "orbital-wg: kill switch armed (iface=$IFACE endpoint=$EPIP)"
}

case "${1:-}" in
  up)
    CONF="${2:-}"
    KS="${3:-}"
    [ -n "$CONF" ] || { echo "orbital-wg: config path required" >&2; exit 2; }
    assert_safe_config "$CONF"
    # Clean slate.
    wg-quick down "$CONF" 2>/dev/null || true
    remove_killswitch
    wg-quick up "$CONF"
    case "$KS" in
      killswitch|ks|1) install_killswitch "$CONF" || echo "orbital-wg: proceeding without kill switch" >&2 ;;
      *) : ;;
    esac
    exit 0
    ;;
  down)
    CONF="${2:-}"
    [ -n "$CONF" ] || { echo "orbital-wg: config path required" >&2; exit 2; }
    # Tearing down is only ever allowed for the caller's own config too, but a
    # missing file must still clear the lockdown rather than strand the machine.
    [ -f "$CONF" ] && assert_safe_config "$CONF"
    remove_killswitch
    wg-quick down "$CONF" 2>/dev/null || true
    exit 0
    ;;
  *)
    echo "usage: orbital-wg up|down <config-path> [killswitch]" >&2
    exit 1
    ;;
esac
