tdagent/tools/rdap.py

import enum

import cachetools
import gradio as gr
import requests
import whois

from tdagent.constants import HttpContentType


# one of domain, ip, autnum, entity etc
_RDAP_URL_TEMPLATE = r"https://rdap.org/{rdap_type}/{rdap_object}"
_CACHE_MAX_SIZE = 4096
_CACHE_TTL_SECONDS = 3600


class RdapTypes(str, enum.Enum):
    """RDAP object types."""

    DOMAIN = "domain"
    IP = "ip"
    AUTNUM = "autnum"
    ENTITY = "entity"


@cachetools.cached(
    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
)
def query_rdap(  # noqa: PLR0911
    url_or_ip: str,
    timeout: int = 30,
) -> dict[str, str | int | float]:
    """Query RDAP to get information about Internet resources.

    The Registration Data Access Protocol (RDAP) is the successor to WHOIS.
    Like WHOIS, RDAP provides access to information about Internet resources
    (domain names, autonomous system numbers, and IP addresses).

    Args:
        url_or_ip: URL, domain or IP to query for RDAP information.
        timeout: Request timeout in seconds. Defaults to 30.

    Returns:
        A JSON formatted string with RDAP information. In there is
        an error, the JSON will contain the key "error" with an
        error message.
    """
    rdap_type = RdapTypes.DOMAIN
    rdap_object = url_or_ip
    if whois.IPV4_OR_V6.match(url_or_ip):
        rdap_type = RdapTypes.IP
    else:
        rdap_object = whois.extract_domain(url_or_ip)

    query_url = _RDAP_URL_TEMPLATE.format(rdap_type=rdap_type, rdap_object=rdap_object)
    response = requests.get(
        query_url,
        timeout=timeout,
        headers={"Accept": HttpContentType.JSON},
    )

    try:
        response.raise_for_status()
    except requests.HTTPError as err:
        if err.response.status_code == 302:
            if "Location" in err.response.headers:
                return {
                    "message": "Follow the location to find RDAP information",
                    "location": err.response.headers["Location"],
                }
            return {
                "error": (
                    "Information not found in RDAP.org but it knows of"
                    " a service which is authoritative for the requested resource."
                ),
            }
        if err.response.status_code == 400:
            return {
                "error": (
                    "Invalid request (malformed path, unsupported object "
                    " type, invalid IP address, etc)"
                ),
            }
        if err.response.status_code == 403:
            return {
                "error": "You've been blocked due to abuse or other misbehavior",
            }
        if err.response.status_code == 404:
            return {
                "error": (
                    "RDAP.org doesn't know of an RDAP service which is"
                    " authoritative for the requested resource. RDAP.org"
                    " only knows about servers that are registered with IANA"
                ),
            }
        return {
            "error": str(err),
        }

    return response.json()


gr_query_rdap = gr.Interface(
    fn=query_rdap,
    inputs=gr.Textbox(label="url or ip"),
    outputs=gr.JSON(label="Report from RDAP"),
    title="Get RDAP information for a given URL.",
    description="Query a RDAP database to gather information about a url or domain.",
    examples=["8.8.8.8", "pastebin.com"],
)