Merge pull request #248 from Dakes/fix-override-cvss-response

Fix cvss and response override
OWASP · Jun 21, 2024 · e96ee5f · e96ee5f
2 parents eae1782 + e12dc89
commit e96ee5f
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 5 deletions.
diff --git a/pytm/pytm.py b/pytm/pytm.py
@@ -859,6 +859,7 @@ def resolve(self):
         elements = defaultdict(list)
         for e in TM._elements:
             if not e.inScope:
+                e.findings = findings
                 continue
 
             override_ids = set(f.threat_id for f in e.overrides)
@@ -2089,6 +2090,8 @@ def encode_threat_data(obj):
         "threat_id",
         "references",
         "condition",
+        "cvss",
+        "response",
     ]
 
     if type(obj) is Finding or (len(obj) != 0 and type(obj[0]) is Finding):
@@ -2104,7 +2107,8 @@ def encode_threat_data(obj):
                 # ignore missing attributes, since this can be called
                 # on both a Finding and a Threat
                 continue
-            setattr(t, a, html.escape(v))
+            if v is not None:
+                setattr(t, a, html.escape(v))
 
         encoded_threat_data.append(t)
 

diff --git a/tests/test_private_func.py b/tests/test_private_func.py
@@ -10,10 +10,12 @@
     Dataflow,
     Datastore,
     DatastoreType,
+    Finding,
     Process,
     Server,
     Threat,
     UIError,
+    encode_threat_data,
 )
 
 
@@ -245,3 +247,41 @@ def test_defaults(self):
                     case["condition"],
                 ),
             )
+
+
+class TestFunction(unittest.TestCase):
+    def test_encode_threat_data(self):
+        findings = [
+            Finding(
+                description="A test description",
+                severity="High",
+                id="1",
+                threat_id="INP01",
+                cvss="9.876",
+                response="A test response",
+            ),
+            Finding(
+                description="An escape test <script>",
+                severity="Medium",
+                id="2",
+                threat_id="INP02",
+                cvss="1.234",
+                response="A test response",
+                assumption=Assumption("Test Assumption", exclude=["INP02"]),
+            )
+        ]
+        encoded_findings = encode_threat_data(findings)
+
+        self.assertEqual(len(encoded_findings), 2)
+        self.assertEqual(encoded_findings[0].description, "A test description")
+        self.assertEqual(encoded_findings[0].severity, "High")
+        self.assertEqual(encoded_findings[0].id, "1")
+        self.assertEqual(encoded_findings[0].threat_id, "INP01")
+        self.assertEqual(encoded_findings[0].cvss, "9.876")
+        self.assertEqual(encoded_findings[0].response, "A test response")
+        self.assertEqual(encoded_findings[1].description, "An escape test &lt;script&gt;")
+        self.assertEqual(encoded_findings[1].severity, "Medium")
+        self.assertEqual(encoded_findings[1].id, "2")
+        self.assertEqual(encoded_findings[1].threat_id, "INP02")
+        self.assertEqual(encoded_findings[1].cvss, "1.234")
+        self.assertEqual(encoded_findings[1].response, "A test response")
diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py
@@ -294,7 +294,11 @@ def test_overrides(self):
         web = Server(
             "Web Server",
             overrides=[
-                Finding(threat_id="Server", response="mitigated by adding TLS"),
+                Finding(
+                    threat_id="Server",
+                    response="mitigated by adding TLS",
+                    cvss="1.234",
+                ),
             ],
         )
         db = Datastore(
@@ -304,6 +308,7 @@ def test_overrides(self):
                 Finding(
                     threat_id="Datastore",
                     response="accepted since inside the trust boundary",
+                    cvss="9.876",
                 ),
             ],
         )
@@ -327,10 +332,18 @@ def test_overrides(self):
         self.assertEqual(
             [f.response for f in web.findings], ["mitigated by adding TLS"]
         )
+        self.assertEqual(
+            [f.cvss for f in web.findings],
+            ["1.234"],
+        )
         self.assertEqual(
             [f.response for f in db.findings],
             ["accepted since inside the trust boundary"],
         )
+        self.assertEqual(
+            [f.cvss for f in db.findings],
+            ["9.876"],
+        )
 
     def test_json_dumps(self):
         random.seed(0)
@@ -435,9 +448,6 @@ def test_report(self):
         self.assertTrue(tm.check())
         output = tm.report("docs/basic_template.md")
 
-        with open(os.path.join(output_path, "output_current.md"), "w") as x:
-            x.write(output)
-
         with open(os.path.join(output_path, "output_current.md"), "w") as x:
             x.write(output)