#!/usr/bin/env python3 from pytm import TM, Server, Datastore, Dataflow, Boundary, Actor, Lambda 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." internet = Boundary("Internet") server_db = Boundary("Server/DB") vpc = Boundary("AWS VPC") user = Actor("User") user.inBoundary = internet web = Server("Web Server") web.OS = "Ubuntu" web.isHardened = True web.sanitizesInput = True web.encodesOutput = True db = Datastore("SQL Database") db.OS = "CentOS" db.isHardened = False db.inBoundary = server_db db.isSQL = True db.inScope = True my_lambda = Lambda("AWS Lambda") my_lambda.hasAccessControl = True my_lambda.inBoundary = vpc
def test_multilevel_dfd(self): random.seed(0) dir_path = os.path.dirname(os.path.realpath(__file__)) with open(os.path.join(dir_path, "dfd_level0.txt")) as x: level_0 = x.read().strip() with open(os.path.join(dir_path, "dfd_level1.txt")) as x: level_1 = 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, 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 = True 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)
#!/usr/bin/env python3 from pytm import ( TM, Actor, Boundary, Classification, Data, Dataflow, Datastore, Lambda, Server, DatastoreType, ) 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 tm.assumptions = [ "Here you can document a list of assumptions about the system", ] internet = Boundary("Internet") server_db = Boundary("Server/DB") server_db.levels = [2] vpc = Boundary("AWS VPC") user = Actor("User")