Skip to content

produce stable diagrams #79

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 4, 2020
Merged

produce stable diagrams #79

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 45 additions & 45 deletions pytm/pytm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
import argparse
import json
import logging
import random
import uuid
from collections import defaultdict
from hashlib import sha224
from os.path import dirname
from re import match, sub
from sys import argv, exit, stderr
from re import match
from sys import exit, stderr
from textwrap import wrap
from weakref import WeakKeyDictionary

Expand All @@ -18,6 +20,9 @@
'''


logger = logging.getLogger(__name__)


class var(object):
''' A descriptor that allows setting a value only once '''
def __init__(self, default, onSet=None):
Expand Down Expand Up @@ -93,19 +98,6 @@ def _setLabel(element):
return "<br/>".join(wrap(element.name, 14))


def _debug(_args, msg):
if _args.debug is True:
stderr.write("DEBUG: {}\n".format(msg))


def _uniq_name(obj_name, obj_uuid):
''' transform name and uuid into a unique string '''
hash_input = '{}{}'.format(obj_name, str(obj_uuid))
h = sha224(hash_input.encode('utf-8')).hexdigest()
hash_without_numbers = sub(r'[0-9]', '', h)
return hash_without_numbers


def _sort(elements, addOrder=False):
ordered = sorted(elements, key=lambda flow: flow.order)
if not addOrder:
Expand Down Expand Up @@ -217,6 +209,14 @@ def __init__(self, name, **kwargs):
self._sf = SuperFormatter()
self._add_threats()

@classmethod
def reset(cls):
cls._BagOfFlows = []
cls._BagOfElements = []
cls._BagOfThreats = []
cls._BagOfFindings = []
cls._BagOfBoundaries = []

def _init_threats(self):
TM._BagOfThreats = []
self._add_threats()
Expand Down Expand Up @@ -264,14 +264,14 @@ def seq(self):
print("@startuml")
for e in TM._BagOfElements:
if isinstance(e, Actor):
print("actor {0} as \"{1}\"".format(_uniq_name(e.name, e.uuid), e.name))
print("actor {0} as \"{1}\"".format(e._uniq_name(), e.name))
elif isinstance(e, Datastore):
print("database {0} as \"{1}\"".format(_uniq_name(e.name, e.uuid), e.name))
print("database {0} as \"{1}\"".format(e._uniq_name(), e.name))
elif not isinstance(e, Dataflow) and not isinstance(e, Boundary):
print("entity {0} as \"{1}\"".format(_uniq_name(e.name, e.uuid), e.name))
print("entity {0} as \"{1}\"".format(e._uniq_name(), e.name))

for e in TM._BagOfFlows:
print("{0} -> {1}: {2}".format(_uniq_name(e.source.name, e.source.uuid), _uniq_name(e.sink.name, e.sink.uuid), e.name))
print("{0} -> {1}: {2}".format(e.source._uniq_name(), e.sink._uniq_name(), e.name))
if e.note != "":
print("note left\n{}\nend note".format(e.note))
print("@enduml")
Expand All @@ -287,6 +287,9 @@ def report(self, *args, **kwargs):
def process(self):
self.check()
result = get_args()
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
if result.debug:
logger.setLevel(logging.DEBUG)
if result.seq is True:
self.seq()
if result.dfd is True:
Expand Down Expand Up @@ -328,16 +331,23 @@ def __init__(self, name, **kwargs):
for key, value in kwargs.items():
setattr(self, key, value)
self.name = name
self.uuid = uuid.uuid4()
self.uuid = uuid.UUID(int=random.getrandbits(128))
self._is_drawn = False
TM._BagOfElements.append(self)

def __repr__(self):
return '<{0}.{1}({2}) at {3}>'.format(
self.__module__, type(self).__name__, self.name, hex(id(self)))
return "<{0}.{1}({2}) at {3}>".format(
self.__module__, type(self).__name__, self.name, hex(id(self))
)

def __str__(self):
return '{0}({1})'.format(type(self).__name__, self.name)
return "{0}({1})".format(type(self).__name__, self.name)

def _uniq_name(self):
''' transform name and uuid into a unique string '''
h = sha224(str(self.uuid).encode('utf-8')).hexdigest()
name = "".join(x for x in self.name if x.isalpha())
return "{0}_{1}_{2}".format(type(self).__name__.lower(), name, h[:10])

def check(self):
return True
Expand All @@ -348,13 +358,11 @@ def check(self):

def dfd(self, **kwargs):
self._is_drawn = True
name = _uniq_name(self.name, self.uuid)
label = _setLabel(self)
print("%s [\n\tshape = square;" % name)
print("%s [\n\tshape = square;" % self._uniq_name())
print('\tlabel = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>{0}</b></td></tr></table>>;'.format(label))
print("]")


class Lambda(Element):
onAWS = varBool(True)
authenticatesSource = varBool(False)
Expand All @@ -375,11 +383,10 @@ def __init__(self, name, **kwargs):

def dfd(self, **kwargs):
self._is_drawn = True
name = _uniq_name(self.name, self.uuid)
color = _setColor(self)
pngpath = dirname(__file__) + "/images/lambda.png"
label = _setLabel(self)
print('{0} [\n\tshape = none\n\tfixedsize=shape\n\timage="{2}"\n\timagescale=true\n\tcolor = {1}'.format(name, color, pngpath))
print('{0} [\n\tshape = none\n\tfixedsize=shape\n\timage="{2}"\n\timagescale=true\n\tcolor = {1}'.format(self._uniq_name(), color, pngpath))
print('\tlabel = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>{}</b></td></tr></table>>;'.format(label))
print("]")

Expand Down Expand Up @@ -423,10 +430,9 @@ def __init__(self, name, **kwargs):

def dfd(self, **kwargs):
self._is_drawn = True
name = _uniq_name(self.name, self.uuid)
color = _setColor(self)
label = _setLabel(self)
print("{0} [\n\tshape = circle\n\tcolor = {1}".format(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("]")

Expand Down Expand Up @@ -471,10 +477,9 @@ def __init__(self, name, **kwargs):

def dfd(self, **kwargs):
self._is_drawn = True
name = _uniq_name(self.name, self.uuid)
color = _setColor(self)
label = _setLabel(self)
print("{0} [\n\tshape = none;\n\tcolor = {1};".format(name, color))
print("{0} [\n\tshape = none;\n\tcolor = {1};".format(self._uniq_name(), color))
print('\tlabel = <<table sides="TB" cellborder="0" cellpadding="2"><tr><td><font color="{1}"><b>{0}</b></font></td></tr></table>>;'.format(label, color))
print("]")

Expand All @@ -487,9 +492,8 @@ def __init__(self, name, **kwargs):

def dfd(self, **kwargs):
self._is_drawn = True
name = _uniq_name(self.name, self.uuid)
label = _setLabel(self)
print("%s [\n\tshape = square;" % name)
print("%s [\n\tshape = square;" % self._uniq_name())
print('\tlabel = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>{0}</b></td></tr></table>>;'.format(label))
print("]")

Expand Down Expand Up @@ -539,10 +543,9 @@ def __init__(self, name, **kwargs):

def dfd(self, **kwargs):
self._is_drawn = True
name = _uniq_name(self.name, self.uuid)
color = _setColor(self)
label = _setLabel(self)
print("{0} [\n\tshape = circle;\n\tcolor = {1};\n".format(name, color))
print("{0} [\n\tshape = circle;\n\tcolor = {1};\n".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("]")

Expand All @@ -553,10 +556,9 @@ def __init__(self, name, **kwargs):

def dfd(self, **kwargs):
self._is_drawn = True
name = _uniq_name(self.name, self.uuid)
color = _setColor(self)
label = _setLabel(self)
print("{0} [\n\tshape = doublecircle;\n\tcolor = {1};\n".format(name, color))
print("{0} [\n\tshape = doublecircle;\n\tcolor = {1};\n".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("]")

Expand Down Expand Up @@ -611,8 +613,8 @@ def dfd(self, mergeResponses=False, **kwargs):
resp_label = "({0}) {1}".format(self.response.order, resp_label)
label += "<br/>" + resp_label
print("\t{0} -> {1} [\n\t\tcolor = {2};\n\t\tdir = {3};\n".format(
_uniq_name(self.source.name, self.source.uuid),
_uniq_name(self.sink.name, self.sink.uuid),
self.source._uniq_name(),
self.sink._uniq_name(),
color,
direction,
))
Expand All @@ -630,16 +632,14 @@ def dfd(self):
if self._is_drawn:
return

result = get_args()
self._is_drawn = True
_debug(result, "Now drawing boundary " + self.name)
name = _uniq_name(self.name, self.uuid)
logger.debug("Now drawing boundary " + self.name)
label = self.name
print("subgraph cluster_{0} {{\n\tgraph [\n\t\tfontsize = 10;\n\t\tfontcolor = firebrick2;\n\t\tstyle = dashed;\n\t\tcolor = firebrick2;\n\t\tlabel = <<i>{1}</i>>;\n\t]\n".format(name, label))
print("subgraph cluster_{0} {{\n\tgraph [\n\t\tfontsize = 10;\n\t\tfontcolor = firebrick2;\n\t\tstyle = dashed;\n\t\tcolor = firebrick2;\n\t\tlabel = <<i>{1}</i>>;\n\t]\n".format(self._uniq_name(), label))
for e in TM._BagOfElements:
if e.inBoundary == self and not e._is_drawn:
# The content to draw can include Boundary objects
_debug(result, "Now drawing content {}".format(e.name))
logger.debug("Now drawing content {}".format(e.name))
e.dfd()
print("\n}\n")

Expand Down
82 changes: 82 additions & 0 deletions tests/dfd.dot
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
digraph tm {
graph [
fontname = Arial;
fontsize = 14;
]
node [
fontname = Arial;
fontsize = 14;
rankdir = lr;
]
edge [
shape = none;
fontname = Arial;
fontsize = 12;
]
labelloc = "t";
fontsize = 20;
nodesep = 1;

subgraph cluster_boundary_Internet_acf3059e70 {
graph [
fontsize = 10;
fontcolor = firebrick2;
style = dashed;
color = firebrick2;
label = <<i>Internet</i>>;
]

actor_User_579e9aae81 [
shape = square;
label = <<table border="0" cellborder="0" cellpadding="2"><tr><td><b>User</b></td></tr></table>>;
]

}

subgraph cluster_boundary_ServerDB_88f2d9c06f {
graph [
fontsize = 10;
fontcolor = firebrick2;
style = dashed;
color = firebrick2;
label = <<i>Server/DB</i>>;
]

datastore_SQLDatabase_d2006ce1bb [
shape = none;
color = black;
label = <<table sides="TB" cellborder="0" cellpadding="2"><tr><td><font color="black"><b>SQL Database</b></font></td></tr></table>>;
]

}

server_WebServer_f2eb7a3ff7 [
shape = circle
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 [
color = black;
dir = forward;

label = <<table border="0" cellborder="0" cellpadding="2"><tr><td><font color="black"><b>User enters<br/>comments (*)</b></font></td></tr></table>>;
]
server_WebServer_f2eb7a3ff7 -> datastore_SQLDatabase_d2006ce1bb [
color = black;
dir = forward;

label = <<table border="0" cellborder="0" cellpadding="2"><tr><td><font color="black"><b>Insert query<br/>with comments</b></font></td></tr></table>>;
]
datastore_SQLDatabase_d2006ce1bb -> server_WebServer_f2eb7a3ff7 [
color = black;
dir = forward;

label = <<table border="0" cellborder="0" cellpadding="2"><tr><td><font color="black"><b>Retrieve<br/>comments</b></font></td></tr></table>>;
]
server_WebServer_f2eb7a3ff7 -> actor_User_579e9aae81 [
color = black;
dir = forward;

label = <<table border="0" cellborder="0" cellpadding="2"><tr><td><font color="black"><b>Show comments<br/>(*)</b></font></td></tr></table>>;
]
}
15 changes: 15 additions & 0 deletions tests/seq.plantuml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
@startuml
actor actor_User_579e9aae81 as "User"
entity server_WebServer_f2eb7a3ff7 as "Web Server"
database datastore_SQLDatabase_d2006ce1bb as "SQL Database"
actor_User_579e9aae81 -> server_WebServer_f2eb7a3ff7: User enters comments (*)
note left
bbb
end note
server_WebServer_f2eb7a3ff7 -> datastore_SQLDatabase_d2006ce1bb: Insert query with comments
note left
ccc
end note
datastore_SQLDatabase_d2006ce1bb -> server_WebServer_f2eb7a3ff7: Retrieve comments
server_WebServer_f2eb7a3ff7 -> actor_User_579e9aae81: Show comments (*)
@enduml
10 changes: 7 additions & 3 deletions tests/test_private_func.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,23 @@
import sys
sys.path.append("..")
import unittest
import random

from pytm.pytm import _uniq_name, Actor, Boundary, Dataflow, Datastore, Server, TM
from pytm.pytm import Actor, Boundary, Dataflow, Datastore, Server, TM


class TestUniqueNames(unittest.TestCase):
def test_duplicate_boundary_names_have_different_unique_names(self):
random.seed(0)
object_1 = Boundary("foo")
object_2 = Boundary("foo")

object_1_uniq_name = _uniq_name(object_1.name, object_1.uuid)
object_2_uniq_name = _uniq_name(object_2.name, object_2.uuid)
object_1_uniq_name = object_1._uniq_name()
object_2_uniq_name = object_2._uniq_name()

self.assertNotEqual(object_1_uniq_name, object_2_uniq_name)
self.assertEqual(object_1_uniq_name, "boundary_foo_acf3059e70")
self.assertEqual(object_2_uniq_name, "boundary_foo_88f2d9c06f")


class TestAttributes(unittest.TestCase):
Expand Down
Loading