diff --git a/src/netsleuth/probes/icmp_win.py b/src/netsleuth/probes/icmp_win.py index 1b242f5..b7a61a7 100644 --- a/src/netsleuth/probes/icmp_win.py +++ b/src/netsleuth/probes/icmp_win.py @@ -1,6 +1,7 @@ from __future__ import annotations import struct +from dataclasses import dataclass from typing import NamedTuple IP_SUCCESS = 0 @@ -22,6 +23,16 @@ class IcmpReply(NamedTuple): ttl: int +@dataclass +class IcmpRequest: + dest: str + ttl: int + timeout_ms: int + payload: bytes = b"netsleuth" + source_ip: str | None = None + df: bool = False + + def classify_status(status: int) -> str: if status == IP_SUCCESS: return "ok" @@ -128,23 +139,16 @@ def _handle(): return iphlpapi, handle -def echo_once( - dest: str, - ttl: int, - timeout_ms: int, - payload: bytes = b"netsleuth", - source_ip: str | None = None, - df: bool = False, -) -> IcmpReply: +def echo_once(req: IcmpRequest) -> IcmpReply: iphlpapi, handle = _handle() try: options = _IpOptionInformation( - Ttl=ttl, Tos=0, Flags=IP_FLAG_DF if df else 0, OptionsSize=0, OptionsData=None + Ttl=req.ttl, Tos=0, Flags=IP_FLAG_DF if req.df else 0, OptionsSize=0, OptionsData=None ) buffer = ctypes.create_string_buffer(_REPLY_BUFFER_SIZE) - address = struct.unpack("= deadline: break - reply = echo_once(address, ttl=ttl, timeout_ms=per_hop_timeout_ms, source_ip=source_ip) + reply = echo_once( + IcmpRequest(dest=address, ttl=ttl, timeout_ms=per_hop_timeout_ms, source_ip=source_ip) + ) hops.append((ttl, reply)) if classify_status(reply.status) == "ok": break diff --git a/src/netsleuth/probes/pmtud.py b/src/netsleuth/probes/pmtud.py index 2c572db..38f9f77 100644 --- a/src/netsleuth/probes/pmtud.py +++ b/src/netsleuth/probes/pmtud.py @@ -69,13 +69,22 @@ def supports_df() -> tuple[bool, str]: async def _probe_windows( resolved_ip: str, size: int, timeout: float, source_ip: str | None ) -> tuple[bool, bool]: - from netsleuth.probes.icmp_win import classify_status, echo_once + from netsleuth.probes.icmp_win import IcmpRequest, classify_status, echo_once payload = b"\x00" * max(0, size - _IPV4_ICMP_OVERHEAD) timeout_ms = int(timeout * 1000) def _run(): - return echo_once(resolved_ip, ttl=64, timeout_ms=timeout_ms, payload=payload, source_ip=source_ip, df=True) + return echo_once( + IcmpRequest( + dest=resolved_ip, + ttl=64, + timeout_ms=timeout_ms, + payload=payload, + source_ip=source_ip, + df=True, + ) + ) reply = await asyncio.to_thread(_run) kind = classify_status(reply.status) diff --git a/tests/test_icmp_win_trace.py b/tests/test_icmp_win_trace.py index 2bbcc2b..17034bd 100644 --- a/tests/test_icmp_win_trace.py +++ b/tests/test_icmp_win_trace.py @@ -3,7 +3,7 @@ from __future__ import annotations import time from unittest.mock import patch -from netsleuth.probes.icmp_win import IP_REQ_TIMED_OUT, IcmpReply, trace_hops_win +from netsleuth.probes.icmp_win import IP_REQ_TIMED_OUT, IcmpReply, IcmpRequest, trace_hops_win def _fake_gethostbyname(host: str) -> str: @@ -19,9 +19,9 @@ def test_trace_hops_win_total_elapsed_is_bounded_by_the_total_budget_not_hops_ti # to the 100ms total budget. call_timeouts: list[int] = [] - def fake_echo_once(dest: str, ttl: int, timeout_ms: int, payload: bytes = b"netsleuth", source_ip=None) -> IcmpReply: - call_timeouts.append(timeout_ms) - time.sleep(timeout_ms / 1000.0) + def fake_echo_once(req: IcmpRequest) -> IcmpReply: + call_timeouts.append(req.timeout_ms) + time.sleep(req.timeout_ms / 1000.0) return IcmpReply(address=None, status=IP_REQ_TIMED_OUT, rtt_ms=None, ttl=0) with ( @@ -43,8 +43,8 @@ def test_trace_hops_win_total_elapsed_is_bounded_by_the_total_budget_not_hops_ti def test_trace_hops_win_per_hop_timeout_is_capped_even_for_a_huge_total_budget(): call_timeouts: list[int] = [] - def fake_echo_once(dest: str, ttl: int, timeout_ms: int, payload: bytes = b"netsleuth", source_ip=None) -> IcmpReply: - call_timeouts.append(timeout_ms) + def fake_echo_once(req: IcmpRequest) -> IcmpReply: + call_timeouts.append(req.timeout_ms) return IcmpReply(address=None, status=IP_REQ_TIMED_OUT, rtt_ms=None, ttl=0) with ( @@ -62,9 +62,9 @@ def test_trace_hops_win_per_hop_timeout_is_capped_even_for_a_huge_total_budget() def test_trace_hops_win_happy_path_hop_count_and_ok_break_are_unchanged(): from netsleuth.probes.icmp_win import IP_SUCCESS - def fake_echo_once(dest: str, ttl: int, timeout_ms: int, payload: bytes = b"netsleuth", source_ip=None) -> IcmpReply: - if ttl < 3: - return IcmpReply(address=f"10.0.0.{ttl}", status=IP_REQ_TIMED_OUT, rtt_ms=None, ttl=0) + def fake_echo_once(req: IcmpRequest) -> IcmpReply: + if req.ttl < 3: + return IcmpReply(address=f"10.0.0.{req.ttl}", status=IP_REQ_TIMED_OUT, rtt_ms=None, ttl=0) return IcmpReply(address="203.0.113.1", status=IP_SUCCESS, rtt_ms=5.0, ttl=64) with (