1 files changed, 153 insertions, 0 deletions
diff --git a/project2/proj2_s4498062/dns/resolver.py b/project2/proj2_s4498062/dns/resolver.py
new file mode 100644
index 0000000..4c09681
--- /dev/null
+++ b/project2/proj2_s4498062/dns/resolver.py
@@ -0,0 +1,153 @@
+#!/usr/bin/env python2
+
+""" DNS Resolver
+
+This module contains a class for resolving hostnames. You will have to implement
+things in this module. This resolver will be both used by the DNS client and the
+DNS server, but with a different list of servers.
+"""
+
+from random import randint
+import socket
+
+from dns.classes import Class
+from dns.types import Type
+
+from dns.cache import RecordCache
+from dns.message import Message, Question, Header
+
+
+class Resolver(object):
+    """ DNS resolver """
+
+    ROOT_SERVERS = [
+        '198.41.0.4',
+        '192.228.79.201',
+        '192.33.4.12',
+        '199.7.91.13',
+        '192.203.230.10',
+        '192.5.5.241',
+        '192.112.36.4',
+        '198.97.190.53',
+        '192.36.148.17',
+        '192.58.128.30',
+        '193.0.14.129',
+        '199.7.83.42',
+        '202.12.27.33'
+        ]
+
+    def __init__(self, caching, ttl, timeout=3):
+        """ Initialize the resolver
+
+        Args:
+            caching (bool): caching is enabled if True
+            ttl (int): ttl of cache entries (if > 0)
+        """
+        self.caching = caching
+        self.ttl = ttl
+        self.timeout = timeout
+
+        if self.caching:
+            self.cache = RecordCache(ttl)
+
+    def do_query(self, hint, hostname, type_, class_=Class.IN, caching=True):
+        """Do a query to a hint"""
+        if self.caching and caching:
+            records = self.cache.lookup(hostname, type_, class_)
+            if records != []:
+                return records
+
+        ident = randint(0, 65535)
+        header = Header(ident, 0, 1, 0, 0, 0)
+        header.qr = 0
+        header.opcode = 0
+        header.rd = 1
+        req = Message(header, [Question(hostname, type_, class_)])
+
+        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        sock.settimeout(self.timeout)
+        sock.sendto(req.to_bytes(), (hint, 53))
+
+        try:
+            data = sock.recv(512)
+            resp = Message.from_bytes(data)
+            if resp.header.ident == ident:
+                if self.caching and caching:
+                    self.cache.add_records_from(resp)
+                return resp.answers + resp.authorities + resp.additionals
+        except socket.timeout:
+            pass
+        return []
+
+    def do_query_to_multiple(
+            self, hints, hostname, type_, class_=Class.IN, caching=True):
+        """Do a query to multiple hints, return the remaining hints"""
+        while hints != []:
+            hint = hints.pop()
+            response = self.do_query(hint, hostname, type_, class_, caching)
+            if response is not None:
+                return hints, response
+        return [], []
+
+    def gethostbyname(self, hostname):
+        """ Translate a host name to IPv4 address.
+
+        Args:
+            hostname (str): the hostname to resolve
+
+        Returns:
+            (str, [str], [str]): (hostname, aliaslist, ipaddrlist)
+        """
+        if self.caching:
+            addrs = self.cache.lookup(hostname, Type.A, Class.IN)
+            cnames = self.cache.lookup(hostname, Type.CNAME, Class.IN)
+            if addrs != []:
+                return hostname, \
+                    [r.rdata.data for r in cnames], \
+                    [r.rdata.data for r in addrs]
+            for cname in cnames:
+                cname, aliases, addrs = self.gethostbyname(cname.rdata.data)
+                if addrs != []:
+                    return str(cname), aliases, addrs
+
+        if hostname == '':
+            return hostname, [], []
+
+        hints = self.ROOT_SERVERS[:]
+        aliases = []
+        while hints != []:
+            hints, info = self.do_query_to_multiple(hints, hostname, Type.A)
+
+            aliases += [
+                r.rdata.data for r in info
+                if r.match(type_=Type.CNAME, class_=Class.IN, name=hostname)]
+
+            # Case 1: answer
+            ips = [
+                r.rdata.data for r in info
+                if r.match(type_=Type.A, class_=Class.IN, name=hostname)]
+            if ips != []:
+                return hostname, aliases, ips
+
+            # Case 2: name servers
+            auths = [
+                r.rdata.data for r in info
+                if r.match(type_=Type.NS, class_=Class.IN)]
+            ips = [
+                add.rdata.data for ns in auths for add in info
+                if add.match(name=ns, type_=Type.A)]
+            if ips != []:
+                hints += ips
+                continue
+            if auths != []:
+                auths = [h for a in auths for h in self.gethostbyname(a)[2]]
+                hints += auths
+                continue
+
+            # Case 3: aliases
+            for alias in aliases:
+                _, extra_aliases, alias_addresses = self.gethostbyname(alias)
+                if alias_addresses != []:
+                    return hostname, aliases + extra_aliases, alias_addresses
+
+        return hostname, aliases, []