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 5, 2024
1 parent 4890300 commit 9d36821
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 2 deletions.
16 changes: 14 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 isinstance(e, (Dataflow, Boundary, EmptyElement)):
participants.append(
'entity {0} as "{1}"'.format(e._uniq_name(), e.display_name())
)
Expand Down Expand Up @@ -1583,6 +1583,18 @@ 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")
# This type is used as a part of manual created Finding
# and is not a component of data flows described by users
# That why it has not be drawn on diagrams
# To do this just mark it as already drawn
self._is_drawn = True # Prevent drawing on diagrams


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

Expand Down
2 changes: 2 additions & 0 deletions tests/test_pytmfunc.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,8 @@ def test_seq(self):
Dataflow(db, web, "Retrieve comments")
Dataflow(web, user, "Show comments (*)")

Finding() # Finding with an empty element

self.assertTrue(tm.check())
output = tm.seq()

Expand Down
6 changes: 6 additions & 0 deletions tm.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
Lambda,
Server,
DatastoreType,
Finding,
)

tm = TM("my test tm")
Expand Down Expand Up @@ -59,6 +60,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 Expand Up @@ -100,6 +105,7 @@
db_to_web.dstPort = 80
db_to_web.data = comment_retrieved
db_to_web.responseTo = web_to_db
db_to_web.overrides = [finding_to_overwrite]

comment_to_show = Data(
"Web server shows comments to the end user", classifcation=Classification.PUBLIC
Expand Down

0 comments on commit 9d36821

Please sign in to comment.