fix: escape quotes and backslashes in DFD and sequence diagram labels

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(),
         )

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), "    "),
         )

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 ``<br/>`` and each line is escaped individually.
+        """
+        return "<br/>".join(
+            html_escape(line, quote=False) for line in self._label_lines()
+        )
 
     def _shape(self) -> str:
         """Get shape for DFD representation."""

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("\\", "<U+005C>").replace('"', "<U+0022>")
+
+
 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(

tests/test_pytmfunc.py

@@ -172,6 +172,84 @@ def test_dfd_escapes_boundary_label_html(self):
         assert "label = <<i>Client &gt; API</i>>;" in output
         assert "label = <<i>Client > API</i>>;" 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 <br/> instead.
+        random.seed(0)
+
+        TM.reset()
+        tm = TM("boundary wrap tm", description="aaa")
+        boundary = Boundary("University of Michigan <Ann Arbor>")
+        Server("Web Server", inBoundary=boundary)
+
+        output = tm.dfd()
+
+        assert (
+            "label = <<i>University of<br/>Michigan &lt;Ann<br/>Arbor&gt;</i>>;"
+            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 <Auth> & Function")
+
+        output = tm.dfd()
+
+        assert (
+            "<b>Very Long<br/>Serverless &lt;Auth&gt;<br/>&amp; Function</b>" 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 <U+0022>admin<U+0022> user"' in output
+        assert 'as "C:<U+005C>Server<U+005C>"' in output
+        assert ": Path C:<U+005C>new<U+005C>tmp<U+005C>\n" in output
+        assert "\nSaid <U+0022>hi<U+0022>\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__))