Skip to content

Commit

Permalink
Introducing EmptyElement to avoid showing "invalid" elements on diagrams
Browse files Browse the repository at this point in the history
 - added EmptyElement
 - skipped generation EmptyElement on DFD and SEQ diagrams
  • Loading branch information
oleh-mykytiuk committed Mar 1, 2024
1 parent 4890300 commit 3d063a3
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 3 deletions.
12 changes: 10 additions & 2 deletions pytm/pytm.py
Original file line number Diff line number Diff line change
Expand Up @@ -685,7 +685,7 @@ def __init__(
if args:
element = args[0]
else:
element = kwargs.pop("element", Element("invalid"))
element = kwargs.pop("element", EmptyElement())

self.target = element.name
self.element = element
Expand Down Expand Up @@ -1016,7 +1016,7 @@ def seq(self):
participants.append(
'database {0} as "{1}"'.format(e._uniq_name(), e.display_name())
)
elif not isinstance(e, Dataflow) and not isinstance(e, Boundary):
elif not any((isinstance(e, Dataflow), isinstance(e, Boundary), isinstance(e, EmptyElement))):
participants.append(
'entity {0} as "{1}"'.format(e._uniq_name(), e.display_name())
)
Expand Down Expand Up @@ -1583,6 +1583,14 @@ def _safeset(self, attr, value):
pass


class EmptyElement(Element):
"""An empty element to avoid generation of elements for standalone Finding"""

def __init__(self):
super().__init__("AutoGenerated", description="Autogenerated element for Finding")
self._is_drawn = True # Prevent drawing on a DFD diagram


class Asset(Element):
"""An asset with outgoing or incoming dataflows"""

Expand Down
6 changes: 5 additions & 1 deletion tm.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
Datastore,
Lambda,
Server,
DatastoreType,
DatastoreType, Finding,
)

tm = TM("my test tm")
Expand Down Expand Up @@ -59,6 +59,10 @@
secretDb.storesPII = True
secretDb.maxClassification = Classification.TOP_SECRET

finding_to_overwrite = Finding(
threat_id="DO01", example="API Gateway is used to check and limit requests",
)

my_lambda = Lambda("AWS Lambda")
my_lambda.controls.hasAccessControl = True
my_lambda.inBoundary = vpc
Expand Down

0 comments on commit 3d063a3

Please sign in to comment.