Skip to content

Commit

Permalink
add test coverage for graphviz
Browse files Browse the repository at this point in the history
  • Loading branch information
@allisonrobbins
allisonrobbins committed Jun 25, 2024
1 parent 34476ad commit 57eaded
Show file tree
Hide file tree
Showing 4 changed files with 143 additions and 7 deletions.
1 change: 1 addition & 0 deletions src/attack_flow/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@ def graphviz(args):
converted = attack_flow.graphviz.convert_attack_tree(flow_bundle)

Check warning on line 100 in src/attack_flow/cli.py

View check run for this annotation

Codecov / codecov/patch

src/attack_flow/cli.py#L97-L100 

Added lines #L97 - L100 were not covered by tests
else:
converted = attack_flow.graphviz.convert_attack_flow(flow_bundle)
break

Check warning on line 103 in src/attack_flow/cli.py

View check run for this annotation

Codecov / codecov/patch

src/attack_flow/cli.py#L102-L103 

Added lines #L102 - L103 were not covered by tests

with open(args.output, "w") as out:
out.write(converted)
Expand Down
12 changes: 6 additions & 6 deletions src/attack_flow/graphviz.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -219,9 +219,9 @@ def _get_attack_tree_action_label(action):
[
'<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5">',
f'<TR><TD BGCOLOR="#B40000" COLSPAN="2"><font color="white"><B>{heading}</B></font></TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{label_escape(action.name)}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{description}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{confidence}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{label_escape(action.name)}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{description}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{confidence}</TD></TR>',
"</TABLE>>",
]
)
Expand Down Expand Up @@ -317,9 +317,9 @@ def _get_operator_label(action, operator_type):
[
'<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5">',
f'<TR><TD BGCOLOR="{color}" COLSPAN="2"><B>{heading}</B></TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{label_escape(action.name)}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{description}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{confidence}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{label_escape(action.name)}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{description}</TD></TR>',
f'<TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">{confidence}</TD></TR>',
"</TABLE>>",
]
)
82 changes: 82 additions & 0 deletions tests/fixtures/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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="[email protected]",
)
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",
)
55 changes: 54 additions & 1 deletion tests/test_graphviz.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down Expand Up @@ -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=<<font point-size="24">My Flow</font><br/><i>(missing description)</i><br/><font point-size="10">Author: Jane Doe &lt;[email protected]&gt;</font><br/><font point-size="10">Created: 2022-08-25 19:26:31</font><br/><font point-size="10">Modified: 2022-08-25 19:26:31</font>>;
\tlabelloc="t";
\t"attack-action--d63857d5-1043-45a4-9397-40ef68db4c5f" [label=<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#B40000" COLSPAN="2"><font color="white"><B>Action</B></font></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Action 1</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Description of action 2</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Very Probable</TD></TR></TABLE>> 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=<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#9CE67E" COLSPAN="2"><B>OR</B></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">My Or Operator</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">this is the description</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Very Probable</TD></TR></TABLE>> 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=<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#B40000" COLSPAN="2"><font color="white"><B>Action</B></font></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Action 2</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Description of action 2</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Very Probable</TD></TR></TABLE>> 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=<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#B40000" COLSPAN="2"><font color="white"><B>Action</B></font></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Action 3</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Description of action 3</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Very Probable</TD></TR></TABLE>> 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=<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#cccccc" COLSPAN="2"><B>Infrastructure</B></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">My Infra</TD></TR></TABLE>> shape=plaintext]
\t"attack-asset--4ae37379-6a11-44c1-b6a8-d11733cfac06" [label=<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#cc99ff" COLSPAN="2"><B>Asset: My Asset</B></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT"></TD></TR></TABLE>> 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(
Expand Down Expand Up @@ -75,3 +106,25 @@ def test_action_label():
attack_flow.graphviz._get_action_label(action)
== '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#99ccff" COLSPAN="2"><B>Action</B></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">My technique</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">This technique has no ID to render in<br/>the header.</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Very Probable</TD></TR></TABLE>>'
)

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")
== '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#99ccff" COLSPAN="2"><B>AND</B></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">My technique</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">This technique has no ID to render in<br/>the header.</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Very Probable</TD></TR></TABLE>>'
)

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)
== '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5"><TR><TD BGCOLOR="#B40000" COLSPAN="2"><font color="white"><B>Action</B></font></TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Name</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">My technique</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Description</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">This technique has no ID to render in<br/>the header.</TD></TR><TR><TD ALIGN="LEFT" BALIGN="LEFT"><B>Confidence</B></TD><TD ALIGN="LEFT" BALIGN="LEFT">Very Probable</TD></TR></TABLE>>'
)

0 comments on commit 57eaded

Please sign in to comment.