tdagent/tools/retrieve_from_mitre_attack.py

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"
    ),
)