def test_multilevel_dfd(self): random.seed(0) dir_path = os.path.dirname(os.path.realpath(__file__)) install_path = os.path.dirname(os.path.realpath(pytm.__file__)) with open(os.path.join(dir_path, "dfd_level0.txt")) as x: level_0 = ( x.read().strip().replace("INSTALL_PATH", os.path.dirname(install_path)) ) with open(os.path.join(dir_path, "dfd_level1.txt")) as x: level_1 = ( x.read().strip().replace("INSTALL_PATH", os.path.dirname(install_path)) ) TM.reset() tm = TM("my test tm", description="aaa") tm.isOrdered = False internet = Boundary("Internet") server_db = Boundary("Server/DB") user = Actor("User", inBoundary=internet, levels=0) 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 (*)") self.assertTrue(tm.check()) output = tm.dfd(levels={0}) with open(os.path.join(dir_path, "0.txt"), "w") as x: x.write(output) self.assertEqual(output, level_0) TM.reset() tm = TM("my test tm", description="aaa") tm.isOrdered = False internet = Boundary("Internet") server_db = Boundary("Server/DB") user = Actor("User", inBoundary=internet, levels=1) 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 (*)") self.assertTrue(tm.check()) output = tm.dfd(levels={1}) with open(os.path.join(dir_path, "1.txt"), "w") as x: x.write(output) self.maxDiff = None self.assertEqual(output, level_1)
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") tm.isOrdered = True 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 (*)") self.assertTrue(tm.check()) output = tm.seq() self.maxDiff = None self.assertEqual(output, expected)
def test_report(self): random.seed(0) dir_path = os.path.dirname(os.path.realpath(__file__)) with open(os.path.join(dir_path, "output.md")) as x: expected = x.read().strip() TM.reset() tm = TM("my test tm", description="aaa", threatsFile="pytm/threatlib/threats.json") tm.isOrdered = True internet = Boundary("Internet") server_db = Boundary("Server/DB") user = Actor("User", inBoundary=internet) web = Server("Web Server") func = Lambda("Lambda func") worker = Process("Task queue worker") db = Datastore("SQL Database", inBoundary=server_db) Dataflow(user, web, "User enters comments (*)", note="bbb", data="auth cookie") Dataflow(web, db, "Insert query with comments", note="ccc") Dataflow(web, func, "Call func") Dataflow(db, web, "Retrieve comments") Dataflow(web, user, "Show comments (*)") Dataflow(worker, db, "Query for tasks") self.assertTrue(tm.check()) output = tm.report("docs/template.md") self.maxDiff = None self.assertEqual(output.strip(), expected.strip())
def test_json_dumps(self): random.seed(0) dir_path = os.path.dirname(os.path.realpath(__file__)) with open(os.path.join(dir_path, "output.json")) as x: expected = x.read().strip() TM.reset() tm = TM("my test tm", description="aaa", threatsFile="pytm/threatlib/threats.json") tm.isOrdered = True internet = Boundary("Internet") server_db = Boundary("Server/DB") user = Actor("User", inBoundary=internet) web = Server("Web Server") func = Lambda("Lambda func") worker = Process("Task queue worker") db = Datastore("SQL Database", inBoundary=server_db) cookie = Data( name="auth cookie", description="auth cookie description", classification=Classification.PUBLIC, ) Dataflow(user, web, "User enters comments (*)", note="bbb", data=cookie) Dataflow(web, db, "Insert query with comments", note="ccc") Dataflow(web, func, "Call func") Dataflow(db, web, "Retrieve comments") Dataflow(web, user, "Show comments (*)") Dataflow(worker, db, "Query for tasks") self.assertTrue(tm.check()) output = json.dumps(tm, default=to_serializable, sort_keys=True, indent=4) with open(os.path.join(dir_path, "output_current.json"), "w") as x: x.write(output) self.maxDiff = None self.assertEqual(output, expected)
from pytm import ( TM, Actor, Boundary, Classification, Data, Dataflow, Datastore, Lambda, Server, ) 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." tm.isOrdered = True tm.mergeResponses = True internet = Boundary("Internet") server_db = Boundary("Server/DB") server_db.levels = [2] vpc = Boundary("AWS VPC") user = Actor("User") user.inBoundary = internet user.levels = [2] web = Server("Web Server") web.OS = "Ubuntu" web.isHardened = True web.sanitizesInput = False
# https://github.com/izar/pytm from pytm import (TM, Server, Dataflow, Boundary, Actor, ExternalEntity, Process) payment_online = TM("stripe") payment_online.description = "stripe payment" payment_online.isOrdered = True payment_online.mergeResponses = True Customer_Client_Web = Boundary("Customer/Internet") Merchant_Web = Boundary("Merchant/Web") Stripe_API = Boundary("Stripe/Web") customer = Actor("Customer") customer_client = ExternalEntity("Customer Client") customer_client.inBoundary = Customer_Client_Web # user.levels = [2] merchant_web = Server("Merchant Web Server") merchant_web.inBoundary = Merchant_Web merchant_web.OS = "Ubuntu" merchant_web.isHardened = True merchant_web.onAWS = True # web.levels = [2] stripe_api = ExternalEntity("Stripe API service") stripe_api.inBoundary = Stripe_API stripe_api.onAWS = False stripe_process = Process("Stripe Payment Service")