#!/bin/sh
# SecuBox DNS Guard - Privacy DNS Manager
# Provides uncensored DNS feeds and smart configuration

. /usr/share/libubox/jshn.sh

# DNS Provider Feed - Uncensored & Privacy-focused
# Format: id|name|primary|secondary|dot_host|doh_url|category|country|description
DNS_PROVIDERS='
fdn|FDN (French Data Network)|80.67.169.12|80.67.169.40|ns0.fdn.fr|https://ns0.fdn.fr/dns-query|privacy|FR|French non-profit, no logs, uncensored
fdn_ipv6|FDN IPv6|2001:910:800::12|2001:910:800::40|ns0.fdn.fr|https://ns0.fdn.fr/dns-query|privacy|FR|FDN IPv6 resolvers
quad9|Quad9|9.9.9.9|149.112.112.112|dns.quad9.net|https://dns.quad9.net/dns-query|security|CH|Malware blocking, Swiss privacy
quad9_unsecured|Quad9 Unsecured|9.9.9.10|149.112.112.10|dns10.quad9.net|https://dns10.quad9.net/dns-query|privacy|CH|No blocking, Swiss privacy
cloudflare|Cloudflare|1.1.1.1|1.0.0.1|cloudflare-dns.com|https://cloudflare-dns.com/dns-query|fast|US|Fast, privacy-focused
cloudflare_family|Cloudflare Family|1.1.1.3|1.0.0.3|family.cloudflare-dns.com|https://family.cloudflare-dns.com/dns-query|family|US|Malware + adult content blocking
mullvad|Mullvad|194.242.2.2|193.19.108.2|dns.mullvad.net|https://dns.mullvad.net/dns-query|privacy|SE|No logs, Swedish privacy
mullvad_adblock|Mullvad Adblock|194.242.2.3|193.19.108.3|adblock.dns.mullvad.net|https://adblock.dns.mullvad.net/dns-query|adblock|SE|Ads + trackers blocking
adguard|AdGuard|94.140.14.14|94.140.15.15|dns.adguard-dns.com|https://dns.adguard-dns.com/dns-query|adblock|CY|Ad blocking DNS
adguard_family|AdGuard Family|94.140.14.15|94.140.15.16|family.adguard-dns.com|https://dns.adguard-dns.com/dns-query|family|CY|Family protection
opendns|OpenDNS|208.67.222.222|208.67.220.220|doh.opendns.com|https://doh.opendns.com/dns-query|security|US|Cisco security DNS
opendns_family|OpenDNS FamilyShield|208.67.222.123|208.67.220.123|doh.familyshield.opendns.com|https://doh.familyshield.opendns.com/dns-query|family|US|Family protection
google|Google|8.8.8.8|8.8.4.4|dns.google|https://dns.google/dns-query|fast|US|Fast, global anycast
cleanbrowsing|CleanBrowsing Security|185.228.168.9|185.228.169.9|security-filter-dns.cleanbrowsing.org|https://doh.cleanbrowsing.org/doh/security-filter|security|US|Security filtering
cleanbrowsing_family|CleanBrowsing Family|185.228.168.168|185.228.169.168|family-filter-dns.cleanbrowsing.org|https://doh.cleanbrowsing.org/doh/family-filter|family|US|Family + adult blocking
controld|Control D|76.76.2.0|76.76.10.0|freedns.controld.com|https://freedns.controld.com/p0|privacy|CA|Canadian, no logs
'

method_status() {
	local current_dns1 current_dns2 provider mode

	# Get current DNS from dnsmasq/network config
	current_dns1=$(uci -q get dhcp.@dnsmasq[0].server 2>/dev/null | awk '{print $1}')
	current_dns2=$(uci -q get dhcp.@dnsmasq[0].server 2>/dev/null | awk '{print $2}')

	# Check if using local resolver (AdGuard Home)
	if netstat -tlnp 2>/dev/null | grep -q ":53.*AdGuard"; then
		mode="adguardhome"
	elif netstat -tlnp 2>/dev/null | grep -q ":53.*dnsmasq"; then
		mode="dnsmasq"
	else
		mode="unknown"
	fi

	# Try to identify provider (grep avoids subshell issue)
	provider="custom"
	if [ -n "$current_dns1" ]; then
		# Search for provider by primary or secondary IP
		matched_line=$(echo "$DNS_PROVIDERS" | grep -E "\|${current_dns1}\||\|${current_dns1}$" | head -1)
		if [ -n "$matched_line" ]; then
			provider=$(echo "$matched_line" | cut -d'|' -f1)
		fi
	fi

	json_init
	json_add_string "mode" "$mode"
	json_add_string "provider" "$provider"
	json_add_string "primary" "${current_dns1:-auto}"
	json_add_string "secondary" "${current_dns2:-}"
	json_add_boolean "dot_enabled" "$(uci -q get dhcp.@dnsmasq[0].dnssec 2>/dev/null | grep -q 1 && echo 1 || echo 0)"
	json_dump
}

method_get_providers() {
	json_init
	json_add_array "providers"

	# Save providers to temp file to avoid subshell issues
	echo "$DNS_PROVIDERS" > /tmp/.dns_providers
	while IFS='|' read id name primary secondary dot_host doh_url category country description; do
		[ -z "$id" ] && continue
		json_add_object
		json_add_string "id" "$id"
		json_add_string "name" "$name"
		json_add_string "primary" "$primary"
		json_add_string "secondary" "$secondary"
		json_add_string "dot_host" "$dot_host"
		json_add_string "doh_url" "$doh_url"
		json_add_string "category" "$category"
		json_add_string "country" "$country"
		json_add_string "description" "$description"
		json_close_object
	done < /tmp/.dns_providers
	rm -f /tmp/.dns_providers

	json_close_array
	json_dump
}

method_get_config() {
	json_init

	# dnsmasq config
	json_add_object "dnsmasq"
	json_add_string "server1" "$(uci -q get dhcp.@dnsmasq[0].server 2>/dev/null | awk '{print $1}')"
	json_add_string "server2" "$(uci -q get dhcp.@dnsmasq[0].server 2>/dev/null | awk '{print $2}')"
	json_add_boolean "noresolv" "$(uci -q get dhcp.@dnsmasq[0].noresolv 2>/dev/null | grep -q 1 && echo 1 || echo 0)"
	json_close_object

	# AdGuard Home config (if present)
	if [ -f /var/lib/adguardhome/AdGuardHome.yaml ]; then
		json_add_object "adguardhome"
		json_add_boolean "installed" 1
		json_add_boolean "running" "$(pgrep -f AdGuardHome >/dev/null && echo 1 || echo 0)"
		json_close_object
	fi

	json_dump
}

method_set_provider() {
	local provider="$1"
	local primary secondary provider_line

	# Find provider in feed using grep (avoids subshell)
	provider_line=$(echo "$DNS_PROVIDERS" | grep "^${provider}|")

	if [ -n "$provider_line" ]; then
		primary=$(echo "$provider_line" | cut -d'|' -f3)
		secondary=$(echo "$provider_line" | cut -d'|' -f4)
	fi

	if [ -z "$primary" ]; then
		json_init
		json_add_boolean "success" 0
		json_add_string "error" "Unknown provider: $provider"
		json_dump
		return
	fi

	# Configure dnsmasq
	uci -q delete dhcp.@dnsmasq[0].server
	uci add_list dhcp.@dnsmasq[0].server="$primary"
	[ -n "$secondary" ] && uci add_list dhcp.@dnsmasq[0].server="$secondary"
	uci set dhcp.@dnsmasq[0].noresolv='1'
	uci commit dhcp

	json_init
	json_add_boolean "success" 1
	json_add_string "provider" "$provider"
	json_add_string "primary" "$primary"
	json_add_string "secondary" "$secondary"
	json_dump
}

method_smart_config() {
	local best_provider best_latency=9999
	local test_providers="fdn quad9 cloudflare mullvad"

	json_init
	json_add_array "results"

	for provider in $test_providers; do
		local primary
		primary=$(echo "$DNS_PROVIDERS" | grep "^${provider}|" | cut -d'|' -f3)
		[ -z "$primary" ] && continue

		# Test latency
		local start end latency
		start=$(date +%s%N 2>/dev/null || date +%s)
		if nslookup -timeout=2 example.com "$primary" >/dev/null 2>&1; then
			end=$(date +%s%N 2>/dev/null || date +%s)
			latency=$(( (end - start) / 1000000 ))
			[ $latency -lt 0 ] && latency=1

			json_add_object
			json_add_string "provider" "$provider"
			json_add_string "server" "$primary"
			json_add_int "latency_ms" "$latency"
			json_add_boolean "reachable" 1
			json_close_object

			if [ "$latency" -lt "$best_latency" ]; then
				best_latency="$latency"
				best_provider="$provider"
			fi
		else
			json_add_object
			json_add_string "provider" "$provider"
			json_add_string "server" "$primary"
			json_add_int "latency_ms" 0
			json_add_boolean "reachable" 0
			json_close_object
		fi
	done

	json_close_array
	json_add_string "recommended" "${best_provider:-fdn}"
	json_add_int "best_latency_ms" "$best_latency"
	json_dump
}

method_test_dns() {
	local server="$1"
	local domain="${2:-example.com}"

	json_init

	local start end latency
	start=$(date +%s%N 2>/dev/null || date +%s)
	if result=$(nslookup -timeout=3 "$domain" "$server" 2>&1); then
		end=$(date +%s%N 2>/dev/null || date +%s)
		latency=$(( (end - start) / 1000000 ))
		[ $latency -lt 0 ] && latency=1

		json_add_boolean "success" 1
		json_add_string "server" "$server"
		json_add_string "domain" "$domain"
		json_add_int "latency_ms" "$latency"
	else
		json_add_boolean "success" 0
		json_add_string "server" "$server"
		json_add_string "error" "DNS query failed"
	fi

	json_dump
}

method_apply() {
	/etc/init.d/dnsmasq restart >/dev/null 2>&1

	json_init
	json_add_boolean "success" 1
	json_add_string "message" "DNS configuration applied"
	json_dump
}

# =============================================================================
# AI DNS Guard Methods (requires secubox-dns-guard package)
# =============================================================================

STATE_DIR="/var/lib/dns-guard"
BLOCKLIST_FILE="/etc/dnsmasq.d/dns-guard-blocklist.conf"

method_guard_status() {
	json_init

	# Check if dns-guard daemon is running
	local running=0
	pgrep -f "dns-guard daemon" >/dev/null 2>&1 && running=1

	json_add_boolean "daemon_running" "$running"

	# Check if dns-guard CLI is available
	if command -v dns-guard >/dev/null 2>&1; then
		json_add_boolean "installed" 1

		# Get config values
		local enabled=$(uci -q get dns-guard.main.enabled)
		local interval=$(uci -q get dns-guard.main.interval)
		local auto_apply=$(uci -q get dns-guard.main.auto_apply_blocks)
		local min_confidence=$(uci -q get dns-guard.main.min_confidence)
		local localai_url=$(uci -q get dns-guard.main.localai_url)

		json_add_boolean "enabled" "${enabled:-0}"
		json_add_int "interval" "${interval:-60}"
		json_add_boolean "auto_apply" "${auto_apply:-0}"
		json_add_int "min_confidence" "${min_confidence:-80}"
		json_add_string "localai_url" "${localai_url:-http://127.0.0.1:8081}"

		# Check LocalAI availability
		if wget -q -O /dev/null --timeout=2 "${localai_url:-http://127.0.0.1:8081}/v1/models" 2>/dev/null; then
			json_add_boolean "localai_online" 1
		else
			json_add_boolean "localai_online" 0
		fi

		# Count alerts
		local alert_count=0
		if [ -f "$STATE_DIR/alerts.json" ]; then
			alert_count=$(jsonfilter -i "$STATE_DIR/alerts.json" -e '@[*]' 2>/dev/null | wc -l)
		fi
		json_add_int "alert_count" "$alert_count"

		# Count pending blocks
		local pending_count=0
		if [ -f "$STATE_DIR/pending_blocks.json" ]; then
			pending_count=$(jsonfilter -i "$STATE_DIR/pending_blocks.json" -e '@[*]' 2>/dev/null | wc -l)
		fi
		json_add_int "pending_count" "$pending_count"

		# Count active blocks
		local active_blocks=0
		if [ -f "$BLOCKLIST_FILE" ]; then
			active_blocks=$(grep -c "^address=" "$BLOCKLIST_FILE" 2>/dev/null || echo "0")
		fi
		json_add_int "active_blocks" "$active_blocks"

		# Last run time
		if [ -f "$STATE_DIR/last_run" ]; then
			json_add_string "last_run" "$(cat "$STATE_DIR/last_run")"
		else
			json_add_string "last_run" ""
		fi

		# Detector status
		json_add_object "detectors"
		for det in dga tunneling rate_anomaly known_bad tld_anomaly; do
			local det_enabled=$(uci -q get "dns-guard.detector_${det}.enabled")
			json_add_boolean "$det" "${det_enabled:-0}"
		done
		json_close_object
	else
		json_add_boolean "installed" 0
	fi

	json_dump
}

method_get_alerts() {
	local limit="${1:-50}"

	json_init
	json_add_array "alerts"

	if [ -f "$STATE_DIR/alerts.json" ]; then
		jsonfilter -i "$STATE_DIR/alerts.json" -e '@[*]' 2>/dev/null | tail -n "$limit" | while read -r alert; do
			local timestamp=$(echo "$alert" | jsonfilter -e '@.timestamp' 2>/dev/null)
			local domain=$(echo "$alert" | jsonfilter -e '@.domain' 2>/dev/null)
			local client=$(echo "$alert" | jsonfilter -e '@.client' 2>/dev/null)
			local type=$(echo "$alert" | jsonfilter -e '@.type' 2>/dev/null)
			local confidence=$(echo "$alert" | jsonfilter -e '@.confidence' 2>/dev/null)
			local reason=$(echo "$alert" | jsonfilter -e '@.reason' 2>/dev/null)

			json_add_object
			json_add_string "timestamp" "$timestamp"
			json_add_string "domain" "$domain"
			json_add_string "client" "$client"
			json_add_string "type" "$type"
			json_add_int "confidence" "${confidence:-0}"
			json_add_string "reason" "$reason"
			json_close_object
		done
	fi

	json_close_array
	json_dump
}

method_get_pending() {
	json_init
	json_add_array "pending"

	if [ -f "$STATE_DIR/pending_blocks.json" ]; then
		jsonfilter -i "$STATE_DIR/pending_blocks.json" -e '@[*]' 2>/dev/null | while read -r block; do
			local id=$(echo "$block" | jsonfilter -e '@.id' 2>/dev/null)
			local domain=$(echo "$block" | jsonfilter -e '@.domain' 2>/dev/null)
			local type=$(echo "$block" | jsonfilter -e '@.type' 2>/dev/null)
			local confidence=$(echo "$block" | jsonfilter -e '@.confidence' 2>/dev/null)
			local reason=$(echo "$block" | jsonfilter -e '@.reason' 2>/dev/null)
			local created=$(echo "$block" | jsonfilter -e '@.created' 2>/dev/null)

			json_add_object
			json_add_string "id" "$id"
			json_add_string "domain" "$domain"
			json_add_string "type" "$type"
			json_add_int "confidence" "${confidence:-0}"
			json_add_string "reason" "$reason"
			json_add_string "created" "$created"
			json_close_object
		done
	fi

	json_close_array
	json_dump
}

method_approve_block() {
	local block_id="$1"

	json_init

	if [ -z "$block_id" ]; then
		json_add_boolean "success" 0
		json_add_string "error" "Missing block_id"
		json_dump
		return
	fi

	if command -v dns-guard >/dev/null 2>&1; then
		if dns-guard approve "$block_id" >/dev/null 2>&1; then
			json_add_boolean "success" 1
			json_add_string "message" "Block approved"
		else
			json_add_boolean "success" 0
			json_add_string "error" "Failed to approve block"
		fi
	else
		json_add_boolean "success" 0
		json_add_string "error" "dns-guard not installed"
	fi

	json_dump
}

method_reject_block() {
	local block_id="$1"

	json_init

	if [ -z "$block_id" ]; then
		json_add_boolean "success" 0
		json_add_string "error" "Missing block_id"
		json_dump
		return
	fi

	if command -v dns-guard >/dev/null 2>&1; then
		if dns-guard reject "$block_id" >/dev/null 2>&1; then
			json_add_boolean "success" 1
			json_add_string "message" "Block rejected"
		else
			json_add_boolean "success" 0
			json_add_string "error" "Failed to reject block"
		fi
	else
		json_add_boolean "success" 0
		json_add_string "error" "dns-guard not installed"
	fi

	json_dump
}

method_approve_all() {
	json_init

	if command -v dns-guard >/dev/null 2>&1; then
		if dns-guard approve-all >/dev/null 2>&1; then
			json_add_boolean "success" 1
			json_add_string "message" "All blocks approved"
		else
			json_add_boolean "success" 0
			json_add_string "error" "Failed to approve blocks"
		fi
	else
		json_add_boolean "success" 0
		json_add_string "error" "dns-guard not installed"
	fi

	json_dump
}

method_ai_check() {
	local domain="$1"

	json_init

	if [ -z "$domain" ]; then
		json_add_boolean "success" 0
		json_add_string "error" "Missing domain"
		json_dump
		return
	fi

	if command -v dns-guard >/dev/null 2>&1; then
		local result=$(dns-guard check "$domain" 2>&1)
		json_add_boolean "success" 1
		json_add_string "domain" "$domain"
		json_add_string "analysis" "$result"
	else
		json_add_boolean "success" 0
		json_add_string "error" "dns-guard not installed"
	fi

	json_dump
}

method_get_blocklist() {
	json_init
	json_add_array "blocked"

	if [ -f "$BLOCKLIST_FILE" ]; then
		grep "^address=" "$BLOCKLIST_FILE" 2>/dev/null | while read -r line; do
			local domain=$(echo "$line" | sed 's/address=\/\([^/]*\)\/.*/\1/')
			json_add_object
			json_add_string "domain" "$domain"
			json_close_object
		done
	fi

	json_close_array
	json_dump
}

method_unblock() {
	local domain="$1"

	json_init

	if [ -z "$domain" ]; then
		json_add_boolean "success" 0
		json_add_string "error" "Missing domain"
		json_dump
		return
	fi

	if [ -f "$BLOCKLIST_FILE" ]; then
		sed -i "/address=\/$domain\//d" "$BLOCKLIST_FILE"
		sed -i "/# Added:.*$domain/d" "$BLOCKLIST_FILE"
		/etc/init.d/dnsmasq restart >/dev/null 2>&1

		json_add_boolean "success" 1
		json_add_string "message" "Domain unblocked: $domain"
	else
		json_add_boolean "success" 0
		json_add_string "error" "Blocklist not found"
	fi

	json_dump
}

method_get_stats() {
	json_init

	local log_file="/var/log/dnsmasq.log"

	if [ -f "$log_file" ]; then
		local total=$(grep -c "query\[" "$log_file" 2>/dev/null || echo "0")
		local unique=$(grep "query\[" "$log_file" 2>/dev/null | sed 's/.*query\[[^]]*\] \([^ ]*\) from.*/\1/' | sort -u | wc -l)
		local a_queries=$(grep -c "query\[A\]" "$log_file" 2>/dev/null || echo "0")
		local aaaa_queries=$(grep -c "query\[AAAA\]" "$log_file" 2>/dev/null || echo "0")
		local txt_queries=$(grep -c "query\[TXT\]" "$log_file" 2>/dev/null || echo "0")

		json_add_int "total_queries" "$total"
		json_add_int "unique_domains" "$unique"
		json_add_int "a_queries" "$a_queries"
		json_add_int "aaaa_queries" "$aaaa_queries"
		json_add_int "txt_queries" "$txt_queries"
	else
		json_add_int "total_queries" 0
		json_add_int "unique_domains" 0
		json_add_int "a_queries" 0
		json_add_int "aaaa_queries" 0
		json_add_int "txt_queries" 0
	fi

	json_dump
}

method_toggle_guard() {
	local enable="$1"

	json_init

	if [ "$enable" = "1" ]; then
		uci set dns-guard.main.enabled='1'
		uci commit dns-guard
		/etc/init.d/dns-guard start >/dev/null 2>&1
		json_add_boolean "success" 1
		json_add_string "message" "DNS Guard enabled"
	else
		uci set dns-guard.main.enabled='0'
		uci commit dns-guard
		/etc/init.d/dns-guard stop >/dev/null 2>&1
		json_add_boolean "success" 1
		json_add_string "message" "DNS Guard disabled"
	fi

	json_dump
}

# RPC interface
case "$1" in
	list)
		cat <<'EOF'
{
	"status": {},
	"get_providers": {},
	"get_config": {},
	"set_provider": { "provider": "string" },
	"smart_config": {},
	"test_dns": { "server": "string", "domain": "string" },
	"apply": {},
	"guard_status": {},
	"get_alerts": { "limit": 50 },
	"get_pending": {},
	"approve_block": { "block_id": "string" },
	"reject_block": { "block_id": "string" },
	"approve_all": {},
	"ai_check": { "domain": "string" },
	"get_blocklist": {},
	"unblock": { "domain": "string" },
	"get_stats": {},
	"toggle_guard": { "enable": 1 }
}
EOF
		;;
	call)
		case "$2" in
			status) method_status ;;
			get_providers) method_get_providers ;;
			get_config) method_get_config ;;
			set_provider)
				read -r input
				provider=$(echo "$input" | jsonfilter -e '@.provider')
				method_set_provider "$provider"
				;;
			smart_config) method_smart_config ;;
			test_dns)
				read -r input
				server=$(echo "$input" | jsonfilter -e '@.server')
				domain=$(echo "$input" | jsonfilter -e '@.domain')
				method_test_dns "$server" "$domain"
				;;
			apply) method_apply ;;
			guard_status) method_guard_status ;;
			get_alerts)
				read -r input
				limit=$(echo "$input" | jsonfilter -e '@.limit' 2>/dev/null)
				method_get_alerts "${limit:-50}"
				;;
			get_pending) method_get_pending ;;
			approve_block)
				read -r input
				block_id=$(echo "$input" | jsonfilter -e '@.block_id')
				method_approve_block "$block_id"
				;;
			reject_block)
				read -r input
				block_id=$(echo "$input" | jsonfilter -e '@.block_id')
				method_reject_block "$block_id"
				;;
			approve_all) method_approve_all ;;
			ai_check)
				read -r input
				domain=$(echo "$input" | jsonfilter -e '@.domain')
				method_ai_check "$domain"
				;;
			get_blocklist) method_get_blocklist ;;
			unblock)
				read -r input
				domain=$(echo "$input" | jsonfilter -e '@.domain')
				method_unblock "$domain"
				;;
			get_stats) method_get_stats ;;
			toggle_guard)
				read -r input
				enable=$(echo "$input" | jsonfilter -e '@.enable')
				method_toggle_guard "$enable"
				;;
		esac
		;;
esac