def test_seq_unused(self): random.seed(0) dir_path = os.path.dirname(os.path.realpath(__file__)) with open(os.path.join(dir_path, "seq_unused.plantuml")) as x: expected = x.read().strip() TM.reset() tm = TM("my test tm", description="aaa", ignoreUnused=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) Lambda("Unused Lambda") 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_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") net = Boundary("Company net") dmz = Boundary("dmz", inBoundary=net) backend = Boundary("backend", inBoundary=net) user = Actor("User", inBoundary=internet) gw = Server("Gateway", inBoundary=dmz) web = Server("Web Server", inBoundary=backend) db = Datastore("SQL Database", inBoundary=backend) Dataflow(user, gw, "User enters comments (*)") Dataflow(gw, web, "Request") Dataflow(web, db, "Insert query with comments") Dataflow(db, web, "Retrieve comments") Dataflow(web, gw, "Response") Dataflow(gw, user, "Show comments (*)") self.assertTrue(tm.check()) output = tm.dfd() self.maxDiff = None self.assertEqual(output, expected)
def test_json_loads(self): random.seed(0) dir_path = os.path.dirname(os.path.realpath(__file__)) with open(os.path.join(dir_path, "input.json")) as x: contents = x.read().strip() TM.reset() tm = loads(contents) self.assertTrue(tm.check()) self.maxDiff = None self.assertEqual([b.name for b in tm._boundaries], ["Internet", "Server/DB"]) self.assertEqual( [e.name for e in tm._elements], [ "Internet", "Server/DB", "User", "Web Server", "SQL Database", "Request", "Insert", "Select", "Response", ], ) self.assertEqual([f.name for f in tm._flows], ["Request", "Insert", "Select", "Response"])
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_resolve(self): 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, inScope=False) web = Server("Web Server") db = Datastore("SQL Database", inBoundary=server_db) req = Dataflow(user, web, "User enters comments (*)") query = Dataflow(web, db, "Insert query with comments") results = Dataflow(db, web, "Retrieve comments") resp = Dataflow(web, user, "Show comments (*)") TM._BagOfThreats = [ Threat(SID=klass, target=klass) for klass in ["Actor", "Server", "Datastore", "Dataflow"] ] tm.resolve() self.maxDiff = None self.assertListEqual([f.id for f in tm.findings], [ 'Server', 'Datastore', 'Dataflow', 'Dataflow', 'Dataflow', 'Dataflow' ]) self.assertListEqual([f.id for f in user.findings], []) self.assertListEqual([f.id for f in web.findings], ["Server"]) self.assertListEqual([f.id for f in db.findings], ["Datastore"]) self.assertListEqual([f.id for f in req.findings], ["Dataflow"]) self.assertListEqual([f.id for f in query.findings], ["Dataflow"]) self.assertListEqual([f.id for f in results.findings], ["Dataflow"]) self.assertListEqual([f.id for f in resp.findings], ["Dataflow"])
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 (*)") tm.check() with captured_output() as (out, err): tm.dfd() output = out.getvalue().strip() self.maxDiff = None self.assertEqual(output, expected)
def test_dfd_duplicates_raise(self): random.seed(0) TM.reset() tm = TM("my test tm", description="aaa", onDuplicates=Action.RESTRICT) 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(user, web, "User views comments") Dataflow(web, db, "Insert query with comments") Dataflow(web, db, "Select query") Dataflow(db, web, "Retrieve comments") Dataflow(web, user, "Show comments (*)") e = re.escape( "Duplicate Dataflow found between Actor(User) " "and Server(Web Server): Dataflow(User enters comments (*)) " "is same as Dataflow(User views comments)" ) with self.assertRaisesRegex(ValueError, e): tm.check()
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_overrides(self): 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, inScope=False) web = Server( "Web Server", overrides=[ Finding(threat_id="Server", response="mitigated by adding TLS"), ], ) db = Datastore( "SQL Database", inBoundary=server_db, overrides=[ Finding( threat_id="Datastore", response="accepted since inside the trust boundary", ), ], ) req = Dataflow(user, web, "User enters comments (*)") query = Dataflow(web, db, "Insert query with comments") results = Dataflow(db, web, "Retrieve comments") resp = Dataflow(web, user, "Show comments (*)") TM._threats = [ Threat(SID="Server", target="Server", condition="False"), Threat(SID="Datastore", target="Datastore"), ] tm.resolve() self.maxDiff = None self.assertEqual( [f.threat_id for f in tm.findings], ["Server", "Datastore"], ) self.assertEqual([f.response for f in web.findings], ["mitigated by adding TLS"]) self.assertEqual( [f.response for f in db.findings], ["accepted since inside the trust boundary"], )
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)
def test_exclude_threats_ignore(self): random.seed(0) TM.reset() excluded_threat = "INP03" remaining_threat = "AA01" TM._threatsExcluded = [excluded_threat] tm = TM("my test tm", description="aaa") web = Server("Web") web.sanitizesInput = False web.encodesOutput = False self.assertTrue(threats[excluded_threat].apply(web)) self.assertTrue(threats[remaining_threat].apply(web)) tm.resolve() self.assertNotIn(excluded_threat, [t.threat_id for t in tm.findings]) self.assertIn(remaining_threat, [t.threat_id for t in tm.findings])