Pickle implemented (but not activated by default). Closes #1
This commit is contained in:
parent
ae00a81f79
commit
04fb72e408
64
ianardap.py
64
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()
|
||||
|
||||
|
@ -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
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user