Skip to content

Commit 64e13e1

Browse files
committed
Introducing EmptyElement to avoid showing "invalid" elements on diagrams
- added EmptyElement - skipped generation EmptyElement on DFD and SEQ diagrams
1 parent 4890300 commit 64e13e1

File tree

3 files changed

+21
-3
lines changed

3 files changed

+21
-3
lines changed

pytm/pytm.py

+14-2
Original file line numberDiff line numberDiff line change
@@ -685,7 +685,7 @@ def __init__(
685685
if args:
686686
element = args[0]
687687
else:
688-
element = kwargs.pop("element", Element("invalid"))
688+
element = kwargs.pop("element", EmptyElement())
689689

690690
self.target = element.name
691691
self.element = element
@@ -1016,7 +1016,7 @@ def seq(self):
10161016
participants.append(
10171017
'database {0} as "{1}"'.format(e._uniq_name(), e.display_name())
10181018
)
1019-
elif not isinstance(e, Dataflow) and not isinstance(e, Boundary):
1019+
elif not isinstance(e, (Dataflow, Boundary, EmptyElement)):
10201020
participants.append(
10211021
'entity {0} as "{1}"'.format(e._uniq_name(), e.display_name())
10221022
)
@@ -1583,6 +1583,18 @@ def _safeset(self, attr, value):
15831583
pass
15841584

15851585

1586+
class EmptyElement(Element):
1587+
"""An empty element to avoid generation of elements for standalone Finding"""
1588+
1589+
def __init__(self):
1590+
super().__init__("AutoGenerated", description="Autogenerated element for Finding")
1591+
# This type is used as a part of manual created Finding
1592+
# and is not a component of data flows described by users
1593+
# That why it has not be drawn on diagrams
1594+
# To do this just mark it as already drawn
1595+
self._is_drawn = True # Prevent drawing on diagrams
1596+
1597+
15861598
class Asset(Element):
15871599
"""An asset with outgoing or incoming dataflows"""
15881600

tests/test_pytmfunc.py

+2
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,8 @@ def test_seq(self):
5858
Dataflow(db, web, "Retrieve comments")
5959
Dataflow(web, user, "Show comments (*)")
6060

61+
Finding() # Finding with an empty element
62+
6163
self.assertTrue(tm.check())
6264
output = tm.seq()
6365

tm.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
Datastore,
1111
Lambda,
1212
Server,
13-
DatastoreType,
13+
DatastoreType, Finding,
1414
)
1515

1616
tm = TM("my test tm")
@@ -59,6 +59,10 @@
5959
secretDb.storesPII = True
6060
secretDb.maxClassification = Classification.TOP_SECRET
6161

62+
finding_to_overwrite = Finding(
63+
threat_id="DO01", example="API Gateway is used to check and limit requests",
64+
)
65+
6266
my_lambda = Lambda("AWS Lambda")
6367
my_lambda.controls.hasAccessControl = True
6468
my_lambda.inBoundary = vpc

0 commit comments

Comments
 (0)