app.py

from pathlib import Path
from typing import NamedTuple

import gradio as gr
import gradio.themes as gr_themes
import markdown

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.retrieve_from_mitre_attack import gr_get_stix_of_attack_id
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 ##


def _read_markdown_body_as_html(path: str = "README.md") -> str:
    with Path(path).open(encoding="utf-8") as f:  # Default mode is "r"
        lines = f.readlines()

    # Skip YAML front matter if present
    if lines and lines[0].strip() == "---":
        for i in range(1, len(lines)):
            if lines[i].strip() == "---":
                lines = lines[i + 1 :]  # skip metadata block
                break

    markdown_body = "".join(lines).strip()
    return markdown.markdown(markdown_body)


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),
    ToolInfo("ATT&CK STIX information", gr_get_stix_of_attack_id),
    ## 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 ##

custom_css = """
.main-header {
    background: linear-gradient(135deg, #00a388 0%, #ffae00 100%);
    padding: 30px;
    border-radius: 5px;
    margin-bottom: 20px;
    text-align: center;
}
"""
with (
    gr.Blocks(
        theme=gr_themes.Origin(
            primary_hue="teal",
            spacing_size="sm",
            font="sans-serif",
        ),
        title="TDAgent",
        fill_height=True,
        fill_width=True,
        css=custom_css,
    ) as gr_app,
):
    gr.HTML(
        """
    <div class="main-header">
        <h1>👩‍💻 TDAgentTools & TDAgent 👨‍💻</h1>
        <p style="font-size: 1.2em; margin: 10px 0 0 0;">
            Empowering Cybersecurity with Agentic AI
        </p>
    </div>
    """,
    )
    with gr.Tabs():
        with gr.TabItem("About"):
            html_content = _read_markdown_body_as_html("README.md")
            gr.Markdown(html_content)
        with gr.TabItem("TDAgentTools"):
            gr.TabbedInterface(
                interface_list=[t_info.interface for t_info in TOOLS],
                tab_names=[t_info.name for t_info in TOOLS],
                title="TDAgentTools",
            )
        with gr.TabItem("Demo"):
            gr.Markdown(
                """
            This is a demo of TDAgentTools, a simple MCP server.
            Be carefull with using well-known urls for malware distribution
            when using the url content extractor tool.
            """,
            )
            gr.HTML(
                """<iframe width="560" height="315" src="https://youtube.com/embed/c7Yg_jOD6J0" frameborder="0" allowfullscreen></iframe>""",  # noqa: E501
            )

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