From ecbb2100581ec75b506036b1f4e47cdfecc49e94 Mon Sep 17 00:00:00 2001
From: KiekerJan <kiekerjan@users.noreply.github.com>
Date: Wed, 26 Oct 2022 23:14:57 +0200
Subject: [PATCH] Add dns time out handling changes

---
 management/dns_update.py    | 26 +++++++++++++++++++-------
 management/status_checks.py | 32 +++++++++++++++++++++-----------
 2 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/management/dns_update.py b/management/dns_update.py
index 9256c0ff..ead03a55 100755
--- a/management/dns_update.py
+++ b/management/dns_update.py
@@ -1030,17 +1030,29 @@ def set_secondary_dns(hostnames, env):
 		resolver = dns.resolver.get_default_resolver()
 		resolver.timeout = 5
 		resolver.lifetime = 5
-		
+				
 		for item in hostnames:
 			if not item.startswith("xfr:"):
 				# Resolve hostname.
-				try:
-					response = resolver.resolve(item, "A")
-				except (dns.resolver.NoNameservers, dns.resolver.NXDOMAIN, dns.resolver.NoAnswer, dns.resolver.Timeout):
+				tries = 2
+				
+				while tries > 0:
+					tries = tries - 1
 					try:
-						response = resolver.resolve(item, "AAAA")
-					except (dns.resolver.NoNameservers, dns.resolver.NXDOMAIN, dns.resolver.NoAnswer, dns.resolver.Timeout):
-						raise ValueError("Could not resolve the IP address of %s." % item)
+						response = resolver.resolve(item, "A")
+						tries = 0
+					except (dns.resolver.NoNameservers, dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
+						try:
+							response = resolver.resolve(item, "AAAA")
+							tries = 0
+						except (dns.resolver.NoNameservers, dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
+							raise ValueError("Could not resolve the IP address of %s." % item)
+						except (dns.resolver.Timeout):
+							if tries < 1:
+								raise ValueError("Could not resolve the IP address of %s due to timeout." % item)
+					except (dns.resolver.Timeout):
+						if tries < 1:
+							raise ValueError("Could not resolve the IP address of %s due to timeout." % item)
 			else:
 				# Validate IP address.
 				try:
diff --git a/management/status_checks.py b/management/status_checks.py
index 675b8f74..29ce2dc1 100755
--- a/management/status_checks.py
+++ b/management/status_checks.py
@@ -303,7 +303,7 @@ def run_network_checks(env, output):
 	# by a spammer, or the user may be deploying on a residential network. We
 	# will not be able to reliably send mail in these cases.
 	rev_ip4 = ".".join(reversed(env['PUBLIC_IP'].split('.')))
-	zen = query_dns(rev_ip4+'.zen.spamhaus.org', 'A', nxdomain=None)
+	zen = query_dns(rev_ip4+'.zen.spamhaus.org', 'A', nxdomain=None, retry = False)
 	if zen is None:
 		output.print_ok("IP address is not blacklisted by zen.spamhaus.org.")
 	elif zen == "[timeout]":
@@ -739,7 +739,7 @@ def check_mail_domain(domain, env, output):
 	# Stop if the domain is listed in the Spamhaus Domain Block List.
 	# The user might have chosen a domain that was previously in use by a spammer
 	# and will not be able to reliably send mail.
-	dbl = query_dns(domain+'.dbl.spamhaus.org', "A", nxdomain=None)
+	dbl = query_dns(domain+'.dbl.spamhaus.org', "A", nxdomain=None, retry=False)
 	if dbl is None:
 		output.print_ok("Domain is not blacklisted by dbl.spamhaus.org.")
 	elif dbl == "[timeout]":
@@ -775,7 +775,7 @@ def check_web_domain(domain, rounded_time, ssl_certificates, env, output):
 	# website for also needs a signed certificate.
 	check_ssl_cert(domain, rounded_time, ssl_certificates, env, output)
 
-def query_dns(qname, rtype, nxdomain='[Not Set]', at=None, as_list=False):
+def query_dns(qname, rtype, nxdomain='[Not Set]', at=None, as_list=False, retry=True):
 	# Make the qname absolute by appending a period. Without this, dns.resolver.query
 	# will fall back a failed lookup to a second query with this machine's hostname
 	# appended. This has been causing some false-positive Spamhaus reports. The
@@ -800,15 +800,25 @@ def query_dns(qname, rtype, nxdomain='[Not Set]', at=None, as_list=False):
 	# lifetime expires a dns.exception.Timeout exception will be raised.
 	resolver.lifetime = 5
 
+	if retry:
+		tries = 2
+	else:
+		tries = 1
+	
 	# Do the query.
-	try:
-		response = resolver.resolve(qname, rtype)
-	except (dns.resolver.NoNameservers, dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
-		# Host did not have an answer for this query; not sure what the
-		# difference is between the two exceptions.
-		return nxdomain
-	except dns.exception.Timeout:
-		return "[timeout]"
+	while tries > 0:
+		tries = tries - 1
+		try:
+			response = resolver.resolve(qname, rtype, search=True)
+			tries = 0
+		except (dns.resolver.NoNameservers, dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
+			# Host did not have an answer for this query; not sure what the
+			# difference is between the two exceptions.
+			if tries < 1:
+				return nxdomain
+		except dns.exception.Timeout:
+			if tries < 1:
+				return "[timeout]"
 
 	# Normalize IP addresses. IP address --- especially IPv6 addresses --- can
 	# be expressed in equivalent string forms. Canonicalize the form before