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__))