diff --git a/pytm/pytm.py b/pytm/pytm.py index a1728d9b..b84bf464 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -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 @@ -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): @@ -93,19 +98,6 @@ def _setLabel(element): return "
".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: @@ -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() @@ -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") @@ -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: @@ -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 @@ -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 = <
{0}
>;'.format(label)) print("]") - class Lambda(Element): onAWS = varBool(True) authenticatesSource = varBool(False) @@ -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 = <
{}
>;'.format(label)) print("]") @@ -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 = <
{}
>;'.format(label)) print("]") @@ -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 = <
{0}
>;'.format(label, color)) print("]") @@ -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 = <
{0}
>;'.format(label)) print("]") @@ -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 = <
{0}
>;'.format(label, color)) print("]") @@ -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 = <
{0}
>;'.format(label, color)) print("]") @@ -611,8 +613,8 @@ def dfd(self, mergeResponses=False, **kwargs): resp_label = "({0}) {1}".format(self.response.order, resp_label) label += "
" + 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, )) @@ -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 = <{1}>;\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 = <{1}>;\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") diff --git a/tests/dfd.dot b/tests/dfd.dot new file mode 100644 index 00000000..f83a261d --- /dev/null +++ b/tests/dfd.dot @@ -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 = <Internet>; + ] + +actor_User_579e9aae81 [ + shape = square; + label = <
User
>; +] + +} + +subgraph cluster_boundary_ServerDB_88f2d9c06f { + graph [ + fontsize = 10; + fontcolor = firebrick2; + style = dashed; + color = firebrick2; + label = <Server/DB>; + ] + +datastore_SQLDatabase_d2006ce1bb [ + shape = none; + color = black; + label = <
SQL Database
>; +] + +} + +server_WebServer_f2eb7a3ff7 [ + shape = circle + color = black + label = <
Web Server
>; +] + actor_User_579e9aae81 -> server_WebServer_f2eb7a3ff7 [ + color = black; + dir = forward; + + label = <
User enters
comments (*)
>; + ] + server_WebServer_f2eb7a3ff7 -> datastore_SQLDatabase_d2006ce1bb [ + color = black; + dir = forward; + + label = <
Insert query
with comments
>; + ] + datastore_SQLDatabase_d2006ce1bb -> server_WebServer_f2eb7a3ff7 [ + color = black; + dir = forward; + + label = <
Retrieve
comments
>; + ] + server_WebServer_f2eb7a3ff7 -> actor_User_579e9aae81 [ + color = black; + dir = forward; + + label = <
Show comments
(*)
>; + ] +} diff --git a/tests/seq.plantuml b/tests/seq.plantuml new file mode 100644 index 00000000..58f572b5 --- /dev/null +++ b/tests/seq.plantuml @@ -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 diff --git a/tests/test_private_func.py b/tests/test_private_func.py index 3c744361..5ee86eb2 100644 --- a/tests/test_private_func.py +++ b/tests/test_private_func.py @@ -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): diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py index b35d48b5..c1f017b3 100644 --- a/tests/test_pytmfunc.py +++ b/tests/test_pytmfunc.py @@ -1,14 +1,89 @@ import sys sys.path.append("..") -import unittest -from pytm import TM, Server, Datastore, Dataflow, Boundary, Actor, Lambda, Process, Threat, ExternalEntity + import json import os +import random +import unittest +from contextlib import contextmanager from os.path import dirname +from io import StringIO + +from pytm import (TM, Actor, Boundary, Dataflow, Datastore, ExternalEntity, + Lambda, Process, Server, Threat) + with open(os.path.abspath(os.path.join(dirname(__file__), '..')) + "/pytm/threatlib/threats.json", "r") as threat_file: threats_json = json.load(threat_file) + +@contextmanager +def captured_output(): + new_out, new_err = StringIO(), StringIO() + old_out, old_err = sys.stdout, sys.stderr + try: + sys.stdout, sys.stderr = new_out, new_err + yield sys.stdout, sys.stderr + finally: + sys.stdout, sys.stderr = old_out, old_err + + +class TestTM(unittest.TestCase): + + def test_seq(self): + random.seed(0) + dir_path = os.path.dirname(os.path.realpath(__file__)) + with open(os.path.join(dir_path, 'seq.plantuml')) as x: + expected = x.read().strip() + + TM.reset() + tm = TM("my test tm", description="aaa") + internet = Boundary("Internet") + server_db = Boundary("Server/DB") + user = Actor("User", inBoundary=internet) + web = Server("Web Server") + db = Datastore("SQL Database", inBoundary=server_db) + + Dataflow(user, web, "User enters comments (*)", note="bbb") + Dataflow(web, db, "Insert query with comments", note="ccc") + Dataflow(db, web, "Retrieve comments") + Dataflow(web, user, "Show comments (*)") + + with captured_output() as (out, err): + tm.seq() + + output = out.getvalue().strip() + self.maxDiff = None + self.assertEqual(output, expected) + + def test_dfd(self): + dir_path = os.path.dirname(os.path.realpath(__file__)) + with open(os.path.join(dir_path, 'dfd.dot')) as x: + expected = x.read().strip() + + random.seed(0) + + TM.reset() + tm = TM("my test tm", description="aaa") + internet = Boundary("Internet") + server_db = Boundary("Server/DB") + user = Actor("User", inBoundary=internet) + web = Server("Web Server") + db = Datastore("SQL Database", inBoundary=server_db) + + Dataflow(user, web, "User enters comments (*)") + Dataflow(web, db, "Insert query with comments") + Dataflow(db, web, "Retrieve comments") + Dataflow(web, user, "Show comments (*)") + + with captured_output() as (out, err): + tm.dfd() + + output = out.getvalue().strip() + self.maxDiff = None + self.assertEqual(output, expected) + + class Testpytm(unittest.TestCase): # Test for all the threats in threats.py - test Threat.apply() function diff --git a/tm.py b/tm.py index 2ed0c2e8..1c1c4b50 100755 --- a/tm.py +++ b/tm.py @@ -1,6 +1,11 @@ #!/usr/bin/env python3 -from pytm import TM, Server, Datastore, Dataflow, Boundary, Actor, Lambda +import random + +from pytm import TM, Actor, Boundary, Dataflow, Datastore, Lambda, Server + +# make sure generated diagrams do not change, makes sense if they're commited +random.seed(0) tm = TM("my test tm") tm.description = "This is a sample threat model of a very simple system - a web-based comment system. The user enters comments and these are added to a database and displayed back to the user. The thought is that it is, though simple, a complete enough example to express meaningful threats." @@ -60,4 +65,6 @@ my_lambda_to_db.dstPort = 3306 my_lambda_to_db.data = "Lamda clears DB every 6 hours" -tm.process() + +if __name__ == "__main__": + tm.process()