From 72420801911ff93c2f6a73a7a78e72aa62cfa997 Mon Sep 17 00:00:00 2001
From: Yorick van Pelt <yorickvanpelt@gmail.com>
Date: Sun, 29 May 2016 18:07:07 +0200
Subject: [PATCH] cache expiration

---
 dns/cache.py    | 24 ++++++++++++++++++------
 dns/resource.py | 12 +++++++++++-
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/dns/cache.py b/dns/cache.py
index ea5b2b0..0348b2f 100644
--- a/dns/cache.py
+++ b/dns/cache.py
@@ -14,6 +14,8 @@ 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
@@ -23,13 +25,16 @@ class ResourceEncoder(json.JSONEncoder):
     """
     def default(self, obj):
         if isinstance(obj, ResourceRecord):
-            return {
+            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)
 
 
@@ -44,14 +49,18 @@ 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)
+    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 gc(self):
+    def sweep(self):
         pass
     def lookup(self, dname, type_, class_):
         return []
@@ -76,8 +85,8 @@ class RecordCache(object):
         self.records = []
         self.ttl = ttl
 
-    def gc(self):
-        pass
+    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
@@ -90,7 +99,7 @@ class RecordCache(object):
             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)
 
@@ -102,6 +111,7 @@ class RecordCache(object):
         """
         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:
@@ -115,8 +125,10 @@ class RecordCache(object):
         """ 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)
 
diff --git a/dns/resource.py b/dns/resource.py
index 083969c..8c98551 100644
--- a/dns/resource.py
+++ b/dns/resource.py
@@ -12,10 +12,11 @@ import struct
 
 from dns.types import Type
 
+from datetime import timedelta, datetime
 
 class ResourceRecord(object):
     """ DNS resource record """
-    def __init__(self, name, type_, class_, ttl, rdata):
+    def __init__(self, name, type_, class_, ttl, rdata, createDate=None):
         """ Create a new resource record
 
         Args:
@@ -29,6 +30,15 @@ class ResourceRecord(object):
         self.class_ = class_
         self.ttl = ttl
         self.rdata = rdata
+        self.createDate = createDate
+
+    def valid(self):
+        if self.createDate is None:
+            return True
+        time = datetime.now()
+        if (self.createDate + timedelta(seconds=self.ttl)) > time:
+            return True
+        return False
 
     def __repr__(self):
         return "<Resource " + Type.by_value[self.type_] + " " + self.name + ">"