diff --git a/pytm/asset.py b/pytm/asset.py index 33a9b61..20412ab 100644 --- a/pytm/asset.py +++ b/pytm/asset.py @@ -182,7 +182,7 @@ def dfd(self, **kwargs) -> str: return self._dfd_template().format( uniq_name=self._uniq_name(), - label=self._label(), + label=self._html_label(), color=color, shape=self._shape(), ) diff --git a/pytm/boundary.py b/pytm/boundary.py index 8601f25..13d8549 100644 --- a/pytm/boundary.py +++ b/pytm/boundary.py @@ -1,6 +1,5 @@ """Boundary model - represents trust boundaries in the threat model.""" -from html import escape as html_escape from typing import List, TYPE_CHECKING from textwrap import indent @@ -64,7 +63,7 @@ def dfd(self, **kwargs) -> str: return self._dfd_template().format( uniq_name=self._uniq_name(), - label=html_escape(self._label(), quote=False), + label=self._html_label(), color=self._color(**kwargs), edges=indent("\n".join(edges), " "), ) diff --git a/pytm/element.py b/pytm/element.py index 6212fdb..f43e18c 100644 --- a/pytm/element.py +++ b/pytm/element.py @@ -4,6 +4,7 @@ import random import uuid as uuid_module from hashlib import sha224 +from html import escape as html_escape from textwrap import wrap from typing import Any, List, Optional, Set, TYPE_CHECKING @@ -29,6 +30,16 @@ def sev_to_color(sev: int) -> str: return "black" +def dot_escape(text: str) -> str: + """Escape *text* for use inside a double-quoted Graphviz string. + + Backslashes are doubled first so that names such as ``C:\\tmp\\`` or + ``\\N`` are drawn literally instead of being read as Graphviz escape + sequences (or escaping the closing quote). + """ + return text.replace("\\", "\\\\").replace('"', '\\"') + + class Element(BaseModel): """A generic element in the threat model. @@ -281,9 +292,23 @@ def display_name(self) -> str: """Get display name for this element.""" return self.name + def _label_lines(self) -> List[str]: + """Get the display name wrapped into lines for DFD labels.""" + return wrap(self.display_name(), 18) + def _label(self) -> str: """Get label for DFD representation.""" - return "\\n".join(wrap(self.display_name(), 18)) + return "\\n".join(dot_escape(line) for line in self._label_lines()) + + def _html_label(self) -> str: + """Get label text for a Graphviz HTML-like label. + + HTML-like labels do not interpret ``\\n`` escapes, so wrapped lines are + joined with ``
`` and each line is escaped individually. + """ + return "
".join( + html_escape(line, quote=False) for line in self._label_lines() + ) def _shape(self) -> str: """Get shape for DFD representation.""" diff --git a/pytm/tm.py b/pytm/tm.py index 82956b3..531690e 100644 --- a/pytm/tm.py +++ b/pytm/tm.py @@ -96,6 +96,17 @@ def __setattr__(cls, name, value): super().__setattr__(name, value) +def plantuml_escape(text: str) -> str: + """Escape *text* for PlantUML sequence diagram names, messages and notes. + + PlantUML turns ``\\n`` and ``\\t`` in text into a newline and a tab, and + reads a trailing ``\\`` as a line continuation (which swallows the next + line, e.g. ``@enduml``). Backslash escaping is not applied consistently, + so backslashes and double quotes are written as PlantUML Unicode escapes. + """ + return text.replace("\\", "").replace('"', "") + + class TM(BaseModel, metaclass=TMModelMetaclass): """Describes the threat model administratively, and holds all details during a run.""" @@ -592,19 +603,22 @@ def seq(self): if isinstance(e, Actor): participants.append( 'actor {0} as "{1}"'.format( - e._uniq_name(), getattr(e, "display_name", lambda: e.name)() + e._uniq_name(), + plantuml_escape(getattr(e, "display_name", lambda: e.name)()), ) ) elif isinstance(e, Datastore): participants.append( 'database {0} as "{1}"'.format( - e._uniq_name(), getattr(e, "display_name", lambda: e.name)() + e._uniq_name(), + plantuml_escape(getattr(e, "display_name", lambda: e.name)()), ) ) elif not isinstance(e, (Dataflow, Boundary)): participants.append( 'entity {0} as "{1}"'.format( - e._uniq_name(), getattr(e, "display_name", lambda: e.name)() + e._uniq_name(), + plantuml_escape(getattr(e, "display_name", lambda: e.name)()), ) ) @@ -613,11 +627,11 @@ def seq(self): message = "{0} -> {1}: {2}".format( e.source._uniq_name(), e.sink._uniq_name(), - getattr(e, "display_name", lambda: e.name)(), + plantuml_escape(getattr(e, "display_name", lambda: e.name)()), ) note = "" if getattr(e, "note", "") != "": - note = "\nnote left\n{}\nend note".format(e.note) + note = "\nnote left\n{}\nend note".format(plantuml_escape(e.note)) messages.append("{}{}".format(message, note)) return self._seq_template().format( diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py index c8b0b2a..00bc7e7 100644 --- a/tests/test_pytmfunc.py +++ b/tests/test_pytmfunc.py @@ -172,6 +172,84 @@ def test_dfd_escapes_boundary_label_html(self): assert "label = <Client > API>;" in output assert "label = <Client > API>;" not in output + def test_dfd_wraps_boundary_label_with_html_line_breaks(self): + # Boundary labels are Graphviz HTML-like labels, which render "\n" + # literally; wrapped lines must be joined with
instead. + random.seed(0) + + TM.reset() + tm = TM("boundary wrap tm", description="aaa") + boundary = Boundary("University of Michigan ") + Server("Web Server", inBoundary=boundary) + + output = tm.dfd() + + assert ( + "label = <University of
Michigan <Ann
Arbor>
>;" + in output + ) + assert "\\n" not in output.split("subgraph cluster_")[1].split("]")[0] + + def test_dfd_wraps_lambda_label_with_html_line_breaks(self): + # Lambda also uses an HTML-like label, so the same rule applies. + random.seed(0) + + TM.reset() + tm = TM("lambda wrap tm", description="aaa") + Lambda("Very Long Serverless & Function") + + output = tm.dfd() + + assert ( + "Very Long
Serverless <Auth>
& Function
" in output + ) + assert "\\n" not in output + + def test_dfd_escapes_quotes_and_backslashes_in_labels(self): + random.seed(0) + + TM.reset() + tm = TM("dfd escape tm", description="aaa") + user = Actor('The "admin" user') + server = Server("C:\\Server\\") + Dataflow(user, server, 'Sends "hello"') + + output = tm.dfd() + + assert 'label = "The \\"admin\\" user";' in output + assert 'label = "C:\\\\Server\\\\";' in output + assert 'label = "Sends \\"hello\\"";' in output + + def test_dfd_escaping_keeps_label_line_breaks(self): + random.seed(0) + + TM.reset() + tm = TM("dfd escape wrap tm", description="aaa") + Server('A "quoted" server with a long name') + + output = tm.dfd() + + assert 'label = "A \\"quoted\\" server\\nwith a long name";' in output + + def test_seq_escapes_quotes_and_backslashes(self): + random.seed(0) + + TM.reset() + tm = TM("seq escape tm", description="aaa") + user = Actor('The "admin" user') + server = Server("C:\\Server\\") + Dataflow(user, server, "Path C:\\new\\tmp\\", note='Said "hi"') + + output = tm.seq() + + assert 'as "The admin user"' in output + assert 'as "C:Server"' in output + assert ": Path C:newtmp\n" in output + assert "\nSaid hi\nend note" in output + # No raw backslash may remain: PlantUML would read "\n"/"\t" as + # control characters and a trailing "\" as a line continuation. + assert "\\" not in output + def test_dfd_duplicates_ignore(self): dir_path = os.path.dirname(os.path.realpath(__file__)) install_path = os.path.dirname(os.path.realpath(pytm.__file__))