From 280007a52e430e6fdb8fe79be2e36cdd98657021 Mon Sep 17 00:00:00 2001 From: Mark Church Date: Mon, 9 Nov 2026 10:49:45 -0500 Subject: [PATCH 1/3] fix: use
line breaks in boundary labels Boundary names are wrapped at 18 characters and joined with "\n", like other element labels. Boundaries use Graphviz HTML-like labels (label = <...>), which do not interpret "\n", so long boundary names were drawn with a literal "\n" in the diagram. Add Element._label_lines() so the wrapping rule is shared, and have Boundary build its HTML label by escaping each line and joining with
. Co-Authored-By: Claude Opus 5.5 --- pytm/boundary.py | 12 +++++++++++- pytm/element.py | 6 +++++- tests/test_pytmfunc.py | 18 ++++++++++++++++++ 3 files changed, 34 insertions(+), 2 deletions(-) diff --git a/pytm/boundary.py b/pytm/boundary.py index 8601f25..8b828aa 100644 --- a/pytm/boundary.py +++ b/pytm/boundary.py @@ -64,11 +64,21 @@ 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), " "), ) + def _html_label(self) -> str: + """Get label 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 _color(self, **kwargs) -> str: """Get color for DFD representation.""" if kwargs.get("colormap", False): diff --git a/pytm/element.py b/pytm/element.py index 6212fdb..e37da98 100644 --- a/pytm/element.py +++ b/pytm/element.py @@ -281,9 +281,13 @@ 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(self._label_lines()) def _shape(self) -> str: """Get shape for DFD representation.""" diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py index c8b0b2a..dc010e8 100644 --- a/tests/test_pytmfunc.py +++ b/tests/test_pytmfunc.py @@ -172,6 +172,24 @@ 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_duplicates_ignore(self): dir_path = os.path.dirname(os.path.realpath(__file__)) install_path = os.path.dirname(os.path.realpath(pytm.__file__)) From 48e05ca5a3c061722f193737600081b49da9a26d Mon Sep 17 00:00:00 2001 From: Mark Church Date: Mon, 9 Nov 2026 11:30:44 -0500 Subject: [PATCH 2/3] fix: use
line breaks in Lambda labels too Lambda is the other element drawn with a Graphviz HTML-like label (
...
), so long Lambda names had the same literal "\n" problem as boundaries, and "<" or "&" in a Lambda name produced an invalid label. Move _html_label() from Boundary to Element and use it for Lambda. Co-Authored-By: Claude Opus 5.5 --- pytm/asset.py | 2 +- pytm/boundary.py | 11 ----------- pytm/element.py | 11 +++++++++++ tests/test_pytmfunc.py | 15 +++++++++++++++ 4 files changed, 27 insertions(+), 12 deletions(-) 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 8b828aa..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 @@ -69,16 +68,6 @@ def dfd(self, **kwargs) -> str: edges=indent("\n".join(edges), " "), ) - def _html_label(self) -> str: - """Get label 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 _color(self, **kwargs) -> str: """Get color for DFD representation.""" if kwargs.get("colormap", False): diff --git a/pytm/element.py b/pytm/element.py index e37da98..9c937c5 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 @@ -289,6 +290,16 @@ def _label(self) -> str: """Get label for DFD representation.""" return "\\n".join(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.""" return "square" diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py index dc010e8..01cea96 100644 --- a/tests/test_pytmfunc.py +++ b/tests/test_pytmfunc.py @@ -190,6 +190,21 @@ def test_dfd_wraps_boundary_label_with_html_line_breaks(self): ) 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_duplicates_ignore(self): dir_path = os.path.dirname(os.path.realpath(__file__)) install_path = os.path.dirname(os.path.realpath(pytm.__file__)) From d1584208e87e2d8242c934b5f08aad99696d9328 Mon Sep 17 00:00:00 2001 From: Mark Church Date: Mon, 9 Nov 2026 11:34:07 -0500 Subject: [PATCH 3/3] fix: escape quotes and backslashes in DFD and sequence diagram labels Element and dataflow names were inserted into diagrams unescaped: - DOT: a '"' in a name ends the quoted label early and dot fails with a syntax error. A trailing '\' escapes the closing quote, and sequences such as '\N' or '\l' are read as Graphviz escapes. - PlantUML: '\n' and '\t' in names, messages or notes are drawn as a newline and a tab, and a trailing '\' is read as a line continuation that swallows '@enduml', so no diagram is produced. Add dot_escape(), applied per wrapped line in Element._label() so the "\n" line breaks between lines are kept, and plantuml_escape(), which writes '\' and '"' as PlantUML Unicode escapes (, ) for participant names, messages and notes. Co-Authored-By: Claude Opus 5.5 --- pytm/element.py | 12 ++++++++++- pytm/tm.py | 24 +++++++++++++++++----- tests/test_pytmfunc.py | 45 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 75 insertions(+), 6 deletions(-) diff --git a/pytm/element.py b/pytm/element.py index 9c937c5..f43e18c 100644 --- a/pytm/element.py +++ b/pytm/element.py @@ -30,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. @@ -288,7 +298,7 @@ def _label_lines(self) -> List[str]: def _label(self) -> str: """Get label for DFD representation.""" - return "\\n".join(self._label_lines()) + 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. diff --git a/pytm/tm.py b/pytm/tm.py index 69df84f..7dfecc5 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.""" @@ -591,19 +602,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)()), ) ) @@ -612,11 +626,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 01cea96..00bc7e7 100644 --- a/tests/test_pytmfunc.py +++ b/tests/test_pytmfunc.py @@ -205,6 +205,51 @@ def test_dfd_wraps_lambda_label_with_html_line_breaks(self): ) 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__))