created clients for core objects to organize code

Author: Cyb3rWard0g

src/attackcti/core/__init__.py

@@ -0,0 +1,10 @@
+"""Core helpers for querying and enriching ATT&CK STIX content."""
+
+from __future__ import annotations
+
+from .query_client import QueryClient
+
+__all__ = [
+    "QueryClient",
+]
+

src/attackcti/core/objects/__init__.py

@@ -0,0 +1,25 @@
+"""Initialization of the attackcti.core.objects module."""
+
+from .analytics import AnalyticsClient
+from .campaigns import CampaignsClient
+from .data_sources import DataSourcesClient
+from .detections import DetectionsClient
+from .groups import GroupsClient
+from .mitigations import MitigationsClient
+from .relationships import RelationshipsClient
+from .software import SoftwareClient
+from .tactics import TacticsClient
+from .techniques import TechniquesClient
+
+__all__ = [
+    "CampaignsClient",
+    "DataSourcesClient",
+    "AnalyticsClient",
+    "DetectionsClient",
+    "GroupsClient",
+    "MitigationsClient",
+    "RelationshipsClient",
+    "SoftwareClient",
+    "TacticsClient",
+    "TechniquesClient",
+]

src/attackcti/core/objects/analytics.py

@@ -0,0 +1,60 @@
+"""Cross-domain analytics query helpers."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Iterable
+
+from stix2 import CompositeDataSource, Filter
+
+from ...models import Analytic as AnalyticModel
+from ...utils.stix import (
+    as_dict,
+    parse_stix_objects,
+    query_stix_objects_by_ids,
+    remove_revoked_deprecated,
+)
+
+
+class AnalyticsClient:
+    """Cross-domain analytics client (COMPOSITE_DS-backed)."""
+
+    def __init__(
+        self,
+        *,
+        data_source: CompositeDataSource,
+        remove_fn: Callable = remove_revoked_deprecated,
+        parse_fn: Callable = parse_stix_objects,
+    ) -> None:
+        """Initialize the client with a data source and helper callbacks."""
+        self._data_source = data_source
+        self._remove_fn = remove_fn
+        self._parse_fn = parse_fn
+
+    def get_analytics(self, *, stix_format: bool = True) -> list[dict[str, Any]]:
+        """Return all analytic objects."""
+        analytics = self._data_source.query(Filter("type", "=", "x-mitre-analytic"))
+        if not stix_format:
+            return self._parse_fn(analytics, AnalyticModel)
+        return [as_dict(a) for a in analytics]
+    
+    def get_analytics_by_ids(self, ids: Iterable[str]) -> dict[str, dict[str, Any]]:
+        """Return analytics keyed by their STIX id."""
+        analytics_dict: dict[str, dict[str, Any]] = {}
+        analytic_ids = {aid for aid in ids if isinstance(aid, str) and aid}
+        if not analytic_ids:
+            return analytics_dict
+
+        analytics = query_stix_objects_by_ids(
+            data_source=self._data_source,
+            stix_type="x-mitre-analytic",
+            ids=analytic_ids
+        )
+        for analytic in analytics:
+            analytic_dict = as_dict(analytic)
+            analytic_id = analytic_dict.get("id")
+            if not isinstance(analytic_id, str) or not analytic_id:
+                continue
+            log_sources = analytic_dict.get("x_mitre_log_source_references") or []
+            analytic_dict["x_attackcti_log_sources"] = [ls for ls in log_sources if isinstance(ls, dict)]
+            analytics_dict[analytic_id] = analytic_dict
+        return analytics_dict

src/attackcti/core/objects/campaigns.py

@@ -0,0 +1,108 @@
+"""Cross-domain campaign query helpers."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Union
+
+from stix2 import CompositeDataSource, Filter
+from stix2.v21.sdo import Campaign as CampaignV21
+
+from ...models import Campaign as CampaignModel
+
+
+class CampaignsClient:
+    """Campaigns query helper class."""
+
+    def __init__(
+        self,
+        *,
+        data_source: CompositeDataSource,
+        remove_fn: Callable | None = None,
+        parse_fn: Callable | None = None,
+    ) -> None:
+        """Initialize the client with a composite data source."""
+        self._data_source = data_source
+        self._remove_fn = remove_fn
+        self._parse_fn = parse_fn
+
+    def get_campaigns(self, *, skip_revoked_deprecated: bool = True, stix_format: bool = True) -> List[Union[CampaignV21, Dict[str, Any]]]:
+        """Return campaigns across ATT&CK matrices.
+
+        Parameters
+        ----------
+        skip_revoked_deprecated
+            When `True`, omit revoked/deprecated campaigns.
+        stix_format
+            When `True`, return STIX objects; when `False`, parse to the
+            `Campaign` Pydantic model.
+
+        Returns
+        -------
+        List[Union[CampaignV21, Dict[str, Any]]]
+            Campaign objects in the requested format.
+        """
+        all_campaigns = self._data_source.query([Filter("type", "=", "campaign")])
+        if skip_revoked_deprecated:
+            all_campaigns = self._remove_fn(all_campaigns)
+        if not stix_format:
+            all_campaigns = self._parse_fn(all_campaigns, CampaignModel)
+        return all_campaigns
+
+
+    def get_campaign_by_alias(self, *, alias: str, case: bool = True, stix_format: bool = True) -> List[Union[CampaignV21, Dict[str, Any]]]:
+        """Return campaigns that match a provided alias.
+
+        Parameters
+        ----------
+        alias
+            Alias to match.
+        case
+            When `True`, perform case-sensitive match; otherwise performs
+            case-insensitive containment match.
+        stix_format
+            When `True`, return STIX objects; when `False`, parse to the
+            `Campaign` Pydantic model.
+
+        Returns
+        -------
+        List[Union[CampaignV21, Dict[str, Any]]]
+            Matching campaign objects in the requested format.
+        """
+        if not case:
+            all_campaigns = self.get_campaigns(stix_format=True)
+            out: list[Any] = []
+            for campaign in all_campaigns:
+                if "aliases" in campaign.keys():
+                    for campaign_alias in campaign["aliases"]:
+                        if alias.lower() in campaign_alias.lower():
+                            out.append(CampaignModel)
+        else:
+            filter_objects = [Filter("type", "=", "campaign"), Filter("aliases", "contains", alias)]
+            out = self._data_source.query(filter_objects)
+
+        if not stix_format:
+            out = self._parse_fn(out, CampaignModel)
+        return out
+
+
+    def get_campaigns_since_time(self, *, timestamp: str, stix_format: bool = True) -> List[Union[CampaignV21, Dict[str, Any]]]:
+        """Return campaigns created after the provided timestamp.
+
+        Parameters
+        ----------
+        timestamp
+            Timestamp string for filtering.
+        stix_format
+            When `True`, return STIX objects; when `False`, parse to the
+            `Campaign` Pydantic model.
+
+        Returns
+        -------
+        List[Union[CampaignV21, Dict[str, Any]]]
+            Campaign objects in the requested format.
+        """
+        filter_objects = [Filter("type", "=", "campaign"), Filter("created", ">", timestamp)]
+        out = self._data_source.query(filter_objects)
+        if not stix_format:
+            out = self._parse_fn(out, CampaignModel)
+        return out

src/attackcti/core/objects/data_sources.py

@@ -0,0 +1,105 @@
+"""Cross-domain data source/component helpers."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Iterable, List
+from warnings import warn
+
+from stix2 import CompositeDataSource, Filter
+
+from ...models import DataComponent as DataComponentModel
+from ...models import DataSource as DataSourceModel
+
+
+class DataSourcesClient:
+    """Data source and data component queries."""
+
+    def __init__(
+        self,
+        *,
+        data_source: CompositeDataSource,
+        remove_fn: Callable | None = None,
+        parse_fn: Callable | None = None,
+    ) -> None:
+        """Initialize the client with a composite data source."""
+        self._data_source = data_source
+        self._remove_fn = remove_fn
+        self._parse_fn = parse_fn
+
+    def get_data_components(self, *, skip_revoked_deprecated: bool = True, stix_format: bool = True) -> List[Dict[str, Any]]:
+        """Return data components across all domains.
+
+        Parameters
+        ----------
+        skip_revoked_deprecated
+            When `True`, omit revoked/deprecated data components.
+        stix_format
+            When `True`, return STIX objects; when `False`, parse to the
+            `DataComponent` Pydantic model.
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            Data component objects in the requested format.
+        """
+        all_data_components = self._data_source.query([Filter("type", "=", "x-mitre-data-component")])
+        if skip_revoked_deprecated:
+            all_data_components = self._remove_fn(all_data_components)
+        if not stix_format:
+            all_data_components = self._parse_fn(all_data_components, DataComponentModel)
+        return all_data_components
+    
+    def get_data_sources(self, *, include_data_components: bool = False, stix_format: bool = True) -> List[Dict[str, Any]]:
+        """Return data sources across all domains.
+
+        Parameters
+        ----------
+        include_data_components
+            When `True`, enrich data sources with related data components
+            (requires Pydantic parsing).
+        stix_format
+            When `True`, return STIX objects; when `False`, parse to the
+            `DataSource` Pydantic model.
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            Data source objects in the requested format.
+        """
+        warn(
+            "Data Sources (`x-mitre-data-source`) are deprecated as of ATT&CK Specification 3.3.0. "
+            "Data Sources are superseded by the Detection Strategy framework..",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        all_data_sources = self._data_source.query([Filter("type", "=", "x-mitre-data-source")])
+        all_data_sources = self._remove_fn(all_data_sources)
+        if include_data_components:
+            all_data_sources = self._parse_fn(all_data_sources, DataSourceModel, include_data_components=True)
+        elif not stix_format:
+            all_data_sources = self._parse_fn(all_data_sources, DataSourceModel)
+        return all_data_sources
+    
+    def get_data_components_by_ids(
+        self,
+        ids: Iterable[str],
+        *,
+        stix_format: bool = True,
+        data_components: list[dict[str, Any]] | None = None,
+        skip_revoked_deprecated: bool = True,
+    ) -> list[dict[str, Any]]:
+        """Return data component objects for the requested ids."""
+        dc_ids = {did for did in ids if isinstance(did, str) and did}
+        if not dc_ids:
+            return []
+
+        if data_components is None:
+            data_components = self.get_data_components(
+                skip_revoked_deprecated=skip_revoked_deprecated,
+                stix_format=True,
+            )
+
+        selected = [dc for dc in data_components if dc.get("id") in dc_ids]
+        if not stix_format:
+            return self._parse_fn(selected, DataComponentModel)
+        return selected