From a941c70b917a0075297887a3b851fa8869894a8b Mon Sep 17 00:00:00 2001
From: RKS <rajesh.sharma@owasp.org>
Date: Sun, 11 Oct 2026 13:39:28 -0400
Subject: [PATCH] fix: preserve assumptions in threat-model JSON

---
 README.md                     |  3 ++
 pytm/pytm.py                  |  3 ++
 tests/test_assumption_json.py | 72 +++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+)
 create mode 100644 tests/test_assumption_json.py

diff --git a/README.md b/README.md
index 316dd5d..962b11c 100644
--- a/README.md
+++ b/README.md
@@ -445,6 +445,9 @@ matching elements connecting to SQL datastores, would be `any(f.sink.oneOf(Datas
 
 With a little bit of Python code it is possible to import a threat model from JSON (notice the special format in the example found in `tests/input.json`). The following example imports the `input.json` example found in tests. Save the following code as `tm2.py`.
 
+JSON exports preserve assumptions as objects with their `name`, `description`, and
+`exclude` threat IDs, including assumptions attached to elements and excluded findings.
+
 ```python
 
 #!/usr/bin/env python3
diff --git a/pytm/pytm.py b/pytm/pytm.py
index 211e3c6..a8b2ff8 100644
--- a/pytm/pytm.py
+++ b/pytm/pytm.py
@@ -532,6 +532,7 @@ def ts_tm(obj):
     return serialize(obj, nested=True)
 
 
+@to_serializable.register(Assumption)
 @to_serializable.register(Controls)
 @to_serializable.register(Data)
 @to_serializable.register(Threat)
@@ -629,6 +630,8 @@ def serialize(obj, nested=False):
 
         if isinstance(value, (Element, Data)):
             value = value.name
+        elif isinstance(value, Assumption):
+            value = to_serializable(value)
         elif hasattr(value, "model_dump") and not isinstance(value, TM):
             value = value.model_dump()
         elif isinstance(obj, Threat) and attr_name == "target":
diff --git a/tests/test_assumption_json.py b/tests/test_assumption_json.py
new file mode 100644
index 0000000..40d50ab
--- /dev/null
+++ b/tests/test_assumption_json.py
@@ -0,0 +1,72 @@
+"""Assumptions must survive threat-model JSON export and import."""
+
+import json
+
+import pytest
+
+from pytm import TM, Assumption, Server, loads
+from pytm.pytm import to_serializable
+
+
+@pytest.fixture(autouse=True)
+def reset_tm():
+    TM.reset()
+    yield
+    TM.reset()
+
+
+@pytest.mark.parametrize("scope", ["model", "element"])
+def test_assumption_round_trip_preserves_fields(scope):
+    assumption = Assumption(
+        "Reviewed input", exclude=["INP01", "INP02"], description="Validated upstream"
+    )
+    tm = TM("Assumptions", assumptions=[assumption] if scope == "model" else [])
+    Server("API", assumptions=[assumption] if scope == "element" else [])
+
+    output = json.dumps(tm, default=to_serializable)
+    TM.reset()
+    restored = loads(output)
+
+    assumptions = (
+        restored.assumptions if scope == "model" else restored._elements[0].assumptions
+    )
+    assert len(assumptions) == 1
+    assert assumptions[0].name == "Reviewed input"
+    assert assumptions[0].description == "Validated upstream"
+    assert assumptions[0].exclude == {"INP01", "INP02"}
+
+
+@pytest.mark.parametrize("scope", ["model", "element"])
+def test_round_trip_preserves_threat_exclusion(tmp_path, scope):
+    threats_file = tmp_path / "threats.json"
+    threats_file.write_text(
+        json.dumps(
+            [
+                {
+                    "SID": "TEST01",
+                    "target": ["Server"],
+                    "description": "Always matches the server",
+                    "condition": "True",
+                }
+            ]
+        )
+    )
+    tm = TM("Exclusions", threatsFile=str(threats_file))
+    server = Server("API")
+    tm.resolve()
+    assert [finding.threat_id for finding in tm.findings] == ["TEST01"]
+
+    owner = tm if scope == "model" else server
+    owner.assumptions = [Assumption("Reviewed", exclude=["TEST01"])]
+    tm.resolve()
+    assert tm.findings == []
+    assert [finding.threat_id for finding in tm.excluded_findings] == ["TEST01"]
+
+    output = json.dumps(tm, default=to_serializable)
+    TM.reset()
+    restored = loads(output)
+    restored.resolve()
+
+    assert restored.findings == []
+    assert [finding.threat_id for finding in restored.excluded_findings] == ["TEST01"]
+    assert restored.excluded_findings[0].assumption.exclude == {"TEST01"}