from textwrap import dedent
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_attack_flow(get_flow_bundle())
assert output == dedent(
"""\
graph TB
classDef action fill:#99ccff
classDef operator fill:#ff9900
classDef condition fill:#99ff99
classDef builtin fill:#cccccc
attack_action__52f2c35a_fa2a_45a4_b84c_46ad9498071f["Action - T1 Action 1: -
Description of action 1 -
Confidence Very Probable"]
class attack_action__52f2c35a_fa2a_45a4_b84c_46ad9498071f action
attack_action__dd3820fa_bae3_4270_8000_5c4642fa780c["Action - Action 2: -
Description of action 2 -
Confidence Very Probable"]
class attack_action__dd3820fa_bae3_4270_8000_5c4642fa780c action
attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17["Action - T3 Action 3: -
Description of action 3 -
Confidence Very Probable"]
class attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17 action
attack_action__7ddab166_c83e_4c79_a701_a0dc2a905dd3["Action - T4 Action 4: -
Description of action 4 -
Confidence Very Probable"]
class attack_action__7ddab166_c83e_4c79_a701_a0dc2a905dd3 action
attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a["Condition: My condition"]
class attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a condition
attack_operator__8932b181_be87_4f81_851a_ab0b4288406a(("OR"))
class attack_operator__8932b181_be87_4f81_851a_ab0b4288406a operator
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__79d21912_36b7_4af9_8958_38949dd0d6de["Infrastructure - Name: My
Infra"]
class infrastructure__79d21912_36b7_4af9_8958_38949dd0d6de builtin
infrastructure__a75c83f7_147e_4695_b173_0981521b2f01["Infrastructure - Name:
Test Infra - Infrastructure
Types: workstation"]
class infrastructure__a75c83f7_147e_4695_b173_0981521b2f01 builtin
attack_action__52f2c35a_fa2a_45a4_b84c_46ad9498071f -->|effect| attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a
attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a -->|on_true| attack_operator__8932b181_be87_4f81_851a_ab0b4288406a
attack_condition__64d5bf0b_6acc_4f43_b0f2_aa93a219897a -->|on_false| attack_action__7ddab166_c83e_4c79_a701_a0dc2a905dd3
attack_operator__8932b181_be87_4f81_851a_ab0b4288406a -->|effect| attack_action__dd3820fa_bae3_4270_8000_5c4642fa780c
attack_operator__8932b181_be87_4f81_851a_ab0b4288406a -->|effect| attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17
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
"""
)