tdagent/tools/whois.py

import json
import shutil

import cachetools
import gradio as gr
import whois

from tdagent.utils.json_utils import TDAgentJsonEncoder


_WHOIS_BINARY = "whois"
_CACHE_MAX_SIZE = 4096
_CACHE_TTL_SECONDS = 3600


@cachetools.cached(
    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
)
def query_whois(url: str) -> str:
    """Query a WHOIS database to gather information about a url or domain.

    WHOIS information includes: domain names, IP address blocks and autonomous
    systems, but it is also used for a wider range of other information.

    Args:
        url: URL to query for WHOIS information.

    Returns:
        A JSON formatted string with the gathered information
    """
    try:
        whois_result = whois.whois(
            url,
            command=shutil.which(_WHOIS_BINARY) is not None,
            executable=_WHOIS_BINARY,
        )
    except whois.parser.PywhoisError as err:
        return json.dumps({"error": str(err)})

    return json.dumps(whois_result, cls=TDAgentJsonEncoder)


gr_query_whois = gr.Interface(
    fn=query_whois,
    inputs=["text"],
    outputs="text",
    title="Get WHOIS information for a given URL.",
    description="Query a WHOIS database to gather information about a url or domain.",
)