Skip to content

Commit 22a6213

Browse files
Cyb3rWard0gCyb3rWard0g
committed
separated sources as modules for local and TAXII collections
1 parent f93d863 commit 22a6213

5 files changed

Lines changed: 374 additions & 0 deletions

File tree

‎src/attackcti/sources/__init__.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
"""Source (transport) implementations for attackcti."""
2+
3+
from .attack_source import MitreAttackSource
4+
from .local_loader import load_local_sources, load_stix_store
5+
from .resolver import load_sources
6+
from .taxii_loader import create_taxii_sources, load_taxii_sources
7+
8+
__all__ = [
9+
"create_taxii_sources",
10+
"load_local_sources",
11+
"MitreAttackSource",
12+
"load_sources",
13+
"load_stix_store",
14+
"load_taxii_sources",
15+
]

‎src/attackcti/sources/attack_source.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
"""Wrapper that encapsulates loaded ATT&CK data sources."""
2+
3+
from __future__ import annotations
4+
5+
from dataclasses import dataclass
6+
from typing import Any
7+
8+
from stix2 import CompositeDataSource
9+
10+
from .resolver import load_sources
11+
12+
13+
@dataclass
14+
class MitreAttackSource:
15+
"""Container for ATT&CK STIX sources.
16+
17+
Attributes
18+
----------
19+
enterprise
20+
Domain-scoped source for Enterprise ATT&CK (or None if not loaded).
21+
mobile
22+
Domain-scoped source for Mobile ATT&CK (or None if not loaded).
23+
ics
24+
Domain-scoped source for ICS ATT&CK (or None if not loaded).
25+
composite
26+
`CompositeDataSource` containing all loaded domain sources.
27+
versions
28+
Mapping of domain -> spec version (or None).
29+
mode
30+
One of local, taxii, mixed, empty.
31+
spec_version
32+
Unified spec version if all domains match, else None.
33+
"""
34+
35+
enterprise: Any | None
36+
mobile: Any | None
37+
ics: Any | None
38+
composite: CompositeDataSource
39+
versions: dict[str, str | None]
40+
mode: str
41+
spec_version: str | None
42+
43+
@classmethod
44+
def load(
45+
cls,
46+
*,
47+
enterprise: str | None = None,
48+
mobile: str | None = None,
49+
ics: str | None = None,
50+
connect_taxii: bool = True,
51+
proxies: dict | None = None,
52+
verify: bool = True,
53+
collection_url: str | None = None,
54+
) -> "MitreAttackSource":
55+
"""Load ATT&CK sources and return a container."""
56+
(ent, mob, ics_src), versions, mode, spec_version = load_sources(
57+
enterprise=enterprise,
58+
mobile=mobile,
59+
ics=ics,
60+
connect_taxii=connect_taxii,
61+
proxies=proxies,
62+
verify=verify,
63+
collection_url=collection_url,
64+
)
65+
composite = CompositeDataSource()
66+
composite.add_data_sources([ds for ds in (ent, mob, ics_src) if ds is not None])
67+
return cls(
68+
enterprise=ent,
69+
mobile=mob,
70+
ics=ics_src,
71+
composite=composite,
72+
versions=versions,
73+
mode=mode,
74+
spec_version=spec_version,
75+
)

‎src/attackcti/sources/local_loader.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
"""Local STIX bundle source helpers.
2+
3+
This module centralizes loading local STIX JSON bundles (files or directories)
4+
into STIX2 data sources via `attackcti.utils.storage.STIXStore`.
5+
"""
6+
7+
from __future__ import annotations
8+
9+
import os
10+
from typing import Any
11+
12+
from ..utils.storage import STIXStore
13+
14+
15+
def load_stix_store(path: str | None) -> tuple[Any | None, str | None]:
16+
"""Load a STIX store from a directory or JSON file path.
17+
18+
Args:
19+
path: Path to a directory of JSON files or a single STIX JSON file. If
20+
`None` or not found on disk, returns `(None, None)`.
21+
22+
Returns
23+
-------
24+
tuple
25+
A `(source, spec_version)` tuple. If `path` is missing, returns
26+
`(None, None)`.
27+
"""
28+
if path and os.path.exists(path):
29+
store = STIXStore(path)
30+
return store.get_store(), store.spec_version
31+
return None, None
32+
33+
34+
def load_local_sources(
35+
*,
36+
enterprise: str | None = None,
37+
mobile: str | None = None,
38+
ics: str | None = None,
39+
) -> tuple[tuple[Any | None, Any | None, Any | None], dict[str, str | None]]:
40+
"""Load local sources for enterprise, mobile, and ICS.
41+
42+
Args:
43+
enterprise: Path to the local enterprise bundle (dir or JSON file).
44+
mobile: Path to the local mobile bundle (dir or JSON file).
45+
46+
Returns
47+
-------
48+
tuple
49+
`((enterprise_source, mobile_source, ics_source), versions)` where
50+
`versions` maps `enterprise|mobile|ics` to detected spec versions (or
51+
`None` when unavailable).
52+
`versions` maps `enterprise|mobile|ics` to detected spec versions (or
53+
`None` when unavailable).
54+
"""
55+
enterprise_source, enterprise_ver = load_stix_store(enterprise)
56+
mobile_source, mobile_ver = load_stix_store(mobile)
57+
ics_source, ics_ver = load_stix_store(ics)
58+
59+
versions = {
60+
"enterprise": enterprise_ver,
61+
"mobile": mobile_ver,
62+
"ics": ics_ver,
63+
}
64+
return (enterprise_source, mobile_source, ics_source), versions

‎src/attackcti/sources/resolver.py‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
"""Source selection helpers.
2+
3+
This module contains the policy for combining multiple source types (local STIX
4+
bundles and TAXII) into the final set of domain sources used by the client.
5+
"""
6+
7+
from __future__ import annotations
8+
9+
from typing import Any
10+
11+
from .local_loader import load_local_sources
12+
from .taxii_loader import load_taxii_sources
13+
14+
15+
def load_sources(
16+
*,
17+
enterprise: str | None = None,
18+
mobile: str | None = None,
19+
ics: str | None = None,
20+
connect_taxii: bool = True,
21+
proxies: dict | None = None,
22+
verify: bool = True,
23+
collection_url: str | None = None,
24+
) -> tuple[tuple[Any | None, Any | None, Any | None], dict[str, str | None], str, str | None]:
25+
"""Load sources using a local-first policy with optional TAXII fallback.
26+
27+
Policy:
28+
- If local sources exist, use them.
29+
- If some local domains are missing and `connect_taxii=True`, fill missing domains from TAXII.
30+
- If no local sources exist:
31+
- If `connect_taxii=True`, load all domains from TAXII.
32+
- If `connect_taxii=False`:
33+
- Raise if the caller provided local paths (they were invalid).
34+
- Otherwise return an empty configuration.
35+
36+
Args:
37+
enterprise: Path to the local enterprise bundle (dir or JSON file).
38+
mobile: Path to the local mobile bundle (dir or JSON file).
39+
ics: Path to the local ICS bundle (dir or JSON file).
40+
connect_taxii: If `True`, allow TAXII fallback/fill behavior.
41+
proxies: Requests proxy configuration for TAXII.
42+
verify: Whether to verify TLS certificates for TAXII.
43+
collection_url: Base TAXII collections URL (ending in `/collections/`).
44+
45+
Returns
46+
-------
47+
tuple
48+
A tuple `(sources, versions, mode, spec_version)` where:
49+
- `sources` is `(enterprise_source, mobile_source, ics_source)`
50+
- `versions` maps `enterprise|mobile|ics` to spec versions (or `None`)
51+
- `mode` is one of `local`, `taxii`, `mixed`, `empty`
52+
- `spec_version` is the unified spec version if known, else `None`
53+
54+
Raises
55+
------
56+
ValueError: If local paths were provided but none were loadable and
57+
`connect_taxii=False`.
58+
"""
59+
local_paths_provided = any((enterprise, mobile, ics))
60+
61+
(enterprise_source, mobile_source, ics_source), versions = load_local_sources(
62+
enterprise=enterprise,
63+
mobile=mobile,
64+
ics=ics,
65+
)
66+
67+
any_local = any((enterprise_source, mobile_source, ics_source))
68+
if not any_local:
69+
if not connect_taxii:
70+
if local_paths_provided:
71+
raise ValueError("No valid local data sources found.")
72+
return (None, None, None), {"enterprise": None, "mobile": None, "ics": None}, "empty", None
73+
74+
(enterprise_source, mobile_source, ics_source), versions = load_taxii_sources(
75+
proxies=proxies,
76+
verify=verify,
77+
collection_url=collection_url,
78+
)
79+
return (enterprise_source, mobile_source, ics_source), versions, "taxii", "2.1"
80+
81+
missing_any = any(ds is None for ds in (enterprise_source, mobile_source, ics_source))
82+
if missing_any and connect_taxii:
83+
(taxii_enterprise, taxii_mobile, taxii_ics), taxii_versions = load_taxii_sources(
84+
proxies=proxies,
85+
verify=verify,
86+
collection_url=collection_url,
87+
)
88+
if enterprise_source is None:
89+
enterprise_source = taxii_enterprise
90+
versions["enterprise"] = taxii_versions["enterprise"]
91+
if mobile_source is None:
92+
mobile_source = taxii_mobile
93+
versions["mobile"] = taxii_versions["mobile"]
94+
if ics_source is None:
95+
ics_source = taxii_ics
96+
versions["ics"] = taxii_versions["ics"]
97+
mode = "mixed"
98+
else:
99+
mode = "local"
100+
101+
non_null_versions = {v for v in versions.values() if v is not None}
102+
spec_version = non_null_versions.pop() if len(non_null_versions) == 1 else None
103+
return (enterprise_source, mobile_source, ics_source), versions, mode, spec_version

‎src/attackcti/sources/taxii_loader.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
"""MITRE ATT&CK TAXII source helpers.
2+
3+
This module builds TAXII 2.1 collection sources for the MITRE ATT&CK datasets.
4+
"""
5+
6+
from __future__ import annotations
7+
8+
from stix2 import TAXIICollectionSource
9+
from taxii2client.v21 import Collection
10+
11+
from ..constants import (
12+
ATTACK_TAXII_COLLECTIONS_URL,
13+
ENTERPRISE_ATTACK_COLLECTION_ID,
14+
ICS_ATTACK_COLLECTION_ID,
15+
MOBILE_ATTACK_COLLECTION_ID,
16+
)
17+
18+
19+
def _normalize_collections_url(collections_url: str) -> str:
20+
"""Normalize a TAXII collections base URL.
21+
22+
Args:
23+
collections_url: Base URL for TAXII collections (typically ends with
24+
/collections/).
25+
26+
Returns
27+
-------
28+
Normalized URL with a trailing /.
29+
30+
Raises
31+
------
32+
ValueError: If collections_url is empty after trimming whitespace.
33+
"""
34+
collections_url = collections_url.strip()
35+
if not collections_url:
36+
raise ValueError("collection_url must be a non-empty string")
37+
if not collections_url.endswith("/"):
38+
collections_url += "/"
39+
return collections_url
40+
41+
42+
def create_taxii_sources(
43+
*,
44+
proxies: dict | None = None,
45+
verify: bool = True,
46+
collection_url: str | None = None,
47+
) -> tuple[TAXIICollectionSource, TAXIICollectionSource, TAXIICollectionSource]:
48+
"""Create TAXII sources for enterprise, mobile, and ICS ATT&CK.
49+
50+
Args:
51+
proxies: Requests proxy configuration passed to taxii2client (and
52+
ultimately requests).
53+
verify: Whether to verify TLS certificates.
54+
collection_url: Base collections URL (ending in /collections/). If
55+
omitted, uses :data:`attackcti.constants.ATTACK_TAXII_COLLECTIONS_URL`.
56+
57+
Returns
58+
-------
59+
A tuple of sources for (enterprise, mobile, ics).
60+
61+
Raises
62+
------
63+
ValueError
64+
If collection_url is provided but empty.
65+
"""
66+
collections_url = _normalize_collections_url(collection_url or ATTACK_TAXII_COLLECTIONS_URL)
67+
68+
enterprise_url = f"{collections_url}{ENTERPRISE_ATTACK_COLLECTION_ID}/"
69+
mobile_url = f"{collections_url}{MOBILE_ATTACK_COLLECTION_ID}/"
70+
ics_url = f"{collections_url}{ICS_ATTACK_COLLECTION_ID}/"
71+
72+
enterprise_collection = Collection(enterprise_url, verify=verify, proxies=proxies)
73+
mobile_collection = Collection(mobile_url, verify=verify, proxies=proxies)
74+
ics_collection = Collection(ics_url, verify=verify, proxies=proxies)
75+
76+
return (
77+
TAXIICollectionSource(enterprise_collection),
78+
TAXIICollectionSource(mobile_collection),
79+
TAXIICollectionSource(ics_collection),
80+
)
81+
82+
83+
def load_taxii_sources(
84+
*,
85+
proxies: dict | None = None,
86+
verify: bool = True,
87+
collection_url: str | None = None,
88+
) -> tuple[tuple[TAXIICollectionSource, TAXIICollectionSource, TAXIICollectionSource], dict[str, str]]:
89+
"""Load TAXII sources for enterprise, mobile, and ICS.
90+
91+
Args:
92+
proxies: Requests proxy configuration passed to taxii2client.
93+
verify: Whether to verify TLS certificates.
94+
collection_url: Base collections URL (ending in /collections/). If
95+
omitted, uses :data:`attackcti.constants.ATTACK_TAXII_COLLECTIONS_URL`.
96+
97+
Returns
98+
-------
99+
`((enterprise_source, mobile_source, ics_source), versions)` where
100+
`versions` maps `enterprise|mobile|ics` to `"2.1"`.
101+
`versions` maps `enterprise|mobile|ics` to `"2.1"`.
102+
103+
Raises
104+
------
105+
ValueError: If collection_url is provided but empty.
106+
"""
107+
sources = create_taxii_sources(
108+
proxies=proxies,
109+
verify=verify,
110+
collection_url=collection_url,
111+
)
112+
versions = {
113+
"enterprise": "2.1",
114+
"mobile": "2.1",
115+
"ics": "2.1",
116+
}
117+
return sources, versions

0 commit comments

Comments
 (0)