From c26db6f7a5687c2d4dc6bde555185473c45662b5 Mon Sep 17 00:00:00 2001 From: Siddharth Gangadhar Date: Fri, 5 Jun 2026 21:31:58 -0400 Subject: [PATCH] feat: replace tcptrace with PyShark in TraceProcessor tcptrace is unmaintained. Replace the subprocess shell-out with pyshark.FileCapture, which wraps Wireshark's actively-maintained dissectors and exposes tcp.analysis.ack_rtt, tcp.analysis.retransmission, and per-stream packet grouping natively in Python. _FlowAccumulator tracks tx_packets, unique_bytes, rexmt_packets, timestamps, and RTT samples per flow direction. _dominant_flow() picks the flow with the most unique bytes. The 12-element metrics list is preserved unchanged for downstream compatibility. Update test_trace.py: replace exe_comm mocks with SimpleNamespace fake packets and pyshark.FileCapture mocks. Add _FlowAccumulator unit tests. 56 tests, 99% coverage. Co-Authored-By: Claude Sonnet 4.6 --- pcapprocessor/trace.py | 167 ++++++++++++++++++---------------- tests/test_trace.py | 198 ++++++++++++++++++++++++++++------------- 2 files changed, 226 insertions(+), 139 deletions(-) diff --git a/pcapprocessor/trace.py b/pcapprocessor/trace.py index 7e2f301..99c3576 100644 --- a/pcapprocessor/trace.py +++ b/pcapprocessor/trace.py @@ -1,10 +1,43 @@ import re -import shlex -import datetime +from collections import defaultdict import numpy as np - -from pcapprocessor import exe_comm +import pyshark + + +class _FlowAccumulator: + """Accumulates per-packet stats for one TCP flow direction.""" + + def __init__(self): + self.tx_packets = 0 + self.unique_bytes = 0 + self.rexmt_packets = 0 + self.first_ts = None + self.last_ts = None + self.rtt_samples = [] + + def add_packet(self, ts: float, payload_len: int, is_retrans: bool, rtt_ms) -> None: + if self.first_ts is None: + self.first_ts = ts + self.last_ts = ts + if payload_len > 0: + self.tx_packets += 1 + if is_retrans: + self.rexmt_packets += 1 + else: + self.unique_bytes += payload_len + if rtt_ms is not None: + self.rtt_samples.append(rtt_ms) + + @property + def tx_time(self) -> float: + if self.first_ts is None: + return 0.0 + return max(self.last_ts - self.first_ts, 0.0) + + @property + def avg_rtt_ms(self) -> float: + return float(np.mean(self.rtt_samples)) if self.rtt_samples else 0.0 class TraceProcessor: @@ -30,84 +63,66 @@ def process(self) -> list: fact_by = self._unit_factor() bn_speed = self._bottleneck_speed(fact_by) - print("Processing ascii trace output") pkt_size = int(self.config.get(self.scenario, "pktSize")) - with open(self.ascii_trace_file, "r") as fl: - lines = list(fl) - axis_y1 = [int(line.strip().split(",")[1]) for line in lines] + with open(self.ascii_trace_file) as fl: + axis_y1 = [int(ln.strip().split(",")[1]) for ln in fl] a = np.array(axis_y1) - queue_mean = a.mean() - queue_variance = a.var() - - trace_cmd = "tcptrace -l -r -n -W --csv " + self.pcap_file - print("Executing tcptrace command. it may take few seconds") - result = exe_comm.exe_comm(shlex.split(trace_cmd)) - print("Processing trace output") - - pcap_lines = result.split("\n") - regex_con = re.compile(r"#([0-9]*) TCP connection traced:") - matches = [ - pcap_lines.index(ln) - for ln in pcap_lines - if re.match(regex_con, ln) - ] + queue_mean = float(a.mean()) + queue_variance = float(a.var()) - if not matches: + flow = self._dominant_flow() + if flow is None: raise ValueError("No TCP connections found in pcap file") - connections = [ - pcap_lines[matches[j]: matches[j + 1]] - for j in range(len(matches) - 1) + tx_time = flow.tx_time + throughput = flow.unique_bytes / tx_time if tx_time > 0 else 0.0 + utilization = throughput * 100.0 / bn_speed if bn_speed > 0 else 0.0 + + return [ + flow.tx_packets, + self.PACKET_OVERHEAD * flow.tx_packets, + round(throughput / fact_by, 3), + round(flow.avg_rtt_ms / 2, 3), + round(throughput * 8 / fact_by, 3), + flow.unique_bytes, + flow.rexmt_packets, + round(utilization, 3), + round(queue_mean, 3), + round(queue_variance, 3), + round(queue_mean * 100.0 / (self.buf_size / pkt_size), 3), + round(tx_time * 1000, 3), ] - connections.append(pcap_lines[matches[-1]:]) - - result_str = "\n" - for i, item in enumerate(connections): - flow_cmp_time = 0 - try: - time_stamp = pcap_lines[matches[i] - 2].split()[-1] - t = datetime.datetime.strptime(time_stamp, "%H:%M:%S.%f") - flow_cmp_time = ( - t.time().hour * 3600 + t.time().minute * 60 + t.time().second - ) * 1000 + t.time().microsecond / 1000 - except Exception: - print("Couldn't parse the flow completion time") - - labels = item[1].split(",") - values = item[3].split(",") - - conn_suffix = "a2b" - if int(values[labels.index("unique_bytes_sent_a2b")]) <= 0: - conn_suffix = "b2a" - - tx_packets = int(values[labels.index("total_packets_" + conn_suffix)]) - rexmt_packets = int(values[labels.index("rexmt_data_pkts_" + conn_suffix)]) - overhead = self.PACKET_OVERHEAD * tx_packets - goodput = 8 * int(values[labels.index("throughput_" + conn_suffix)]) - unique_bytes = int(values[labels.index("unique_bytes_sent_" + conn_suffix)]) - tx_time = (1.0 * unique_bytes) / goodput - throughput = ( - int(values[labels.index("actual_data_bytes_" + conn_suffix)]) / tx_time - ) - rtt = float(values[labels.index("RTT_avg_" + conn_suffix)]) / 2 - utilization = throughput * 100.0 / bn_speed - - result_str = [ - tx_packets, - overhead, - round(throughput / fact_by, 3), - round(rtt, 3), - round(goodput / fact_by, 3), - unique_bytes, - rexmt_packets, - round(utilization, 3), - round(queue_mean, 3), - round(queue_variance, 3), - round(queue_mean * 100.0 / (self.buf_size / pkt_size), 3), - round(flow_cmp_time, 3), - ] - - return result_str + + def _dominant_flow(self): + """Return the flow accumulator with the most unique bytes, or None.""" + flows = defaultdict(_FlowAccumulator) + cap = pyshark.FileCapture( + self.pcap_file, + display_filter="tcp", + keep_packets=False, + ) + try: + for pkt in cap: + self._process_packet(pkt, flows) + finally: + cap.close() + if not flows: + return None + best = max(flows.values(), key=lambda f: f.unique_bytes) + return best if best.unique_bytes > 0 else None + + @staticmethod + def _process_packet(pkt, flows) -> None: + try: + tcp = pkt.tcp + key = (int(tcp.stream), pkt.ip.src, tcp.srcport, pkt.ip.dst, tcp.dstport) + ts = float(pkt.sniff_timestamp) + payload_len = int(tcp.len) if hasattr(tcp, "len") else 0 + is_retrans = hasattr(tcp, "analysis_retransmission") + rtt_ms = float(tcp.analysis_ack_rtt) * 1000 if hasattr(tcp, "analysis_ack_rtt") else None + flows[key].add_packet(ts, payload_len, is_retrans, rtt_ms) + except AttributeError: + pass def _unit_factor(self) -> float: first, second = list(self.unit) diff --git a/tests/test_trace.py b/tests/test_trace.py index d97d91b..0fbe32a 100644 --- a/tests/test_trace.py +++ b/tests/test_trace.py @@ -1,31 +1,9 @@ import pytest from configparser import ConfigParser -from unittest.mock import patch - -from pcapprocessor.trace import TraceProcessor - - -_A2B_OUTPUT = "\n".join([ - "timestamp line 00:01:00.000000", - "filler line", - "#1 TCP connection traced:", - "unique_bytes_sent_a2b,total_packets_a2b,rexmt_data_pkts_a2b," - "throughput_a2b,actual_data_bytes_a2b,RTT_avg_a2b", - "extra line", - "1000,10,0,1000,500,5.0", - "", -]) - -_B2A_OUTPUT = "\n".join([ - "timestamp line 00:01:00.000000", - "filler line", - "#1 TCP connection traced:", - "unique_bytes_sent_a2b,total_packets_b2a,rexmt_data_pkts_b2a," - "throughput_b2a,unique_bytes_sent_b2a,actual_data_bytes_b2a,RTT_avg_b2a", - "extra line", - "0,10,0,1000,1000,500,5.0", - "", -]) +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from pcapprocessor.trace import TraceProcessor, _FlowAccumulator def _make_config(scenario="myScenario"): @@ -47,76 +25,170 @@ def _make_processor(tmp_path, scenario="myScenario", unit="MB", buf_size=1000): ) +def _fake_pkt(stream="0", src="1.1.1.1", srcport="1234", + dst="2.2.2.2", dstport="80", + payload_len=1000, ts=0.0, + is_retrans=False, rtt_ms=None): + """SimpleNamespace fake pyshark packet — hasattr behaves naturally.""" + tcp = SimpleNamespace(stream=stream, srcport=srcport, dstport=dstport, len=str(payload_len)) + if is_retrans: + tcp.analysis_retransmission = "1" + if rtt_ms is not None: + tcp.analysis_ack_rtt = str(rtt_ms / 1000) + ip = SimpleNamespace(src=src, dst=dst) + pkt = SimpleNamespace(tcp=tcp, ip=ip) + pkt.sniff_timestamp = str(ts) + return pkt + + +def _mock_cap(packets): + cap = MagicMock() + cap.__iter__ = MagicMock(return_value=iter(packets)) + return cap + + +# --- _FlowAccumulator --- + +def test_flow_accumulator_tracks_unique_bytes(): + acc = _FlowAccumulator() + acc.add_packet(0.0, 500, False, None) + acc.add_packet(1.0, 300, False, 10.0) + assert acc.unique_bytes == 800 + assert acc.tx_packets == 2 + + +def test_flow_accumulator_separates_retransmissions(): + acc = _FlowAccumulator() + acc.add_packet(0.0, 1000, False, None) + acc.add_packet(0.5, 500, True, None) + assert acc.unique_bytes == 1000 + assert acc.rexmt_packets == 1 + assert acc.tx_packets == 2 + + +def test_flow_accumulator_tx_time(): + acc = _FlowAccumulator() + acc.add_packet(1.0, 100, False, None) + acc.add_packet(3.5, 100, False, None) + assert acc.tx_time == pytest.approx(2.5) + + +def test_flow_accumulator_avg_rtt_ms(): + acc = _FlowAccumulator() + acc.add_packet(0.0, 100, False, 10.0) + acc.add_packet(1.0, 100, False, 20.0) + assert acc.avg_rtt_ms == pytest.approx(15.0) + + +def test_flow_accumulator_avg_rtt_ms_empty(): + assert _FlowAccumulator().avg_rtt_ms == 0.0 + + +def test_flow_accumulator_tx_time_no_packets(): + assert _FlowAccumulator().tx_time == 0.0 + + +# --- _unit_factor --- + def test_unit_factor_megabytes(): - proc = TraceProcessor("f.pcap", "MB", None, "s", "t.txt", 100) - assert proc._unit_factor() == 8_000_000 + assert TraceProcessor("f.pcap", "MB", None, "s", "t.txt", 100)._unit_factor() == 8_000_000 def test_unit_factor_megabits(): - proc = TraceProcessor("f.pcap", "Mb", None, "s", "t.txt", 100) - assert proc._unit_factor() == 1_000_000 + assert TraceProcessor("f.pcap", "Mb", None, "s", "t.txt", 100)._unit_factor() == 1_000_000 def test_unit_factor_kilobytes(): - proc = TraceProcessor("f.pcap", "KB", None, "s", "t.txt", 100) - assert proc._unit_factor() == 8_000 + assert TraceProcessor("f.pcap", "KB", None, "s", "t.txt", 100)._unit_factor() == 8_000 def test_unit_factor_gigabits(): - proc = TraceProcessor("f.pcap", "Gb", None, "s", "t.txt", 100) - assert proc._unit_factor() == 1_000_000_000 + assert TraceProcessor("f.pcap", "Gb", None, "s", "t.txt", 100)._unit_factor() == 1_000_000_000 +# --- _bottleneck_speed --- + def test_bottleneck_speed_extracts_numeric(tmp_path): proc = _make_processor(tmp_path) - # bottleneckSpeed="10Mbps", fact_by=1 → 10*1 = 10 assert proc._bottleneck_speed(1.0) == 10.0 +# --- _dominant_flow --- + +def test_dominant_flow_returns_none_when_no_packets(tmp_path): + proc = _make_processor(tmp_path) + with patch("pcapprocessor.trace.pyshark.FileCapture", return_value=_mock_cap([])): + assert proc._dominant_flow() is None + + +def test_dominant_flow_selects_flow_with_most_bytes(tmp_path): + proc = _make_processor(tmp_path) + pkts = [ + _fake_pkt(src="1.1.1.1", srcport="1234", dst="2.2.2.2", dstport="80", + payload_len=1000, ts=0.0), + _fake_pkt(src="2.2.2.2", srcport="80", dst="1.1.1.1", dstport="1234", + payload_len=50, ts=0.1), + ] + with patch("pcapprocessor.trace.pyshark.FileCapture", return_value=_mock_cap(pkts)): + flow = proc._dominant_flow() + assert flow.unique_bytes == 1000 + + +def test_dominant_flow_counts_retransmissions(tmp_path): + proc = _make_processor(tmp_path) + pkts = [ + _fake_pkt(payload_len=1000, ts=0.0), + _fake_pkt(payload_len=500, ts=0.1, is_retrans=True), + ] + with patch("pcapprocessor.trace.pyshark.FileCapture", return_value=_mock_cap(pkts)): + flow = proc._dominant_flow() + assert flow.rexmt_packets == 1 + assert flow.unique_bytes == 1000 + + +def test_dominant_flow_collects_rtt_samples(tmp_path): + proc = _make_processor(tmp_path) + pkts = [ + _fake_pkt(payload_len=1000, ts=0.0, rtt_ms=10.0), + _fake_pkt(payload_len=500, ts=0.5, rtt_ms=20.0), + ] + with patch("pcapprocessor.trace.pyshark.FileCapture", return_value=_mock_cap(pkts)): + flow = proc._dominant_flow() + assert flow.avg_rtt_ms == pytest.approx(15.0) + + +# --- process --- + def test_process_returns_12_metrics(tmp_path): proc = _make_processor(tmp_path) - with patch("pcapprocessor.exe_comm.exe_comm", return_value=_A2B_OUTPUT): + pkts = [ + _fake_pkt(payload_len=1000, ts=0.0, rtt_ms=10.0), + _fake_pkt(payload_len=1000, ts=1.0), + ] + with patch("pcapprocessor.trace.pyshark.FileCapture", return_value=_mock_cap(pkts)): result = proc.process() assert isinstance(result, list) assert len(result) == 12 -def test_process_a2b_tx_packets(tmp_path): +def test_process_overhead_is_packet_overhead_times_tx_packets(tmp_path): proc = _make_processor(tmp_path) - with patch("pcapprocessor.exe_comm.exe_comm", return_value=_A2B_OUTPUT): + pkts = [_fake_pkt(payload_len=500, ts=0.0), _fake_pkt(payload_len=500, ts=1.0)] + with patch("pcapprocessor.trace.pyshark.FileCapture", return_value=_mock_cap(pkts)): result = proc.process() - assert result[0] == 10 # tx_packets - assert result[1] == 320 # overhead = PACKET_OVERHEAD * tx_packets + assert result[1] == result[0] * TraceProcessor.PACKET_OVERHEAD -def test_process_b2a_path_returns_12_metrics(tmp_path): +def test_process_flow_completion_time_in_ms(tmp_path): proc = _make_processor(tmp_path) - with patch("pcapprocessor.exe_comm.exe_comm", return_value=_B2A_OUTPUT): + pkts = [_fake_pkt(payload_len=1000, ts=0.0), _fake_pkt(payload_len=500, ts=1.0)] + with patch("pcapprocessor.trace.pyshark.FileCapture", return_value=_mock_cap(pkts)): result = proc.process() - assert isinstance(result, list) - assert len(result) == 12 + assert result[-1] == pytest.approx(1000.0) def test_process_raises_on_no_connections(tmp_path): proc = _make_processor(tmp_path) - with patch("pcapprocessor.exe_comm.exe_comm", return_value="no tcp data here\n"): + with patch("pcapprocessor.trace.pyshark.FileCapture", return_value=_mock_cap([])): with pytest.raises(ValueError, match="No TCP connections found"): proc.process() - - -def test_process_flow_cmp_time_fallback_on_bad_timestamp(tmp_path): - # Timestamp that can't be parsed → flow_cmp_time stays 0 - bad_ts_output = "\n".join([ - "no timestamp here", - "filler line", - "#1 TCP connection traced:", - "unique_bytes_sent_a2b,total_packets_a2b,rexmt_data_pkts_a2b," - "throughput_a2b,actual_data_bytes_a2b,RTT_avg_a2b", - "extra line", - "1000,10,0,1000,500,5.0", - "", - ]) - proc = _make_processor(tmp_path) - with patch("pcapprocessor.exe_comm.exe_comm", return_value=bad_ts_output): - result = proc.process() - assert result[-1] == 0 # flow_cmp_time defaults to 0