From 3d063a3bbf29dc64af39db9b96ce7b2d137c0b26 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 | 12 ++++++++++--
 tm.py        |  6 +++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/pytm/pytm.py b/pytm/pytm.py
index 9906648..0adc6c5 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 any((isinstance(e, Dataflow), isinstance(e, Boundary), isinstance(e, EmptyElement))):
                 participants.append(
                     'entity {0} as "{1}"'.format(e._uniq_name(), e.display_name())
                 )
@@ -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"""
 
diff --git a/tm.py b/tm.py
index 68a4197..5b7141a 100755
--- a/tm.py
+++ b/tm.py
@@ -10,7 +10,7 @@
     Datastore,
     Lambda,
     Server,
-    DatastoreType,
+    DatastoreType, Finding,
 )
 
 tm = TM("my test tm")
@@ -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