Fix WHOIS timeout

tdagent/tools/virus_total.py

@@ -6,11 +6,16 @@ import gradio as gr
 import vt
 
 
+_CACHE_MAX_SIZE = 4096
+_CACHE_TTL_SECONDS = 3600
+
 # Get API key from environment variable
 API_KEY = os.getenv("VT_API_KEY")
 
 
-@cachetools.cached(cache=cachetools.TTLCache(maxsize=100, ttl=3600))
+@cachetools.cached(
+    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
+)
 def get_virus_total_url_info(url: str) -> str:
     """Get URL Info from VirusTotal URL Scanner. Scan URL is not available."""
     if not API_KEY:

tdagent/tools/whois.py

@@ -1,11 +1,21 @@
 import json
+import shutil
 
+import cachetools
 import gradio as gr
 import whois
 
 from tdagent.utils.json_utils import TDAgentJsonEncoder
 
 
+_WHOIS_BINARY = "whois"
+_CACHE_MAX_SIZE = 4096
+_CACHE_TTL_SECONDS = 3600
+
+
+@cachetools.cached(
+    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
+)
 def query_whois(url: str) -> str:
     """Query a WHOIS database to gather information about a url or domain.
 
@@ -19,7 +29,11 @@ def query_whois(url: str) -> str:
         A JSON formatted string with the gathered information
     """
     try:
-        whois_result = whois.whois(url)
+        whois_result = whois.whois(
+            url,
+            command=shutil.which(_WHOIS_BINARY) is not None,
+            executable=_WHOIS_BINARY,
+        )
     except whois.parser.PywhoisError as err:
         return json.dumps({"error": str(err)})