From 94e1504ff769fcc2b7e90932f9d37e832de25320 Mon Sep 17 00:00:00 2001 From: Raphael Ahrens Date: Wed, 8 Oct 2025 09:50:54 +0200 Subject: [PATCH 1/2] Fix for #278 missing html escaping for dot It was possible to inject html markup in the label of a dot node. This lead to the error observed in #278. This fix is currently only for the label attribute. Other attribute might be affected as well. --- pytm/pytm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytm/pytm.py b/pytm/pytm.py index a711687..545704f 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -1496,7 +1496,7 @@ def display_name(self): return self.name def _label(self): - return "\\n".join(wrap(self.display_name(), 18)) + return "\\n".join(wrap(html.escape(self.display_name()), 18)) def _shape(self): return "square" From 2b9d5588a8476ed14b54357dd79f4b227bca3fd5 Mon Sep 17 00:00:00 2001 From: Raphael Ahrens Date: Sat, 29 Mar 2025 11:38:06 +0100 Subject: [PATCH 2/2] Refactor of the `serialize` and `to_serializable` function This commit tries to tackle #268 and rewrite the `serialize` function to handle less class specific cases and move it into the single dispatch function `to_serializable`. --- pytm/pytm.py | 106 ++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 75 insertions(+), 31 deletions(-) diff --git a/pytm/pytm.py b/pytm/pytm.py index 545704f..662cc02 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -1269,7 +1269,7 @@ def get_table(db, klass): for e in TM._threats + TM._data + TM._elements + self.findings + [self]: table = get_table(db, e.__class__) row = {} - for k, v in serialize(e).items(): + for k, v in to_serializable(e).items(): if k == "id": k = "SID" row[k] = ", ".join(str(i) for i in v) if isinstance(v, list) else v @@ -1278,7 +1278,6 @@ def get_table(db, klass): db.close() - class Controls: """Controls implemented by/on and Element""" @@ -2003,54 +2002,99 @@ def to_serializable(val): @to_serializable.register(TM) -def ts_tm(obj): - return serialize(obj, nested=True) +def _(obj): + ignore = ("_sf", "_duplicate_ignored_attrs", "_threats", "_elements") + result = {} + for attr_name in dir(obj): + if ( + attr_name.startswith("__") + or callable(getattr(obj.__class__, attr_name, {})) + or attr_name in ignore + ): + # ignore all functions and __atrributes and ignore attributes + continue + value = getattr(obj, attr_name) + if isinstance(value, (Element, Data)): + value = value.name + result[attr_name.lstrip("_")] = value + result["elements"] = [e for e in obj._elements if isinstance(e, (Actor, Asset))] + return result @to_serializable.register(Controls) @to_serializable.register(Data) +@to_serializable.register(Finding) +def _(obj): + return serialize(obj) + + @to_serializable.register(Threat) -@to_serializable.register(Element) +def _(obj): + return serialize(obj, replace={"target": [v.__name__ for v in obj.target]}) + + @to_serializable.register(Finding) +def _(obj): + return serialize(obj, ignore=["element"]) + + +@to_serializable.register(Element) def ts_element(obj): - return serialize(obj, nested=False) + return serialize(obj, + ignore=("_is_drawn", "uuid"), + replace={ + "levels": list(obj.levels), + "sourceFiles": list(obj.sourceFiles), + "findings": [v.id for v in obj.findings], + }) + + +@to_serializable.register(Actor) +@to_serializable.register(Asset) +def _(obj): + # Note that we use the ts_element function defined for the Element class + result = ts_element(obj) + result["__class__"] = obj.__class__.__name__ + return result + + +def serialize(obj, ignore=None, replace=None): + """ + Serialize an object into a dictionary. + Keyword arguments: + ignore -- a collection of attribute names, which are not included in the serialized dictionary + replace -- dictionary keyed with attribute names which should be replaced by the diven value. + """ + if ignore is None: + ignore = [] + if replace is None: + replace = {} -def serialize(obj, nested=False): - """Used if *obj* is an instance of TM, Element, Threat or Finding.""" - klass = obj.__class__ result = {} - if isinstance(obj, (Actor, Asset)): - result["__class__"] = klass.__name__ - for i in dir(obj): + for attr_name in dir(obj): if ( - i.startswith("__") - or callable(getattr(klass, i, {})) - or ( - isinstance(obj, TM) - and i in ("_sf", "_duplicate_ignored_attrs", "_threats") - ) - or (isinstance(obj, Element) and i in ("_is_drawn", "uuid")) - or (isinstance(obj, Finding) and i == "element") + attr_name.startswith("__") + or callable(getattr(obj.__class__, attr_name, {})) + or attr_name in ignore ): + # ignore all functions and __atrributes and ignore attributes continue - value = getattr(obj, i) - if isinstance(obj, TM) and i == "_elements": - value = [e for e in value if isinstance(e, (Actor, Asset))] + try: + result[attr_name] = replace[attr_name] + continue + except KeyError: + pass + value = getattr(obj, attr_name) if value is not None: if isinstance(value, (Element, Data)): value = value.name - elif isinstance(obj, Threat) and i == "target": - value = [v.__name__ for v in value] - elif i in ("levels", "sourceFiles", "assumptions"): - value = list(value) elif ( - not nested - and not isinstance(value, str) + not isinstance(value, str) and isinstance(value, Iterable) ): - value = [v.id if isinstance(v, Finding) else v.name for v in value] - result[i.lstrip("_")] = value + value = [v.name for v in value] + result[attr_name.lstrip("_")] = value return result