diff options
Diffstat (limited to 'project2/proj2_s4498062/dns/cache.py')
-rw-r--r-- | project2/proj2_s4498062/dns/cache.py | 55 |
1 files changed, 43 insertions, 12 deletions
diff --git a/project2/proj2_s4498062/dns/cache.py b/project2/proj2_s4498062/dns/cache.py index 3ef14b3..9cde66f 100644 --- a/project2/proj2_s4498062/dns/cache.py +++ b/project2/proj2_s4498062/dns/cache.py @@ -7,6 +7,7 @@ It is highly recommended to use these. """ import json +import time from dns.resource import ResourceRecord, RecordData from dns.types import Type @@ -22,11 +23,12 @@ class ResourceEncoder(json.JSONEncoder): 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 + "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) @@ -42,20 +44,33 @@ def resource_from_json(dct): class_ = Class.from_string(dct["class"]) ttl = dct["ttl"] rdata = RecordData.create(type_, dct["rdata"]) - return ResourceRecord(name, type_, class_, ttl, rdata) + timestamp = dct["timestamp"] + return ResourceRecord(name, type_, class_, ttl, rdata, timestamp) class RecordCache(object): """ Cache for ResourceRecords """ - def __init__(self, ttl): + FILE = '.dns.cache' + + def __init__(self): """ Initialize the RecordCache Args: ttl (int): TTL of cached entries (if > 0) """ self.records = [] - self.ttl = ttl + self.read_cache_file() + + def __del__(self): + self.write_cache_file() + + def remove_old(self): + """Remove entries for which the TTL has expired""" + now = int(time.clock()) + for record in reversed(self.records): + if record.ttl + record.timestamp < now: + self.records.remove(record) def lookup(self, dname, type_, class_): """ Lookup resource records in cache @@ -68,7 +83,16 @@ class RecordCache(object): type_ (Type): type class_ (Class): class """ - pass + 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 @@ -76,12 +100,19 @@ class RecordCache(object): Args: record (ResourceRecord): the record added to the cache """ - pass + self.records.append(record) def read_cache_file(self): """ Read the cache file from disk """ - pass + 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 """ - pass + self.remove_old() + with open(self.FILE, 'w') as jsonfile: + json.dump(self.records, jsonfile, cls=ResourceEncoder, indent=4) |