virustotal.py

"""VirusTotal lookups (online, opt-in via API key).

Looks up an IP or domain against the VirusTotal v3 API and returns its
detection counts (how many engines flagged it as malicious/suspicious).

The API key is **optional**. Supply it either in ``config.ini``::

    [virustotal]
    api_key = <your key>

or via the ``VIRUSTOTAL_API_KEY`` environment variable. If no key is configured
the UI simply shows "API token not added" — nothing breaks, and nothing is sent
to VirusTotal.

Note: a lookup sends the IP/domain to VirusTotal, so it is an *online* action
(like the AbuseIPDB/RDAP enrichment). Results are cached for the process
lifetime to avoid burning the free-tier rate limit.
"""

from __future__ import annotations

import threading

import requests

import config

API = "https://www.virustotal.com/api/v3"

NOT_CONFIGURED = (
    "VirusTotal API token not added — add api_key under [virustotal] in "
    "config.ini (copy config.ini.example) or set VIRUSTOTAL_API_KEY."
)

_cache: dict[str, dict] = {}
_cache_lock = threading.Lock()


def api_key() -> str:
    return config.get("virustotal", "api_key", env="VIRUSTOTAL_API_KEY", default="")


def _timeout() -> float:
    try:
        return float(config.get("virustotal", "timeout", env="VT_TIMEOUT", default="8") or 8)
    except ValueError:
        return 8.0


def available() -> bool:
    return bool(api_key())


def scan_max() -> int:
    """Max indicators a single 'Find Malicious IoCs' scan will query (keeps the
    free-tier rate limit / daily quota in check). Configurable via
    [virustotal] scan_max or VT_SCAN_MAX."""
    try:
        return max(1, int(config.get("virustotal", "scan_max", env="VT_SCAN_MAX", default="40") or 40))
    except ValueError:
        return 40


def lookup(kind: str, value: str) -> dict:
    """Look up an IP ('ip') or domain ('domain') on VirusTotal.

    Returns a dict the frontend renders directly. When no key is configured it
    returns ``{"available": False, "note": ...}`` without making any request."""
    key = api_key()
    if not key:
        return {"available": False, "note": NOT_CONFIGURED}

    endpoint = "domains" if kind == "domain" else "ip_addresses"
    gui = "domain" if kind == "domain" else "ip-address"
    cache_key = f"{endpoint}:{value}"
    with _cache_lock:
        if cache_key in _cache:
            return _cache[cache_key]

    out: dict = {"available": True, "value": value}
    try:
        r = requests.get(
            f"{API}/{endpoint}/{value}",
            headers={"x-apikey": key, "Accept": "application/json"},
            timeout=_timeout(),
        )
        if r.status_code == 404:
            out.update({"found": False, "note": "No VirusTotal record for this indicator."})
        elif r.status_code in (401, 403):
            out = {"available": False,
                   "note": f"VirusTotal rejected the API token (HTTP {r.status_code})."}
        elif r.status_code == 429:
            out["error"] = "VirusTotal rate limit reached (HTTP 429). Try again shortly."
        elif r.ok:
            attrs = (r.json().get("data") or {}).get("attributes", {}) or {}
            stats = attrs.get("last_analysis_stats", {}) or {}
            malicious = int(stats.get("malicious", 0) or 0)
            suspicious = int(stats.get("suspicious", 0) or 0)
            total = sum(int(v) for v in stats.values() if isinstance(v, (int, float)))
            out.update({
                "found": True,
                "malicious": malicious,
                "suspicious": suspicious,
                "harmless": int(stats.get("harmless", 0) or 0),
                "undetected": int(stats.get("undetected", 0) or 0),
                "total": total,
                "reputation": attrs.get("reputation"),
                "permalink": f"https://www.virustotal.com/gui/{gui}/{value}",
            })
        else:
            out["error"] = f"VirusTotal returned HTTP {r.status_code}."
    except Exception as exc:  # noqa: BLE001
        out["error"] = str(exc)

    # Only cache definitive answers (found True/False) — never transient errors.
    if "error" not in out and out.get("available"):
        with _cache_lock:
            _cache[cache_key] = out
    return out