Project 2: simple DNS resolver

1 files changed, 34 insertions, 6 deletions

diff --git a/project2/proj2_s4498062/dns/message.py b/project2/proj2_s4498062/dns/message.py
index 9f8ead4..a8ad621 100644
--- a/project2/proj2_s4498062/dns/message.py
+++ b/project2/proj2_s4498062/dns/message.py
@@ -8,14 +8,15 @@ import struct
 
 from dns.domainname import Parser, Composer
 from dns.resource import ResourceRecord
+from dns.types import Type
 
 
 class Message(object):
     """ DNS message """
 
     def __init__(
-            self, header, questions=[], answers=[], authorities=[],
-            additionals=[]):
+            self, header, questions=None, answers=None, authorities=None,
+            additionals=None):
         """ Create a new DNS message
 
         Args:
@@ -26,10 +27,10 @@ class Message(object):
             additionals ([ResourceRecord]): the additional section
         """
         self.header = header
-        self.questions = questions
-        self.answers = answers
-        self.authorities = authorities
-        self.additionals = additionals
+        self.questions = questions if questions != None else []
+        self.answers = answers if answers != None else []
+        self.authorities = authorities if authorities != None else []
+        self.additionals = additionals if additionals != None else []
 
     @property
     def resources(self):
@@ -103,6 +104,33 @@ class Message(object):
 
         return cls(header, questions, answers, authorities, additionals)
 
+    def get_addresses(self):
+        """Get the addresses from a response"""
+        return self.get_data('answers', Type.A)
+
+    def get_hints(self):
+        """Get the nameservers from a response"""
+        hints = self.get_data('authorities', Type.NS)
+        results = []
+        for hint in hints:
+            addits = [
+                a.rdata.data for a in self.additionals if
+                a.name == hint and a.type_ == Type.A]
+            if addits == []:
+                continue
+            results.append((hint, addits))
+        return results
+
+    def get_aliases(self):
+        """Get the aliases from a response"""
+        return self.get_data('additionals', Type.CNAME)
+
+    def get_data(self, key, type_=None):
+        """Yield from a message section, possibly filtering by type"""
+        for k in getattr(self, key):
+            if type_ == None or k.type_ == type_:
+                yield k.rdata.data
+
 
 class Header(object):
     """ The header section of a DNS message