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.
    """
    try:
        lift = attack_client()
        return lift.get_object_by_attack_id(
            object_type=object_type.strip(),
            attack_id=attack_id.strip(),
            stix_format=False,
        )[0]
    except Exception as e:  # noqa: BLE001
        return {"Exception": str(e)}


gr_get_stix_of_attack_id = gr.Interface(
    fn=get_stix_object_of_attack_id,
    inputs=[
        gr.Textbox(label="Mitre technique ID"),
        gr.Textbox(label="Mitre object type"),
    ],
    outputs=gr.JSON(label="Mitre report"),
    title="MITRE ATT&CK STIX information",
    description=(
        "Retrieves a specific STIX object identified by an ATT&CK ID across all ATT&CK"
        " matrices"
    ),
    examples=[
        ["T1568.002", "attack-pattern"],
        ["M1042", "course-of-action"],
    ],
)