diff --git a/pytm/pytm.py b/pytm/pytm.py index 211e3c6..0f31bfa 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -288,6 +288,12 @@ def _add_data(container, value): container.append(item) +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 class _FlowDefaultsBuilder: """Collect and apply default relationships for data flows.""" @@ -400,37 +406,16 @@ 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) + _fill(flow, "srcPort", flow.source, "port") + _fill(flow.controls, "isEncrypted", flow.source.controls) def _apply_forward_defaults(self, flow: Dataflow) -> None: - flow._safeset( - "protocol", getattr(flow.sink, "protocol", getattr(flow, "protocol", "")) - ) - flow._safeset( - "dstPort", getattr(flow.sink, "port", getattr(flow, "dstPort", -1)) - ) - 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), - ) + _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 7f7b06e..fc74796 100644 --- a/tests/test_private_func.py +++ b/tests/test_private_func.py @@ -10,6 +10,7 @@ Dataflow, Datastore, DatastoreType, + Element, Finding, Process, Server, @@ -73,6 +74,37 @@ 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_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")