From 9d36821620c8c5cf28092e31360f0aad77b54721 Mon Sep 17 00:00:00 2001
From: Oleh Mykytiuk <oleh_mykytiuk@epam.com>
Date: Fri, 1 Mar 2024 15:52:53 +0200
Subject: [PATCH] Introducing EmptyElement to avoid showing "invalid" elements
 on diagrams

 - added EmptyElement
 - skipped generation EmptyElement on DFD and SEQ diagrams
---
 pytm/pytm.py           | 16 ++++++++++++++--
 tests/test_pytmfunc.py |  2 ++
 tm.py                  |  6 ++++++
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/pytm/pytm.py b/pytm/pytm.py
index 9906648..fe00231 100644
--- a/pytm/pytm.py
+++ b/pytm/pytm.py
@@ -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
@@ -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())
                 )
@@ -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"""
 
diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py
index f5129d5..9192216 100644
--- a/tests/test_pytmfunc.py
+++ b/tests/test_pytmfunc.py
@@ -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()
 
diff --git a/tm.py b/tm.py
index 68a4197..b92f4c7 100755
--- a/tm.py
+++ b/tm.py
@@ -11,6 +11,7 @@
     Lambda,
     Server,
     DatastoreType,
+    Finding,
 )
 
 tm = TM("my test tm")
@@ -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
@@ -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