tdagent/tools/virus_total.py

import gradio as gr
import vt
import os
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')


@lru_cache(maxsize=100)
def get_url_info_cached(url: str, timestamp: Optional[int] = None) -> str:
    """
    Get URL Info from VirusTotal URL Scanner. Scan URL is not available
    """
    try:
        # 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}

Analysis Statistics:
- Harmless: {last_analysis_stats['harmless']}
- Malicious: {last_analysis_stats['malicious']}
- Suspicious: {last_analysis_stats['suspicious']}
- Undetected: {last_analysis_stats['undetected']}
- Timeout: {last_analysis_stats['timeout']}

Reputation Score: {url_analysis.reputation}
Times Submitted: {url_analysis.times_submitted}

Cache Status: Hit
            """

            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"],
    outputs="text",
    title="VirusTotal URL Scanner",
    description="Get URL Info from VirusTotal URL Scanner. Scan URL is not available",
)