net-dnsserver/dns/cache.py

#!/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 """
        try:
            with open("cache.json") as f:
                self.records = json.load(f, object_hook=resource_from_json)
        except Exception, e:
            print "error loading cache", repr(e)
        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)