cache_and_increase_timeout (#8)

- build: update requirements files (6012e73d5a715309ec953641f276c2c7c5654cfe)
- feat: cache results and increase timeout (6b950bc1385d603ea7c076aea72210f0d5e72c03)
- Merge branch 'main' into pr/8 (6553ce6e6f2ddb20ed280971c8b44d107cd97334)
- fix: repited tag section in README (160aed774ab0d62460121d694a92ae12ac34e58f)


Co-authored-by: Rodrigo Domínguez Sanz <[email protected]>

README.md

@@ -11,8 +11,6 @@ license: apache-2.0
 tags:
  - mcp-server-track
 short_description: tdb
-tags:
-- mcp-server-track
 ---
 
 # TDA Agent Tools

requirements-dev.txt

@@ -1,5 +1,5 @@
 # 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

requirements.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv export --format requirements.txt --no-hashes --no-dev -o requirements.txt
+#    uv export --format requirements-txt --no-hashes --no-dev -o requirements.txt
 aiofiles==24.1.0
     # via
     #   gradio

tdagent/tools/get_domain_information.py

@@ -4,6 +4,7 @@ from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Any
 
+import cachetools
 import gradio as gr
 import requests
 import urllib3
@@ -30,7 +31,13 @@ _DNS_RECORD_TYPES = [
 
 _COMMON_SUBDOMAINS_TXT_PATH = Path("./subdomains/subdomains.txt")
 
+_CACHE_MAX_SIZE = 4096
+_CACHE_TTL_SECONDS = 3600
 
+
+@cachetools.cached(
+    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
+)
 def get_geolocation(ip: str) -> dict[str, Any] | str:
     """Get location information from an ip address.
 
@@ -65,13 +72,16 @@ def get_geolocation(ip: str) -> dict[str, Any] | str:
     try:
         return requests.get(
             f"https://geolocation-db.com/json/{ip}",
-            timeout=0.5,
+            timeout=1,
         ).json()
     except Exception as e:  # noqa: BLE001
         return str(e)
 
-
-def _request_dns_record(domain: str, record_type: str) -> list[str]:
+def _request_dns_record(  # noqa: D417
+    domain: str,
+    record_type: str,
+    timeout: float = 0.5,
+    ) -> list[str]:
     """Utility to build dns resolve requests that do not use port 53.
 
     Args:
@@ -90,7 +100,7 @@ def _request_dns_record(domain: str, record_type: str) -> list[str]:
         },
         data=q.to_wire(),
         verify=True,
-        timeout=0.2,
+        timeout=timeout,
     )
     dns_message = message.from_wire(response.content)
     return [str(rdata) for rdata in dns_message.answer[0]] if dns_message.answer else []
@@ -98,6 +108,9 @@ def _request_dns_record(domain: str, record_type: str) -> list[str]:
 
 # see: https://thepythoncode.com/article/dns-enumeration-with-python
 # https://dnspython.readthedocs.io
+@cachetools.cached(
+    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
+)
 def enumerate_dns(domain_name: str) -> dict[str, Any] | None:
     r"""Enumerates information about a specific domain's DNS configuration.
 
@@ -160,14 +173,13 @@ def enumerate_dns(domain_name: str) -> dict[str, Any] | None:
     enumeration = {}
     for record_type in _DNS_RECORD_TYPES:
         try:
-            record = _request_dns_record(domain_name, record_type)
+            record = _request_dns_record(domain_name, record_type, timeout=1)
             if record:
                 enumeration[record_type] = record
         except Exception as e:  # noqa: BLE001, PERF203
             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.
 
@@ -180,12 +192,14 @@ def resolve_subdomain_ipv4(domain: str) -> str | None:
             is returned.
     """
     try:
-        _request_dns_record(domain, "A")
+        _request_dns_record(domain, "A", timeout=0.6)
         return domain  # noqa: TRY300
     except Exception:  # noqa: BLE001
         return None
 
-
+@cachetools.cached(
+    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
+)
 def scrap_subdomains_for_domain(domain_name: str) -> list[str]:
     """Retrieves subdomains associated to a domain if any.
 
@@ -223,11 +237,14 @@ def scrap_subdomains_for_domain(domain_name: str) -> list[str]:
         return []
 
     potential_subdomains = [f"{subdomain}.{domain_name}" for subdomain in subdomains]
-    with ThreadPoolExecutor(max_workers=5) as executor:
+    with ThreadPoolExecutor(max_workers=None) as executor:
         results = executor.map(resolve_subdomain_ipv4, potential_subdomains)
         return [domain for domain in results if domain]
 
 
+@cachetools.cached(
+    cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
+)
 def retrieve_ioc_from_threatfox(potentially_ioc: str) -> str:
     r"""Retrieves information about a potential IoC from ThreatFox.
 
@@ -284,6 +301,7 @@ def retrieve_ioc_from_threatfox(potentially_ioc: str) -> str:
         port=443,
         maxsize=50,
         headers=headers,
+        timeout=5,
     )
     data = {
         "query": "search_ioc",
@@ -300,7 +318,7 @@ def retrieve_ioc_from_threatfox(potentially_ioc: str) -> str:
 geo_location_tool = gr.Interface(
     fn=get_geolocation,
     inputs=["text"],
-    outputs=["text"],
+    outputs="json",
     title="Domain Associated Geolocation Finder",
     description="Retrieves the geolocation associated to an input ip address",
     theme="default",
@@ -309,7 +327,7 @@ geo_location_tool = gr.Interface(
 dns_enumeration_tool = gr.Interface(
     fn=enumerate_dns,
     inputs=["text"],
-    outputs=["text"],
+    outputs="json",
     title="DNS record enumerator of domains",
     description="Retrieves several dns record types for the input domain names",
     theme="default",
@@ -318,7 +336,7 @@ dns_enumeration_tool = gr.Interface(
 scrap_subdomains_tool = gr.Interface(
     fn=scrap_subdomains_for_domain,
     inputs=["text"],
-    outputs=["text"],
+    outputs="json",
     title="Subdomains Extractor of domains",
     description="Retrieves the subdomains for the input domain if they are common",
     theme="default",
@@ -327,7 +345,7 @@ scrap_subdomains_tool = gr.Interface(
 extractor_of_ioc_from_threatfox_tool = gr.Interface(
     fn=retrieve_ioc_from_threatfox,
     inputs=["text"],
-    outputs=["text"],
+    outputs="text",
     title="IoC information extractor associated to particular entities",
     description=(
         "If information as an Indicator of Compromise (IoC) exists"