From 6336dbbff74b937d6d72bb725094560637bba1c0 Mon Sep 17 00:00:00 2001
From: KiekerJan <kiekerjan@users.noreply.github.com>
Date: Sun, 23 Oct 2022 23:32:47 +0200
Subject: [PATCH] add dns query handling changes

---
 management/dns_update.py    | 26 ++++++++++++++++++++------
 management/status_checks.py |  9 +++++++--
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/management/dns_update.py b/management/dns_update.py
index 1567371a..d7c1f617 100755
--- a/management/dns_update.py
+++ b/management/dns_update.py
@@ -983,6 +983,7 @@ def set_custom_dns_record(qname, rtype, value, action, env):
 def get_secondary_dns(custom_dns, mode=None):
 	resolver = dns.resolver.get_default_resolver()
 	resolver.timeout = 10
+	resolver.lifetime = 10
 
 	values = []
 	for qname, rtype, value in custom_dns:
@@ -1000,10 +1001,17 @@ def get_secondary_dns(custom_dns, mode=None):
 			# doesn't.
 			if not hostname.startswith("xfr:"):
 				if mode == "xfr":
-					response = dns.resolver.resolve(hostname+'.', "A", raise_on_no_answer=False)
-					values.extend(map(str, response))
-					response = dns.resolver.resolve(hostname+'.', "AAAA", raise_on_no_answer=False)
-					values.extend(map(str, response))
+					try:
+						response = resolver.resolve(hostname+'.', "A", raise_on_no_answer=False)
+						values.extend(map(str, response))
+					except dns.exception.DNSException:
+						pass
+						
+					try:
+						response = resolver.resolve(hostname+'.', "AAAA", raise_on_no_answer=False)
+						values.extend(map(str, response))
+					except dns.exception.DNSException:
+						pass
 					continue
 				values.append(hostname)
 
@@ -1021,6 +1029,8 @@ def set_secondary_dns(hostnames, env):
 		# Validate that all hostnames are valid and that all zone-xfer IP addresses are valid.
 		resolver = dns.resolver.get_default_resolver()
 		resolver.timeout = 5
+		resolver.lifetime = 5
+		
 		for item in hostnames:
 			if not item.startswith("xfr:"):
 				# Resolve hostname.
@@ -1030,7 +1040,11 @@ def set_secondary_dns(hostnames, env):
 					try:
 						response = resolver.resolve(item, "AAAA")
 					except (dns.resolver.NoNameservers, dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
-						raise ValueError("Could not resolve the IP address of %s." % item)
+						pass
+					except (dns.resolver.Timeout):
+						raise ValueError("Could not resolve the IP address of %s due to timeout." % item)
+				except (dns.resolver.Timeout):
+					raise ValueError("Could not resolve the IP address of %s due to timeout." % item)
 			else:
 				# Validate IP address.
 				try:
@@ -1062,7 +1076,7 @@ def get_custom_dns_records(custom_dns, qname, rtype):
 def build_recommended_dns(env):
 	ret = []
 	for (domain, zonefile, records) in build_zones(env):
-		# remove records that we don't dislay
+		# remove records that we don't display
 		records = [r for r in records if r[3] is not False]
 
 		# put Required at the top, then Recommended, then everythiing else
diff --git a/management/status_checks.py b/management/status_checks.py
index 0d555441..675b8f74 100755
--- a/management/status_checks.py
+++ b/management/status_checks.py
@@ -541,7 +541,7 @@ def check_dns_zone(domain, env, output, dns_zonefiles):
 		for ns in custom_secondary_ns:
 			# We must first resolve the nameserver to an IP address so we can query it.
 			ns_ips = query_dns(ns, "A")
-			if not ns_ips:
+			if not ns_ips or ns_ips in {'[Not Set]', '[timeout]'}:
 				output.print_error("Secondary nameserver %s is not valid (it doesn't resolve to an IP address)." % ns)
 				continue
 			# Choose the first IP if nameserver returns multiple
@@ -788,12 +788,17 @@ def query_dns(qname, rtype, nxdomain='[Not Set]', at=None, as_list=False):
 	# running bind server), or if the 'at' argument is specified, use that host
 	# as the nameserver.
 	resolver = dns.resolver.get_default_resolver()
-	if at:
+	
+	# Make sure at is not a string that cannot be used as a nameserver
+	if at and at not in {'[Not set]', '[timeout]'}:
 		resolver = dns.resolver.Resolver()
 		resolver.nameservers = [at]
 
 	# Set a timeout so that a non-responsive server doesn't hold us back.
 	resolver.timeout = 5
+	# The number of seconds to spend trying to get an answer to the question. If the
+	# lifetime expires a dns.exception.Timeout exception will be raised.
+	resolver.lifetime = 5
 
 	# Do the query.
 	try: