136 lines
4.1 KiB
Python
136 lines
4.1 KiB
Python
#!/usr/bin/env python2
|
|
|
|
"""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
|
|
|
|
from dns.resource import ResourceRecord, RecordData
|
|
from dns.types import Type
|
|
from dns.classes import Class
|
|
|
|
import datetime
|
|
|
|
|
|
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):
|
|
res = {
|
|
"name": obj.name,
|
|
"type": Type.to_string(obj.type_),
|
|
"class": Class.to_string(obj.class_),
|
|
"ttl": obj.ttl,
|
|
"rdata": obj.rdata.data
|
|
}
|
|
if obj.createDate is not None:
|
|
res['createDate'] = obj.createDate.isoformat()
|
|
return res
|
|
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"])
|
|
try:
|
|
createDate = datetime.datetime.strptime(dct["createDate"], "%Y-%m-%dT%H:%M:%S.%f")
|
|
except KeyError, e:
|
|
createDate = None
|
|
return ResourceRecord(name, type_, class_, ttl, rdata, createDate)
|
|
|
|
class DummyCache(object):
|
|
""" Cache for ResourceRecords """
|
|
|
|
def __init__(self, ttl):
|
|
pass
|
|
def sweep(self):
|
|
pass
|
|
def lookup(self, dname, type_, class_):
|
|
return []
|
|
def add_record(self, record):
|
|
pass
|
|
def read_cache_file(self):
|
|
""" Read the cache file from disk """
|
|
pass
|
|
def write_cache_file(self):
|
|
""" Write the cache file to disk """
|
|
pass
|
|
|
|
class RecordCache(object):
|
|
""" Cache for ResourceRecords """
|
|
|
|
def __init__(self, ttl):
|
|
""" Initialize the RecordCache
|
|
|
|
Args:
|
|
ttl (int): TTL of cached entries (if > 0)
|
|
"""
|
|
self.records = []
|
|
self.ttl = ttl
|
|
|
|
def sweep(self):
|
|
self.records = [rec for rec in self.records if rec.valid()]
|
|
|
|
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.sweep()
|
|
matches = lambda rec: rec.type_ == type_ and rec.class_ == class_ and rec.name.lower() == dname.lower()
|
|
return filter(matches, self.records)
|
|
|
|
def add_record(self, record):
|
|
""" Add a new Record to the cache
|
|
|
|
Args:
|
|
record (ResourceRecord): the record added to the cache
|
|
"""
|
|
if self.ttl > 0:
|
|
record.ttl = self.ttl
|
|
record.createDate = datetime.datetime.now()
|
|
for rec in self.records:
|
|
if rec.type_ == record.type_ and rec.class_ == record.class_ and rec.name.lower() == record.name.lower() \
|
|
and rec.rdata.data == record.rdata.data:
|
|
rec.ttl = record.ttl
|
|
# update last seen time
|
|
break
|
|
else:
|
|
self.records.append(record)
|
|
|
|
def read_cache_file(self):
|
|
""" Read the cache file from disk """
|
|
with open("cache.json") as f:
|
|
self.records = json.load(f, object_hook=resource_from_json)
|
|
self.sweep()
|
|
def write_cache_file(self):
|
|
""" Write the cache file to disk """
|
|
self.sweep()
|
|
with open("cache.json", 'w') as f:
|
|
json.dump(self.records, f, cls=ResourceEncoder, indent=4)
|
|
|
|
|