app.py

from typing import NamedTuple

import gradio as gr

from tdagent.tools.get_domain_information import (
    dns_enumeration_tool,
    extractor_of_ioc_from_threatfox_tool,
    geo_location_tool,
    scrap_subdomains_tool,
)
from tdagent.tools.get_url_content import gr_get_url_http_content
from tdagent.tools.internal_company_user_search import gr_internal_company
from tdagent.tools.lookup_company_cloud_account_information import (
    gr_lookup_company_cloud_account_information,
)
from tdagent.tools.query_abuse_ip_db import gr_query_abuseipdb
from tdagent.tools.rdap import gr_query_rdap
from tdagent.tools.send_email import gr_send_email
from tdagent.tools.virus_total import gr_virus_total_url_info


# from tdagent.tools.whois import gr_query_whois


## Tools to load into the application interface ##


class ToolInfo(NamedTuple):
    """Gradio MCP tool info."""

    name: str
    interface: gr.Interface


TOOLS = (
    ToolInfo("Get URL Content", gr_get_url_http_content),
    ToolInfo("Query AbuseIPDB", gr_query_abuseipdb),
    # Whois does not work from Spaces (port 43 blocked)
    # ToolInfo("Query WHOIS", gr_query_whois),
    ToolInfo("Query RDAP", gr_query_rdap),
    ToolInfo("Virus Total URL info", gr_virus_total_url_info),
    ToolInfo("Get IP's Location", geo_location_tool),
    ToolInfo("DNS Enumerator", dns_enumeration_tool),
    ToolInfo("Subdomain Retriever", scrap_subdomains_tool),
    ToolInfo("Extractor of IoCs", extractor_of_ioc_from_threatfox_tool),
    ## Fake tools
    ToolInfo("Fake company directory", gr_internal_company),
    ToolInfo(
        "Fake company cloud accounts",
        gr_lookup_company_cloud_account_information,
    ),
    ToolInfo("Send email", gr_send_email),
)

## Application Interface ##
gr_app = gr.TabbedInterface(
    interface_list=[t_info.interface for t_info in TOOLS],
    tab_names=[t_info.name for t_info in TOOLS],
    title="TDAgentTools",
)

if __name__ == "__main__":
    gr_app.launch(mcp_server=True)