#!/bin/bash

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


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

dnsleak_protection_up() 
{

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

	if [ -z "$interface_name" ]
	then
   		echo "\$interface_name is empty"
   		return 0
	fi

	# disable DNS traffic on interface != utun420
	local rules=$'*filter\n'
	printf -v rules '%s:windscribe_dnsleaks - [0:0]\n' "$rules"
	printf -v rules '%s-A OUTPUT -j windscribe_dnsleaks %s\n' "$rules" "$marker"

	printf -v rules '%s-A windscribe_dnsleaks -o lo -j ACCEPT %s\n' "$rules" "$marker"
	printf -v rules '%s-A windscribe_dnsleaks ! -o %s -p udp --dport 53 -j DROP %s\n' "$rules" "$interface_name" "$marker"
	printf -v rules '%s-A windscribe_dnsleaks ! -o %s -p tcp --dport 53 -j DROP %s\n' "$rules" "$interface_name" "$marker"

	printf -v rules '%s-A windscribe_dnsleaks -j ACCEPT %s\n' "$rules" "$marker"
	printf -v rules '%sCOMMIT\n' "$rules"

	echo -n "$rules" | cmd "/sbin/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 < <(/sbin/iptables-save 2>/dev/null)
	
	if [[ $found -ne 0 ]]; then
		echo -n "$restore" | cmd "/sbin/iptables-restore" -n
	fi
}


main() {

	if [ "$#" -gt 2 ]; then
    	echo "Usage: dns-leak-protect up|down [interface_name]"
    	return 1
	fi

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


main "$@"
