fixed Virus total dependencies

tdagent/tools/virus_total.py

@@ -2,16 +2,24 @@ import gradio as gr
 import vt
 import os
 from datetime import datetime
+from functools import lru_cache
+from typing import Optional
+import time
 
 # Get API key from environment variable
 API_KEY = os.getenv('VT_API_KEY')
 
+# Initialize the VT client globally
+client = vt.Client(API_KEY)
 
-def get_url_info(url):
-    try:
-        # Initialize the client
-        client = vt.Client(API_KEY)
 
+@lru_cache(maxsize=100)
+def get_url_info_cached(url: str, timestamp: Optional[int] = None) -> str:
+    """
+    Cached version of URL info retrieval.
+    timestamp parameter is used to force cache invalidation when needed
+    """
+    try:
         # URL ID is created by computing the base64-encoded SHA-256 of the URL
         url_id = vt.url_id(url)
 
@@ -20,12 +28,20 @@ def get_url_info(url):
 
         # Format the results
         last_analysis_stats = url_analysis.last_analysis_stats
-        last_analysis_date = datetime.utcfromtimestamp(url_analysis.last_analysis_date).strftime(
-            '%Y-%m-%d %H:%M:%S UTC')
+
+        # Fix: Convert the datetime attribute properly
+        last_analysis_date = url_analysis.last_analysis_date
+        if last_analysis_date:
+            # Convert to datetime if it's a timestamp
+            if isinstance(last_analysis_date, (int, float)):
+                last_analysis_date = datetime.utcfromtimestamp(last_analysis_date)
+            date_str = last_analysis_date.strftime('%Y-%m-%d %H:%M:%S UTC')
+        else:
+            date_str = "Not available"
 
         result = f"""
 URL: {url}
-Last Analysis Date: {last_analysis_date}
+Last Analysis Date: {date_str}
 
 Analysis Statistics:
 - Harmless: {last_analysis_stats['harmless']}
@@ -36,15 +52,29 @@ Analysis Statistics:
 
 Reputation Score: {url_analysis.reputation}
 Times Submitted: {url_analysis.times_submitted}
+
+Cache Status: Hit
         """
 
-        client.close()
         return result
 
     except Exception as e:
         return f"Error: {str(e)}"
 
 
+def get_url_info(url: str) -> str:
+    """
+    Wrapper function to handle the cached URL info retrieval
+    """
+    # Clean the URL to ensure consistent caching
+    url = url.strip().lower()
+
+    # Use current timestamp rounded to nearest hour to maintain cache for an hour
+    timestamp = int(time.time()) // 3600
+
+    return get_url_info_cached(url, timestamp)
+
+
 gr_virus_total = gr.Interface(
     fn=get_url_info,
     inputs=["text"],