#!/bin/bash
# ==============================================================================
# 06-configure-dns.sh
#
# Configure dnsmasq in sys-fips so .fips queries are forwarded to local
# FIPS DNS resolver on 127.0.0.1:5354, and all other DNS goes upstream.
#
# Also applies Qubes-specific nftables DNS interception so AppVM DNS queries
# are forced to local dnsmasq in sys-fips.
#
# Usage:
#   sudo bash 06-configure-dns.sh
# ==============================================================================
set -euo pipefail

echo "=== sys-fips DNS Configuration ==="
echo ""

if [ "$(id -u)" -ne 0 ]; then
    echo "✗ This script must be run as root (sudo)"
    exit 1
fi

ensure_nft_rule() {
    local table="$1"
    local chain="$2"
    local rule="$3"

    if ! nft list chain "$table" "$chain" >/dev/null 2>&1; then
        echo "  WARN: nft chain $table $chain not found; skipping: $rule"
        return 0
    fi

    if nft list chain "$table" "$chain" | grep -F -- "$rule" >/dev/null 2>&1; then
        return 0
    fi

    nft add rule "$table" "$chain" $rule
}

systemctl stop dnsmasq 2>/dev/null || true
killall dnsmasq 2>/dev/null || true

UPSTREAM_DNS=$(grep '^nameserver' /etc/resolv.conf | awk '{print $2}' | head -1)
UPSTREAM_DNS="${UPSTREAM_DNS:-10.139.1.1}"

echo "  Upstream DNS: $UPSTREAM_DNS"

mkdir -p /etc/dnsmasq.d
cat > /etc/dnsmasq.d/fips.conf << EOF
# sys-fips DNS resolution
port=53
listen-address=0.0.0.0
bind-interfaces

# .fips domains -> FIPS resolver
server=/fips/127.0.0.1#5354

# everything else -> upstream
server=$UPSTREAM_DNS

no-resolv
no-hosts
# log-queries
EOF

if [ ! -f /etc/dnsmasq.conf ]; then
    touch /etc/dnsmasq.conf
fi
if ! grep -q 'conf-dir=/etc/dnsmasq.d/,*.conf' /etc/dnsmasq.conf; then
    echo 'conf-dir=/etc/dnsmasq.d/,*.conf' >> /etc/dnsmasq.conf
fi

if ! dnsmasq --test >/dev/null 2>&1; then
    echo "✗ dnsmasq config failed validation"
    exit 1
fi

# Qubes AppVMs often use /32 point-to-point addressing.
# Debian's dnsmasq systemd-helper adds --local-service by default,
# which can reject DNS queries from downstream AppVMs in this topology.
DNSMASQ_HELPER="/usr/share/dnsmasq/systemd-helper"
if [ -f "$DNSMASQ_HELPER" ] && grep -q -- '--local-service' "$DNSMASQ_HELPER"; then
    if [ ! -f "${DNSMASQ_HELPER}.bak" ]; then
        cp "$DNSMASQ_HELPER" "${DNSMASQ_HELPER}.bak"
    fi
    sed -i 's/--local-service//g' "$DNSMASQ_HELPER"
    echo "  Patched dnsmasq helper: removed --local-service for Qubes /32 compatibility"
fi

if command -v systemctl >/dev/null 2>&1; then
    systemctl restart dnsmasq || true
fi

if ! pgrep -x dnsmasq >/dev/null 2>&1; then
    dnsmasq
fi

echo "Applying Qubes nftables DNS interception rules..."
ensure_nft_rule ip qubes dnat-dns 'iifname "vif*" udp dport 53 dnat to 10.139.1.1:53'
ensure_nft_rule ip qubes dnat-dns 'iifname "vif*" tcp dport 53 dnat to 10.139.1.1:53'
ensure_nft_rule ip qubes custom-input 'iifname "vif*" ip daddr 10.139.1.1 udp dport 53 accept'
ensure_nft_rule ip qubes custom-input 'iifname "vif*" ip daddr 10.139.1.1 tcp dport 53 accept'
echo "  nftables DNS interception ensured"

RCLOCAL="/rw/config/rc.local"
MARK_START="# === SYS-FIPS DNS START ==="
MARK_END="# === SYS-FIPS DNS END ==="

if [ ! -f "$RCLOCAL" ]; then
    mkdir -p /rw/config
    cat > "$RCLOCAL" << 'EOF'
#!/bin/sh
EOF
fi

sed -i "/$MARK_START/,/$MARK_END/d" "$RCLOCAL"

cat >> "$RCLOCAL" << 'EOF'
# === SYS-FIPS DNS START ===
# Keep dnsmasq alive with .fips forwarding
if [ -f /usr/share/dnsmasq/systemd-helper ]; then
  sed -i 's/--local-service//g' /usr/share/dnsmasq/systemd-helper
fi
systemctl restart dnsmasq >/dev/null 2>&1 || true
pgrep -x dnsmasq >/dev/null 2>&1 || dnsmasq >/dev/null 2>&1 || true

# Qubes DNS interception for downstream AppVMs
nft list chain ip qubes dnat-dns >/dev/null 2>&1 && \
  nft list chain ip qubes dnat-dns | grep -F 'iifname "vif*" udp dport 53 dnat to 10.139.1.1:53' >/dev/null 2>&1 || \
  nft add rule ip qubes dnat-dns iifname "vif*" udp dport 53 dnat to 10.139.1.1:53 >/dev/null 2>&1 || true

nft list chain ip qubes dnat-dns >/dev/null 2>&1 && \
  nft list chain ip qubes dnat-dns | grep -F 'iifname "vif*" tcp dport 53 dnat to 10.139.1.1:53' >/dev/null 2>&1 || \
  nft add rule ip qubes dnat-dns iifname "vif*" tcp dport 53 dnat to 10.139.1.1:53 >/dev/null 2>&1 || true

nft list chain ip qubes custom-input >/dev/null 2>&1 && \
  nft list chain ip qubes custom-input | grep -F 'iifname "vif*" ip daddr 10.139.1.1 udp dport 53 accept' >/dev/null 2>&1 || \
  nft add rule ip qubes custom-input iifname "vif*" ip daddr 10.139.1.1 udp dport 53 accept >/dev/null 2>&1 || true

nft list chain ip qubes custom-input >/dev/null 2>&1 && \
  nft list chain ip qubes custom-input | grep -F 'iifname "vif*" ip daddr 10.139.1.1 tcp dport 53 accept' >/dev/null 2>&1 || \
  nft add rule ip qubes custom-input iifname "vif*" ip daddr 10.139.1.1 tcp dport 53 accept >/dev/null 2>&1 || true
# === SYS-FIPS DNS END ===
EOF

chmod +x "$RCLOCAL"

echo ""
echo "Verifying DNS..."
if dig @127.0.0.1 google.com A +short +time=2 >/dev/null 2>&1; then
    echo "  Upstream DNS: OK"
else
    echo "  Upstream DNS: WARN (may be expected in isolated setup)"
fi

STATUS=$(dig @127.0.0.1 test.fips AAAA +time=2 2>&1 | grep 'status:' || true)
if echo "$STATUS" | grep -qE 'NXDOMAIN|SERVFAIL|NOERROR'; then
    echo "  .fips forwarding: OK"
else
    echo "  .fips forwarding: WARN"
fi

echo ""
echo "✓ DNS configured"