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