fixed Virus total dependencies

tdagent/tools/virus_total.py

@@ -5,13 +5,11 @@ from datetime import datetime
 from functools import lru_cache
 from typing import Optional
 import time
+import asyncio
 
 # Get API key from environment variable
 API_KEY = os.getenv('VT_API_KEY')
 
-# Initialize the VT client globally
-client = vt.Client(API_KEY)
-
 
 @lru_cache(maxsize=100)
 def get_url_info_cached(url: str, timestamp: Optional[int] = None) -> str:
@@ -20,26 +18,28 @@ def get_url_info_cached(url: str, timestamp: Optional[int] = None) -> str:
     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)
-
-        # Get the URL analysis
-        url_analysis = client.get_object(f"/urls/{url_id}")
-
-        # Format the results
-        last_analysis_stats = url_analysis.last_analysis_stats
-
-        # 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"""
+        # Create a new client for each request (thread-safe)
+        with vt.Client(API_KEY) as client:
+            # URL ID is created by computing the base64-encoded SHA-256 of the URL
+            url_id = vt.url_id(url)
+
+            # Get the URL analysis
+            url_analysis = client.get_object(f"/urls/{url_id}")
+
+            # Format the results
+            last_analysis_stats = url_analysis.last_analysis_stats
+
+            # 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: {date_str}
 
@@ -54,9 +54,9 @@ Reputation Score: {url_analysis.reputation}
 Times Submitted: {url_analysis.times_submitted}
 
 Cache Status: Hit
-        """
+            """
 
-        return result
+            return result
 
     except Exception as e:
         return f"Error: {str(e)}"