From ea8e9ac0b6a79ce8f9ee4604c8cf402c29df9bae Mon Sep 17 00:00:00 2001
From: Eskil Andreen <eskil@ndrn.se>
Date: Mon, 25 Jan 2021 19:35:59 +0100
Subject: [PATCH] Fix typo in DS06 which made it always True

---
 pytm/threatlib/threats.json |  2 +-
 tests/test_pytmfunc.py      | 42 +++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/pytm/threatlib/threats.json b/pytm/threatlib/threats.json
index 31ab062..f6061a4 100644
--- a/pytm/threatlib/threats.json
+++ b/pytm/threatlib/threats.json
@@ -1306,7 +1306,7 @@
     "Likelihood Of Attack": "High",
     "severity": "Very High",
     "prerequisites": "",
-    "condition": "target.hasDataLeaks",
+    "condition": "target.hasDataLeaks()",
     "mitigations": "All data should be encrypted in transit. All PII and restricted data must be encrypted at rest. If a service is storing credentials used to authenticate users or incoming connections, it must only store hashes of them created using cryptographic functions, so it is only possible to compare them against user input, without fully decoding them. If a client is storing credentials in either files or other data store, access to them must be as restrictive as possible, including using proper file permissions, database users with restricted access or separate storage.",
     "example": "An application, which connects to a database without TLS, performs a database query in which it compares the password to a stored hash, instead of fetching the hash and comparing it locally.",
     "references": "https://cwe.mitre.org/data/definitions/311.html, https://cwe.mitre.org/data/definitions/312.html, https://cwe.mitre.org/data/definitions/916.html, https://cwe.mitre.org/data/definitions/653.html"
diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py
index a293ef3..75bfe10 100644
--- a/tests/test_pytmfunc.py
+++ b/tests/test_pytmfunc.py
@@ -10,6 +10,7 @@
     Action,
     Actor,
     Boundary,
+    Classification,
     Data,
     Dataflow,
     Datastore,
@@ -887,6 +888,47 @@ def test_DS05(self):
         threat = threats["DS05"]
         self.assertTrue(threat.apply(web))
 
+    def test_DS06(self):
+        threat = threats["DS06"]
+
+        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_
+
+        with self.subTest("Doesn't apply unless dataflow has data defined"):
+            dataflow = create_dataflow(define_data=False)
+            self.assertFalse(threat.apply(dataflow))
+
+        with self.subTest("Data classification equals sink, source and dataflow"):
+            dataflow = create_dataflow()
+            self.assertFalse(threat.apply(dataflow))
+
+        with self.subTest("Data classification is less than sink, source and dataflow"):
+            dataflow = create_dataflow(data=Classification.PUBLIC)
+            self.assertFalse(threat.apply(dataflow))
+
+        with self.subTest("Data classification exceeds source"):
+            dataflow = create_dataflow(source=Classification.PUBLIC)
+            self.assertTrue(threat.apply(dataflow))
+
+        with self.subTest("Data classification exceeds sink"):
+            dataflow = create_dataflow(sink=Classification.PUBLIC)
+            self.assertTrue(threat.apply(dataflow))
+
+        with self.subTest("Data classification exceeds dataflow"):
+            dataflow = create_dataflow(dataflow=Classification.PUBLIC)
+            self.assertTrue(threat.apply(dataflow))
+
     def test_SC05(self):
         web = Server("Web Server")
         web.providesIntegrity = False