From 7c7eb834607d45d72e3f1e2b5e6c8a30e2e9f339 Mon Sep 17 00:00:00 2001 From: svastikkka Date: Thu, 29 Oct 2026 10:50:55 +0530 Subject: [PATCH 1/2] fix: keep explicit dataflow values when check() applies defaults --- pytm/pytm.py | 46 ++++++++++++++++---------------------- tests/test_private_func.py | 20 +++++++++++++++++ 2 files changed, 39 insertions(+), 27 deletions(-) diff --git a/pytm/pytm.py b/pytm/pytm.py index 211e3c69..74ed2de5 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -288,6 +288,11 @@ def _add_data(container, value): container.append(item) +def _fill(obj, attr, value): + if attr not in obj.model_fields_set: + obj._safeset(attr, value) + + @dataclass class _FlowDefaultsBuilder: """Collect and apply default relationships for data flows.""" @@ -400,37 +405,24 @@ def _merge_overrides(self, flow: Dataflow) -> None: pass def _apply_response_defaults(self, flow: Dataflow) -> None: - flow._safeset( - "protocol", getattr(flow.source, "protocol", getattr(flow, "protocol", "")) - ) - flow._safeset( - "srcPort", getattr(flow.source, "port", getattr(flow, "srcPort", -1)) - ) - if hasattr(flow.source, "controls"): - flow.controls._safeset( - "isEncrypted", getattr(flow.source.controls, "isEncrypted", False) - ) + _fill(flow, "protocol", flow.source.protocol) + _fill(flow, "srcPort", flow.source.port) + _fill(flow.controls, "isEncrypted", flow.source.controls.isEncrypted) def _apply_forward_defaults(self, flow: Dataflow) -> None: - flow._safeset( - "protocol", getattr(flow.sink, "protocol", getattr(flow, "protocol", "")) + _fill(flow, "protocol", flow.sink.protocol) + _fill(flow, "dstPort", flow.sink.port) + _fill(flow.controls, "isEncrypted", flow.sink.controls.isEncrypted) + _fill( + flow.controls, + "authenticatesDestination", + flow.source.controls.authenticatesDestination, ) - flow._safeset( - "dstPort", getattr(flow.sink, "port", getattr(flow, "dstPort", -1)) + _fill( + flow.controls, + "checksDestinationRevocation", + flow.source.controls.checksDestinationRevocation, ) - if hasattr(flow.sink, "controls"): - flow.controls._safeset( - "isEncrypted", getattr(flow.sink.controls, "isEncrypted", False) - ) - if hasattr(flow.source, "controls"): - flow.controls._safeset( - "authenticatesDestination", - getattr(flow.source.controls, "authenticatesDestination", False), - ) - flow.controls._safeset( - "checksDestinationRevocation", - getattr(flow.source.controls, "checksDestinationRevocation", False), - ) def _enrich_data_attributes(self, flow: Dataflow) -> None: for datum in list(getattr(flow, "data", [])): diff --git a/tests/test_private_func.py b/tests/test_private_func.py index 7f7b06e0..cd21f41f 100644 --- a/tests/test_private_func.py +++ b/tests/test_private_func.py @@ -73,6 +73,26 @@ def test_responses(self): assert insert.response is None assert insert.isResponse is False + def test_explicit_dataflow_values_survive_check(self): + tm = TM("TM") + user = Actor("User") + web = Server("Web", port=8080, protocol="HTTP") + login = Dataflow(user, web, "Login") + login.protocol = "HTTPS" + login.dstPort = 443 + login.controls.isEncrypted = True + reply = Dataflow(web, user, "Reply", isResponse=True) + reply.protocol = "HTTPS" + reply.srcPort = 443 + reply.controls.isEncrypted = True + assert tm.check() + assert login.protocol == "HTTPS" + assert login.dstPort == 443 + assert login.controls.isEncrypted is True + assert reply.protocol == "HTTPS" + assert reply.srcPort == 443 + assert reply.controls.isEncrypted is True + def test_defaults(self): tm = TM("TM") user_data = Data("HTTP") From 41c68949262ea1bf3c64ea204f8f35d24edd59e3 Mon Sep 17 00:00:00 2001 From: svastikkka Date: Sat, 14 Nov 2026 10:55:25 +0530 Subject: [PATCH 2/2] fix: skip protocol and port defaults for endpoints that do not have them --- pytm/pytm.py | 31 ++++++++++++------------------- tests/test_private_func.py | 12 ++++++++++++ 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/pytm/pytm.py b/pytm/pytm.py index 74ed2de5..0f31bfa5 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -288,9 +288,10 @@ def _add_data(container, value): container.append(item) -def _fill(obj, attr, value): - if attr not in obj.model_fields_set: - obj._safeset(attr, value) +def _fill(obj, attr, source, source_attr=None): + source_attr = source_attr or attr + if attr not in obj.model_fields_set and hasattr(source, source_attr): + obj._safeset(attr, getattr(source, source_attr)) @dataclass @@ -405,24 +406,16 @@ def _merge_overrides(self, flow: Dataflow) -> None: pass def _apply_response_defaults(self, flow: Dataflow) -> None: - _fill(flow, "protocol", flow.source.protocol) - _fill(flow, "srcPort", flow.source.port) - _fill(flow.controls, "isEncrypted", flow.source.controls.isEncrypted) + _fill(flow, "protocol", flow.source) + _fill(flow, "srcPort", flow.source, "port") + _fill(flow.controls, "isEncrypted", flow.source.controls) def _apply_forward_defaults(self, flow: Dataflow) -> None: - _fill(flow, "protocol", flow.sink.protocol) - _fill(flow, "dstPort", flow.sink.port) - _fill(flow.controls, "isEncrypted", flow.sink.controls.isEncrypted) - _fill( - flow.controls, - "authenticatesDestination", - flow.source.controls.authenticatesDestination, - ) - _fill( - flow.controls, - "checksDestinationRevocation", - flow.source.controls.checksDestinationRevocation, - ) + _fill(flow, "protocol", flow.sink) + _fill(flow, "dstPort", flow.sink, "port") + _fill(flow.controls, "isEncrypted", flow.sink.controls) + _fill(flow.controls, "authenticatesDestination", flow.source.controls) + _fill(flow.controls, "checksDestinationRevocation", flow.source.controls) def _enrich_data_attributes(self, flow: Dataflow) -> None: for datum in list(getattr(flow, "data", [])): diff --git a/tests/test_private_func.py b/tests/test_private_func.py index cd21f41f..fc74796e 100644 --- a/tests/test_private_func.py +++ b/tests/test_private_func.py @@ -10,6 +10,7 @@ Dataflow, Datastore, DatastoreType, + Element, Finding, Process, Server, @@ -93,6 +94,17 @@ def test_explicit_dataflow_values_survive_check(self): assert reply.srcPort == 443 assert reply.controls.isEncrypted is True + def test_check_with_element_and_boundary_endpoints(self): + tm = TM("TM") + user = Actor("User") + element = Element("Element") + zone = Boundary("Zone") + Dataflow(user, element, "To element") + Dataflow(element, user, "From element", isResponse=True) + Dataflow(user, zone, "To boundary") + Dataflow(zone, user, "From boundary", isResponse=True) + assert tm.check() + def test_defaults(self): tm = TM("TM") user_data = Data("HTTP")