From 6a685adfe1ef9ef0d65b6fdebe9624fb60bfd782 Mon Sep 17 00:00:00 2001
From: Nick Le Mouton <nick@noodles.net.nz>
Date: Thu, 16 Apr 2026 10:26:46 +1200
Subject: [PATCH] Fix likelihood regression and added tests

---
 pytm/finding.py        |  6 +++++-
 pytm/tm.py             |  1 +
 tests/test_pytmfunc.py | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/pytm/finding.py b/pytm/finding.py
index 75d201f..9ab5559 100644
--- a/pytm/finding.py
+++ b/pytm/finding.py
@@ -27,6 +27,7 @@ class Finding(BaseModel):
         assumption (Assumption): The assumption that caused this finding to be excluded
         response (str): Describes how this threat matching this particular asset or dataflow is being handled. Can be one of: mitigated, transferred, avoided, accepted
         cvss (str): The CVSS score and/or vector
+        likelihood (str): Likelihood of the threat
     """
 
     model_config = ConfigDict(
@@ -57,6 +58,7 @@ class Finding(BaseModel):
         description="Describes how this threat matching this particular asset or dataflow is being handled. Can be one of: mitigated, transferred, avoided, accepted",
     )
     cvss: str = Field(default="", description="The CVSS score and/or vector")
+    likelihood: str = Field(default="", description="Likelihood of the threat")
 
     def __init__(self, *args, **kwargs):
         """Initialize a Finding.
@@ -66,7 +68,7 @@ def __init__(self, *args, **kwargs):
             **kwargs: Finding properties:
                 - element (Element): Element this finding applies to
                 - target (str): Name of the element this finding applies to
-                - threat (Threat): Threat object to copy attributes from (description, details, severity, mitigations, example, references, condition)
+                - threat (Threat): Threat object to copy attributes from (description, details, severity, mitigations, example, references, condition, likelihood)
                 - description (str): Threat description
                 - details (str): Threat details
                 - severity (str): Threat severity
@@ -79,6 +81,7 @@ def __init__(self, *args, **kwargs):
                 - assumption (Assumption): The assumption that caused this finding to be excluded
                 - response (str): Describes how this threat matching this particular asset or dataflow is being handled. Can be one of: mitigated, transferred, avoided, accepted
                 - cvss (str): The CVSS score and/or vector
+                - likelihood (str): Likelihood of the threat
         """
         # Handle positional element argument
         if args:
@@ -105,6 +108,7 @@ def __init__(self, *args, **kwargs):
                 "example",
                 "references",
                 "condition",
+                "likelihood",
             ]
             for attr in threat_attrs:
                 if attr not in kwargs:  # Don't override explicit values
diff --git a/pytm/tm.py b/pytm/tm.py
index fb75e0c..66fdeff 100644
--- a/pytm/tm.py
+++ b/pytm/tm.py
@@ -656,6 +656,7 @@ def encode_threat_data(obj):
                 "condition",
                 "cvss",
                 "response",
+                "likelihood",
             ]
 
             items = obj if isinstance(obj, list) else [obj]
diff --git a/tests/test_pytmfunc.py b/tests/test_pytmfunc.py
index c0696be..e44cd83 100644
--- a/tests/test_pytmfunc.py
+++ b/tests/test_pytmfunc.py
@@ -1728,3 +1728,44 @@ def test_override_finding_applied_through_resolve(self):
         assert len(server.findings) == 1
         assert server.findings[0].response == "accepted"
         assert server.findings[0].cvss == "5.0"
+
+    def test_likelihood_copied_from_threat_to_finding(self):
+        """likelihood is propagated from Threat to resolved Finding."""
+        TM.reset()
+        tm = TM("test tm", description="aaa")
+        Server("Web Server")
+        TM._threats = [Threat(SID="T01", target="Server", severity="High", likelihood="Medium")]
+        tm.resolve()
+
+        server = next(e for e in TM._elements if e.name == "Web Server")
+        assert len(server.findings) == 1
+        assert server.findings[0].likelihood == "Medium"
+
+    def test_override_finding_likelihood_not_overwritten(self):
+        """An explicit likelihood on a Finding override is preserved after resolve."""
+        TM.reset()
+        tm = TM("test tm", description="aaa")
+        Server(
+            "Web Server",
+            overrides=[
+                Finding(threat_id="T01", likelihood="High"),
+            ],
+        )
+        TM._threats = [Threat(SID="T01", target="Server", severity="High", likelihood="Low")]
+        tm.resolve()
+
+        server = next(e for e in TM._elements if e.name == "Web Server")
+        assert len(server.findings) == 1
+        assert server.findings[0].likelihood == "High"
+
+    def test_finding_likelihood_defaults_to_empty(self):
+        """likelihood defaults to empty string when the threat has none."""
+        TM.reset()
+        tm = TM("test tm", description="aaa")
+        Server("Web Server")
+        TM._threats = [Threat(SID="T01", target="Server", severity="High")]
+        tm.resolve()
+
+        server = next(e for e in TM._elements if e.name == "Web Server")
+        assert len(server.findings) == 1
+        assert server.findings[0].likelihood == ""