from textwrap import dedent import attack_flow.graphviz from attack_flow.model import ( AttackAction, AttackCondition, ) from .fixtures import get_flow_bundle, get_tree_bundle def test_convert_attack_flow_to_graphviz(): output = attack_flow.graphviz.convert_attack_flow(get_flow_bundle()) assert output == dedent( """\ digraph { \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--52f2c35a-fa2a-45a4-b84c-46ad9498071f" [label=<
Action: T1
NameAction 1
DescriptionDescription of action 1
ConfidenceVery Probable
> shape=plaintext] \t"attack-action--52f2c35a-fa2a-45a4-b84c-46ad9498071f" -> "attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a" [label=effect] \t"attack-action--dd3820fa-bae3-4270-8000-5c4642fa780c" [label=<
Action
NameAction 2
DescriptionDescription of action 2
ConfidenceVery Probable
> shape=plaintext] \t"attack-action--dd3820fa-bae3-4270-8000-5c4642fa780c" -> "attack-asset--4ae37379-6a11-44c1-b6a8-d11733cfac06" [label=asset] \t"attack-action--a0847849-a533-4b1f-a94a-720bbd25fc17" [label=<
Action: T3
NameAction 3
DescriptionDescription of action 3
ConfidenceVery Probable
> shape=plaintext] \t"attack-action--7ddab166-c83e-4c79-a701-a0dc2a905dd3" [label=<
Action: T4
NameAction 4
DescriptionDescription of action 4
ConfidenceVery Probable
> shape=plaintext] \t"attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a" [label=<
Condition
DescriptionMy condition
> shape=plaintext] \t"attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a" -> "attack-operator--8932b181-be87-4f81-851a-ab0b4288406a" [label=on_true] \t"attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a" -> "attack-action--7ddab166-c83e-4c79-a701-a0dc2a905dd3" [label=on_false] \t"attack-operator--8932b181-be87-4f81-851a-ab0b4288406a" [label=OR fillcolor="#ff9900" shape=circle style=filled] \t"attack-operator--8932b181-be87-4f81-851a-ab0b4288406a" -> "attack-action--dd3820fa-bae3-4270-8000-5c4642fa780c" [label=effect] \t"attack-operator--8932b181-be87-4f81-851a-ab0b4288406a" -> "attack-action--a0847849-a533-4b1f-a94a-720bbd25fc17" [label=effect] \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--79d21912-36b7-4af9-8958-38949dd0d6de" [label=<
Infrastructure
NameMy Infra
> shape=plaintext] \t"infrastructure--a75c83f7-147e-4695-b173-0981521b2f01" [label=<
Infrastructure
NameTest Infra
Infrastructure Typesworkstation
> shape=plaintext] \t"attack-action--dd3820fa-bae3-4270-8000-5c4642fa780c" -> "infrastructure--a75c83f7-147e-4695-b173-0981521b2f01" [label="related-to"] } """ ) 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( id="attack-action--2f375dbd-4d6e-4036-9efa-d67f7fc93d1e", technique_id="T1", name="Action 1", description="Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.", ) assert ( attack_flow.graphviz._get_action_label(action) == '<
Action: T1
NameAction 1
DescriptionLorem ipsum dolor sit amet, consectetur
adipiscing elit, sed do eiusmod tempor
incididunt ut labore et dolore magna
aliqua.
ConfidenceVery Probable
>' ) def test_wrap_condition_description(): """Long descriptions should be wrapped.""" condition = AttackCondition( id="attack-condition--64d5bf0b-6acc-4f43-b0f2-aa93a219897a", description="Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.", on_true_refs=[], on_false_refs=[], ) assert ( attack_flow.graphviz._get_condition_label(condition) == '<
Condition
DescriptionLorem ipsum dolor sit amet, consectetur
adipiscing elit, sed do eiusmod tempor
incididunt ut labore et dolore magna
aliqua.
>' ) def test_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_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
>' )