diff --git a/src/attack_flow/cli.py b/src/attack_flow/cli.py index da8c9cfb..5b4ec315 100644 --- a/src/attack_flow/cli.py +++ b/src/attack_flow/cli.py @@ -91,7 +91,15 @@ def graphviz(args): """ path = Path(args.attack_flow) flow_bundle = attack_flow.model.load_attack_flow_bundle(path) - converted = attack_flow.graphviz.convert(flow_bundle) + + if ( + flow_bundle.get("objects", "") + and attack_flow.model.get_flow_object(flow_bundle).scope == "attack-tree" + ): + converted = attack_flow.graphviz.convert_attack_tree(flow_bundle) + else: + converted = attack_flow.graphviz.convert_attack_flow(flow_bundle) + with open(args.output, "w") as out: out.write(converted) return 0 @@ -106,7 +114,14 @@ def mermaid(args): """ path = Path(args.attack_flow) flow_bundle = attack_flow.model.load_attack_flow_bundle(path) - converted = attack_flow.mermaid.convert(flow_bundle) + if ( + flow_bundle.get("objects", "") + and attack_flow.model.get_flow_object(flow_bundle).scope == "attack-tree" + ): + converted = attack_flow.mermaid.convert_attack_tree(flow_bundle) + else: + converted = attack_flow.mermaid.convert_attack_flow(flow_bundle) + with open(args.output, "w") as out: out.write(converted) return 0 diff --git a/src/attack_flow/graphviz.py b/src/attack_flow/graphviz.py index c7205dde..9d2f66fc 100644 --- a/src/attack_flow/graphviz.py +++ b/src/attack_flow/graphviz.py @@ -19,7 +19,7 @@ def label_escape(text): return graphviz.escape(html.escape(text)) -def convert(bundle): +def convert_attack_flow(bundle): """ Convert an Attack Flow STIX bundle into Graphviz format. @@ -71,6 +71,87 @@ def convert(bundle): return gv.source +def convert_attack_tree(bundle): + """ + Convert an Attack Flow STIX bundle into Graphviz format. + + :param stix2.Bundle flow: + :rtype: str + """ + + gv = graphviz.Digraph(graph_attr={"rankdir": "BT"}) + gv.body = _get_body_label(bundle) + ignored_ids = get_viz_ignored_ids(bundle) + + objects = bundle.objects + + id_to_remove = [] + ids = [] + for i, o in enumerate(objects): + if o.type == "attack-operator": + id_to_remove.append( + { + "id": o.id, + "prev_id": objects[i - 1].id, + "next_id": o.effect_refs[0], + "type": o.operator, + } + ) + + ids = [i["id"] for i in id_to_remove] + objects = [item for item in objects if item.id not in ids] + new_operator_ids = [i["next_id"] for i in id_to_remove] + for operator in id_to_remove: + for i, o in enumerate(objects): + if o.type == "relationship" and o.source_ref == operator["id"]: + o.source_ref = operator.prev_id + if o.type == "relationship" and o.target_ref == operator["id"]: + o.target_ref = operator.next_id + if o.get("effect_refs") and operator["id"] in o.effect_refs: + for i, j in enumerate(o.effect_refs): + if j == operator["id"]: + o.effect_refs[i] = operator["next_id"] + + for o in objects: + logger.debug("Processing object id=%s", o.id) + if o.type == "attack-action": + if o.id in new_operator_ids: + operator_type = [ + item["type"] for item in id_to_remove if item["next_id"] == o.id + ][0] + gv.node( + o.id, + label=_get_operator_label(o, operator_type), + shape="plaintext", + ) + else: + gv.node( + o.id, + _get_attack_tree_action_label(o), + shape="plaintext", + ) + for ref in o.get("asset_refs", []): + gv.edge(o.id, ref) + for ref in o.get("effect_refs", []): + gv.edge(o.id, ref) + elif o.type == "attack-asset": + gv.node(o.id, _get_asset_label(o), shape="plaintext") + if object_ref := o.get("object_ref"): + gv.edge(o.id, object_ref, "object") + elif o.type == "attack-condition": + gv.node(o.id, _get_condition_label(o), shape="plaintext") + for ref in o.get("on_true_refs", []): + gv.edge(o.id, ref, "on_true") + for ref in o.get("on_false_refs", []): + gv.edge(o.id, ref, "on_false") + elif o.type == "relationship": + gv.edge(o.source_ref, o.target_ref, o.relationship_type) + elif o.id not in ignored_ids: + gv.node(o.id, _get_builtin_label(o), shape="plaintext") + + return gv.source + + def _get_body_label(bundle): flow = get_flow_object(bundle) author = bundle.get_obj(flow.created_by_ref)[0] @@ -119,6 +200,33 @@ def _get_action_label(action): ) +def _get_attack_tree_action_label(action): + """ + Generate the GraphViz label for an action node as a table. + + :param action: + :rtype: str + """ + if tid := action.get("technique_id", None): + heading = f"Action: {tid}" + else: + heading = "Action" + description = "
".join( + textwrap.wrap(label_escape(action.get("description", "")), width=40) + ) + confidence = confidence_num_to_label(action.get("confidence", 95)) + return "".join( + [ + '<', + f'', + f'', + f'', + f'', + "
{heading}
Name{label_escape(action.name)}
Description{description}
Confidence{confidence}
>", + ] + ) + + def _get_asset_label(asset): """ Generate the GraphViz label for an asset node as a table. @@ -184,3 +292,34 @@ def _get_condition_label(condition): ">", ] ) + + +def _get_operator_label(action, operator_type): + """ + Generate the GraphViz label for an action node as a table. + + :param action: + :rtype: str + """ + if tid := action.get("technique_id", None): + heading = f"{operator_type} {tid}" + else: + heading = f"{operator_type}" + description = "
".join( + textwrap.wrap(label_escape(action.get("description", "")), width=40) + ) + confidence = confidence_num_to_label(action.get("confidence", 95)) + if operator_type == "AND": + color = "#99ccff" + else: + color = "#9CE67E" + return "".join( + [ + '<', + f'', + f'', + f'', + f'', + "
{heading}
Name{label_escape(action.name)}
Description{description}
Confidence{confidence}
>", + ] + ) diff --git a/src/attack_flow/mermaid.py b/src/attack_flow/mermaid.py index a0b121a9..b352fa10 100644 --- a/src/attack_flow/mermaid.py +++ b/src/attack_flow/mermaid.py @@ -20,6 +20,7 @@ def __init__(self): self.classes = dict() self.nodes = list() self.edges = list() + self.direction = "" def add_class(self, class_, shape, style): self.classes[class_] = (shape, style) @@ -33,7 +34,10 @@ def add_edge(self, src_id, target_id, text): def render(self): # Mermaid can't handle IDs with hyphens in them: convert_id = lambda id_: id_.replace("-", "_") - lines = ["graph TB"] + if self.direction: + lines = [f"graph {self.direction}"] + else: + lines = ["graph TB"] for class_, (_, style) in self.classes.items(): lines.append(f" classDef {class_} {style}") @@ -46,6 +50,9 @@ def render(self): if self.classes[node_class][0] == "circle": shape_start = "((" shape_end = "))" + elif self.classes[node_class][0] == "trap": + shape_start = "[/" + shape_end = "\]" else: shape_start = "[" shape_end = "]" @@ -64,7 +71,7 @@ def render(self): return "\n".join(lines) -def convert(bundle): +def convert_attack_flow(bundle): """ Convert an Attack Flow STIX bundle into Mermaid format. @@ -119,3 +126,94 @@ def convert(bundle): graph.add_node(o.id, "builtin", " - ".join(label_lines)) return graph.render() + + +def convert_attack_tree(bundle): + + """ + Convert an Attack Flow STIX bundle into Mermaid format. + + :param stix2.Bundle flow: + :rtype: str + """ + graph = MermaidGraph() + graph.direction = "BT" + graph.add_class("action", "rect", "fill:#B40000, color:white") + graph.add_class("AND", "rect", "fill:#99ccff") + graph.add_class("OR", "trap", "fill:#9CE67E") + graph.add_class("condition", "rect", "fill:#99ff99") + graph.add_class("builtin", "rect", "fill:#cccccc") + ignored_ids = get_viz_ignored_ids(bundle) + + objects = bundle.objects + id_to_remove = [] + ids = [] + + for i, o in enumerate(objects): + if o.type == "attack-operator": + id_to_remove.append( + { + "id": o.id, + "prev_id": objects[i - 1].id, + "next_id": o.effect_refs[0], + "type": o.operator, + } + ) + ids = [i["id"] for i in id_to_remove] + objects = [item for item in objects if item.id not in ids] + new_operator_ids = [i["next_id"] for i in id_to_remove] + for operator in id_to_remove: + for i, o in enumerate(objects): + if o.type == "relationship" and o.source_ref == operator["id"]: + o.source_ref = operator.prev_id + if o.type == "relationship" and o.target_ref == operator["id"]: + o.target_ref = operator.next_id + if o.get("effect_refs") and operator["id"] in o.effect_refs: + for i, j in enumerate(o.effect_refs): + if j == operator["id"]: + o.effect_refs[i] = operator["next_id"] + + for o in bundle.objects: + if o.type == "attack-action": + if tid := o.get("technique_id", None): + name = f"{tid} {o.name}" + else: + name = o.name + if o.id in new_operator_ids: + operator_type = [ + item["type"] for item in id_to_remove if item["next_id"] == o.id + ][0] + label_lines = [ + f"{operator_type}", + f"{name}", + ] + graph.add_node(o.id, operator_type, " - ".join(label_lines)) + else: + label_lines = [ + "Action", + f"{name}", + ] + graph.add_node(o.id, "action", " - ".join(label_lines)) + for ref in o.get("effect_refs", []): + graph.add_edge(o.id, ref, " ") + elif o.type == "attack-condition": + graph.add_node(o.id, "condition", f"Condition: {o.description}") + for ref in o.get("on_true_refs", []): + graph.add_edge(o.id, ref, "on_true") + for ref in o.get("on_false_refs", []): + graph.add_edge(o.id, ref, "on_false") + elif o.type == "relationship": + graph.add_edge(o.source_ref, o.target_ref, o.relationship_type) + elif o.id not in ignored_ids and o.id not in ids: + type_ = o.type.replace("-", " ").title() + label_lines = [f"{type_}"] + for key, value in o.items(): + if key in VIZ_IGNORE_COMMON_PROPERTIES: + continue + key = key.replace("_", " ").title() + if isinstance(value, list): + value = ", ".join(str(v) for v in value) + label_lines.append(f"{key}: {value}") + graph.add_node(o.id, "builtin", " - ".join(label_lines)) + + return graph.render() diff --git a/src/attack_flow_builder/src/assets/configuration/builder.config.ts b/src/attack_flow_builder/src/assets/configuration/builder.config.ts index dd91a4e5..1cab2754 100644 --- a/src/attack_flow_builder/src/assets/configuration/builder.config.ts +++ b/src/attack_flow_builder/src/assets/configuration/builder.config.ts @@ -93,6 +93,7 @@ const config: AppConfiguration = { ["campaign", "Campaign"], ["threat-actor", "Threat Actor"], ["malware", "Malware"], + ["attack-tree", "Attack Tree"], ["other", "Other"] ] }, diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py index 0292e217..b6deaf1c 100644 --- a/tests/fixtures/__init__.py +++ b/tests/fixtures/__init__.py @@ -136,3 +136,112 @@ def get_flow_bundle(): extension_creator, id="bundle--06cf9129-8d0d-4d58-9484-b5323caf09ad", ) + +def get_tree_bundle(): + asset_obj = stix2.Infrastructure( + id="infrastructure--79d21912-36b7-4af9-8958-38949dd0d6de", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="My Infra", + ) + asset = AttackAsset( + id="attack-asset--4ae37379-6a11-44c1-b6a8-d11733cfac06", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="My Asset", + object_ref=asset_obj.id, + ) + action3 = AttackAction( + id="attack-action--a0847849-a533-4b1f-a94a-720bbd25fc17", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="Action 3", + technique_id="T3", + description="Description of action 3", + asset_refs=[asset.id], + ) + or_action = AttackAction( + id="attack-action--1994e9f2-11f1-489a-a5e7-3ad4cfd8890a", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="My Or Operator", + technique_id="T3", + description="this is the description", + effect_refs=[action3.id] + ) + or_operator = AttackOperator( + id="attack-operator--8932b181-be87-4f81-851a-ab0b4288406a", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + operator="OR", + effect_refs=[or_action.id], + ) + action1 = AttackAction( + id="attack-action--d63857d5-1043-45a4-9397-40ef68db4c5f", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="Action 1", + description="Description of action 2", + effect_refs=[or_operator.id], + ) + action2 = AttackAction( + id="attack-action--24fc6003-33f6-4dd7-a929-b6031927940f", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="Action 2", + description="Description of action 2", + effect_refs=[or_operator.id], + ) + infra = stix2.Infrastructure( + id="infrastructure--a75c83f7-147e-4695-b173-0981521b2f01", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="Test Infra", + infrastructure_types=["workstation"], + ) + infra_rel = stix2.Relationship( + id="relationship--5286c903-9afc-4e29-ab42-644976d3aae7", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + source_ref=action2.id, + target_ref=infra.id, + relationship_type="related-to", + ) + author = stix2.Identity( + id="identity--bbe39bd7-9c12-41de-b5c0-dcd3fb98b360", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="Jane Doe", + contact_information="jdoe@example.com", + ) + flow = AttackFlow( + id="attack-flow--7cabcb58-6930-47b9-b15c-3be2f3a5fce1", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="My Flow", + start_refs=[action1.id, action2.id], + created_by_ref=author.id, + ) + condition = AttackCondition( + id="attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + description="My condition", + on_true_refs=[action1.id], + on_false_refs=[action2.id], + ) + return stix2.Bundle( + flow, + author, + action1, + or_action, + action2, + or_operator, + action3, + asset_obj, + asset, + infra, + infra_rel, + condition, + id="bundle--06cf9129-8d0d-4d58-9484-b5323caf09ad", + ) \ No newline at end of file diff --git a/tests/test_cli.py b/tests/test_cli.py index 0fbf9702..2d1bf531 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -16,6 +16,8 @@ import stix2 import attack_flow.schema +from attack_flow.model import AttackFlow +from datetime import datetime @patch("sys.exit") @@ -95,12 +97,12 @@ def test_doc_schema(schema_mock, generate_mock, insert_mock, exit_mock): @patch("sys.exit") -@patch("attack_flow.graphviz.convert") +@patch("attack_flow.graphviz.convert_attack_flow") @patch("attack_flow.model.load_attack_flow_bundle") -def test_graphviz(load_mock, convert_mock, exit_mock): +def test_graphviz_attack_flow(load_mock, convert_mock, exit_mock): """ Test that the script parses a JSON file and passes the resulting object - to convert(). + to convert_attack_flow(). """ convert_mock.return_value = dedent( r"""\ @@ -111,6 +113,7 @@ def test_graphviz(load_mock, convert_mock, exit_mock): ) bundle = stix2.Bundle() load_mock.return_value = bundle + with NamedTemporaryFile() as flow, NamedTemporaryFile() as graphviz: sys.argv = ["af", "graphviz", flow.name, graphviz.name] runpy.run_module("attack_flow.cli", run_name="__main__") @@ -119,11 +122,46 @@ def test_graphviz(load_mock, convert_mock, exit_mock): convert_mock.assert_called_with(bundle) exit_mock.assert_called_with(0) +@patch("sys.exit") +@patch("attack_flow.graphviz.convert_attack_tree") +@patch("attack_flow.model.load_attack_flow_bundle") +def test_graphviz_attack_tree(load_mock, convert_mock, exit_mock): + """ + Test that the script parses a JSON file and passes the resulting object + to convert_attack_flow(). + """ + convert_mock.return_value = dedent( + r"""\ + graph { + "node1" -> "node2"; + } + """ + ) + + flow = AttackFlow( + id="attack-flow--7cabcb58-6930-47b9-b15c-3be2f3a5fce1", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="My Flow", + start_refs=[], + created_by_ref="identity--bbe39bd7-9c12-41de-b5c0-dcd3fb98b360", + scope="attack-tree" + ) + bundle = stix2.Bundle(flow) + load_mock.return_value = bundle + + with NamedTemporaryFile() as flow, NamedTemporaryFile() as graphviz: + sys.argv = ["af", "graphviz", flow.name, graphviz.name] + runpy.run_module("attack_flow.cli", run_name="__main__") + load_mock.assert_called() + assert str(load_mock.call_args[0][0]) == flow.name + convert_mock.assert_called_with(bundle) + exit_mock.assert_called_with(0) @patch("sys.exit") -@patch("attack_flow.mermaid.convert") +@patch("attack_flow.mermaid.convert_attack_flow") @patch("attack_flow.model.load_attack_flow_bundle") -def test_mermaid(load_mock, convert_mock, exit_mock): +def test_mermaid_attack_flow(load_mock, convert_mock, exit_mock): """ Test that the script parses a JSON file and passes the resulting object to convert(). @@ -144,6 +182,38 @@ def test_mermaid(load_mock, convert_mock, exit_mock): convert_mock.assert_called_with(bundle) exit_mock.assert_called_with(0) +@patch("sys.exit") +@patch("attack_flow.mermaid.convert_attack_tree") +@patch("attack_flow.model.load_attack_flow_bundle") +def test_mermaid_attack_tree(load_mock, convert_mock, exit_mock): + """ + Test that the script parses a JSON file and passes the resulting object + to convert(). + """ + convert_mock.return_value = dedent( + r"""\ + graph TB + node1 ---> node2 + """ + ) + flow = AttackFlow( + id="attack-flow--7cabcb58-6930-47b9-b15c-3be2f3a5fce1", + created=datetime(2022, 8, 25, 19, 26, 31), + modified=datetime(2022, 8, 25, 19, 26, 31), + name="My Flow", + start_refs=[], + created_by_ref="identity--bbe39bd7-9c12-41de-b5c0-dcd3fb98b360", + scope="attack-tree" + ) + bundle = stix2.Bundle(flow) + load_mock.return_value = bundle + with NamedTemporaryFile() as flow, NamedTemporaryFile() as graphviz: + sys.argv = ["af", "mermaid", flow.name, graphviz.name] + runpy.run_module("attack_flow.cli", run_name="__main__") + load_mock.assert_called() + assert str(load_mock.call_args[0][0]) == flow.name + convert_mock.assert_called_with(bundle) + exit_mock.assert_called_with(0) @patch("sys.exit") @patch("attack_flow.matrix.render") diff --git a/tests/test_graphviz.py b/tests/test_graphviz.py index 088825cc..55bdeaaa 100644 --- a/tests/test_graphviz.py +++ b/tests/test_graphviz.py @@ -1,15 +1,14 @@ from textwrap import dedent - import attack_flow.graphviz from attack_flow.model import ( AttackAction, AttackCondition, ) -from .fixtures import get_flow_bundle +from .fixtures import get_flow_bundle, get_tree_bundle def test_convert_attack_flow_to_graphviz(): - output = attack_flow.graphviz.convert(get_flow_bundle()) + output = attack_flow.graphviz.convert_attack_flow(get_flow_bundle()) assert output == dedent( """\ digraph { @@ -37,6 +36,35 @@ def test_convert_attack_flow_to_graphviz(): ) +def test_convert_attack_tree_to_graphviz(): + output = attack_flow.graphviz.convert_attack_tree(get_tree_bundle()) + assert output == dedent( + """\ + digraph { + \tgraph [rankdir=BT] + \tlabel=<My Flow
(missing description)
Author: Jane Doe <jdoe@example.com>
Created: 2022-08-25 19:26:31
Modified: 2022-08-25 19:26:31>; + \tlabelloc="t"; + \t"attack-action--d63857d5-1043-45a4-9397-40ef68db4c5f" [label=<
Action
NameAction 1
DescriptionDescription of action 2
ConfidenceVery Probable
> shape=plaintext] + \t"attack-action--d63857d5-1043-45a4-9397-40ef68db4c5f" -> "attack-action--1994e9f2-11f1-489a-a5e7-3ad4cfd8890a" + \t"attack-action--1994e9f2-11f1-489a-a5e7-3ad4cfd8890a" [label=<
OR T3
NameMy Or Operator
Descriptionthis is the description
ConfidenceVery Probable
> shape=plaintext] + \t"attack-action--1994e9f2-11f1-489a-a5e7-3ad4cfd8890a" -> "attack-action--a0847849-a533-4b1f-a94a-720bbd25fc17" + \t"attack-action--24fc6003-33f6-4dd7-a929-b6031927940f" [label=<
Action
NameAction 2
DescriptionDescription of action 2
ConfidenceVery Probable
> shape=plaintext] + \t"attack-action--24fc6003-33f6-4dd7-a929-b6031927940f" -> "attack-action--1994e9f2-11f1-489a-a5e7-3ad4cfd8890a" + \t"attack-action--a0847849-a533-4b1f-a94a-720bbd25fc17" [label=<
Action: T3
NameAction 3
DescriptionDescription of action 3
ConfidenceVery Probable
> shape=plaintext] + \t"attack-action--a0847849-a533-4b1f-a94a-720bbd25fc17" -> "attack-asset--4ae37379-6a11-44c1-b6a8-d11733cfac06" + \t"infrastructure--79d21912-36b7-4af9-8958-38949dd0d6de" [label=<
Infrastructure
NameMy Infra
> shape=plaintext] + \t"attack-asset--4ae37379-6a11-44c1-b6a8-d11733cfac06" [label=<
Asset: My Asset
Description
> shape=plaintext] + \t"attack-asset--4ae37379-6a11-44c1-b6a8-d11733cfac06" -> "infrastructure--79d21912-36b7-4af9-8958-38949dd0d6de" [label=object] + \t"infrastructure--a75c83f7-147e-4695-b173-0981521b2f01" [label=<
Infrastructure
NameTest Infra
Infrastructure Typesworkstation
> shape=plaintext] + \t"attack-action--24fc6003-33f6-4dd7-a929-b6031927940f" -> "infrastructure--a75c83f7-147e-4695-b173-0981521b2f01" [label="related-to"] + \t"attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a" [label=<
Condition
DescriptionMy condition
> shape=plaintext] + \t"attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a" -> "attack-action--d63857d5-1043-45a4-9397-40ef68db4c5f" [label=on_true] + \t"attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a" -> "attack-action--24fc6003-33f6-4dd7-a929-b6031927940f" [label=on_false] + } + """ + ) + + def test_wrap_action_description(): """Long descriptions should be wrapped.""" action = AttackAction( @@ -75,3 +103,25 @@ def test_action_label(): attack_flow.graphviz._get_action_label(action) == '<
Action
NameMy technique
DescriptionThis technique has no ID to render in
the header.
ConfidenceVery Probable
>' ) + +def test_get_operator_label(): + action = AttackAction( + id="attack-action--b5696498-66e8-41b6-87e1-19d2657ac48b", + name="My technique", + description="This technique has no ID to render in the header.", + ) + assert ( + attack_flow.graphviz._get_operator_label(action, operator_type="AND") + == '<
AND
NameMy technique
DescriptionThis technique has no ID to render in
the header.
ConfidenceVery Probable
>' + ) + +def test_get_attack_tree_action_label(): + action = AttackAction( + id="attack-action--b5696498-66e8-41b6-87e1-19d2657ac48b", + name="My technique", + description="This technique has no ID to render in the header.", + ) + assert ( + attack_flow.graphviz._get_attack_tree_action_label(action) + == '<
Action
NameMy technique
DescriptionThis technique has no ID to render in
the header.
ConfidenceVery Probable
>' + ) diff --git a/tests/test_mermaid.py b/tests/test_mermaid.py index 74784015..99cc13c3 100644 --- a/tests/test_mermaid.py +++ b/tests/test_mermaid.py @@ -1,11 +1,10 @@ from textwrap import dedent -from .fixtures import get_flow_bundle +from .fixtures import get_flow_bundle, get_tree_bundle import attack_flow.mermaid - def test_convert_attack_flow_to_mermaid(): - output = attack_flow.mermaid.convert(get_flow_bundle()) + output = attack_flow.mermaid.convert_attack_flow(get_flow_bundle()) assert output == dedent( """\ graph TB @@ -41,3 +40,40 @@ class infrastructure__a75c83f7_147e_4695_b173_0981521b2f01 builtin attack_action__dd3820fa_bae3_4270_8000_5c4642fa780c -->|related-to| infrastructure__a75c83f7_147e_4695_b173_0981521b2f01 """ ) + +def test_convert_attack_tree_to_mermaid(): + output = attack_flow.mermaid.convert_attack_tree(get_tree_bundle()) + assert output == dedent( + """\ + graph BT + classDef action fill:#B40000, color:white + classDef AND fill:#99ccff + classDef OR fill:#9CE67E + classDef condition fill:#99ff99 + classDef builtin fill:#cccccc + + attack_action__d63857d5_1043_45a4_9397_40ef68db4c5f["Action - Action 1"] + class attack_action__d63857d5_1043_45a4_9397_40ef68db4c5f action + attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a[/"OR - T3 My Or Operator"\\] + class attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a OR + attack_action__24fc6003_33f6_4dd7_a929_b6031927940f["Action - Action 2"] + class attack_action__24fc6003_33f6_4dd7_a929_b6031927940f action + attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17["Action - T3 Action 3"] + class attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17 action + infrastructure__79d21912_36b7_4af9_8958_38949dd0d6de["Infrastructure - Name: My
Infra"] + class infrastructure__79d21912_36b7_4af9_8958_38949dd0d6de builtin + attack_asset__4ae37379_6a11_44c1_b6a8_d11733cfac06["Attack Asset - Name: My
Asset - Object Ref:
infrastructure--
79d21912-36b7-4af9-8958-38949dd0d6de"] + class attack_asset__4ae37379_6a11_44c1_b6a8_d11733cfac06 builtin + infrastructure__a75c83f7_147e_4695_b173_0981521b2f01["Infrastructure - Name:
Test Infra - Infrastructure
Types: workstation"] + class infrastructure__a75c83f7_147e_4695_b173_0981521b2f01 builtin + attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a["Condition: My condition"] + class attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a condition + + attack_action__d63857d5_1043_45a4_9397_40ef68db4c5f -->| | attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a + attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a -->| | attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17 + attack_action__24fc6003_33f6_4dd7_a929_b6031927940f -->| | attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a + attack_action__24fc6003_33f6_4dd7_a929_b6031927940f -->|related-to| infrastructure__a75c83f7_147e_4695_b173_0981521b2f01 + attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a -->|on_true| attack_action__d63857d5_1043_45a4_9397_40ef68db4c5f + attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a -->|on_false| attack_action__24fc6003_33f6_4dd7_a929_b6031927940f + """ + )