tools: Add a caching layer to WHOIS

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>

diff --git a/tools/base.py b/tools/base.py
index 7ee3978f479c34c3819973f668cc4249010757ea..0b8a068c8bb270ac902685f753d50e078eeed8ef 100644 (file)
--- a/tools/base.py
+++ b/tools/base.py
@@ -53,7 +53,7 @@ class RIR(object):
        whois_server = None
 
        def __init__(self):
-               pass
+               self.db = self._open_database(".cache.db")
 
        def __str__(self):
                if self.name:
@@ -61,6 +61,20 @@ class RIR(object):
 
                return self.__class__.__name__
 
+       def _open_database(self, path=None):
+               db = sqlite3.connect(path or ":memory:")
+               db.set_trace_callback(logging.debug)
+
+               # Create tables
+               db.executescript("""
+                       CREATE TABLE IF NOT EXISTS whois_query_cache(query TEXT, response TEXT,
+                               fetched_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP);
+                       CREATE UNIQUE INDEX IF NOT EXISTS whois_query_cache_query
+                               ON whois_query_cache(query);
+               """)
+
+               return db
+
        @property
        def parser(self):
                return RIRParser
@@ -85,7 +99,7 @@ class RIR(object):
                # Write the database to disk
                p.export_database(directory)
 
-       def whois(self, query):
+       def _whois(self, query):
                command = [
                        "whois", query,
                ]
@@ -104,6 +118,34 @@ class RIR(object):
 
                return output.decode(errors="ignore")
 
+       def whois(self, query):
+               # Try fetching a response from the cache
+               with self.db as c:
+                       res = c.execute("SELECT response, fetched_at FROM whois_query_cache \
+                               WHERE query = ?", (query,))
+
+                       # Return any results
+                       for row in res:
+                               response, fetched_at = row
+
+                               logging.debug("Fetched response for %s from cache (%s)"
+                                       % (query, fetched_at))
+
+                               return response
+
+               # If we could not find anything, we will have to contact the whois server
+               response = self._whois(query)
+
+               # Store the response in the database
+               with self.db as c:
+                       c.execute("INSERT INTO whois_query_cache(query, response) \
+                               VALUES(?, ?)", (query, response))
+
+               # Commit changes to disk
+               self.db.commit()
+
+               return response
+
        def get_name_for_asn(self, asn):
                result = self.whois("AS%s" % asn)