#!/bin/bash

# Copyright (c) 2020-2024, Windscribe Limited. All rights reserved.
# Usage: dns-leak-protect up|down [interface_name]

PATH="$PATH:/usr/local/sbin:/usr/sbin:/sbin"

cmd() {
    echo "[#] $*" >&2
    "$@"
}

get_os_default_dnsservers_resolvectl()
{
    mapfile -t dns_array < <( resolvectl dns | awk -F: '$1 !~ /\(utun420|\(tun/ {st = index($0,":"); print substr($0,st+1)}' |
                  awk '$1 {print $1}' | sed 's/^[ \t]*//;s/[ \t]*$//' | grep '^[0-9]\+\.[0-9]\+\.[0-9]\+\.[0-9]\+$' )
}


get_os_default_dnsservers_NMCLI()
{
    mapfile -t dns_array < <( (nmcli dev list || nmcli dev show ) 2>/dev/null | grep IP4.DNS |
                  awk -F: '{ st = index($0,":");print substr($0,st+1)}' | awk '$1 {print $1}' | sed 's/^[ \t]*//;s/[ \t]*$//' )
}

dnsleak_protection_up()
{
    # get dns default DNS-servers and put it into dns_array variable based on resolvectl or nmcli methods.
    if ! command -v resolvectl &> /dev/null; then
        get_os_default_dnsservers_NMCLI
    else
        get_os_default_dnsservers_resolvectl
        # try the nmcli method if resolvectl failed
                if [ -z "$dns_array" ]; then
            get_os_default_dnsservers_NMCLI
        fi
    fi

    local marker="-m comment --comment \"Windscribe client dns leak protection\""

    local rules=$'*filter\n'
    printf -v rules '%s:windscribe_dnsleaks - [0:0]\n' "$rules"
    printf -v rules '%s-I OUTPUT -j windscribe_dnsleaks %s\n' "$rules" "$marker"

    # disallow all OS-default DNS servers
    for dnsIp in "${dns_array[@]}"; do
        for allowed in "$@"; do
            if [[ $dnsIp == "$allowed" ]]; then
                continue 2
            fi
        done
        printf -v rules '%s-A windscribe_dnsleaks -d %s -p udp --dport 53 -j DROP %s\n' "$rules" "$dnsIp" "$marker"
        printf -v rules '%s-A windscribe_dnsleaks -d %s -p tcp --dport 53 -j DROP %s\n' "$rules" "$dnsIp" "$marker"
    done

    printf -v rules '%sCOMMIT\n' "$rules"

    echo -n "$rules" | cmd "iptables-restore" -n
}


dnsleak_protection_down()
{

    local line iptables found restore in_filter_section
    restore="" found=0 is_filter_section=0
    while read -r line; do
        if [[ $line == "*filter"* ]]; then
            is_filter_section=1
        else
            [[ $line == "*"* ]] && is_filter_section=0
        fi

        [[ $line == "*"* || $line == COMMIT || $line == "-A "*"-m comment --comment \"Windscribe client dns leak protection\""* ]] || continue
        [[ $line == "-A"* ]] && found=1

        if [[ $is_filter_section -ne 0 && $line == COMMIT ]]; then
            printf -v restore '%s-X windscribe_dnsleaks\n' "$restore"
        fi
        printf -v restore '%s%s\n' "$restore" "${line/#-A/-D}"

    done < <(iptables-save 2>/dev/null)

    if [[ $found -ne 0 ]]; then
        echo -n "$restore" | cmd "iptables-restore" -n
    fi
}

main()
{
    local action="$1"

    if [[ $action == "up" ]]; then
        dnsleak_protection_up $@
    elif [[ $action == "down" ]]; then
        dnsleak_protection_down
    else
        echo "Usage: dns-leak-protect up|down [servers_to_allow]"
        return 1
    fi
}


main "$@"
