Add RType AAAA

yasasa · Sep 27, 2019 · af0dc95 · af0dc95
1 parent 8028b01
commit af0dc95
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 13 deletions.
diff --git a/dns/message.py b/dns/message.py
@@ -10,6 +10,7 @@
 QUERY_TYPE_NS = 2
 
 QUERY_TYPE_CNAME = 5
+QUERY_TYPE_AAAA = 28
 
 
 def check_name(name_fields):
@@ -121,6 +122,14 @@ def seek_response(response_start, full_msg):
         answer, _ = decode_name(full_msg, rdata)
     elif rtype == QUERY_TYPE_NS:
         answer, _ = decode_name(full_msg, rdata)
+    elif rtype == QUERY_TYPE_AAAA:
+        ip = rdata[:rdlen * 8]
+        answer = "{:04X}:{:04X}:{:04X}:{:04X}:{:04X}:{:04X}:{:04X}:{:04X}".format(
+            ip[:16].uint, ip[16:32].uint, ip[32:48].uint, ip[48:64].uint,
+            ip[64:80].uint, ip[80:96].uint, ip[96:112].uint, ip[112:128].uint)
+    else:
+        raise ResponsePacketError(
+            "Found RType {} in response packet.".format(rtype))
 
     resp = rdata[rdlen * 8:]
 
@@ -252,6 +261,9 @@ def _pp_response_list(self, response_list):
                 msg = "NS\t{:s}".format(answer)
             elif rtype == QUERY_TYPE_CNAME:
                 msg = "CNAME\t{:s}".format(answer)
+            elif rtype == QUERY_TYPE_AAAA:
+                msg = "IP\t{:s}".format(answer)
+
             auth_type = "auth" if self.aa else "nonauth"
             msg = "{:s}\t{:d}\t{:s}".format(msg, ttl, auth_type)
             msgs.append(msg)

diff --git a/dns/resolver.py b/dns/resolver.py
@@ -89,7 +89,8 @@ def query(self, query, server, timeout=1., port=53, max_retries=5):
                 retries += 1
 
         if max_retries == retries:
-            raise socket.timeout("{} retries, {:.2f} seconds".format(retries, time.time() - start))
-
+            raise socket.timeout("{} retries, {:.2f} seconds".format(
+                retries,
+                time.time() - start))
 
         return resp, retries, end - start
diff --git a/main.py b/main.py
@@ -25,11 +25,8 @@
     )
     parser.add_argument("--port", type=int, default=53, help="UDP port number")
     parser.add_argument(
-        "@server",
-        type=str,
-        help="IPv4 address for DNS server.")
-    parser.add_argument(
-        "name", type=str, help="Domain name to query for.")
+        "@server", type=str, help="IPv4 address for DNS server.")
+    parser.add_argument("name", type=str, help="Domain name to query for.")
 
     group = parser.add_mutually_exclusive_group()
     group.add_argument("-mx", action="store_true")
@@ -42,16 +39,23 @@
     with Resolver() as resolver:
         print("DNS client sending request for {}".format(args.name))
         print("Server: {}".format(server))
-        print("Request Type: {}".format("NS" if args.ns else "MX" if args.mx else "A"))
+        print("Request Type: {}".format("NS" if args.ns else "MX"
+                                        if args.mx else "A"))
         query = Query(names=[(args.name, query_type)])
         try:
-            response, attempts, time = resolver.query(query, server, timeout=args.timeout, port=args.port, max_retries=args.max_retries)
-            print("Response received after {:.2f} seconds ({} retries)".format(time, attempts))
+            response, attempts, time = resolver.query(
+                query,
+                server,
+                timeout=args.timeout,
+                port=args.port,
+                max_retries=args.max_retries)
+            print("Response received after {:.2f} seconds ({} retries)".format(
+                time, attempts))
             print(response)
         except dns.InvalidNameError:
             print("NOTFOUND")
         except socket.timeout as e:
-            print("ERROR\tRequest Timed Out:",e)
+            print("ERROR\tRequest Timed Out:", e)
         except dns.FormatError as e:
             print("ERROR\tIncorrect Input Format:", e)
         except dns.ServerRefusedError:
@@ -62,5 +66,3 @@
             print("ERROR\tUnexpected Response:", e)
         except dns.ServerError:
             print("ERRPR\tServer Error")
-
-