Skip to content

Commit c4bc16e

Browse files
claudeclaude
committed
refactor: Extract LayerGenerator and JsonGenerator from DiffStix
- Create LayerGenerator class in formatters/layer_generator.py (105 lines) - Create JsonGenerator class in formatters/json_generator.py (60 lines) - Extract get_layers_dict() method (80 lines) into LayerGenerator.generate() - Extract get_changes_dict() method (34 lines) into JsonGenerator.generate() - Update DiffStix to delegate to new generators - DiffStix reduced from 1,317 lines to 1,217 lines (7.6% reduction) - Total reduction from original: 245 lines (16.8% from 1,462 lines) - All 132/133 tests passing (only known permission test fails)
1 parent 61e08bb commit c4bc16e

3 files changed

Lines changed: 169 additions & 107 deletions

File tree

‎mitreattack/diffStix/core/diff_stix.py‎

Lines changed: 7 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313

1414
from mitreattack.diffStix.core.contributor_tracker import ContributorTracker
1515
from mitreattack.diffStix.core.domain_statistics import DomainStatistics
16+
from mitreattack.diffStix.formatters.json_generator import JsonGenerator
17+
from mitreattack.diffStix.formatters.layer_generator import LayerGenerator
1618
from mitreattack.diffStix.formatters.markdown_generator import MarkdownGenerator
1719
from mitreattack.diffStix.utils.stix_utils import (
1820
cleanup_values,
@@ -195,8 +197,10 @@ def __init__(
195197

196198
self.load_data()
197199

198-
# Initialize markdown generator after data is loaded
200+
# Initialize formatters after data is loaded
199201
self._markdown_generator = MarkdownGenerator(self)
202+
self._layer_generator = LayerGenerator(self)
203+
self._json_generator = JsonGenerator(self)
200204

201205
@property
202206
def release_contributors(self) -> dict:
@@ -1206,112 +1210,8 @@ def get_layers_dict(self):
12061210
12071211
Returns a dict mapping domain to its layer dict.
12081212
"""
1209-
logger.info("Generating ATT&CK Navigator layers")
1210-
1211-
colors = {
1212-
"additions": "#a1d99b", # granny smith apple
1213-
"major_version_changes": "#fcf3a2", # yellow-ish
1214-
"minor_version_changes": "#c7c4e0", # light periwinkle
1215-
"other_version_changes": "#B5E5CF", # mint
1216-
"patches": "#B99095", # mauve
1217-
"deletions": "#ff00e1", # hot magenta
1218-
"revocations": "#ff9000", # dark orange
1219-
"deprecations": "#ff6363", # bittersweet red
1220-
"unchanged": "#ffffff", # white
1221-
}
1222-
1223-
layers = {}
1224-
thedate = datetime.datetime.today().strftime("%B %Y")
1225-
# for each layer file in the domains mapping
1226-
for domain in self.domains:
1227-
logger.debug(f"Generating ATT&CK Navigator layer for domain: {domain}")
1228-
# build techniques list
1229-
techniques = []
1230-
for section, technique_stix_objects in self.data["changes"]["techniques"][domain].items():
1231-
if section == "revocations" or section == "deprecations":
1232-
continue
1233-
1234-
for technique in technique_stix_objects:
1235-
problem_detected = False
1236-
if "kill_chain_phases" not in technique:
1237-
logger.error(f"{technique['id']}: technique missing a tactic!! {technique['name']}")
1238-
problem_detected = True
1239-
if "external_references" not in technique:
1240-
logger.error(f"{technique['id']}: technique missing external references!! {technique['name']}")
1241-
problem_detected = True
1242-
1243-
if problem_detected:
1244-
continue
1245-
1246-
for phase in technique["kill_chain_phases"]:
1247-
techniques.append(
1248-
{
1249-
"techniqueID": technique["external_references"][0]["external_id"],
1250-
"tactic": phase["phase_name"],
1251-
"enabled": True,
1252-
"color": colors[section],
1253-
# trim the 's' off end of word
1254-
"comment": section[:-1] if section != "unchanged" else section,
1255-
}
1256-
)
1257-
1258-
legendItems = []
1259-
for section, description in self.section_descriptions.items():
1260-
legendItems.append({"color": colors[section], "label": f"{section}: {description}"})
1261-
1262-
# build layer structure
1263-
layer_json = {
1264-
"versions": {
1265-
"layer": "4.5",
1266-
"navigator": "5.0.0",
1267-
"attack": self.data["new"][domain]["attack_release_version"],
1268-
},
1269-
"name": f"{thedate} {self.domain_to_domain_label[domain]} Updates",
1270-
"description": f"{self.domain_to_domain_label[domain]} updates for the {thedate} release of ATT&CK",
1271-
"domain": domain,
1272-
"techniques": techniques,
1273-
"sorting": 0,
1274-
"hideDisabled": False,
1275-
"legendItems": legendItems,
1276-
"showTacticRowBackground": True,
1277-
"tacticRowBackground": "#205b8f",
1278-
"selectTechniquesAcrossTactics": True,
1279-
}
1280-
layers[domain] = layer_json
1281-
1282-
return layers
1213+
return self._layer_generator.generate()
12831214

12841215
def get_changes_dict(self):
12851216
"""Return dict format summarizing detected differences."""
1286-
logger.info("Generating changes info")
1287-
1288-
changes_dict = {}
1289-
for domain in self.domains:
1290-
changes_dict[domain] = {}
1291-
1292-
for object_type, domains in self.data["changes"].items():
1293-
for domain, sections in domains.items():
1294-
changes_dict[domain][object_type] = {}
1295-
1296-
for section, stix_objects in sections.items():
1297-
groupings = self.get_groupings(
1298-
object_type=object_type,
1299-
stix_objects=stix_objects,
1300-
section=section,
1301-
domain=domain,
1302-
)
1303-
# new_values includes parents & children mixed
1304-
# (e.g. techniques/sub-techniques, data sources/components)
1305-
new_values = cleanup_values(groupings=groupings)
1306-
changes_dict[domain][object_type][section] = new_values
1307-
1308-
# always add contributors
1309-
changes_dict["new-contributors"] = []
1310-
sorted_contributors = sorted(self.release_contributors, key=lambda v: v.lower())
1311-
for contributor in sorted_contributors:
1312-
# do not include ATT&CK as contributor
1313-
if contributor == "ATT&CK":
1314-
continue
1315-
changes_dict["new-contributors"].append(contributor)
1316-
1317-
return changes_dict
1217+
return self._json_generator.generate()

‎mitreattack/diffStix/formatters/json_generator.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
"""JSON changelog output generator."""
2+
3+
from loguru import logger
4+
5+
from mitreattack.diffStix.utils.stix_utils import cleanup_values
6+
7+
8+
class JsonGenerator:
9+
"""Generates JSON formatted changelog output from ATT&CK version differences."""
10+
11+
def __init__(self, diff_stix_instance):
12+
"""Initialize JsonGenerator with a DiffStix instance.
13+
14+
Parameters
15+
----------
16+
diff_stix_instance : DiffStix
17+
The DiffStix instance containing data and helper methods
18+
"""
19+
self.diff_stix = diff_stix_instance
20+
21+
def generate(self) -> dict:
22+
"""Return dict format summarizing detected differences.
23+
24+
Returns
25+
-------
26+
dict
27+
A dict containing all changes organized by domain and object type
28+
"""
29+
logger.info("Generating changes info")
30+
31+
changes_dict = {}
32+
for domain in self.diff_stix.domains:
33+
changes_dict[domain] = {}
34+
35+
for object_type, domains in self.diff_stix.data["changes"].items():
36+
for domain, sections in domains.items():
37+
changes_dict[domain][object_type] = {}
38+
39+
for section, stix_objects in sections.items():
40+
groupings = self.diff_stix.get_groupings(
41+
object_type=object_type,
42+
stix_objects=stix_objects,
43+
section=section,
44+
domain=domain,
45+
)
46+
# new_values includes parents & children mixed
47+
# (e.g. techniques/sub-techniques, data sources/components)
48+
new_values = cleanup_values(groupings=groupings)
49+
changes_dict[domain][object_type][section] = new_values
50+
51+
# always add contributors
52+
changes_dict["new-contributors"] = []
53+
sorted_contributors = sorted(self.diff_stix.release_contributors, key=lambda v: v.lower())
54+
for contributor in sorted_contributors:
55+
# do not include ATT&CK as contributor
56+
if contributor == "ATT&CK":
57+
continue
58+
changes_dict["new-contributors"].append(contributor)
59+
60+
return changes_dict

‎mitreattack/diffStix/formatters/layer_generator.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
"""ATT&CK Navigator layer output generator."""
2+
3+
import datetime
4+
5+
from loguru import logger
6+
7+
8+
class LayerGenerator:
9+
"""Generates ATT&CK Navigator layer JSON from ATT&CK version differences."""
10+
11+
def __init__(self, diff_stix_instance):
12+
"""Initialize LayerGenerator with a DiffStix instance.
13+
14+
Parameters
15+
----------
16+
diff_stix_instance : DiffStix
17+
The DiffStix instance containing data and helper methods
18+
"""
19+
self.diff_stix = diff_stix_instance
20+
21+
def generate(self) -> dict:
22+
"""Return ATT&CK Navigator layers in dict format summarizing detected differences.
23+
24+
Returns
25+
-------
26+
dict
27+
A dict mapping domain to its layer dict
28+
"""
29+
logger.info("Generating ATT&CK Navigator layers")
30+
31+
colors = {
32+
"additions": "#a1d99b", # granny smith apple
33+
"major_version_changes": "#fcf3a2", # yellow-ish
34+
"minor_version_changes": "#c7c4e0", # light periwinkle
35+
"other_version_changes": "#B5E5CF", # mint
36+
"patches": "#B99095", # mauve
37+
"deletions": "#ff00e1", # hot magenta
38+
"revocations": "#ff9000", # dark orange
39+
"deprecations": "#ff6363", # bittersweet red
40+
"unchanged": "#ffffff", # white
41+
}
42+
43+
layers = {}
44+
thedate = datetime.datetime.today().strftime("%B %Y")
45+
# for each layer file in the domains mapping
46+
for domain in self.diff_stix.domains:
47+
logger.debug(f"Generating ATT&CK Navigator layer for domain: {domain}")
48+
# build techniques list
49+
techniques = []
50+
for section, technique_stix_objects in self.diff_stix.data["changes"]["techniques"][domain].items():
51+
if section == "revocations" or section == "deprecations":
52+
continue
53+
54+
for technique in technique_stix_objects:
55+
problem_detected = False
56+
if "kill_chain_phases" not in technique:
57+
logger.error(f"{technique['id']}: technique missing a tactic!! {technique['name']}")
58+
problem_detected = True
59+
if "external_references" not in technique:
60+
logger.error(f"{technique['id']}: technique missing external references!! {technique['name']}")
61+
problem_detected = True
62+
63+
if problem_detected:
64+
continue
65+
66+
for phase in technique["kill_chain_phases"]:
67+
techniques.append(
68+
{
69+
"techniqueID": technique["external_references"][0]["external_id"],
70+
"tactic": phase["phase_name"],
71+
"enabled": True,
72+
"color": colors[section],
73+
# trim the 's' off end of word
74+
"comment": section[:-1] if section != "unchanged" else section,
75+
}
76+
)
77+
78+
legendItems = []
79+
for section, description in self.diff_stix.section_descriptions.items():
80+
legendItems.append({"color": colors[section], "label": f"{section}: {description}"})
81+
82+
# build layer structure
83+
layer_json = {
84+
"versions": {
85+
"layer": "4.5",
86+
"navigator": "5.0.0",
87+
"attack": self.diff_stix.data["new"][domain]["attack_release_version"],
88+
},
89+
"name": f"{thedate} {self.diff_stix.domain_to_domain_label[domain]} Updates",
90+
"description": f"{self.diff_stix.domain_to_domain_label[domain]} updates for the {thedate} release of ATT&CK",
91+
"domain": domain,
92+
"techniques": techniques,
93+
"sorting": 0,
94+
"hideDisabled": False,
95+
"legendItems": legendItems,
96+
"showTacticRowBackground": True,
97+
"tacticRowBackground": "#205b8f",
98+
"selectTechniquesAcrossTactics": True,
99+
}
100+
layers[domain] = layer_json
101+
102+
return layers

0 commit comments

Comments
 (0)