def test_SC01(self):
process1 = Process("Process1")
process1.implementsNonce = False
json = Data(name="JSON", description="some JSON data", format="JSON")
process1.data = json
threat = threats["SC01"]
self.assertTrue(threat.apply(process1))
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,
isEncryptedAtRest=True)
comment = Data("Comment", isStored=True)
Dataflow(user, gw, "User enters comments (*)")
Dataflow(gw, web, "Request")
Dataflow(web, db, "Insert query with comments", data=[comment])
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_DR01(self):
web = Server("Web Server")
db = Datastore("Database")
insert = Dataflow(web, db, "Insert query")
insert.data = Data("ssn", isPII=True, isStored=True)
insert.isEncrypted = False
threat = threats["DR01"]
self.assertTrue(threat.apply(insert))
def test_DO03(self):
user = Actor("User")
web = Server("Web Server")
user_to_web = Dataflow(user, web, "User enters comments (*)")
user_to_web.protocol = "HTTP"
xml = Data(name="user to web data", description="textual", format="XML")
user_to_web.data = xml
threat = threats["DO03"]
self.assertTrue(threat.apply(user_to_web))
def test_AC04(self):
user = Actor("User")
web = Server("Web Server")
user_to_web = Dataflow(user, web, "User enters comments (*)")
xml = Data(name="user to web data", description="textual", format="XML")
user_to_web.data = xml
user_to_web.authorizesSource = False
threat = threats["AC04"]
self.assertTrue(threat.apply(user_to_web))
def test_AC22(self):
user = Actor("User")
web = Server("Web Server")
user_to_web = Dataflow(user, web, "User enters comments (*)")
user_to_web.data = Data("password",
isCredentials=True,
credentialsLife=Lifetime.HARDCODED)
user_to_web.protocol = "HTTPS"
user_to_web.isEncrypted = True
threat = threats["AC22"]
self.assertTrue(threat.apply(user_to_web))
def create_dataflow(
source=Classification.RESTRICTED,
sink=Classification.RESTRICTED,
dataflow=Classification.RESTRICTED,
data=Classification.RESTRICTED,
define_data=True,
):
source_ = Server("Source", maxClassification=source)
sink_ = Datastore("Sink", maxClassification=sink)
flow_ = Dataflow(source_, sink_, "Flow", maxClassification=dataflow)
if define_data:
flow_.data = Data("Data", classification=data)
return flow_
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)
secretDb = Datastore("Real Identity Database")
secretDb.OS = "CentOS"
secretDb.sourceFiles = ["pytm/pytm.py"]
secretDb.isHardened = True
secretDb.inBoundary = server_db
secretDb.isSQL = True
secretDb.inScope = True
secretDb.storesPII = True
secretDb.maxClassification = Classification.TOP_SECRET
my_lambda = Lambda("AWS Lambda")
my_lambda.hasAccessControl = True
my_lambda.inBoundary = vpc
my_lambda.levels = [1, 2]
token_user_identity = Data("Token verifying user identity",
classification=Classification.SECRET)
db_to_secretDb = Dataflow(db, secretDb, "Database verify real user identity")
db_to_secretDb.protocol = "RDA-TCP"
db_to_secretDb.dstPort = 40234
db_to_secretDb.data = token_user_identity
db_to_secretDb.note = "Verifying that the user is who they say they are."
db_to_secretDb.maxClassification = Classification.SECRET
comments_in_text = Data("Comments in HTML or Markdown",
classification=Classification.PUBLIC)
user_to_web = Dataflow(user, web, "User enters comments (*)")
user_to_web.protocol = "HTTP"
user_to_web.dstPort = 80
user_to_web.data = comments_in_text
user_to_web.note = "This is a simple web app\nthat stores and retrieves user comments."
web_to_db.data = 'MySQL insert statement, all literals'
web_to_db.note = "Web server inserts user comments\ninto it's SQL query and stores them in the DB."
db_to_web = Dataflow(db, web, "Retrieve comments")
db_to_web.protocol = "MySQL"
db_to_web.dstPort = 80
db_to_web.data = 'Web server retrieves comments from DB'
db_to_web.responseTo = web_to_db
web_to_user = Dataflow(web, user, "Show comments (*)")
web_to_user.protocol = "HTTP"
web_to_user.data = 'Web server shows comments to the end user'
web_to_user.responseTo = user_to_web
my_lambda_to_db = Dataflow(my_lambda, db, "Lambda periodically cleans DB")
my_lambda_to_db.protocol = "MySQL"
my_lambda_to_db.dstPort = 3306
my_lambda_to_db.data = "Lambda clears DB every 6 hours"
userIdToken = Data(
name="User ID Token",
description=
"Some unique token that represents the user real data in the secret database",
classification=Classification.TOP_SECRET,
traverses=[user_to_web, db_to_secretDb],
processedBy=[db, secretDb],
)
if __name__ == "__main__":
tm.process()