summaryrefslogtreecommitdiff
path: root/project2/proj2_s4498062/dns/cache.py
blob: a8b62be8dfcb7b7716f0f3de0b9902d9296d87c1 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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)