#!/bin/bash
# /etc/NetworkManager/dispatcher.d/90-nft-wg
# Automatically toggle DNS leak protection for wireguard

INTERFACE="$1"
ACTION="$2"

NFT_BIN="$(command -v nft || true)"
FAMILY="inet"
TABLE="filter"
CHAIN="output"

log() { logger -t nft-wg "$*"; }


add_dns_protection() {
    $NFT_BIN add rule inet filter output udp dport 53 oifname "!= \"$INTERFACE\"" drop\; 2>/tmp/nft-wg-add.err
    if [ $? -ne 0 ]; then
        log "failed to add udp rule for $INTERFACE: $(cat /tmp/nft-wg-add.err)"
        exit 1
    fi

    $NFT_BIN add rule inet filter output tcp dport 53 oifname "!= \"$INTERFACE\"" drop\; 2>/tmp/nft-wg-add.err
    if [ $? -ne 0 ]; then
        log "failed to add tcp rule for $INTERFACE: $(cat /tmp/nft-wg-add.err)"
        exit 1
    fi

    rm -f /tmp/nft-wg-add.err
    log "Added DNS leak protection for $INTERFACE"
}

remove_dns_protection() {
    $NFT_BIN --handle --numeric list chain inet filter output \
        | awk -v iface="$INTERFACE" '$0 ~ iface && /drop/ {print $NF}' \
        | xargs -r -I {} sudo $NFT_BIN delete rule inet filter output handle {}
    log "Removed DNS leak protection for $INTERFACE"
}

# Only run for WireGuard interfaces (start with "wg-")
if [[ "$INTERFACE" == wg-* ]]; then
    case "$ACTION" in
        up)
            add_dns_protection
            ;;
        down)
            remove_dns_protection
            ;;
    esac
fi

exit 0