fix: rodrigo's laziness second try

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(

tdagent/tools/get_domain_information.py

@@ -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:
@@ -164,7 +164,7 @@ 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
 
 

tdagent/tools/retrieve_from_mitre_attack.py

@@ -0,0 +1,41 @@
+from typing import Any
+
+import gradio as gr
+from attackcti import attack_client
+
+
+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"
+    ),
+)

tdagent/tools/retrieve_from_mitre_attack_information

@@ -1,19 +0,0 @@
-from typing import Any
-
-from attackcti import attack_client
-
-
-def get_stix_object_of_attack_id(attack_id: str, object_type: str = "attack-pattern") -> dict[str, Any]:
-    """Retrieves a specific STIX object identified by an ATT&CK ID across 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:
-        List: 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)[0]