From b23a3ce081762684889f79a6b85f1d3e5fcafc45 Mon Sep 17 00:00:00 2001
From: Babak Farrokhi <babak@farrokhi.net>
Date: Wed, 15 Jun 2016 17:14:10 +0430
Subject: [PATCH] Fix looking up NS record from root server

- also resolve hostname to IP if name is given as dns server
- raise timeout to 2 seconds
---
 dnstraceroute.py | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/dnstraceroute.py b/dnstraceroute.py
index 4db23ac..21e304a 100755
--- a/dnstraceroute.py
+++ b/dnstraceroute.py
@@ -38,7 +38,6 @@ import time
 import dns.query
 import dns.rdatatype
 import dns.resolver
-
 from cymruwhois import cymruwhois
 
 __author__ = 'Babak Farrokhi (babak@farrokhi.net)'
@@ -181,7 +180,7 @@ def ping(resolver, hostname, dnsrecord, ttl):
     _ttl = ttl
 
     try:
-        resolver.query(hostname, dnsrecord)
+        resolver.query(hostname, dnsrecord, raise_on_no_answer=False)
 
     except dns.resolver.NoNameservers as e:
         if not quiet:
@@ -199,8 +198,8 @@ def ping(resolver, hostname, dnsrecord, ttl):
         pass
     except SystemExit:
         pass
-    except:
-        print("unxpected error: ", sys.exc_info()[0])
+    except Exception as e:
+        print("unxpected error: ", e)
         sys.exit(1)
     else:
         reached = True
@@ -222,7 +221,7 @@ def main():
 
     dnsrecord = 'A'
     count = 30
-    timeout = 1
+    timeout = 2
     dnsserver = dns.resolver.get_default_resolver().nameservers[0]
     dest_port = 53
     hops = 0
@@ -273,6 +272,16 @@ def main():
 
     color = Colors(color_mode)
 
+    # check if we have a valid dns server address
+    try:
+        ipaddress.ip_address(dnsserver)
+    except ValueError:  # so it is not a valid IPv4 or IPv6 address, so try to resolve host name
+        try:
+            dnsserver = socket.getaddrinfo(dnsserver, port=None, family=socket.AF_INET)[1][4][0]
+        except OSError:
+            print('Error: cannot resolve hostname:', dnsserver)
+            sys.exit(1)
+
     resolver = dns.resolver.Resolver()
     resolver.nameservers = [dnsserver]
     resolver.timeout = timeout
@@ -286,7 +295,8 @@ def main():
     trace_path = []
 
     if not quiet:
-        print("%s DNS: %s:%d, hostname: %s, rdatatype: %s" % (__PROGNAME__, dnsserver, dest_port, hostname, dnsrecord))
+        print("%s DNS: %s:%d, hostname: %s, rdatatype: %s" % (__PROGNAME__, dnsserver, dest_port, hostname, dnsrecord),
+              flush=True)
 
     while True:
         if shutdown: