From 04fb72e408a2e16c320a6287d09f1d94ea3a746a Mon Sep 17 00:00:00 2001
From: Stephane Bortzmeyer <stephane+framagit@bortzmeyer.org>
Date: Wed, 7 Jul 2021 15:50:47 +0200
Subject: [PATCH] Pickle implemented (but not activated by default). Closes #1

---
 ianardap.py      | 64 +++++++++++++++++++++++++++++++++---------------
 test_ianardap.py |  9 ++++++-
 2 files changed, 52 insertions(+), 21 deletions(-)

diff --git a/ianardap.py b/ianardap.py
index c659849..6ba0da9 100755
--- a/ianardap.py
+++ b/ianardap.py
@@ -14,43 +14,64 @@ import os
 import sys
 import time
 import fcntl
+import pickle
 
 IANABASE = "https://data.iana.org/rdap/dns.json"
-CACHE = os.environ["HOME"] + "/.ianardapcache.json" 
+CACHE = os.environ["HOME"] + "/.ianardapcache" 
 MAXAGE = 24 # Hours
 IANATIMEOUT = 10 # Seconds
 MAXTESTS = 3 # Maximum attempts to get the database
 
 class IanaRDAPDatabase():
 
-    def __init__(self, maxage=MAXAGE, cachefile=CACHE):
-        """ Retrieves the IANA databse, if not already cached. maxage is in hours. """
+    def __init__(self, maxage=MAXAGE, cachefile=CACHE, pickleformat=False):
+        """Retrieves the IANA database, if not already cached. maxage is in
+hours. The cache file argument should not have an extension (it will
+be added automatically). pickleformat is not the default because it is
+not really faster *and* it introduces security risks if someone can
+write in the file (see the documentation of the module)."""
         cache_valid = False
-        self.cachefile = cachefile
+        if pickleformat:
+            self.cachefile = cachefile + ".pickle"
+        else:
+            self.cachefile = cachefile + ".json"
         self.lockname = self.cachefile + ".lock"
         loaded = False
         tests = 0
         errmsg = "No error"
         while not loaded and tests < MAXTESTS:
             self.lock()
-            if os.path.exists(cachefile) and \
-               datetime.datetime.fromtimestamp(os.path.getmtime(cachefile)) > \
+            if os.path.exists(self.cachefile) and \
+               datetime.datetime.fromtimestamp(os.path.getmtime(self.cachefile)) > \
                (datetime.datetime.now() - datetime.timedelta(hours = maxage)):
-                cache = open(cachefile, "rb")
+                cache = open(self.cachefile, "rb")
                 content = cache.read()
                 cache.close()
                 self.unlock()
-                try:
-                    database = json.loads(content)
-                    loaded = True
-                    self.retrieved = datetime.datetime.fromtimestamp(os.path.getmtime(cachefile))
-                    cache_valid = True
-                except json.decoder.JSONDecodeError:
-                    tests += 1
-                    errmsg = "Invalid JSON content in %s" % cachefile
-                    # Delete it without mercy
-                    os.remove(cachefile)
-                    continue
+                if pickleformat:
+                    try:
+                        database = pickle.loads(content)
+                        loaded = True
+                        self.retrieved = datetime.datetime.fromtimestamp(os.path.getmtime(self.cachefile))
+                        cache_valid = True
+                    except (pickle.UnpicklingError, EOFError):
+                        tests += 1
+                        errmsg = "Invalid pickle content in %s" % self.cachefile
+                        # Delete it without mercy
+                        os.remove(self.cachefile)
+                        continue                        
+                else:
+                    try:
+                        database = json.loads(content)
+                        loaded = True
+                        self.retrieved = datetime.datetime.fromtimestamp(os.path.getmtime(self.cachefile))
+                        cache_valid = True
+                    except json.decoder.JSONDecodeError:
+                        tests += 1
+                        errmsg = "Invalid JSON content in %s" % self.cachefile
+                        # Delete it without mercy
+                        os.remove(self.cachefile)
+                        continue
             else:
                 self.unlock()
                 response = requests.get(IANABASE, timeout=IANATIMEOUT)
@@ -81,8 +102,11 @@ class IanaRDAPDatabase():
                     self.services[tld] = server
         if not cache_valid:
             self.lock()
-            cache = open(cachefile, "wb")
-            cache.write(content)
+            cache = open(self.cachefile, "wb")
+            if pickleformat:
+                cache.write(pickle.dumps(database))
+            else:
+                cache.write(content)
             cache.close()
             self.unlock()
         
diff --git a/test_ianardap.py b/test_ianardap.py
index 7a2f449..26d1fb2 100644
--- a/test_ianardap.py
+++ b/test_ianardap.py
@@ -23,7 +23,7 @@ def test_alternative_cache():
         datetime.datetime.fromtimestamp(os.path.getmtime(tmpfile.name)) > \
         (datetime.datetime.now() - datetime.timedelta(minutes=1))
     os.remove(tmpfile.name)
-    os.remove(tmpfile.name + ".lock")
+    os.remove(tmpfile.name + ".json.lock")
 
 def test_refresh():
     # Force a resfresh
@@ -41,3 +41,10 @@ def test_find_not_exists():
     database = ianardap.IanaRDAPDatabase()
     server = database.find("www.foobar.example")
     assert server is None
+
+def test_pickle():
+    database = ianardap.IanaRDAPDatabase(pickleformat=True)
+    assert database.description.startswith("RDAP bootstrap file") and database.version == "1.0" and \
+        len(database.services) > 1000
+    
+