tdagent/tools/virus_total.py

import os
from datetime import datetime, timezone

import cachetools
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=_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:
        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=gr.Textbox(label="url"),
    outputs=gr.Text(label="VirusTotal report"),
    title="VirusTotal URL Scanner",
    description="Get URL Info from VirusTotal URL Scanner. Scan URL is not available",
    examples=["https://advertipros.com//?u=script", "https://google.com"],
    example_labels=["👾 Malicious URL", "🧑‍💻 Benign URL"],
)