def test_DS02(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.environment = 'Production'
lambda1.environment = 'Production'
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "DS02"))
self.assertTrue(ThreatObj.apply(process1))
self.assertTrue(ThreatObj.apply(lambda1))
def test_API01(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.implementsAPI = True
lambda1.implementsAPI = True
threat = threats["API01"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_DS02(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.environment = "Production"
lambda1.environment = "Production"
threat = threats["DS02"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_API01(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.implementsAPI = True
lambda1.implementsAPI = True
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "API01"))
self.assertTrue(ThreatObj.apply(process1))
self.assertTrue(ThreatObj.apply(lambda1))
def test_INP13(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.validatesInput = False
lambda1.validatesInput = False
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "INP13"))
self.assertTrue(ThreatObj.apply(process1))
self.assertTrue(ThreatObj.apply(lambda1))
def test_INP13(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.validatesInput = False
lambda1.validatesInput = False
threat = threats["INP13"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_API02(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.implementsAPI = True
process1.validatesInput = False
lambda1.implementsAPI = True
lambda1.validatesInput = False
threat = threats["API02"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_INP26(self):
process1 = Process("Process")
lambda1 = Lambda("lambda")
process1.validatesInput = False
process1.sanitizesInput = False
lambda1.validatesInput = False
lambda1.sanitizesInput = False
threat = threats["INP26"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_INP24(self):
process1 = Process("Process")
lambda1 = Lambda("lambda")
process1.checksInputBounds = False
process1.validatesInput = False
lambda1.checksInputBounds = False
lambda1.validatesInput = False
threat = threats["INP24"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_INP14(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
web = Server("Web Server")
process1.validatesInput = False
lambda1.validatesInput = False
web.validatesInput = False
threat = threats["INP14"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
self.assertTrue(threat.apply(web))
def test_INP24(self):
process1 = Process("Process")
lambda1 = Lambda("lambda")
process1.checksInputBounds = False
process1.validatesInput = False
lambda1.checksInputBounds = False
lambda1.validatesInput = False
ThreatObj = Threat(
next(item for item in threats_json if item["SID"] == "INP24"))
self.assertTrue(ThreatObj.apply(process1))
self.assertTrue(ThreatObj.apply(lambda1))
def test_INP01(self):
lambda1 = Lambda("mylambda")
process1 = Process("myprocess")
lambda1.usesEnvironmentVariables = True
lambda1.sanitizesInput = False
lambda1.checksInputBounds = False
process1.usesEnvironmentVariables = True
process1.sanitizesInput = False
process1.checksInputBounds = False
threat = threats["INP01"]
self.assertTrue(threat.apply(lambda1))
self.assertTrue(threat.apply(process1))
def test_LB01(self):
process1 = Process("Process1")
process1.implementsAPI = True
process1.validatesInput = False
process1.sanitizesInput = False
lambda1 = Lambda("Lambda1")
lambda1.implementsAPI = True
lambda1.validatesInput = False
lambda1.sanitizesInput = False
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "LB01"))
self.assertTrue(ThreatObj.apply(process1))
self.assertTrue(ThreatObj.apply(lambda1))
def test_INP01(self):
lambda1 = Lambda('mylambda')
process1 = Process('myprocess')
lambda1.usesEnvironmentVariables = True
lambda1.sanitizesInput = False
lambda1.checksInputBounds = False
process1.usesEnvironmentVariables = True
process1.sanitizesInput = False
process1.checksInputBounds = False
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "INP01"))
self.assertTrue(ThreatObj.apply(lambda1))
self.assertTrue(ThreatObj.apply(process1))
def test_LB01(self):
process1 = Process("Process1")
process1.implementsAPI = True
process1.controls.validatesInput = False
process1.controls.sanitizesInput = False
lambda1 = Lambda("Lambda1")
lambda1.implementsAPI = True
lambda1.controls.validatesInput = False
lambda1.controls.sanitizesInput = False
threat = threats["LB01"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_DO02(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
web = Server("Web Server")
db = Datastore("DB")
process1.handlesResourceConsumption = False
lambda1.handlesResourceConsumption = False
web.handlesResourceConsumption = False
db.handlesResourceConsumption = False
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "DO02"))
self.assertTrue(ThreatObj.apply(process1))
self.assertTrue(ThreatObj.apply(lambda1))
self.assertTrue(ThreatObj.apply(web))
self.assertTrue(ThreatObj.apply(db))
def test_AC03(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.usesEnvironmentVariables = True
process1.controls.implementsAuthenticationScheme = False
process1.controls.validatesInput = False
process1.controls.authorizesSource = False
lambda1.usesEnvironmentVariables = True
lambda1.controls.implementsAuthenticationScheme = False
lambda1.controls.validatesInput = False
lambda1.controls.authorizesSource = False
threat = threats["AC03"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_DO02(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
web = Server("Web Server")
db = Datastore("DB")
process1.handlesResourceConsumption = False
lambda1.handlesResourceConsumption = False
web.handlesResourceConsumption = False
db.handlesResourceConsumption = False
threat = threats["DO02"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
self.assertTrue(threat.apply(web))
self.assertTrue(threat.apply(db))
def test_INP08(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
web = Server("Web Server")
process1.validatesInput = False
process1.sanitizesInput = False
lambda1.validatesInput = False
lambda1.sanitizesInput = False
web.validatesInput = False
web.sanitizesInput = False
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "INP08"))
self.assertTrue(ThreatObj.apply(process1))
self.assertTrue(ThreatObj.apply(lambda1))
self.assertTrue(ThreatObj.apply(web))
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_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_INP12(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.controls.checksInputBounds = False
process1.controls.validatesInput = False
lambda1.controls.checksInputBounds = False
lambda1.controls.validatesInput = False
threat = threats["INP12"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
def test_INP08(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
web = Server("Web Server")
process1.controls.validatesInput = False
process1.controls.sanitizesInput = False
lambda1.controls.validatesInput = False
lambda1.controls.sanitizesInput = False
web.controls.validatesInput = False
web.controls.sanitizesInput = False
threat = threats["INP08"]
self.assertTrue(threat.apply(process1))
self.assertTrue(threat.apply(lambda1))
self.assertTrue(threat.apply(web))
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_AC03(self):
process1 = Process("Process1")
lambda1 = Lambda("Lambda1")
process1.usesEnvironmentVariables = True
process1.implementsAuthenticationScheme = False
process1.validatesInput = False
process1.authorizesSource = False
lambda1.usesEnvironmentVariables = True
lambda1.implementsAuthenticationScheme = False
lambda1.validatesInput = False
lambda1.authorizesSource = False
ThreatObj = Threat(next(item for item in threats_json if item["SID"] == "AC03"))
self.assertTrue(ThreatObj.apply(process1))
self.assertTrue(ThreatObj.apply(lambda1))
db.isSQL = True
db.inScope = True
db.maxClassification = Classification.RESTRICTED
db.levels = [2]
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)
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")
my_lambda_to_db.protocol = "SQL"
my_lambda_to_db.dstPort = 3306
db.type = DatastoreType.SQL
db.inScope = True
db.maxClassification = Classification.RESTRICTED
db.levels = [2]
secretDb = Datastore("Real Identity Database")
secretDb.OS = "CentOS"
secretDb.sourceFiles = ["pytm/pytm.py"]
secretDb.controls.isHardened = True
secretDb.inBoundary = server_db
secretDb.type = DatastoreType.SQL
secretDb.inScope = True
secretDb.storesPII = True
secretDb.maxClassification = Classification.TOP_SECRET
my_lambda = Lambda("AWS Lambda")
my_lambda.controls.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)
