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_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_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_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)
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
web.encodesOutput = True
web.authorizesSource = False
web.sourceFiles = ["pytm/json.py", "docs/template.md"]
#!/usr/bin/env python3
from pytm import TM, Server, Datastore, Dataflow, Boundary, Actor
tm = TM("my test tm")
tm.description = "another test tm"
tm.isOrdered = True
User_Web = Boundary("User/Web")
Web_DB = Boundary("Web/DB")
#!/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."
User_Web = Boundary("User/Web")
Web_DB = Boundary("Web/DB")
VPC = Boundary("AWS VPC")
user = Actor("User")
user.inBoundary = User_Web
web = Server("Web Server")
web.OS = "CloudOS"
web.isHardened = True
my_lambda = Lambda("cleanDBevery6hours")
my_lambda.hasAccessControl = True
my_lambda.inBoundary = Web_DB
#my_lambda.inBoundary = VPC # TODO: need multiple boundaries capability for these situations
db = Datastore("SQL Database")
db.OS = "CentOS"
db.isHardened = False
db.inBoundary = Web_DB
db.isSQL = True
db.inScope = False
my_lambda_to_db = Dataflow(my_lambda, db, "(λ)Periodically cleans DB")
#!/usr/bin/env python3
from pytm import TM, Server, Datastore, Dataflow, Boundary, Actor, Element
tm = TM("Apps")
tm.description = "Apps threat modeling"
internet = Boundary("Internet")
machine = Boundary("User's machine")
apps_vpc = Boundary("Apps VPC")
rds_boundary = Boundary("RDS security group")
cache_boundary = Boundary("ElastiCache security group")
user = Actor("User/Browser")
user.inBoundary = machine
apigee = Element("Apigee")
apigee.inBoundary = internet
apigee.isHardened = True
apigee = Element("Apigee")
apigee.inBoundary = internet
apigee.isHardened = True
server = Server("Apps Server")
server.inBoundary = apps_vpc
server.isHardened = True
server.hasAccessControl = True
server.encodesOutput = True
db = Datastore("MySQL DB")
db.isHardened = True
# 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")