From af41c8e142a9e1a58773891297e520ee6b4b0172 Mon Sep 17 00:00:00 2001 From: arobbins Date: Fri, 31 May 2024 13:47:06 -0400 Subject: [PATCH 01/11] configure build to allow attack tree rendering --- src/attack_flow/cli.py | 8 +- src/attack_flow/graphviz.py | 132 +++++++++++++++++- .../assets/configuration/builder.config.ts | 1 + 3 files changed, 138 insertions(+), 3 deletions(-) diff --git a/src/attack_flow/cli.py b/src/attack_flow/cli.py index da8c9cfb..dfffad7d 100644 --- a/src/attack_flow/cli.py +++ b/src/attack_flow/cli.py @@ -91,7 +91,13 @@ 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) + for object in flow_bundle.objects: + if object.type == "attack-flow": + if object.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 diff --git a/src/attack_flow/graphviz.py b/src/attack_flow/graphviz.py index c7205dde..188e12d3 100644 --- a/src/attack_flow/graphviz.py +++ b/src/attack_flow/graphviz.py @@ -18,8 +18,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. @@ -70,6 +69,81 @@ 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() + 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, "asset") + for ref in o.get("effect_refs", []): + gv.edge(o.id, ref, "effect") + 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) @@ -117,7 +191,31 @@ 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): """ @@ -184,3 +282,33 @@ 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}
>", + ] + ) \ No newline at end of file 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 f245c93e..21c328ad 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", "ATT&CK Tree"], ["other", "Other"] ] }, From a78aa7d45caf0d60c29256845ca0e547e35f151f Mon Sep 17 00:00:00 2001 From: arobbins Date: Mon, 3 Jun 2024 15:46:00 -0400 Subject: [PATCH 02/11] render Attack Tree nodes bottom to top, remove labels from edges, fix typo in config --- src/attack_flow/graphviz.py | 6 +++--- .../src/assets/configuration/builder.config.ts | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/attack_flow/graphviz.py b/src/attack_flow/graphviz.py index 188e12d3..238cef61 100644 --- a/src/attack_flow/graphviz.py +++ b/src/attack_flow/graphviz.py @@ -77,7 +77,7 @@ def convert_attack_tree(bundle): :rtype: str """ - gv = graphviz.Digraph() + gv = graphviz.Digraph(graph_attr={'rankdir':'BT'}) gv.body = _get_body_label(bundle) ignored_ids = get_viz_ignored_ids(bundle) @@ -124,9 +124,9 @@ def convert_attack_tree(bundle): shape="plaintext", ) for ref in o.get("asset_refs", []): - gv.edge(o.id, ref, "asset") + gv.edge(o.id, ref) for ref in o.get("effect_refs", []): - gv.edge(o.id, ref, "effect") + 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"): 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 49f6d1b7..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,7 +93,7 @@ const config: AppConfiguration = { ["campaign", "Campaign"], ["threat-actor", "Threat Actor"], ["malware", "Malware"], - ["attack-tree", "ATT&CK Tree"], + ["attack-tree", "Attack Tree"], ["other", "Other"] ] }, From 2f32e5d26b7f1ed5e9ecdc09013ba5c84b2816ac Mon Sep 17 00:00:00 2001 From: arobbins Date: Wed, 19 Jun 2024 14:33:11 -0400 Subject: [PATCH 03/11] lint code and fix unit tests --- src/attack_flow/cli.py | 2 +- src/attack_flow/graphviz.py | 41 +++++++++++++++++++++++-------------- tests/test_cli.py | 5 +++-- tests/test_graphviz.py | 2 +- 4 files changed, 31 insertions(+), 19 deletions(-) diff --git a/src/attack_flow/cli.py b/src/attack_flow/cli.py index dfffad7d..09109759 100644 --- a/src/attack_flow/cli.py +++ b/src/attack_flow/cli.py @@ -97,7 +97,7 @@ def graphviz(args): 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 diff --git a/src/attack_flow/graphviz.py b/src/attack_flow/graphviz.py index 238cef61..08901a3d 100644 --- a/src/attack_flow/graphviz.py +++ b/src/attack_flow/graphviz.py @@ -18,6 +18,7 @@ def label_escape(text): return graphviz.escape(html.escape(text)) + def convert_attack_flow(bundle): """ Convert an Attack Flow STIX bundle into Graphviz format. @@ -69,6 +70,7 @@ def convert_attack_flow(bundle): return gv.source + def convert_attack_tree(bundle): """ Convert an Attack Flow STIX bundle into Graphviz format. @@ -77,7 +79,7 @@ def convert_attack_tree(bundle): :rtype: str """ - gv = graphviz.Digraph(graph_attr={'rankdir':'BT'}) + gv = graphviz.Digraph(graph_attr={"rankdir": "BT"}) gv.body = _get_body_label(bundle) ignored_ids = get_viz_ignored_ids(bundle) @@ -85,10 +87,14 @@ def convert_attack_tree(bundle): id_to_remove = [] ids = [] - for i,o in enumerate(objects): + 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 + { + "id": o.id, + "prev_id": objects[i - 1].id, + "next_id": o.effect_refs[0], + "type": o.operator, } ) @@ -96,22 +102,23 @@ def convert_attack_tree(bundle): 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 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] + 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), @@ -191,6 +198,8 @@ def _get_action_label(action): ">", ] ) + + def _get_attack_tree_action_label(action): """ Generate the GraphViz label for an action node as a table. @@ -217,6 +226,7 @@ def _get_attack_tree_action_label(action): ] ) + def _get_asset_label(asset): """ Generate the GraphViz label for an asset node as a table. @@ -283,6 +293,7 @@ def _get_condition_label(condition): ] ) + def _get_operator_label(action, operator_type): """ Generate the GraphViz label for an action node as a table. @@ -311,4 +322,4 @@ def _get_operator_label(action, operator_type): f'Confidence{confidence}', ">", ] - ) \ No newline at end of file + ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 0fbf9702..83f434f1 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -95,12 +95,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): """ 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 +111,7 @@ def test_graphviz(load_mock, convert_mock, exit_mock): ) bundle = stix2.Bundle() load_mock.return_value = bundle + print("printing resp bundle ", 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__") diff --git a/tests/test_graphviz.py b/tests/test_graphviz.py index 088825cc..26341449 100644 --- a/tests/test_graphviz.py +++ b/tests/test_graphviz.py @@ -9,7 +9,7 @@ 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 { From 34476adff6fa81804ffa37227ec68dd5931c7129 Mon Sep 17 00:00:00 2001 From: arobbins Date: Wed, 19 Jun 2024 14:47:40 -0400 Subject: [PATCH 04/11] fix unit test by adding check for objects in graph bundle --- src/attack_flow/cli.py | 15 +++++++++------ tests/test_cli.py | 2 +- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/attack_flow/cli.py b/src/attack_flow/cli.py index 09109759..8c45f4df 100644 --- a/src/attack_flow/cli.py +++ b/src/attack_flow/cli.py @@ -91,12 +91,15 @@ def graphviz(args): """ path = Path(args.attack_flow) flow_bundle = attack_flow.model.load_attack_flow_bundle(path) - for object in flow_bundle.objects: - if object.type == "attack-flow": - if object.scope == "attack-tree": - converted = attack_flow.graphviz.convert_attack_tree(flow_bundle) - else: - converted = attack_flow.graphviz.convert_attack_flow(flow_bundle) + if not flow_bundle.get("objects", ""): + converted = attack_flow.graphviz.convert_attack_flow(flow_bundle) + else: + for object in flow_bundle.objects: + if object.type == "attack-flow": + if object.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) diff --git a/tests/test_cli.py b/tests/test_cli.py index 83f434f1..14868380 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -111,7 +111,7 @@ def test_graphviz(load_mock, convert_mock, exit_mock): ) bundle = stix2.Bundle() load_mock.return_value = bundle - print("printing resp bundle ", 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__") From 57eadedf5d97615c432d2f9961bdd64143162a7f Mon Sep 17 00:00:00 2001 From: arobbins Date: Tue, 25 Jun 2024 15:55:54 -0400 Subject: [PATCH 05/11] add test coverage for graphviz --- src/attack_flow/cli.py | 1 + src/attack_flow/graphviz.py | 12 +++--- tests/fixtures/__init__.py | 82 +++++++++++++++++++++++++++++++++++++ tests/test_graphviz.py | 55 ++++++++++++++++++++++++- 4 files changed, 143 insertions(+), 7 deletions(-) diff --git a/src/attack_flow/cli.py b/src/attack_flow/cli.py index 8c45f4df..2d20d1d1 100644 --- a/src/attack_flow/cli.py +++ b/src/attack_flow/cli.py @@ -100,6 +100,7 @@ def graphviz(args): converted = attack_flow.graphviz.convert_attack_tree(flow_bundle) else: converted = attack_flow.graphviz.convert_attack_flow(flow_bundle) + break with open(args.output, "w") as out: out.write(converted) diff --git a/src/attack_flow/graphviz.py b/src/attack_flow/graphviz.py index 08901a3d..9d2f66fc 100644 --- a/src/attack_flow/graphviz.py +++ b/src/attack_flow/graphviz.py @@ -219,9 +219,9 @@ def _get_attack_tree_action_label(action): [ '<', f'', - f'', - f'', - f'', + f'', + f'', + f'', "
{heading}
Name{label_escape(action.name)}
Description{description}
Confidence{confidence}
Name{label_escape(action.name)}
Description{description}
Confidence{confidence}
>", ] ) @@ -317,9 +317,9 @@ def _get_operator_label(action, operator_type): [ '<', f'', - f'', - f'', - f'', + f'', + f'', + f'', "
{heading}
Name{label_escape(action.name)}
Description{description}
Confidence{confidence}
Name{label_escape(action.name)}
Description{description}
Confidence{confidence}
>", ] ) diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py index 0292e217..1dc943a7 100644 --- a/tests/fixtures/__init__.py +++ b/tests/fixtures/__init__.py @@ -136,3 +136,85 @@ 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", + 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", + 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], + ) + + 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, + ) + return stix2.Bundle( + flow, + author, + action1, + or_action, + action2, + or_operator, + action3, + asset_obj, + asset, + id="bundle--06cf9129-8d0d-4d58-9484-b5323caf09ad", + ) \ No newline at end of file diff --git a/tests/test_graphviz.py b/tests/test_graphviz.py index 26341449..069893ce 100644 --- a/tests/test_graphviz.py +++ b/tests/test_graphviz.py @@ -5,7 +5,8 @@ AttackAction, AttackCondition, ) -from .fixtures import get_flow_bundle +from .fixtures import get_flow_bundle, get_tree_bundle +import json def test_convert_attack_flow_to_graphviz(): @@ -37,6 +38,36 @@ def test_convert_attack_flow_to_graphviz(): ) +def test_convert_attack_tree_to_graphviz(): + output = attack_flow.graphviz.convert_attack_tree(get_tree_bundle()) + # Serializing json + json_object = json.dumps(output, indent=4) + + # Writing to sample.json + with open("sample.json", "w") as outfile: + outfile.write(json_object) + 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
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
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] + } + """ + ) + + def test_wrap_action_description(): """Long descriptions should be wrapped.""" action = AttackAction( @@ -75,3 +106,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
>' + ) From f6501dffecc89bd6107f71b86686345be89b43db Mon Sep 17 00:00:00 2001 From: arobbins Date: Wed, 26 Jun 2024 12:31:23 -0400 Subject: [PATCH 06/11] add attack tree code for mermaid --- src/attack_flow/cli.py | 18 ++++--- src/attack_flow/mermaid.py | 102 ++++++++++++++++++++++++++++++++++++- 2 files changed, 110 insertions(+), 10 deletions(-) diff --git a/src/attack_flow/cli.py b/src/attack_flow/cli.py index 2d20d1d1..600229a6 100644 --- a/src/attack_flow/cli.py +++ b/src/attack_flow/cli.py @@ -93,14 +93,10 @@ def graphviz(args): flow_bundle = attack_flow.model.load_attack_flow_bundle(path) if not flow_bundle.get("objects", ""): converted = attack_flow.graphviz.convert_attack_flow(flow_bundle) + elif attack_flow.model.get_flow_object(flow_bundle).scope == "attack-tree": + converted = attack_flow.graphviz.convert_attack_tree(flow_bundle) else: - for object in flow_bundle.objects: - if object.type == "attack-flow": - if object.scope == "attack-tree": - converted = attack_flow.graphviz.convert_attack_tree(flow_bundle) - else: - converted = attack_flow.graphviz.convert_attack_flow(flow_bundle) - break + converted = attack_flow.graphviz.convert_attack_flow(flow_bundle) with open(args.output, "w") as out: out.write(converted) @@ -116,7 +112,13 @@ 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 not flow_bundle.get("objects", ""): + converted = attack_flow.mermaid.convert_attack_flow(flow_bundle) + elif 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/mermaid.py b/src/attack_flow/mermaid.py index a0b121a9..19a257b9 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() From dcdc920f5851b4cdabbb481b89b331c4c36964e8 Mon Sep 17 00:00:00 2001 From: arobbins Date: Wed, 26 Jun 2024 12:55:52 -0400 Subject: [PATCH 07/11] add tests for mermaid --- tests/test_cli.py | 2 +- tests/test_graphviz.py | 8 +------- tests/test_mermaid.py | 35 ++++++++++++++++++++++++++++++++--- 3 files changed, 34 insertions(+), 11 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index 14868380..f7586579 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -122,7 +122,7 @@ def test_graphviz(load_mock, convert_mock, exit_mock): @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): """ diff --git a/tests/test_graphviz.py b/tests/test_graphviz.py index 069893ce..2c75a024 100644 --- a/tests/test_graphviz.py +++ b/tests/test_graphviz.py @@ -6,7 +6,6 @@ AttackCondition, ) from .fixtures import get_flow_bundle, get_tree_bundle -import json def test_convert_attack_flow_to_graphviz(): @@ -40,12 +39,7 @@ def test_convert_attack_flow_to_graphviz(): def test_convert_attack_tree_to_graphviz(): output = attack_flow.graphviz.convert_attack_tree(get_tree_bundle()) - # Serializing json - json_object = json.dumps(output, indent=4) - - # Writing to sample.json - with open("sample.json", "w") as outfile: - outfile.write(json_object) + assert output == dedent( """\ digraph { diff --git a/tests/test_mermaid.py b/tests/test_mermaid.py index 74784015..d4559c43 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,33 @@ 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 - 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 - 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 + + 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 + """ + ) From 0cd10622b2298cdb4fe5746ea25e7c47a933682a Mon Sep 17 00:00:00 2001 From: arobbins Date: Wed, 26 Jun 2024 12:57:43 -0400 Subject: [PATCH 08/11] linting fix --- src/attack_flow/mermaid.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/attack_flow/mermaid.py b/src/attack_flow/mermaid.py index 19a257b9..b352fa10 100644 --- a/src/attack_flow/mermaid.py +++ b/src/attack_flow/mermaid.py @@ -34,7 +34,7 @@ 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("-", "_") - if self.direction: + if self.direction: lines = [f"graph {self.direction}"] else: lines = ["graph TB"] @@ -50,7 +50,7 @@ def render(self): if self.classes[node_class][0] == "circle": shape_start = "((" shape_end = "))" - elif self.classes[node_class][0] =="trap": + elif self.classes[node_class][0] == "trap": shape_start = "[/" shape_end = "\]" else: @@ -126,7 +126,8 @@ def convert_attack_flow(bundle): graph.add_node(o.id, "builtin", " - ".join(label_lines)) return graph.render() - + + def convert_attack_tree(bundle): """ @@ -147,7 +148,7 @@ def convert_attack_tree(bundle): objects = bundle.objects id_to_remove = [] ids = [] - + for i, o in enumerate(objects): if o.type == "attack-operator": id_to_remove.append( @@ -172,7 +173,6 @@ def convert_attack_tree(bundle): 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): From f1c081e25056ccb0000e85990075a94d9c5fd8be Mon Sep 17 00:00:00 2001 From: arobbins Date: Fri, 28 Jun 2024 14:26:03 -0400 Subject: [PATCH 09/11] add up to 98% test coverage --- src/attack_flow/cli.py | 9 ++--- tests/fixtures/__init__.py | 20 ++++++++++- tests/test_cli.py | 73 ++++++++++++++++++++++++++++++++++++-- tests/test_graphviz.py | 7 ++-- tests/test_mermaid.py | 47 ++++++++++++------------ 5 files changed, 122 insertions(+), 34 deletions(-) diff --git a/src/attack_flow/cli.py b/src/attack_flow/cli.py index 600229a6..1bac46ee 100644 --- a/src/attack_flow/cli.py +++ b/src/attack_flow/cli.py @@ -91,9 +91,8 @@ def graphviz(args): """ path = Path(args.attack_flow) flow_bundle = attack_flow.model.load_attack_flow_bundle(path) - if not flow_bundle.get("objects", ""): - converted = attack_flow.graphviz.convert_attack_flow(flow_bundle) - elif attack_flow.model.get_flow_object(flow_bundle).scope == "attack-tree": + + 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) @@ -112,9 +111,7 @@ def mermaid(args): """ path = Path(args.attack_flow) flow_bundle = attack_flow.model.load_attack_flow_bundle(path) - if not flow_bundle.get("objects", ""): - converted = attack_flow.mermaid.convert_attack_flow(flow_bundle) - elif attack_flow.model.get_flow_object(flow_bundle).scope == "attack-tree": + 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) diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py index 1dc943a7..05cd045a 100644 --- a/tests/fixtures/__init__.py +++ b/tests/fixtures/__init__.py @@ -156,6 +156,7 @@ def get_tree_bundle(): 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], ) @@ -164,6 +165,7 @@ def get_tree_bundle(): 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] ) @@ -190,7 +192,21 @@ def get_tree_bundle(): 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), @@ -216,5 +232,7 @@ def get_tree_bundle(): action3, asset_obj, asset, + infra, + infra_rel, 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 f7586579..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") @@ -97,7 +99,7 @@ def test_doc_schema(schema_mock, generate_mock, insert_mock, exit_mock): @patch("sys.exit") @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_attack_flow(). @@ -120,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_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(). @@ -145,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 2c75a024..e5728f01 100644 --- a/tests/test_graphviz.py +++ b/tests/test_graphviz.py @@ -1,5 +1,4 @@ from textwrap import dedent - import attack_flow.graphviz from attack_flow.model import ( AttackAction, @@ -48,15 +47,17 @@ def test_convert_attack_tree_to_graphviz(): \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
NameMy Or Operator
Descriptionthis is the description
ConfidenceVery Probable
> shape=plaintext] + \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
NameAction 3
DescriptionDescription of action 3
ConfidenceVery Probable
> shape=plaintext] + \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"] } """ ) diff --git a/tests/test_mermaid.py b/tests/test_mermaid.py index d4559c43..0203d170 100644 --- a/tests/test_mermaid.py +++ b/tests/test_mermaid.py @@ -45,28 +45,31 @@ 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 + 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 - 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 - 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 + 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_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__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 + """ ) From 6cae3b4a82e160a0b741cbb8217c546c5d3e2dfc Mon Sep 17 00:00:00 2001 From: arobbins Date: Fri, 28 Jun 2024 14:27:30 -0400 Subject: [PATCH 10/11] linting --- src/attack_flow/cli.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/attack_flow/cli.py b/src/attack_flow/cli.py index 1bac46ee..5b4ec315 100644 --- a/src/attack_flow/cli.py +++ b/src/attack_flow/cli.py @@ -91,8 +91,11 @@ def graphviz(args): """ path = Path(args.attack_flow) flow_bundle = attack_flow.model.load_attack_flow_bundle(path) - - if flow_bundle.get("objects", "") and attack_flow.model.get_flow_object(flow_bundle).scope == "attack-tree": + + 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) @@ -111,7 +114,10 @@ def mermaid(args): """ path = Path(args.attack_flow) flow_bundle = attack_flow.model.load_attack_flow_bundle(path) - if flow_bundle.get("objects", "") and attack_flow.model.get_flow_object(flow_bundle).scope == "attack-tree": + 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) From 8571e32755fd20eb7e1e30a5907f46ddafb950f7 Mon Sep 17 00:00:00 2001 From: arobbins Date: Fri, 28 Jun 2024 15:13:32 -0400 Subject: [PATCH 11/11] add last 1% for test coverage --- tests/fixtures/__init__.py | 9 +++++++++ tests/test_graphviz.py | 4 +++- tests/test_mermaid.py | 4 ++++ 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py index 05cd045a..b6deaf1c 100644 --- a/tests/fixtures/__init__.py +++ b/tests/fixtures/__init__.py @@ -222,6 +222,14 @@ def get_tree_bundle(): 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, @@ -234,5 +242,6 @@ def get_tree_bundle(): asset, infra, infra_rel, + condition, id="bundle--06cf9129-8d0d-4d58-9484-b5323caf09ad", ) \ No newline at end of file diff --git a/tests/test_graphviz.py b/tests/test_graphviz.py index e5728f01..55bdeaaa 100644 --- a/tests/test_graphviz.py +++ b/tests/test_graphviz.py @@ -38,7 +38,6 @@ 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 { @@ -58,6 +57,9 @@ def test_convert_attack_tree_to_graphviz(): \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] } """ ) diff --git a/tests/test_mermaid.py b/tests/test_mermaid.py index 0203d170..99cc13c3 100644 --- a/tests/test_mermaid.py +++ b/tests/test_mermaid.py @@ -66,10 +66,14 @@ class infrastructure__79d21912_36b7_4af9_8958_38949dd0d6de builtin 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 """ )