WHOIS and fixes

app.py

@@ -1,23 +1,47 @@
+from typing import NamedTuple
+
 import gradio as gr
 
 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.letter_counter import gr_letter_counter
-from tdagent.tools.lookup_company_cloud_account_information import gr_lookup_company_cloud_account_information
+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.send_email import gr_send_email
-from tdagent.tools.virus_total import gr_virus_total
+from tdagent.tools.virus_total import gr_virus_total_url_info
+from tdagent.tools.whois import gr_query_whois
 
-gr_app = gr.TabbedInterface(
-    [
-        gr_get_url_http_content,
-        gr_letter_counter,
-        gr_query_abuseipdb,
-        gr_virus_total,
-        gr_internal_company,
+
+## 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),
+    ToolInfo("Query WHOIS", gr_query_whois),
+    ToolInfo("Virus Total URL info", gr_virus_total_url_info),
+    ## Fake tools
+    ToolInfo("Fake company directory", gr_internal_company),
+    ToolInfo(
+        "Fake company cloud accounts",
         gr_lookup_company_cloud_account_information,
-        gr_send_email,
-    ],
+    ),
+    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__":

packages.txt

@@ -0,0 +1 @@
+whois

pyproject.toml

@@ -12,7 +12,13 @@ authors = [
 requires-python = ">=3.10,<4"
 readme = "README.md"
 license = ""
-dependencies = ["gradio[mcp]>=5.32.1", "requests>=2.32.3", "vt-py~=0.21.0"]
+dependencies = [
+    "cachetools>=6.0.0",
+    "gradio[mcp]>=5.32.1",
+    "python-whois>=0.9.5",
+    "requests>=2.32.3",
+    "vt-py~=0.21.0",
+]
 
 [project.scripts]
 
@@ -98,7 +104,7 @@ line-length = 88
 
 [tool.ruff.lint]
 select = ["ALL"]
-ignore = ["D100", "D104", "D107", "D401", "EM102", "ERA001", "TRY003"]
+ignore = ["D100", "D104", "D107", "D401", "EM102", "ERA001", "TRY003", "UP038"]
 
 [tool.ruff.lint.flake8-quotes]
 inline-quotes = "double"

requirements-dev.txt

@@ -1,301 +1,109 @@
 # This file was autogenerated by uv via the following command:
 #    uv export --format requirements-txt --no-hashes --group dev --group test -o requirements-dev.txt
 aiofiles==24.1.0
-    # via
-    #   gradio
-    #   vt-py
 aiohappyeyeballs==2.6.1
-    # via aiohttp
 aiohttp==3.12.8
-    # via vt-py
 aiosignal==1.3.2
-    # via aiohttp
 annotated-types==0.7.0
-    # via pydantic
 anyio==4.9.0
-    # via
-    #   gradio
-    #   httpx
-    #   mcp
-    #   sse-starlette
-    #   starlette
 async-timeout==5.0.1 ; python_full_version < '3.11'
-    # via aiohttp
 attrs==25.3.0
-    # via aiohttp
 audioop-lts==0.2.1 ; python_full_version >= '3.13'
-    # via gradio
 boolean-py==5.0
-    # via license-expression
 cachecontrol==0.14.3
-    # via pip-audit
+cachetools==6.0.0
 certifi==2025.4.26
-    # via
-    #   httpcore
-    #   httpx
-    #   requests
 cfgv==3.4.0
-    # via pre-commit
 charset-normalizer==3.4.2
-    # via requests
 click==8.2.1 ; sys_platform != 'emscripten'
-    # via
-    #   typer
-    #   uvicorn
 colorama==0.4.6 ; sys_platform == 'win32'
-    # via
-    #   click
-    #   pytest
-    #   tqdm
 coverage==7.8.2
-    # via pytest-cov
 cyclonedx-python-lib==9.1.0
-    # via pip-audit
 defusedxml==0.7.1
-    # via py-serializable
 distlib==0.3.9
-    # via virtualenv
 exceptiongroup==1.3.0 ; python_full_version < '3.11'
-    # via
-    #   anyio
-    #   pytest
 fastapi==0.115.12
-    # via gradio
 ffmpy==0.6.0
-    # via gradio
 filelock==3.18.0
-    # via
-    #   cachecontrol
-    #   huggingface-hub
-    #   virtualenv
 frozenlist==1.6.2
-    # via
-    #   aiohttp
-    #   aiosignal
 fsspec==2025.5.1
-    # via
-    #   gradio-client
-    #   huggingface-hub
 gradio==5.32.1
-    # via tdagent
 gradio-client==1.10.2
-    # via gradio
 groovy==0.1.2
-    # via gradio
 h11==0.16.0
-    # via
-    #   httpcore
-    #   uvicorn
 hf-xet==1.1.2 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
-    # via huggingface-hub
 httpcore==1.0.9
-    # via httpx
 httpx==0.28.1
-    # via
-    #   gradio
-    #   gradio-client
-    #   mcp
-    #   safehttpx
 httpx-sse==0.4.0
-    # via mcp
 huggingface-hub==0.32.4
-    # via
-    #   gradio
-    #   gradio-client
 identify==2.6.12
-    # via pre-commit
 idna==3.10
-    # via
-    #   anyio
-    #   httpx
-    #   requests
-    #   yarl
 iniconfig==2.1.0
-    # via pytest
 jinja2==3.1.6
-    # via gradio
 license-expression==30.4.1
-    # via cyclonedx-python-lib
 markdown-it-py==3.0.0
-    # via rich
 markupsafe==3.0.2
-    # via
-    #   gradio
-    #   jinja2
 mcp==1.9.0
-    # via gradio
 mdurl==0.1.2
-    # via markdown-it-py
 msgpack==1.1.0
-    # via cachecontrol
 multidict==6.4.4
-    # via
-    #   aiohttp
-    #   yarl
 mypy==1.16.0
 mypy-extensions==1.1.0
-    # via mypy
 nodeenv==1.9.1
-    # via pre-commit
 numpy==2.2.6
-    # via
-    #   gradio
-    #   pandas
 orjson==3.10.18
-    # via gradio
 packageurl-python==0.16.0
-    # via cyclonedx-python-lib
 packaging==25.0
-    # via
-    #   gradio
-    #   gradio-client
-    #   huggingface-hub
-    #   pip-audit
-    #   pip-requirements-parser
-    #   pytest
 pandas==2.2.3
-    # via gradio
 pathspec==0.12.1
-    # via mypy
 pillow==11.2.1
-    # via gradio
 pip==25.1.1
-    # via pip-api
 pip-api==0.0.34
-    # via pip-audit
 pip-audit==2.9.0
 pip-requirements-parser==32.0.1
-    # via pip-audit
 platformdirs==4.3.8
-    # via
-    #   pip-audit
-    #   virtualenv
 pluggy==1.6.0
-    # via pytest
 pre-commit==3.8.0
 propcache==0.3.1
-    # via
-    #   aiohttp
-    #   yarl
 py-serializable==2.0.0
-    # via cyclonedx-python-lib
 pydantic==2.11.5
-    # via
-    #   fastapi
-    #   gradio
-    #   mcp
-    #   pydantic-settings
 pydantic-core==2.33.2
-    # via pydantic
 pydantic-settings==2.9.1
-    # via mcp
 pydub==0.25.1
-    # via gradio
 pygments==2.19.1
-    # via rich
 pyparsing==3.2.3
-    # via pip-requirements-parser
 pytest==7.4.4
-    # via
-    #   pytest-cov
-    #   pytest-randomly
 pytest-cov==4.1.0
 pytest-randomly==3.16.0
 python-dateutil==2.9.0.post0
-    # via pandas
 python-dotenv==1.1.0
-    # via pydantic-settings
 python-multipart==0.0.20
-    # via
-    #   gradio
-    #   mcp
+python-whois==0.9.5
 pytz==2025.2
-    # via pandas
 pyyaml==6.0.2
-    # via
-    #   gradio
-    #   huggingface-hub
-    #   pre-commit
 requests==2.32.3
-    # via
-    #   cachecontrol
-    #   huggingface-hub
-    #   pip-audit
-    #   tdagent
 rich==14.0.0
-    # via
-    #   pip-audit
-    #   typer
 ruff==0.11.12
-    # via gradio
 safehttpx==0.1.6
-    # via gradio
 semantic-version==2.10.0
-    # via gradio
 shellingham==1.5.4 ; sys_platform != 'emscripten'
-    # via typer
 six==1.17.0
-    # via python-dateutil
 sniffio==1.3.1
-    # via anyio
 sortedcontainers==2.4.0
-    # via cyclonedx-python-lib
 sse-starlette==2.3.6
-    # via mcp
 starlette==0.46.2
-    # via
-    #   fastapi
-    #   gradio
-    #   mcp
 toml==0.10.2
-    # via pip-audit
 tomli==2.2.1 ; python_full_version <= '3.11'
-    # via
-    #   coverage
-    #   mypy
-    #   pytest
 tomlkit==0.13.2
-    # via gradio
 tqdm==4.67.1
-    # via huggingface-hub
 typer==0.16.0 ; sys_platform != 'emscripten'
-    # via gradio
 typing-extensions==4.14.0
-    # via
-    #   anyio
-    #   exceptiongroup
-    #   fastapi
-    #   gradio
-    #   gradio-client
-    #   huggingface-hub
-    #   multidict
-    #   mypy
-    #   pydantic
-    #   pydantic-core
-    #   rich
-    #   typer
-    #   typing-inspection
-    #   uvicorn
 typing-inspection==0.4.1
-    # via
-    #   pydantic
-    #   pydantic-settings
 tzdata==2025.2
-    # via pandas
 urllib3==2.4.0
-    # via
-    #   gradio
-    #   requests
 uvicorn==0.34.3 ; sys_platform != 'emscripten'
-    # via
-    #   gradio
-    #   mcp
 virtualenv==20.31.2
-    # via pre-commit
 vt-py==0.21.0
-    # via tdagent
 websockets==15.0.1
-    # via gradio-client
 xdoctest==1.2.0
 yarl==1.20.0
-    # via aiohttp

requirements.txt

@@ -1,211 +1,84 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile pyproject.toml -o requirements.txt
+#    uv export --format requirements-txt --no-hashes --no-dev -o requirements.txt
 aiofiles==24.1.0
-    # via
-    #   gradio
-    #   vt-py
 aiohappyeyeballs==2.6.1
-    # via aiohttp
 aiohttp==3.12.8
-    # via vt-py
 aiosignal==1.3.2
-    # via aiohttp
 annotated-types==0.7.0
-    # via pydantic
 anyio==4.9.0
-    # via
-    #   gradio
-    #   httpx
-    #   mcp
-    #   sse-starlette
-    #   starlette
-async-timeout==5.0.1
-    # via aiohttp
+async-timeout==5.0.1 ; python_full_version < '3.11'
 attrs==25.3.0
-    # via aiohttp
+audioop-lts==0.2.1 ; python_full_version >= '3.13'
+cachetools==6.0.0
 certifi==2025.4.26
-    # via
-    #   httpcore
-    #   httpx
-    #   requests
 charset-normalizer==3.4.2
-    # via requests
-click==8.2.1
-    # via
-    #   typer
-    #   uvicorn
-exceptiongroup==1.3.0
-    # via anyio
+click==8.2.1 ; sys_platform != 'emscripten'
+colorama==0.4.6 ; sys_platform == 'win32'
+coverage==7.8.2
+exceptiongroup==1.3.0 ; python_full_version < '3.11'
 fastapi==0.115.12
-    # via gradio
 ffmpy==0.6.0
-    # via gradio
 filelock==3.18.0
-    # via huggingface-hub
 frozenlist==1.6.2
-    # via
-    #   aiohttp
-    #   aiosignal
 fsspec==2025.5.1
-    # via
-    #   gradio-client
-    #   huggingface-hub
 gradio==5.32.1
-    # via tdagent (pyproject.toml)
 gradio-client==1.10.2
-    # via gradio
 groovy==0.1.2
-    # via gradio
 h11==0.16.0
-    # via
-    #   httpcore
-    #   uvicorn
-hf-xet==1.1.2
-    # via huggingface-hub
+hf-xet==1.1.2 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
 httpcore==1.0.9
-    # via httpx
 httpx==0.28.1
-    # via
-    #   gradio
-    #   gradio-client
-    #   mcp
-    #   safehttpx
 httpx-sse==0.4.0
-    # via mcp
 huggingface-hub==0.32.4
-    # via
-    #   gradio
-    #   gradio-client
 idna==3.10
-    # via
-    #   anyio
-    #   httpx
-    #   requests
-    #   yarl
+iniconfig==2.1.0
 jinja2==3.1.6
-    # via gradio
-markdown-it-py==3.0.0
-    # via rich
+markdown-it-py==3.0.0 ; sys_platform != 'emscripten'
 markupsafe==3.0.2
-    # via
-    #   gradio
-    #   jinja2
 mcp==1.9.0
-    # via gradio
-mdurl==0.1.2
-    # via markdown-it-py
+mdurl==0.1.2 ; sys_platform != 'emscripten'
 multidict==6.4.4
-    # via
-    #   aiohttp
-    #   yarl
 numpy==2.2.6
-    # via
-    #   gradio
-    #   pandas
 orjson==3.10.18
-    # via gradio
 packaging==25.0
-    # via
-    #   gradio
-    #   gradio-client
-    #   huggingface-hub
 pandas==2.2.3
-    # via gradio
 pillow==11.2.1
-    # via gradio
+pluggy==1.6.0
 propcache==0.3.1
-    # via
-    #   aiohttp
-    #   yarl
 pydantic==2.11.5
-    # via
-    #   fastapi
-    #   gradio
-    #   mcp
-    #   pydantic-settings
 pydantic-core==2.33.2
-    # via pydantic
 pydantic-settings==2.9.1
-    # via mcp
 pydub==0.25.1
-    # via gradio
-pygments==2.19.1
-    # via rich
+pygments==2.19.1 ; sys_platform != 'emscripten'
+pytest==7.4.4
+pytest-cov==4.1.0
+pytest-randomly==3.16.0
 python-dateutil==2.9.0.post0
-    # via pandas
 python-dotenv==1.1.0
-    # via pydantic-settings
 python-multipart==0.0.20
-    # via
-    #   gradio
-    #   mcp
+python-whois==0.9.5
 pytz==2025.2
-    # via pandas
 pyyaml==6.0.2
-    # via
-    #   gradio
-    #   huggingface-hub
 requests==2.32.3
-    # via
-    #   tdagent (pyproject.toml)
-    #   huggingface-hub
-rich==14.0.0
-    # via typer
-ruff==0.11.12
-    # via gradio
+rich==14.0.0 ; sys_platform != 'emscripten'
+ruff==0.11.12 ; sys_platform != 'emscripten'
 safehttpx==0.1.6
-    # via gradio
 semantic-version==2.10.0
-    # via gradio
-shellingham==1.5.4
-    # via typer
+shellingham==1.5.4 ; sys_platform != 'emscripten'
 six==1.17.0
-    # via python-dateutil
 sniffio==1.3.1
-    # via anyio
 sse-starlette==2.3.6
-    # via mcp
 starlette==0.46.2
-    # via
-    #   fastapi
-    #   gradio
-    #   mcp
+tomli==2.2.1 ; python_full_version <= '3.11'
 tomlkit==0.13.2
-    # via gradio
 tqdm==4.67.1
-    # via huggingface-hub
-typer==0.16.0
-    # via gradio
+typer==0.16.0 ; sys_platform != 'emscripten'
 typing-extensions==4.14.0
-    # via
-    #   anyio
-    #   exceptiongroup
-    #   fastapi
-    #   gradio
-    #   gradio-client
-    #   huggingface-hub
-    #   multidict
-    #   pydantic
-    #   pydantic-core
-    #   rich
-    #   typer
-    #   typing-inspection
-    #   uvicorn
 typing-inspection==0.4.1
-    # via
-    #   pydantic
-    #   pydantic-settings
 tzdata==2025.2
-    # via pandas
 urllib3==2.4.0
-    # via requests
-uvicorn==0.34.3
-    # via
-    #   gradio
-    #   mcp
+uvicorn==0.34.3 ; sys_platform != 'emscripten'
 vt-py==0.21.0
-    # via tdagent (pyproject.toml)
 websockets==15.0.1
-    # via gradio-client
+xdoctest==1.2.0
 yarl==1.20.0
-    # via aiohttp

tdagent/tools/get_url_content.py

@@ -37,11 +37,15 @@ def get_url_http_content(
 
     if content_type:
         headers["Accept"] = ",".join(content_type)
-    response = requests.get(
-        url,
-        headers=headers,
-        timeout=timeout,
-    )
+
+    try:
+        response = requests.get(
+            url,
+            headers=headers,
+            timeout=timeout,
+        )
+    except requests.exceptions.MissingSchema as err:
+        return "", str(err)
 
     try:
         response.raise_for_status()

tdagent/tools/internal_company_user_search.py

@@ -1,5 +1,6 @@
 import gradio as gr
 
+
 # Fake user database
 users_db = {
     "jsmith": {
@@ -9,7 +10,7 @@ users_db = {
         "user_id": "US789456",
         "jobtitle": "Software Engineer",
         "department": "Engineering",
-        "country": "United States"
+        "country": "United States",
     },
     "mhacker": {
         "username": "mhacker",
@@ -18,16 +19,19 @@ users_db = {
         "user_id": "US123789",
         "jobtitle": "Security Specialist",
         "department": "Pentests",
-        "country": "Germany"
-    }
+        "country": "Germany",
+    },
 }
 
 
-def lookup_user(username):
-    """
-    Function to lookup user information.
+def lookup_user(username: str) -> str:
+    """Function to lookup user information.
+
     Company User Lookup System. Enter a username to get user details.
-    Returns a formatted string with user details if found, otherwise returns error message
+
+    Returns:
+        A formatted string with user details if found, otherwise an
+        error message.
     """
     if username in users_db:
         user = users_db[username]
@@ -39,8 +43,8 @@ User ID: {user['user_id']}
 Job Title: {user['jobtitle']}
 Department: {user['department']}
 Country: {user['country']}"""
-    else:
-        return """User Not Found
+
+    return """User Not Found
 Username: Not found
 Email: N/A
 Name: N/A
@@ -56,6 +60,6 @@ gr_internal_company = gr.Interface(
     inputs=["text"],
     outputs=["text"],
     title="Company User Lookup System",
-    description="Company User Lookup System. Enter a username to get user details",
-    theme="default"
+    description="Company User Lookup System.",
+    theme="default",
 )

tdagent/tools/lookup_company_cloud_account_information.py

@@ -1,5 +1,6 @@
 import gradio as gr
 
+
 # Fake cloud accounts database
 cloud_accounts = {
     ("AWS", "123456789012"): {
@@ -8,7 +9,7 @@ cloud_accounts = {
         "cloud_account_name": "Production-Main",
         "owner_user_id": "AWS001",
         "owner_email": "cloud.admin@company.com",
-        "deployed_applications": ["ERP System", "Customer Portal", "Data Lake"]
+        "deployed_applications": ["ERP System", "Customer Portal", "Data Lake"],
     },
     ("AWS", "098765432109"): {
         "public_cloud_provider": "AWS",
@@ -16,7 +17,7 @@ cloud_accounts = {
         "cloud_account_name": "Development-Team1",
         "owner_user_id": "AWS002",
         "owner_email": "dev.lead@company.com",
-        "deployed_applications": ["Test Environment", "CI/CD Pipeline"]
+        "deployed_applications": ["Test Environment", "CI/CD Pipeline"],
     },
     ("Azure", "sub-abc-123-def-456"): {
         "public_cloud_provider": "Azure",
@@ -24,7 +25,7 @@ cloud_accounts = {
         "cloud_account_name": "Enterprise-Solutions",
         "owner_user_id": "AZ001",
         "owner_email": "azure.admin@company.com",
-        "deployed_applications": ["Microsoft 365 Integration", "Azure AD Connect"]
+        "deployed_applications": ["Microsoft 365 Integration", "Azure AD Connect"],
     },
     ("Azure", "sub-xyz-789-uvw-321"): {
         "public_cloud_provider": "Azure",
@@ -32,7 +33,7 @@ cloud_accounts = {
         "cloud_account_name": "Research-Division",
         "owner_user_id": "AZ002",
         "owner_email": "research.lead@company.com",
-        "deployed_applications": ["ML Platform", "Research Portal"]
+        "deployed_applications": ["ML Platform", "Research Portal"],
     },
     ("GCP", "project-id-123456"): {
         "public_cloud_provider": "GCP",
@@ -40,7 +41,7 @@ cloud_accounts = {
         "cloud_account_name": "Analytics-Platform",
         "owner_user_id": "GCP001",
         "owner_email": "gcp.admin@company.com",
-        "deployed_applications": ["BigQuery Data Warehouse", "Kubernetes Cluster"]
+        "deployed_applications": ["BigQuery Data Warehouse", "Kubernetes Cluster"],
     },
     ("GCP", "project-id-789012"): {
         "public_cloud_provider": "GCP",
@@ -48,45 +49,49 @@ cloud_accounts = {
         "cloud_account_name": "ML-Operations",
         "owner_user_id": "GCP002",
         "owner_email": "ml.ops@company.com",
-        "deployed_applications": ["TensorFlow Training", "Model Serving Platform"]
-    }
+        "deployed_applications": ["TensorFlow Training", "Model Serving Platform"],
+    },
 }
 
 
-def lookup_cloud_account(public_cloud_provider, account_id):
-    """
-    Function to lookup cloud account information
-    Returns a formatted string with account details if found, otherwise returns error message
+def lookup_cloud_account(public_cloud_provider: str, account_id: str) -> dict[str, str]:
+    """Function to lookup cloud account information.
+
+    Returns a formatted string with account details if
+    found, otherwise returns error message.
     """
     account_key = (public_cloud_provider, account_id)
 
     if account_key in cloud_accounts:
         account = cloud_accounts[account_key]
-        return f"""{{
-    "public_cloud_provider": "{account['public_cloud_provider']}",
-    "account_id": "{account['account_id']}",
-    "cloud_account_name": "{account['cloud_account_name']}",
-    "owner_user_id": "{account['owner_user_id']}",
-    "owner_email": "{account['owner_email']}",
-    "deployed_applications": {str(account['deployed_applications'])}
-}}"""
-    else:
-        return f"""{{
-    "error": "Account not found",
-    "public_cloud_provider": "{public_cloud_provider}",
-    "account_id": "{account_id}",
-    "message": "No cloud account information found"
-}}"""
+        return {
+            "public_cloud_provider": f"{account['public_cloud_provider']}",
+            "account_id": f"{account['account_id']}",
+            "cloud_account_name": f"{account['cloud_account_name']}",
+            "owner_user_id": f"{account['owner_user_id']}",
+            "owner_email": f"{account['owner_email']}",
+            "deployed_applications": f"{account['deployed_applications']!s}",
+        }
+
+    return {
+        "error": "Account not found",
+        "public_cloud_provider": f"{public_cloud_provider}",
+        "account_id": f"{account_id}",
+        "message": "No cloud account information found",
+    }
 
 
 # Create Gradio Interface
 gr_lookup_company_cloud_account_information = gr.Interface(
     fn=lookup_cloud_account,
-    inputs=[gr.Dropdown(
-        choices=["AWS", "Azure", "GCP"],
-        label="Cloud Provider",
-        info="Select the cloud provider"
-    ), "text"],
+    inputs=[
+        gr.Dropdown(
+            choices=["AWS", "Azure", "GCP"],
+            label="Cloud Provider",
+            info="Select the cloud provider",
+        ),
+        "text",
+    ],
     outputs="text",
     title="Company Cloud Account Lookup System",
     description="""
@@ -97,5 +102,5 @@ gr_lookup_company_cloud_account_information = gr.Interface(
     Azure: sub-abc-123-def-456, sub-xyz-789-uvw-321
     GCP: project-id-123456, project-id-789012
     """,
-    theme="default"
+    theme="default",
 )

tdagent/tools/query_abuse_ip_db.py

@@ -1,16 +1,19 @@
-import os
+from __future__ import annotations
 
-import requests
-from datetime import datetime
+import os
 from dataclasses import dataclass
-from typing import List, Optional, Dict, Any, Tuple, Union
+from datetime import datetime
+from typing import Any
+
 import gradio as gr
+import requests
 
 
 # API docs: https://docs.abuseipdb.com/#check-endpoint
 
+
 @dataclass
-class AbuseReport:
+class AbuseReport:  # noqa: D101
     date: str
     categories: str
     comment: str
@@ -18,7 +21,7 @@ class AbuseReport:
 
 
 @dataclass
-class IPCheckResult:
+class IPCheckResult:  # noqa: D101
     ip_address: str
     abuse_confidence_score: int
     total_reports: int
@@ -26,10 +29,10 @@ class IPCheckResult:
     domain: str
     isp: str
     last_reported: str
-    reports: List[AbuseReport]
+    reports: list[AbuseReport]
 
     def format_summary(self) -> str:
-        """Format a summary of the IP check result"""
+        """Format a summary of the IP check result."""
         return f"""
         IP Address: {self.ip_address}
         Abuse Confidence Score: {self.abuse_confidence_score}%
@@ -41,9 +44,8 @@ class IPCheckResult:
         """
 
 
-def check_ip(ip_address: str, api_key: str, days: str = "30") -> Dict[str, Any]:
-    """
-    Query the AbuseIPDB API to check if an IP address has been reported.
+def check_ip(ip_address: str, api_key: str, days: str = "30") -> dict[str, Any]:
+    """Query the AbuseIPDB API to check if an IP address has been reported.
 
     Args:
         ip_address: The IP address to check
@@ -53,30 +55,33 @@ def check_ip(ip_address: str, api_key: str, days: str = "30") -> Dict[str, Any]:
     Returns:
         API response data as a dictionary
     """
-    url = 'https://api.abuseipdb.com/api/v2/check'
+    url = "https://api.abuseipdb.com/api/v2/check"
 
-    headers = {
-        'Accept': 'application/json',
-        'Key': api_key
-    }
+    headers = {"Accept": "application/json", "Key": api_key}
 
     params = {
-        'ipAddress': ip_address,
-        'maxAgeInDays': days,
-        'verbose': True
+        "ipAddress": ip_address,
+        "maxAgeInDays": days,
+        "verbose": str(True),
     }
 
     try:
-        response = requests.get(url, headers=headers, params=params)
+        response = requests.get(
+            url,
+            headers=headers,
+            params=params,
+            timeout=30,
+        )
         response.raise_for_status()  # Raise exception for HTTP errors
         return response.json()
     except requests.exceptions.RequestException as e:
         return {"error": str(e)}
 
 
-def parse_response(response: Dict[str, Any]) -> Tuple[Optional[IPCheckResult], Optional[str]]:
-    """
-    Parse the API response into a dataclass
+def parse_response(
+    response: dict[str, Any],
+) -> tuple[IPCheckResult | None, str]:
+    """Parse the API response into a dataclass.
 
     Args:
         response: The API response dictionary
@@ -93,18 +98,21 @@ def parse_response(response: Dict[str, Any]) -> Tuple[Optional[IPCheckResult], O
 
     # Create a list of AbuseReport objects
     reports = []
-    if "reports" in data and data["reports"]:
+    if data.get("reports"):
         for report in data["reports"]:
-            reported_at = datetime.fromisoformat(report["reportedAt"].replace("Z", "+00:00")).strftime(
-                "%Y-%m-%d %H:%M:%S")
+            reported_at = datetime.fromisoformat(
+                report["reportedAt"].replace("Z", "+00:00"),
+            ).strftime("%Y-%m-%d %H:%M:%S")
             categories = ", ".join([str(cat) for cat in report.get("categories", [])])
 
-            reports.append(AbuseReport(
-                date=reported_at,
-                categories=categories,
-                comment=report.get("comment", ""),
-                reporter=str(report.get("reporterId", "Anonymous"))
-            ))
+            reports.append(
+                AbuseReport(
+                    date=reported_at,
+                    categories=categories,
+                    comment=report.get("comment", ""),
+                    reporter=str(report.get("reporterId", "Anonymous")),
+                ),
+            )
 
     # Create the main result object
     result = IPCheckResult(
@@ -115,15 +123,14 @@ def parse_response(response: Dict[str, Any]) -> Tuple[Optional[IPCheckResult], O
         domain=data.get("domain", "N/A"),
         isp=data.get("isp", "N/A"),
         last_reported=data.get("lastReportedAt", "Never"),
-        reports=reports
+        reports=reports,
     )
 
-    return result, None
+    return result, ""
 
 
 def query_abuseipdb(ip_address: str, days: int = 30) -> str:
-    """
-    Main function to query AbuseIPDB and format the response for Gradio
+    """Query AbuseIP to find if an IP has been reported for abusive behavior.
 
     Args:
         ip_address: The IP address to check
@@ -141,10 +148,10 @@ def query_abuseipdb(ip_address: str, days: int = 30) -> str:
     response = check_ip(ip_address, api_key, str(days))
     result, error = parse_response(response)
 
-    if error:
-        return error
+    if result:
+        return result.format_summary()
 
-    return result.format_summary()
+    return error
 
 
 gr_query_abuseipdb = gr.Interface(
@@ -152,5 +159,8 @@ gr_query_abuseipdb = gr.Interface(
     inputs=["text"],
     outputs="text",
     title="AbuseIPDB IP Checker",
-    description="Check if an IP address has been reported for abusive behavior using AbuseIP DB API",
+    description=(
+        "Check if an IP address has been reported for abusive behavior"
+        " using AbuseIP DB API"
+    ),
 )

tdagent/tools/send_email.py

@@ -1,53 +1,48 @@
-import gradio as gr
 import datetime
 import random
 
+import gradio as gr
+
 
-def send_email(recipient, subject, message):
-    """
-    Simulates sending an email from cert@company.com and prints it to the console in a well-formatted way.
+def send_email(recipient: str, subject: str, message: str) -> str:
+    """Simulates sending an email from cert@company.com.
 
     This function takes email details, formats them into a standard email structure
     with headers and body, prints the formatted email to the console, and returns
     a success message. The sender is always set to cert@company.com.
 
-    Parameters:
-    -----------
-    recipient : str
-        The email address of the recipient (To field)
-    subject : str
-        The subject line of the email
-    message : str
-        The main body content of the email
+    Args:
+        recipient: The email address of the recipient (To field)
+        subject: The subject line of the email
+        message: The main body content of the email
 
     Returns:
-    --------
-    str
         A success message indicating that the email was sent, including
         the recipient address and the current time
 
     Example:
-    --------
-    >>> send_email("jane@example.com", "Security Alert", "Please update your password.")
-    ------ EMAIL SENT ------
-    From: cert@company.com
-    To: jane@example.com
-    Subject: Security Alert
-    Date: Wed, 04 Jun 2025 16:40:58 +0000
-    Message-ID: <123456.7890123456@mail-server>
-
-    Please update your password.
-    ------------------------
-    'Email successfully sent to jane@example.com at 16:40:58'
-    """
+        >>> send_email("jane@example.com", "Security Alert", "Please update your password.")
+        ------ EMAIL SENT ------
+        From: cert@company.com
+        To: jane@example.com
+        Subject: Security Alert
+        Date: Wed, 04 Jun 2025 16:40:58 +0000
+        Message-ID: <123456.7890123456@mail-server>
+
+        Please update your password.
+        ------------------------
+        'Email successfully sent to jane@example.com at 16:40:58'
+    """  # noqa: E501
     # Fixed sender email
     sender = "cert@company.com"
 
     # Generate a random message ID
-    message_id = f"<{random.randint(100000, 999999)}.{random.randint(1000000000, 9999999999)}@mail-server>"
+    message_id = f"<{random.randint(100000, 999999)}.{random.randint(1000000000, 9999999999)}@mail-server>"  # noqa: E501, S311
 
     # Get current timestamp
-    timestamp = datetime.datetime.now().strftime("%a, %d %b %Y %H:%M:%S +0000")
+    timestamp = datetime.datetime.now(datetime.timezone.utc).strftime(
+        "%a, %d %b %Y %H:%M:%S +0000",
+    )
 
     # Format the email
     email_format = f"""
@@ -63,10 +58,13 @@ Message-ID: {message_id}
 """
 
     # Print the email to console
-    print(email_format)
+    print(email_format)  # noqa: T201
 
     # Return a success message
-    return f"Email successfully sent from {sender} to {recipient} at {datetime.datetime.now().strftime('%H:%M:%S')}"
+    return (
+        f"Email successfully sent from {sender} to {recipient} at "
+        + datetime.datetime.now(datetime.timezone.utc).strftime("%H:%M:%S")
+    )
 
 
 # Create Gradio Interface
@@ -75,5 +73,5 @@ gr_send_email = gr.Interface(
     inputs=["text", "text", "text"],
     outputs="text",
     title="Email Sender Simulator",
-    description="This tool simulates sending an email by formatting and printing it to the console."
+    description="This tool simulates sending an email.",
 )

tdagent/tools/virus_total.py

@@ -1,21 +1,21 @@
+import os
+from datetime import datetime, timezone
+
+import cachetools
 import gradio as gr
 import vt
-import os
-from datetime import datetime
-from functools import lru_cache
-from typing import Optional
-import time
-import asyncio
+
 
 # Get API key from environment variable
-API_KEY = os.getenv('VT_API_KEY')
+API_KEY = os.getenv("VT_API_KEY")
+
 
+@cachetools.cached(cache=cachetools.TTLCache(maxsize=100, ttl=3600))
+def get_virus_total_url_info(url: str) -> str:
+    """Get URL Info from VirusTotal URL Scanner. Scan URL is not available."""
+    if not API_KEY:
+        return "Error: Virus total API key not configured."
 
-@lru_cache(maxsize=100)
-def get_url_info_cached(url: str, timestamp: Optional[int] = None) -> str:
-    """
-    Get URL Info from VirusTotal URL Scanner. Scan URL is not available
-    """
     try:
         # Create a new client for each request (thread-safe)
         with vt.Client(API_KEY) as client:
@@ -33,12 +33,15 @@ def get_url_info_cached(url: str, timestamp: Optional[int] = None) -> str:
             if last_analysis_date:
                 # Convert to datetime if it's a timestamp
                 if isinstance(last_analysis_date, (int, float)):
-                    last_analysis_date = datetime.utcfromtimestamp(last_analysis_date)
-                date_str = last_analysis_date.strftime('%Y-%m-%d %H:%M:%S UTC')
+                    last_analysis_date = datetime.fromtimestamp(
+                        last_analysis_date,
+                        timezone.utc,
+                    )
+                date_str = last_analysis_date.strftime("%Y-%m-%d %H:%M:%S UTC")
             else:
                 date_str = "Not available"
 
-            result = f"""
+            return f"""
 URL: {url}
 Last Analysis Date: {date_str}
 
@@ -53,29 +56,14 @@ Reputation Score: {url_analysis.reputation}
 Times Submitted: {url_analysis.times_submitted}
 
 Cache Status: Hit
-            """
-
-            return result
-
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-
-def get_url_info(url: str) -> str:
-    """
-    Wrapper function to handle the cached URL info retrieval
-    """
-    # Clean the URL to ensure consistent caching
-    url = url.strip().lower()
-
-    # Use current timestamp rounded to nearest hour to maintain cache for an hour
-    timestamp = int(time.time()) // 3600
+            """.strip()
 
-    return get_url_info_cached(url, timestamp)
+    except Exception as err:  # noqa: BLE001
+        return f"Error: {err}"
 
 
-gr_virus_total = gr.Interface(
-    fn=get_url_info,
+gr_virus_total_url_info = gr.Interface(
+    fn=get_virus_total_url_info,
     inputs=["text"],
     outputs="text",
     title="VirusTotal URL Scanner",

tdagent/tools/whois.py

@@ -0,0 +1,35 @@
+import json
+
+import gradio as gr
+import whois
+
+from tdagent.utils.json_utils import TDAgentJsonEncoder
+
+
+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)
+    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.",
+)

tdagent/utils/json_utils.py

@@ -0,0 +1,14 @@
+import datetime
+import json
+
+
+class TDAgentJsonEncoder(json.JSONEncoder):
+    """Extend JSON encoder with known types."""
+
+    def default(self, o: object) -> object:  # noqa: D102
+        if isinstance(o, datetime.datetime):
+            return {"__type__": "datetime", "value": o.isoformat()}
+        if isinstance(o, datetime.date):
+            return {"__type__": "date", "value": o.isoformat()}
+
+        return super().default(o)