ensure all items has correct color (based on scope)

OWASP · Mar 26, 2020 · 665f663 · 665f663
1 parent d6d0afe
commit 665f663
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/pytm/pytm.py b/pytm/pytm.py
@@ -432,8 +432,9 @@ def check(self):
 
     def dfd(self, **kwargs):
         self._is_drawn = True
+        color = _setColor(self)
         label = _setLabel(self)
-        print("%s [\n\tshape = square;" % self._uniq_name())
+        print("{0} [\n\tshape = square;\n\tcolor = {1};".format(self._uniq_name(), color))
         print('\tlabel = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>{0}</b></td></tr></table>>;'.format(label))
         print("]")
 
@@ -568,7 +569,7 @@ def dfd(self, **kwargs):
         self._is_drawn = True
         color = _setColor(self)
         label = _setLabel(self)
-        print("{0} [\n\tshape = circle\n\tcolor = {1}".format(self._uniq_name(), color))
+        print("{0} [\n\tshape = circle\n\tcolor = {1};".format(self._uniq_name(), color))
         print('\tlabel = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>{}</b></td></tr></table>>;'.format(label))
         print("]")
 
@@ -632,8 +633,9 @@ def __init__(self, name, **kwargs):
 
     def dfd(self, **kwargs):
         self._is_drawn = True
+        color = _setColor(self)
         label = _setLabel(self)
-        print("%s [\n\tshape = square;" % self._uniq_name())
+        print("{0} [\n\tshape = square;\n\tcolor = {1};".format(self._uniq_name(), color))
         print('\tlabel = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>{0}</b></td></tr></table>>;'.format(label))
         print("]")
 
@@ -685,7 +687,7 @@ def dfd(self, **kwargs):
         self._is_drawn = True
         color = _setColor(self)
         label = _setLabel(self)
-        print("{0} [\n\tshape = circle;\n\tcolor = {1};\n".format(self._uniq_name(), color))
+        print("{0} [\n\tshape = circle;\n\tcolor = {1};".format(self._uniq_name(), color))
         print('\tlabel = <<table border="0" cellborder="0" cellpadding="2"><tr><td><font color="{1}"><b>{0}</b></font></td></tr></table>>;'.format(label, color))
         print("]")
 
@@ -698,7 +700,7 @@ def dfd(self, **kwargs):
         self._is_drawn = True
         color = _setColor(self)
         label = _setLabel(self)
-        print("{0} [\n\tshape = doublecircle;\n\tcolor = {1};\n".format(self._uniq_name(), color))
+        print("{0} [\n\tshape = doublecircle;\n\tcolor = {1};".format(self._uniq_name(), color))
         print('\tlabel = <<table border="0" cellborder="0" cellpadding="2"><tr><td><font color="{1}"><b>{0}</b></font></td></tr></table>>;'.format(label, color))
         print("]")
 

diff --git a/tests/dfd.dot b/tests/dfd.dot
@@ -28,6 +28,7 @@ subgraph cluster_boundary_Internet_acf3059e70 {
 
 actor_User_579e9aae81 [
 	shape = square;
+	color = black;
 	label = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>User</b></td></tr></table>>;
 ]
 
@@ -52,7 +53,7 @@ datastore_SQLDatabase_d2006ce1bb [
 
 server_WebServer_f2eb7a3ff7 [
 	shape = circle
-	color = black
+	color = black;
 	label = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>Web Server</b></td></tr></table>>;
 ]
 	actor_User_579e9aae81 -> server_WebServer_f2eb7a3ff7 [