From 724378945c52e665bc0c065d191a4aa692e9accf Mon Sep 17 00:00:00 2001
From: tke <tobias.kessels@certbw.de>
Date: Fri, 5 May 2023 11:13:11 +0200
Subject: [PATCH] Changed Diskcache to sqlite with adjustable timeout

---
 config/visidatarc | 61 +++++++++++++++++++++++++++++------------------
 1 file changed, 38 insertions(+), 23 deletions(-)

diff --git a/config/visidatarc b/config/visidatarc
index 6b06b71..c5407ba 100644
--- a/config/visidatarc
+++ b/config/visidatarc
@@ -10,24 +10,39 @@ import json
 from urllib.parse import unquote_plus
 
 import os.path
-import shelve
+import pickle
+import time
+import sqlite3
 
-cache_path = os.path.expanduser('~/.visidata_cache')
+cache_path = os.path.expanduser('~/.visidata_cache.db')
 
-def disk_cache_decorator(func):
-    @functools.lru_cache(maxsize=1000)
-    def get_from_shelve(*args, **kwargs):
-        with shelve.open(cache_path) as cache:
+def init_cache_db():
+    with sqlite3.connect(cache_path) as conn:
+        conn.execute('''CREATE TABLE IF NOT EXISTS cache
+                     (key TEXT PRIMARY KEY, value BLOB, timestamp INTEGER)''')
+
+init_cache_db()
+
+def disk_cache_decorator(max_age=None, lru_cache_size=1000):
+    def decorator(func):
+        @functools.lru_cache(maxsize=lru_cache_size)
+        def get_from_sqlite(*args, **kwargs):
             key = f"{func.__name__}:{str(args)}:{str(kwargs)}"
-            if key in cache:
-                return cache[key]
-            else:
-                result = func(*args, **kwargs)
-                cache[key] = result
-                return result
-    return get_from_shelve
-
-
+            with sqlite3.connect(cache_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute('SELECT value, timestamp FROM cache WHERE key=?', (key,))
+                row = cursor.fetchone()
+                current_time = int(time.time())
+                if row and (max_age is None or current_time - row[1] <= max_age):
+                    return pickle.loads(row[0])
+                else:
+                    result = func(*args, **kwargs)
+                    serialized_value = pickle.dumps(result)
+                    cursor.execute('INSERT OR REPLACE INTO cache (key, value, timestamp) VALUES (?, ?, ?)', (key, serialized_value, current_time))
+                    conn.commit()
+                    return result
+        return get_from_sqlite
+    return decorator
 def decode_url_safe(url_safe_string):
     utf8_string = unquote_plus(url_safe_string)
     return utf8_string
@@ -92,7 +107,7 @@ def sym_time(val):
     return datetime.fromtimestamp(b)
 
 
-@disk_cache_decorator
+@functools.lru_cache(maxsize=1000)
 def vendor(mac):
     try:
         from mac_vendor_lookup import InvalidMacError, MacLookup as mlu
@@ -115,7 +130,7 @@ def _get_vt():
     except:
         return None
 
-@disk_cache_decorator
+@disk_cache_decorator()
 def vt_ip(ip):
     vt = _get_vt()
     if vt is None:
@@ -124,7 +139,7 @@ def vt_ip(ip):
     return response
 
 
-@disk_cache_decorator
+@disk_cache_decorator()
 def vt_file(hash):
     vt = _get_vt()
     if vt is None:
@@ -133,7 +148,7 @@ def vt_file(hash):
     return response
 
 
-@disk_cache_decorator
+@disk_cache_decorator()
 def dns_lookup(domain, record='A'):
     if len(domain.split(",")) > 1:
         return ",".join([dns_lookup(x, record) for x in domain.split(",")])
@@ -150,7 +165,7 @@ def dns_lookup(domain, record='A'):
     except ModuleNotFoundError:
         return "module not available"
 
-@disk_cache_decorator
+@disk_cache_decorator()
 def _asn(ip):
     from bs4 import BeautifulSoup
     import requests
@@ -170,7 +185,7 @@ def _asn(ip):
         res['name']=" ".join(name_split[:-1])
     return res
 
-@disk_cache_decorator
+@disk_cache_decorator()
 def asn(ip, type="asn"):
     if len(ip.split(",")) > 1:
         return ",".join([_asn(x, type) for x in ip.split(",")])
@@ -179,7 +194,7 @@ def asn(ip, type="asn"):
     except:
         return ""
 
-@disk_cache_decorator
+@disk_cache_decorator()
 def _ipinfo(ip):
     try:
         import requests
@@ -220,7 +235,7 @@ def mx_lookup(domain):
         return str(e)
 
 
-@disk_cache_decorator
+@disk_cache_decorator(max_age=60*60*24)
 def _grab_banner(ip, port=25):
     try:
         import socket