tdagent/tools/virus_total.py

import os
from datetime import datetime, timezone

import cachetools
import gradio as gr
import vt


# Get API key from environment variable
API_KEY = os.getenv("VT_API_KEY")


@cachetools.cached(cache=cachetools.TTLCache(maxsize=100, ttl=3600))
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:
        return "Error: Virus total API key not configured."

    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.fromtimestamp(
                        last_analysis_date,
                        timezone.utc,
                    )
                date_str = last_analysis_date.strftime("%Y-%m-%d %H:%M:%S UTC")
            else:
                date_str = "Not available"

            return 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
            """.strip()

    except Exception as err:  # noqa: BLE001
        return f"Error: {err}"


gr_virus_total_url_info = gr.Interface(
    fn=get_virus_total_url_info,
    inputs=["text"],
    outputs="text",
    title="VirusTotal URL Scanner",
    description="Get URL Info from VirusTotal URL Scanner. Scan URL is not available",
)