#!/bin/sh
# PID 1 of a session VM: firewalls egress to the public internet (or only the
# proxy, or only the WireGuard server and its tunnel), then runs knotchy-agent as
# the knotchy user with no capabilities, so nothing in the browser can undo it.
# Fails closed: a missing or malformed setting stops the VM. Needs NET_ADMIN.
# Env: KNOTCHY_ROUTE (direct|proxy|wireguard), KNOTCHY_BLOCK_CIDRS, KNOTCHY_PROXY_ADDR,
# KNOTCHY_PROXY_URL*, KNOTCHY_WG_SETCONF*, KNOTCHY_WG_ENDPOINT, KNOTCHY_WG_ADDRESS,
# KNOTCHY_WG_DNS, KNOTCHY_WG_MTU. *Secrets, normally in /run/knotchy-secrets/route.
set -eu

fail() {
  echo "knotchy-init: $*" >&2
  exit 1
}

blocked="0.0.0.0/8, 10.0.0.0/8, 100.64.0.0/10, 127.0.0.0/8, 169.254.0.0/16,
  172.16.0.0/12, 192.0.0.0/24, 192.0.2.0/24, 192.88.99.0/24, 192.168.0.0/16,
  198.18.0.0/15, 198.51.100.0/24, 203.0.113.0/24, 224.0.0.0/4, 240.0.0.0/4"

is_ipv4() {
  case "$1" in *[!0-9.]* | "" | *..* | .* | *.) return 1 ;; esac
  old_ifs=$IFS; IFS=.
  # shellcheck disable=SC2086  # splitting on dots is the point
  set -- $1
  IFS=$old_ifs
  [ $# -eq 4 ] || return 1
  for octet in "$@"; do
    [ "${#octet}" -le 3 ] && [ "$octet" -le 255 ] || return 1
  done
}

is_port() {
  case "$1" in *[!0-9]* | "") return 1 ;; esac
  [ "${#1}" -le 5 ] && [ "$1" -ge 1 ] && [ "$1" -le 65535 ]
}

is_cidr() {
  case "$1" in */*) ;; *) return 1 ;; esac
  prefix=${1##*/}
  case "$prefix" in *[!0-9]* | "") return 1 ;; esac
  is_ipv4 "${1%/*}" && [ "${#prefix}" -le 2 ] && [ "$prefix" -le 32 ]
}

set -f
for cidr in ${KNOTCHY_BLOCK_CIDRS:-}; do
  case "$cidr" in
    *[!0-9./]* | "") fail "malformed CIDR '$cidr' in KNOTCHY_BLOCK_CIDRS" ;;
    *) blocked="$blocked, $cidr" ;;
  esac
done

route=${KNOTCHY_ROUTE:-direct}
case $route in
  direct | proxy | wireguard) ;;
  *) fail "unknown KNOTCHY_ROUTE '$route'" ;;
esac

# Secrets come on a share because `container` keeps -e values on disk. Delete
# the file (on the Mac too) and unmount before anything else runs.
proxy_url=${KNOTCHY_PROXY_URL:-}
wg_setconf=${KNOTCHY_WG_SETCONF:-}
secrets=/run/knotchy-secrets
if [ -e "$secrets/route" ]; then
  while IFS= read -r line; do
    case "$line" in
      KNOTCHY_PROXY_URL=*) proxy_url=${line#*=} ;;
      KNOTCHY_WG_SETCONF=*) wg_setconf=${line#*=} ;;
    esac
  done < "$secrets/route"
  rm -f "$secrets/route" || fail "could not delete the secrets file"
fi
if grep -q " $secrets " /proc/mounts; then
  umount "$secrets" || fail "could not unmount the secrets share"
fi

case $route in
  direct)
    [ -z "${KNOTCHY_PROXY_ADDR:-}${KNOTCHY_WG_ENDPOINT:-}$proxy_url$wg_setconf" ] ||
      fail "proxy or VPN settings given for a direct route"
    egress="ip daddr @blocked4 reject with icmpx type admin-prohibited" ;;
  proxy)
    proxy_addr=${KNOTCHY_PROXY_ADDR:-}
    proxy_ip=${proxy_addr%:*}
    proxy_port=${proxy_addr##*:}
    is_ipv4 "$proxy_ip" && is_port "$proxy_port" || fail "malformed KNOTCHY_PROXY_ADDR"
    case "$proxy_url" in
      http://*@"$proxy_addr" | http://"$proxy_addr" | socks5://*@"$proxy_addr" | socks5://"$proxy_addr") ;;
      *) fail "malformed or missing KNOTCHY_PROXY_URL" ;;
    esac
    egress="ip daddr $proxy_ip tcp dport $proxy_port accept
    reject with icmpx type admin-prohibited" ;;
  wireguard)
    [ -n "$wg_setconf" ] || fail "KNOTCHY_WG_SETCONF is missing"
    wg_endpoint=${KNOTCHY_WG_ENDPOINT:-}
    wg_ip=${wg_endpoint%:*}
    wg_port=${wg_endpoint##*:}
    is_ipv4 "$wg_ip" && is_port "$wg_port" || fail "malformed KNOTCHY_WG_ENDPOINT"
    [ -n "${KNOTCHY_WG_ADDRESS:-}" ] || fail "KNOTCHY_WG_ADDRESS is empty"
    for cidr in $KNOTCHY_WG_ADDRESS; do is_cidr "$cidr" || fail "malformed KNOTCHY_WG_ADDRESS"; done
    wg_dns=${KNOTCHY_WG_DNS:-1.1.1.1}
    [ -n "$wg_dns" ] || wg_dns=1.1.1.1
    for ns in $wg_dns; do is_ipv4 "$ns" || fail "malformed KNOTCHY_WG_DNS"; done
    wg_mtu=${KNOTCHY_WG_MTU:-1420}
    case "$wg_mtu" in *[!0-9]* | "") fail "malformed KNOTCHY_WG_MTU" ;; esac
    [ "${#wg_mtu}" -le 4 ] && [ "$wg_mtu" -ge 576 ] || fail "malformed KNOTCHY_WG_MTU"
    wg_dns_set=$(echo $wg_dns | sed 's/ /, /g')
    # DNS only to the VPN's (often private) resolvers; outside the tunnel, only wireguard-go.
    egress="oifname \"wg0\" ip daddr { $wg_dns_set } meta l4proto { tcp, udp } th dport 53 accept
    oifname \"wg0\" ip daddr @blocked4 reject with icmpx type admin-prohibited
    oifname \"wg0\" accept
    ip daddr $wg_ip udp dport $wg_port meta skuid 0 accept
    reject with icmpx type admin-prohibited" ;;
esac

# No DNS exception in the firewall: a --dns resolver in a blocked range can't answer.
dns=$(awk '$1 == "nameserver" && $2 ~ /^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$/ {
  printf "%s%s", sep, $2; sep = ", " }' /etc/resolv.conf)
bad_ns=$(awk '$1 == "nameserver" && $2 ~ /^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$/ {
  split($2, o, "."); a = o[1] + 0; b = o[2] + 0
  if (a == 10 || a == 127 || (a == 169 && b == 254) || (a == 172 && b >= 16 && b <= 31) ||
      (a == 192 && b == 168) || (a == 100 && b >= 64 && b <= 127)) { print $2; exit }
}' /etc/resolv.conf)
[ -z "$bad_ns" ] || fail "resolver $bad_ns is in a blocked range; it cannot be reached through the egress firewall"

nft -f - <<EOF || fail "could not apply the egress firewall"
flush ruleset
table inet knotchy {
  set blocked4 {
    type ipv4_addr
    flags interval
    auto-merge
    elements = { $blocked }
  }

  chain output {
    type filter hook output priority filter; policy accept;
    oifname "lo" accept
    meta nfproto ipv6 reject with icmpx type admin-prohibited
    $egress
  }

  chain input {
    type filter hook input priority filter; policy drop;
    iifname "lo" accept
    ct state established,related accept
  }

  chain forward {
    type filter hook forward priority filter; policy drop;
  }
}
EOF

echo 1 > /proc/sys/net/ipv6/conf/all/disable_ipv6 2>/dev/null || true

nft list table inet knotchy >/dev/null 2>&1 || fail "egress firewall is not in place"

case $route in
  direct)
    echo "knotchy-init: egress locked to the public internet${dns:+ (DNS via $dns)}" ;;
  proxy)
    echo "knotchy-init: egress locked to the proxy at $KNOTCHY_PROXY_ADDR"
    export http_proxy="$proxy_url" https_proxy="$proxy_url" \
      HTTP_PROXY="$proxy_url" HTTPS_PROXY="$proxy_url" no_proxy="" ;;
  wireguard)
    # The firewall is already in place, so a half-built tunnel leaks nothing.
    old_umask=$(umask)
    umask 077
    printf '%s' "$wg_setconf" | base64 -d > /run/knotchy-wg.conf 2>/dev/null ||
      fail "malformed KNOTCHY_WG_SETCONF"
    # Without SYS_ADMIN, which init only needed to unmount the secrets share.
    WG_I_PREFER_BUGGY_USERSPACE_TO_POLISHED_KMOD=1 setpriv --bounding-set -sys_admin \
      wireguard-go wg0 >/dev/null 2>&1 ||
      fail "could not start wireguard-go"
    wg setconf wg0 /run/knotchy-wg.conf || fail "the WireGuard settings were rejected"
    rm -f /run/knotchy-wg.conf
    umask "$old_umask"
    for cidr in $KNOTCHY_WG_ADDRESS; do
      ip -4 address add "$cidr" dev wg0 || fail "could not give the tunnel its address"
    done
    ip link set dev wg0 mtu "$wg_mtu" up || fail "could not bring the tunnel up"
    via=$(ip -4 route get "$wg_ip" | awk '{ for (i = 1; i < NF; i++) if ($i == "via") { print $(i + 1); exit } }')
    dev=$(ip -4 route get "$wg_ip" | awk '{ for (i = 1; i < NF; i++) if ($i == "dev") { print $(i + 1); exit } }')
    [ -n "$dev" ] || fail "no route to the VPN server"
    ip -4 route add "$wg_ip/32" ${via:+via "$via"} dev "$dev" || fail "could not route to the VPN server"
    ip -4 route replace default dev wg0 || fail "could not route through the tunnel"
    printf 'nameserver %s\n' $wg_dns > /etc/resolv.conf || fail "could not set the VPN's DNS"
    echo "knotchy-init: egress locked to the WireGuard tunnel via $KNOTCHY_WG_ENDPOINT" ;;
esac

# The browser has no business knowing the Mac's networks or the VPN's keys.
unset KNOTCHY_ROUTE KNOTCHY_BLOCK_CIDRS KNOTCHY_PROXY_ADDR KNOTCHY_PROXY_URL KNOTCHY_WG_SETCONF \
  KNOTCHY_WG_ENDPOINT KNOTCHY_WG_ADDRESS KNOTCHY_WG_DNS KNOTCHY_WG_MTU
export HOME=/home/knotchy USER=knotchy LOGNAME=knotchy
exec setpriv --reuid=knotchy --regid=knotchy --init-groups \
  --inh-caps=-all --bounding-set=-all --no-new-privs \
  /usr/local/bin/knotchy-agent