#!/bin/sh

# SPDX-FileCopyrightText: 2023 - 2024 Benjamin Grande M. S. <ben.grande.b@gmail.com>
#
# SPDX-License-Identifier: AGPL-3.0-or-later

set -eu

wg_conf="/etc/wireguard/wireguard.conf"
nft_conf="/var/run/wireguard/dnat.nft"

mkdir -p -- "${nft_conf%/*}"
rm -f -- "${nft_conf}"
touch -- "${nft_conf}"

set_nft(){
  printf '%s\n' "${*}" | tee -a -- "${nft_conf}" >/dev/null
}

set_nft_dnat(){
  ipv="${1}" # empty(4), 6
  proto="${2}" # tcp, udp
  dns_host="${3}"
  rule_prefix="insert rule ip${ipv} qubes custom-dnat iifgroup 2 ${proto}"
  rule_suffix="dport 53 dnat to ${dns_host}"
  set_nft "${rule_prefix} ${rule_suffix}"
}

dns="$(grep -s -e "^\s*DNS\s*=\s*\S\+" -- "${wg_conf}" |
  sed -e "s/.*=//;s/ //g")"

if test -z "${dns}"; then
  set_nft "insert rule ip  qubes custom-dnat drop"
  set_nft "insert rule ip6 qubes custom-dnat drop"
  exit
fi

dns_primary="$(printf '%s\n' "${dns}" | cut -d "," -f 1)"
dns_secondary="$(printf '%s\n' "${dns}" | cut -d "," -f 2)"

dns_primary_ipv=""
if printf '%s\n' "${dns_primary}" | grep -qF -e ":"; then
  dns_primary_ipv=6
fi

dns_secondary_ipv=""
if printf '%s\n' "${dns_secondary}" | grep -qF -e ":"; then
  dns_secondary_ipv=6
fi

if test -n "${dns}"; then
  set_nft_dnat "${dns_primary_ipv}" udp "${dns_primary}"
  set_nft_dnat "${dns_primary_ipv}" tcp "${dns_primary}"
  if printf '%s\n' "${dns}" | grep -qF -e ","; then
    set_nft_dnat "${dns_secondary_ipv}" udp "${dns_secondary}"
    set_nft_dnat "${dns_secondary_ipv}" tcp "${dns_secondary}"
  fi
fi

ln -sf -- /run/resolvconf/resolv.conf /etc/resolv.conf