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
>' + )