diff options
Diffstat (limited to 'project2/proj2_s4498062/dns/cache.py')
-rw-r--r-- | project2/proj2_s4498062/dns/cache.py | 119 |
1 files changed, 119 insertions, 0 deletions
diff --git a/project2/proj2_s4498062/dns/cache.py b/project2/proj2_s4498062/dns/cache.py new file mode 100644 index 0000000..a8b62be --- /dev/null +++ b/project2/proj2_s4498062/dns/cache.py @@ -0,0 +1,119 @@ +"""A cache for resource records + +This module contains a class which implements a cache for DNS resource records, +you still have to do most of the implementation. The module also provides a +class and a function for converting ResourceRecords from and to JSON strings. +It is highly recommended to use these. +""" + +import json +import time + +from dns.resource import ResourceRecord, RecordData +from dns.types import Type +from dns.classes import Class + + +class ResourceEncoder(json.JSONEncoder): + """ Conver ResourceRecord to JSON + + Usage: + string = json.dumps(records, cls=ResourceEncoder, indent=4) + """ + def default(self, obj): + if isinstance(obj, ResourceRecord): + return { + "name": obj.name, + "type": Type.to_string(obj.type_), + "class": Class.to_string(obj.class_), + "ttl": obj.ttl, + "rdata": obj.rdata.data, + "timestamp": obj.timestamp + } + return json.JSONEncoder.default(self, obj) + + +def resource_from_json(dct): + """ Convert JSON object to ResourceRecord + + Usage: + records = json.loads(string, object_hook=resource_from_json) + """ + name = dct["name"] + type_ = Type.from_string(dct["type"]) + class_ = Class.from_string(dct["class"]) + ttl = dct["ttl"] + rdata = RecordData.create(type_, dct["rdata"]) + timestamp = dct["timestamp"] + return ResourceRecord(name, type_, class_, ttl, rdata, timestamp) + + +class RecordCache(object): + """ Cache for ResourceRecords """ + + FILE = '.dns.cache' + + def __init__(self, ttl): + """ Initialize the RecordCache + + Args: + ttl (int): TTL of cached entries (if > 0) + """ + self.records = [] + self.read_cache_file() + self.ttl = ttl + + def __del__(self): + self.write_cache_file() + + def remove_old(self): + """Remove entries for which the TTL has expired""" + now = int(time.time()) + for record in reversed(self.records): + if min(self.ttl, record.ttl) + record.timestamp < now: + self.records.remove(record) + + def lookup(self, dname, type_, class_): + """ Lookup resource records in cache + + Lookup for the resource records for a domain name with a specific type + and class. + + Args: + dname (str): domain name + type_ (Type): type + class_ (Class): class + """ + self.remove_old() + return [ + r for r in self.records + if r.match(name=dname, type_=type_, class_=class_)] + + def add_records_from(self, msg): + for record in msg.answers + msg.authorities + msg.additionals: + if record.type_ in [Type.A, Type.AAAA, Type.CNAME, Type.NS] and \ + record.class_ == Class.IN: + self.add_record(record) + + def add_record(self, record): + """ Add a new Record to the cache + + Args: + record (ResourceRecord): the record added to the cache + """ + self.records.append(record) + + def read_cache_file(self): + """ Read the cache file from disk """ + try: + with open(self.FILE, 'r') as jsonfile: + self.records = json.load( + jsonfile, object_hook=resource_from_json) + except IOError: + pass + + def write_cache_file(self): + """ Write the cache file to disk """ + self.remove_old() + with open(self.FILE, 'w') as jsonfile: + json.dump(self.records, jsonfile, cls=ResourceEncoder, indent=4) |