diff --git a/CHANGELOG.md b/CHANGELOG.md index ae817261..5739dddc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ ## New features +- Use numbered dataflow labels in sequence diagram [#94](https://github.com/izar/pytm/pull/94) - Move authenticateDestination to base Element [#88](https://github.com/izar/pytm/pull/88) - Assign inputs and outputs to all elements [#89](https://github.com/izar/pytm/pull/89) - Allow detecting and/or hiding duplicate dataflows by setting `TM.onDuplicates` [#100](https://github.com/izar/pytm/pull/100) diff --git a/pytm/pytm.py b/pytm/pytm.py index 5314771e..f8cd2b7b 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -18,7 +18,9 @@ ''' Helper functions ''' -''' The base for this (descriptors instead of properties) has been shamelessly lifted from https://nbviewer.jupyter.org/urls/gist.github.com/ChrisBeaumont/5758381/raw/descriptor_writeup.ipynb +''' The base for this (descriptors instead of properties) has been + shamelessly lifted from + https://nbviewer.jupyter.org/urls/gist.github.com/ChrisBeaumont/5758381/raw/descriptor_writeup.ipynb By Chris Beaumont ''' @@ -284,7 +286,8 @@ class Threat(): id = varString("", required=True) description = varString("") - condition = varString("", doc="a Python expression that should evaluate to a boolean True or False") + condition = varString("", doc="""a Python expression that should evaluate +to a boolean True or False""") details = varString("") severity = varString("") mitigations = varString("") @@ -323,7 +326,8 @@ def apply(self, target): class Finding(): - """Represents a Finding - the element in question and a description of the finding """ + """Represents a Finding - the element in question +and a description of the finding""" element = varElement(None, required=True, doc="Element this finding applies to") target = varString("", doc="Name of the element this finding applies to") @@ -359,7 +363,8 @@ def __init__( class TM(): - """Describes the threat model administratively, and holds all details during a run""" + """Describes the threat model administratively, +and holds all details during a run""" _BagOfFlows = [] _BagOfElements = [] @@ -377,7 +382,8 @@ class TM(): mergeResponses = varBool(False, doc="Merge response edges in DFDs") ignoreUnused = varBool(False, doc="Ignore elements not used in any Dataflow") findings = varFindings([], doc="threats found for elements of this model") - onDuplicates = varAction(Action.NO_ACTION, doc="How to handle duplicate Dataflow with same properties, except name and notes") + onDuplicates = varAction(Action.NO_ACTION, doc="""How to handle duplicate Dataflow +with same properties, except name and notes""") def __init__(self, name, **kwargs): for key, value in kwargs.items(): @@ -415,7 +421,16 @@ def resolve(self): for t in TM._BagOfThreats: if not t.apply(e): continue - f = Finding(e, t.description, t.details, t.severity, t.mitigations, t.example, t.id, t.references) + f = Finding( + e, + t.description, + t.details, + t.severity, + t.mitigations, + t.example, + t.id, + t.references, + ) findings.append(f) elements[e].append(f) self.findings = findings @@ -424,12 +439,15 @@ def resolve(self): def check(self): if self.description is None: - raise ValueError("Every threat model should have at least a brief description of the system being modeled.") + raise ValueError("""Every threat model should have at least +a brief description of the system being modeled.""") TM._BagOfFlows = _match_responses(_sort(TM._BagOfFlows, self.isOrdered)) self._check_duplicates(TM._BagOfFlows) _apply_defaults(TM._BagOfFlows) if self.ignoreUnused: - TM._BagOfElements, TM._BagOfBoundaries = _get_elements_and_boundaries(TM._BagOfFlows) + TM._BagOfElements, TM._BagOfBoundaries = _get_elements_and_boundaries( + TM._BagOfFlows + ) result = True for e in (TM._BagOfElements): if not e.check(): @@ -514,15 +532,23 @@ def seq(self): participants = [] for e in TM._BagOfElements: if isinstance(e, Actor): - participants.append("actor {0} as \"{1}\"".format(e._uniq_name(), e.name)) + participants.append( + 'actor {0} as "{1}"'.format(e._uniq_name(), e.display_name()) + ) elif isinstance(e, Datastore): - participants.append("database {0} as \"{1}\"".format(e._uniq_name(), e.name)) + participants.append( + 'database {0} as "{1}"'.format(e._uniq_name(), e.display_name()) + ) elif not isinstance(e, Dataflow) and not isinstance(e, Boundary): - participants.append("entity {0} as \"{1}\"".format(e._uniq_name(), e.name)) + participants.append( + 'entity {0} as "{1}"'.format(e._uniq_name(), e.display_name()) + ) messages = [] for e in TM._BagOfFlows: - message = "{0} -> {1}: {2}".format(e.source._uniq_name(), e.sink._uniq_name(), e.name) + message = "{0} -> {1}: {2}".format( + e.source._uniq_name(), e.sink._uniq_name(), e.display_name() + ) note = "" if e.note != "": note = "\nnote left\n{}\nend note".format(e.note) @@ -712,7 +738,6 @@ def inside(self, *boundaries): return True return False - def _attr_values(self): klass = self.__class__ result = {} @@ -755,7 +780,7 @@ def __init__(self, name, **kwargs): def _dfd_template(self): return """{uniq_name} [ - shape = none; + shape = {shape}; fixedsize = shape; image = "{image}"; imagescale = true; @@ -775,9 +800,13 @@ def dfd(self, **kwargs): uniq_name=self._uniq_name(), label=self._label(), color=self._color(), + shape=self._shape(), image=os.path.join(os.path.dirname(__file__), "images", "lambda.png"), ) + def _shape(self): + return "none" + class Server(Element): """An entity processing data""" @@ -877,7 +906,7 @@ def __init__(self, name, **kwargs): def _dfd_template(self): return """{uniq_name} [ - shape = none; + shape = {shape}; color = {color}; fontcolor = {color}; label = < @@ -888,6 +917,9 @@ def _dfd_template(self): ] """ + def _shape(self): + return "none" + class Actor(Element): """An entity usually initiating actions""" @@ -1043,9 +1075,9 @@ def _dfd_template(self): return """subgraph cluster_{uniq_name} {{ graph [ fontsize = 10; - fontcolor = firebrick2; + fontcolor = {color}; style = dashed; - color = firebrick2; + color = {color}; label = <{label}>; ] @@ -1069,19 +1101,31 @@ def dfd(self): return self._dfd_template().format( uniq_name=self._uniq_name(), label=self._label(), + color=self._color(), edges=indent("\n".join(edges), " "), ) + def _color(self): + return "firebrick2" + def get_args(): _parser = argparse.ArgumentParser() - _parser.add_argument('--debug', action='store_true', help='print debug messages') - _parser.add_argument('--dfd', action='store_true', help='output DFD') - _parser.add_argument('--report', help='output report using the named template file (sample template file is under docs/template.md)') - _parser.add_argument('--exclude', help='specify threat IDs to be ignored') - _parser.add_argument('--seq', action='store_true', help='output sequential diagram') - _parser.add_argument('--list', action='store_true', help='list all available threats') - _parser.add_argument('--describe', help='describe the properties available for a given element') + _parser.add_argument("--debug", action="store_true", help="print debug messages") + _parser.add_argument("--dfd", action="store_true", help="output DFD") + _parser.add_argument( + "--report", + help="""output report using the named template file +(sample template file is under docs/template.md)""", + ) + _parser.add_argument("--exclude", help="specify threat IDs to be ignored") + _parser.add_argument("--seq", action="store_true", help="output sequential diagram") + _parser.add_argument( + "--list", action="store_true", help="list all available threats" + ) + _parser.add_argument( + "--describe", help="describe the properties available for a given element" + ) _args = _parser.parse_args() return _args diff --git a/tests/seq.plantuml b/tests/seq.plantuml index e444a8ed..72eb3374 100644 --- a/tests/seq.plantuml +++ b/tests/seq.plantuml @@ -3,14 +3,14 @@ actor actor_User_579e9aae81 as "User" entity server_WebServer_f2eb7a3ff7 as "Web Server" database datastore_SQLDatabase_d2006ce1bb as "SQL Database" -actor_User_579e9aae81 -> server_WebServer_f2eb7a3ff7: User enters comments (*) +actor_User_579e9aae81 -> server_WebServer_f2eb7a3ff7: (1) User enters comments (*) note left bbb end note -server_WebServer_f2eb7a3ff7 -> datastore_SQLDatabase_d2006ce1bb: Insert query with comments +server_WebServer_f2eb7a3ff7 -> datastore_SQLDatabase_d2006ce1bb: (2) Insert query with comments note left ccc end note -datastore_SQLDatabase_d2006ce1bb -> server_WebServer_f2eb7a3ff7: Retrieve comments -server_WebServer_f2eb7a3ff7 -> actor_User_579e9aae81: Show comments (*) +datastore_SQLDatabase_d2006ce1bb -> server_WebServer_f2eb7a3ff7: (3) Retrieve comments +server_WebServer_f2eb7a3ff7 -> actor_User_579e9aae81: (4) Show comments (*) @enduml diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py index deaf003a..f81b7edb 100644 --- a/tests/test_pytmfunc.py +++ b/tests/test_pytmfunc.py @@ -36,6 +36,7 @@ def test_seq(self): TM.reset() tm = TM("my test tm", description="aaa") + tm.isOrdered = True internet = Boundary("Internet") server_db = Boundary("Server/DB") user = Actor("User", inBoundary=internet) @@ -47,7 +48,7 @@ def test_seq(self): Dataflow(db, web, "Retrieve comments") Dataflow(web, user, "Show comments (*)") - tm.check() + self.assertTrue(tm.check()) output = tm.seq() self.maxDiff = None @@ -73,7 +74,7 @@ def test_seq_unused(self): Dataflow(db, web, "Retrieve comments") Dataflow(web, user, "Show comments (*)") - tm.check() + self.assertTrue(tm.check()) output = tm.seq() self.maxDiff = None @@ -99,7 +100,7 @@ def test_dfd(self): Dataflow(db, web, "Retrieve comments") Dataflow(web, user, "Show comments (*)") - tm.check() + self.assertTrue(tm.check()) output = tm.dfd() self.maxDiff = None