add_mitre_attack_information_retrieval (#7)

- feat: add stix object retrieval for an attack id (8e93ae4498ecb4cdd0db9728e834c9e7446fe841)
- Merge branch 'main' of https://huggingface.co/spaces/Agents-MCP-Hackathon/TDAgentTools (c59ae19eb829e398b16e5091760a9b6c99e9e480)
- fix: update requirements files and add missing dependency (b1c435fae63c7e8efa1cef272312cfdb163dfee9)
- fix: rodrigo's laziness second try (dd8af64b17133f1a7d87725c06ebbbfe6760ede6)
- feat: add cache to mittre retrieval (b0da92f1a0715fc63217c181f7323cf3fa3a2cd7)

app.py

@@ -15,6 +15,7 @@ from tdagent.tools.lookup_company_cloud_account_information import (
 )
 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
 
@@ -43,6 +44,7 @@ TOOLS = (
     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(

pyproject.toml

@@ -13,6 +13,9 @@ requires-python = ">=3.10,<4"
 readme = "README.md"
 license = ""
 dependencies = [
+    "attackcti>=0.5.4",
+    "audioop-lts>=0.2.1 ; python_full_version >= '3.13'",
+    "black>=25.1.0",
     "cachetools>=6.0.0",
     "dnspython>=2.7.0",
     "gradio[mcp]>=5.32.1",

requirements-dev.txt

@@ -1,17 +1,19 @@
 # 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
+#    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
+aiohttp==3.12.9
     # via vt-py
 aiosignal==1.3.2
     # via aiohttp
 annotated-types==0.7.0
     # via pydantic
+antlr4-python3-runtime==4.9.3
+    # via stix2-patterns
 anyio==4.9.0
     # via
     #   gradio
@@ -21,10 +23,16 @@ anyio==4.9.0
     #   starlette
 async-timeout==5.0.1 ; python_full_version < '3.11'
     # via aiohttp
+attackcti==0.5.4
+    # via tdagent
 attrs==25.3.0
     # via aiohttp
 audioop-lts==0.2.1 ; python_full_version >= '3.13'
-    # via gradio
+    # via
+    #   gradio
+    #   tdagent
+black==25.1.0
+    # via tdagent
 boolean-py==5.0
     # via license-expression
 cachecontrol==0.14.3
@@ -40,8 +48,9 @@ cfgv==3.4.0
     # via pre-commit
 charset-normalizer==3.4.2
     # via requests
-click==8.2.1 ; sys_platform != 'emscripten'
+click==8.2.1
     # via
+    #   black
     #   typer
     #   uvicorn
 colorama==0.4.6 ; sys_platform == 'win32'
@@ -80,7 +89,7 @@ fsspec==2025.5.1
     # via
     #   gradio-client
     #   huggingface-hub
-gradio==5.32.1
+gradio==5.33.0
     # via tdagent
 gradio-client==1.10.2
     # via gradio
@@ -90,7 +99,7 @@ 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'
+hf-xet==1.1.3 ; 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
@@ -138,7 +147,9 @@ multidict==6.4.4
     #   yarl
 mypy==1.16.0
 mypy-extensions==1.1.0
-    # via mypy
+    # via
+    #   black
+    #   mypy
 nodeenv==1.9.1
     # via pre-commit
 numpy==2.2.6
@@ -147,20 +158,23 @@ numpy==2.2.6
     #   pandas
 orjson==3.10.18
     # via gradio
-packageurl-python==0.16.0
+packageurl-python==0.17.0
     # via cyclonedx-python-lib
 packaging==25.0
     # via
+    #   black
     #   gradio
     #   gradio-client
     #   huggingface-hub
     #   pip-audit
     #   pip-requirements-parser
     #   pytest
-pandas==2.2.3
+pandas==2.3.0
     # via gradio
 pathspec==0.12.1
-    # via mypy
+    # via
+    #   black
+    #   mypy
 pillow==11.2.1
     # via gradio
 pip==25.1.1
@@ -172,6 +186,7 @@ pip-requirements-parser==32.0.1
     # via pip-audit
 platformdirs==4.3.8
     # via
+    #   black
     #   pip-audit
     #   virtualenv
 pluggy==1.6.0
@@ -185,6 +200,7 @@ py-serializable==2.0.0
     # via cyclonedx-python-lib
 pydantic==2.11.5
     # via
+    #   attackcti
     #   fastapi
     #   gradio
     #   mcp
@@ -218,7 +234,10 @@ python-multipart==0.0.20
 python-whois==0.9.5
     # via tdagent
 pytz==2025.2
-    # via pandas
+    # via
+    #   pandas
+    #   stix2
+    #   taxii2-client
 pyyaml==6.0.2
     # via
     #   gradio
@@ -229,6 +248,8 @@ requests==2.32.3
     #   cachecontrol
     #   huggingface-hub
     #   pip-audit
+    #   stix2
+    #   taxii2-client
     #   tdagent
 rich==14.0.0
     # via
@@ -242,8 +263,13 @@ semantic-version==2.10.0
     # via gradio
 shellingham==1.5.4 ; sys_platform != 'emscripten'
     # via typer
+simplejson==3.20.1
+    # via stix2
 six==1.17.0
-    # via python-dateutil
+    # via
+    #   python-dateutil
+    #   stix2-patterns
+    #   taxii2-client
 sniffio==1.3.1
     # via anyio
 sortedcontainers==2.4.0
@@ -255,14 +281,21 @@ starlette==0.46.2
     #   fastapi
     #   gradio
     #   mcp
+stix2==3.0.1
+    # via attackcti
+stix2-patterns==2.0.0
+    # via stix2
+taxii2-client==2.3.0
+    # via attackcti
 toml==0.10.2
     # via pip-audit
 tomli==2.2.1 ; python_full_version <= '3.11'
     # via
+    #   black
     #   coverage
     #   mypy
     #   pytest
-tomlkit==0.13.2
+tomlkit==0.13.3
     # via gradio
 tqdm==4.67.1
     # via huggingface-hub
@@ -271,6 +304,7 @@ typer==0.16.0 ; sys_platform != 'emscripten'
 typing-extensions==4.14.0
     # via
     #   anyio
+    #   black
     #   exceptiongroup
     #   fastapi
     #   gradio

requirements.txt

@@ -1,17 +1,19 @@
 # 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
+aiohttp==3.12.9
     # via vt-py
 aiosignal==1.3.2
     # via aiohttp
 annotated-types==0.7.0
     # via pydantic
+antlr4-python3-runtime==4.9.3
+    # via stix2-patterns
 anyio==4.9.0
     # via
     #   gradio
@@ -19,12 +21,20 @@ anyio==4.9.0
     #   mcp
     #   sse-starlette
     #   starlette
+async-timeout==5.0.1 ; python_full_version < '3.11'
+    # via aiohttp
+attackcti==0.5.4
+    # via tdagent
 attrs==25.3.0
     # via aiohttp
-audioop-lts==0.2.1
-    # via gradio
+audioop-lts==0.2.1 ; python_full_version >= '3.13'
+    # via
+    #   gradio
+    #   tdagent
+black==25.1.0
+    # via tdagent
 cachetools==6.0.0
-    # via tdagent (pyproject.toml)
+    # via tdagent
 certifi==2025.4.26
     # via
     #   httpcore
@@ -34,10 +44,22 @@ charset-normalizer==3.4.2
     # via requests
 click==8.2.1
     # via
+    #   black
     #   typer
     #   uvicorn
+colorama==0.4.6 ; sys_platform == 'win32'
+    # via
+    #   click
+    #   pytest
+    #   tqdm
+coverage==7.8.2
+    # via pytest-cov
 dnspython==2.7.0
-    # via tdagent (pyproject.toml)
+    # via tdagent
+exceptiongroup==1.3.0 ; python_full_version < '3.11'
+    # via
+    #   anyio
+    #   pytest
 fastapi==0.115.12
     # via gradio
 ffmpy==0.6.0
@@ -52,8 +74,8 @@ fsspec==2025.5.1
     # via
     #   gradio-client
     #   huggingface-hub
-gradio==5.32.1
-    # via tdagent (pyproject.toml)
+gradio==5.33.0
+    # via tdagent
 gradio-client==1.10.2
     # via gradio
 groovy==0.1.2
@@ -62,7 +84,7 @@ h11==0.16.0
     # via
     #   httpcore
     #   uvicorn
-hf-xet==1.1.2
+hf-xet==1.1.3 ; 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
@@ -84,9 +106,11 @@ idna==3.10
     #   httpx
     #   requests
     #   yarl
+iniconfig==2.1.0
+    # via pytest
 jinja2==3.1.6
     # via gradio
-markdown-it-py==3.0.0
+markdown-it-py==3.0.0 ; sys_platform != 'emscripten'
     # via rich
 markupsafe==3.0.2
     # via
@@ -94,12 +118,14 @@ markupsafe==3.0.2
     #   jinja2
 mcp==1.9.0
     # via gradio
-mdurl==0.1.2
+mdurl==0.1.2 ; sys_platform != 'emscripten'
     # via markdown-it-py
 multidict==6.4.4
     # via
     #   aiohttp
     #   yarl
+mypy-extensions==1.1.0
+    # via black
 numpy==2.2.6
     # via
     #   gradio
@@ -108,19 +134,28 @@ orjson==3.10.18
     # via gradio
 packaging==25.0
     # via
+    #   black
     #   gradio
     #   gradio-client
     #   huggingface-hub
-pandas==2.2.3
+    #   pytest
+pandas==2.3.0
     # via gradio
+pathspec==0.12.1
+    # via black
 pillow==11.2.1
     # via gradio
+platformdirs==4.3.8
+    # via black
+pluggy==1.6.0
+    # via pytest
 propcache==0.3.1
     # via
     #   aiohttp
     #   yarl
 pydantic==2.11.5
     # via
+    #   attackcti
     #   fastapi
     #   gradio
     #   mcp
@@ -131,8 +166,14 @@ pydantic-settings==2.9.1
     # via mcp
 pydub==0.25.1
     # via gradio
-pygments==2.19.1
+pygments==2.19.1 ; sys_platform != 'emscripten'
     # via rich
+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
@@ -144,29 +185,39 @@ python-multipart==0.0.20
     #   gradio
     #   mcp
 python-whois==0.9.5
-    # via tdagent (pyproject.toml)
+    # via tdagent
 pytz==2025.2
-    # via pandas
+    # via
+    #   pandas
+    #   stix2
+    #   taxii2-client
 pyyaml==6.0.2
     # via
     #   gradio
     #   huggingface-hub
 requests==2.32.3
     # via
-    #   tdagent (pyproject.toml)
     #   huggingface-hub
-rich==14.0.0
+    #   stix2
+    #   taxii2-client
+    #   tdagent
+rich==14.0.0 ; sys_platform != 'emscripten'
     # via typer
-ruff==0.11.12
+ruff==0.11.12 ; sys_platform != 'emscripten'
     # via gradio
 safehttpx==0.1.6
     # via gradio
 semantic-version==2.10.0
     # via gradio
-shellingham==1.5.4
+shellingham==1.5.4 ; sys_platform != 'emscripten'
     # via typer
+simplejson==3.20.1
+    # via stix2
 six==1.17.0
-    # via python-dateutil
+    # via
+    #   python-dateutil
+    #   stix2-patterns
+    #   taxii2-client
 sniffio==1.3.1
     # via anyio
 sse-starlette==2.3.6
@@ -176,22 +227,39 @@ starlette==0.46.2
     #   fastapi
     #   gradio
     #   mcp
-tomlkit==0.13.2
+stix2==3.0.1
+    # via attackcti
+stix2-patterns==2.0.0
+    # via stix2
+taxii2-client==2.3.0
+    # via attackcti
+tomli==2.2.1 ; python_full_version <= '3.11'
+    # via
+    #   black
+    #   coverage
+    #   pytest
+tomlkit==0.13.3
     # via gradio
 tqdm==4.67.1
     # via huggingface-hub
-typer==0.16.0
+typer==0.16.0 ; sys_platform != 'emscripten'
     # via gradio
 typing-extensions==4.14.0
     # via
+    #   anyio
+    #   black
+    #   exceptiongroup
     #   fastapi
     #   gradio
     #   gradio-client
     #   huggingface-hub
+    #   multidict
     #   pydantic
     #   pydantic-core
+    #   rich
     #   typer
     #   typing-inspection
+    #   uvicorn
 typing-inspection==0.4.1
     # via
     #   pydantic
@@ -199,14 +267,17 @@ typing-inspection==0.4.1
 tzdata==2025.2
     # via pandas
 urllib3==2.4.0
-    # via requests
-uvicorn==0.34.3
+    # via
+    #   gradio
+    #   requests
+uvicorn==0.34.3 ; sys_platform != 'emscripten'
     # via
     #   gradio
     #   mcp
 vt-py==0.21.0
-    # via tdagent (pyproject.toml)
+    # via tdagent
 websockets==15.0.1
     # via gradio-client
+xdoctest==1.2.0
 yarl==1.20.0
     # via aiohttp

tdagent/tools/get_domain_information.py

@@ -10,7 +10,7 @@ import urllib3
 from dns import message
 
 
-_DNS_SERVER  = "https://dns.google/dns-query" # can use others
+_DNS_SERVER = "https://dns.google/dns-query"  # can use others
 _DNS_RECORD_TYPES = [
     "A",
     "AAAA",
@@ -71,7 +71,7 @@ def get_geolocation(ip: str) -> dict[str, Any] | str:
         return str(e)
 
 
-def _request_dns_record(domain: str, record_type: str) -> str:
+def _request_dns_record(domain: str, record_type: str) -> list[str]:
     """Utility to build dns resolve requests that do not use port 53.
 
     Args:
@@ -95,6 +95,7 @@ def _request_dns_record(domain: str, record_type: str) -> str:
     dns_message = message.from_wire(response.content)
     return [str(rdata) for rdata in dns_message.answer[0]] if dns_message.answer else []
 
+
 # see: https://thepythoncode.com/article/dns-enumeration-with-python
 # https://dnspython.readthedocs.io
 def enumerate_dns(domain_name: str) -> dict[str, Any] | None:
@@ -163,9 +164,10 @@ def enumerate_dns(domain_name: str) -> dict[str, Any] | None:
             if record:
                 enumeration[record_type] = record
         except Exception as e:  # noqa: BLE001, PERF203
-            enumeration[record_type] = str(e)
+            enumeration[record_type] = [str(e)]
     return enumeration if enumeration else None
 
+
 def resolve_subdomain_ipv4(domain: str) -> str | None:
     """Resolve the IPv4 address of a domain.
 
@@ -225,6 +227,7 @@ def scrap_subdomains_for_domain(domain_name: str) -> list[str]:
         results = executor.map(resolve_subdomain_ipv4, potential_subdomains)
         return [domain for domain in results if domain]
 
+
 def retrieve_ioc_from_threatfox(potentially_ioc: str) -> str:
     r"""Retrieves information about a potential IoC from ThreatFox.
 

tdagent/tools/retrieve_from_mitre_attack.py

@@ -0,0 +1,49 @@
+from typing import Any
+
+import cachetools
+import gradio as gr
+from attackcti import attack_client
+
+
+_CACHE_MAX_SIZE = 4096
+_CACHE_TTL_SECONDS = 3600
+
+
+@cachetools.cached(
+    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
+)
+def get_stix_object_of_attack_id(
+    attack_id: str,
+    object_type: str = "attack-pattern",
+) -> dict[str, Any]:
+    """Retrieves a STIX object identified by an ATT&CK ID in all ATT&CK matrices.
+
+    Args:
+        attack_id (str): The ATT&CK ID (e.g., 'T1234') of the STIX object to retrieve.
+        object_type (str): The type of STIX object to retrieve, such as
+            'attack-pattern', 'course-of-action', 'intrusion-set',
+            'malware', 'tool', or 'x-mitre-data-component'. Default is 'attack-pattern'
+
+    Returns:
+        A list containing the matched STIX object, either in its raw STIX format
+        or as a custom dictionary following the structure defined by the relevant
+        Pydantic model, depending on the 'stix_format' flag.
+    """
+    lift = attack_client()
+    return lift.get_object_by_attack_id(
+        object_type=object_type,
+        attack_id=attack_id,
+        stix_format=False,
+    )[0]
+
+
+gr_get_stix_of_attack_id = gr.Interface(
+    fn=get_stix_object_of_attack_id,
+    inputs=["text", "text"],
+    outputs="json",
+    title="MITRE ATT&CK STIX information",
+    description=(
+        "Retrieves a specific STIX object identified by an ATT&CK ID across all ATT&CK"
+        " matrices"
+    ),
+)