Skip to content

Commit 343f343

Browse files
ccie18643ccie18643
committed
Refactored IP fragment handling code.
1 parent eb56471 commit 343f343

File tree

4 files changed

+108
-41
lines changed

4 files changed

+108
-41
lines changed

‎pytcp/lib/ip_frag.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#!/usr/bin/env python3
2+
3+
################################################################################
4+
## ##
5+
## PyTCP - Python TCP/IP stack ##
6+
## Copyright (C) 2020-present Sebastian Majewski ##
7+
## ##
8+
## This program is free software: you can redistribute it and/or modify ##
9+
## it under the terms of the GNU General Public License as published by ##
10+
## the Free Software Foundation, either version 3 of the License, or ##
11+
## (at your option) any later version. ##
12+
## ##
13+
## This program is distributed in the hope that it will be useful, ##
14+
## but WITHOUT ANY WARRANTY; without even the implied warranty of ##
15+
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the ##
16+
## GNU General Public License for more details. ##
17+
## ##
18+
## You should have received a copy of the GNU General Public License ##
19+
## along with this program. If not, see <https://www.gnu.org/licenses/>. ##
20+
## ##
21+
## Author's email: ccie18643@gmail.com ##
22+
## Github repository: https://github.com/ccie18643/PyTCP ##
23+
## ##
24+
################################################################################
25+
26+
27+
"""
28+
Module contains class FragFlow used for IPv4/IPv6 packet fragmentation and reassembly.
29+
30+
pytcp/lib/frag_flow.py
31+
32+
ver 3.0.2
33+
"""
34+
35+
36+
from __future__ import annotations
37+
38+
import time
39+
from dataclasses import dataclass, field
40+
41+
from net_addr.ip_address import IpAddress
42+
43+
44+
@dataclass(kw_only=True, frozen=True)
45+
class IpFragFlowId:
46+
"""
47+
Class stores IPv4/IPv6 packet fragmentation flow ID.
48+
"""
49+
50+
src: IpAddress
51+
dst: IpAddress
52+
id: int
53+
54+
55+
@dataclass(kw_only=True)
56+
class IpFragData:
57+
"""
58+
Class stores IPv4/IPv6 packet fragmentation data.
59+
"""
60+
61+
timestamp: float = field(default_factory=time.time, init=False)
62+
header: bytes
63+
last: bool = field(default=False, init=False)
64+
payload: dict[int, bytes]

‎pytcp/stack/packet_handler/__init__.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
Ip6Network,
4949
MacAddress,
5050
)
51+
from pytcp.lib.ip_frag import IpFragData, IpFragFlowId
5152
from pytcp.lib.logger import log
5253
from pytcp.lib.packet_stats import PacketStatsRx, PacketStatsTx
5354
from pytcp.protocols.dhcp4__legacy.client import Dhcp4Client
@@ -171,8 +172,8 @@ def __init__(
171172
self.ip6_id: int = 0
172173

173174
# Used to defragment IPv4 and IPv6 packets
174-
self.ip4_frag_flows: dict[tuple[Ip4Address, Ip4Address, int], dict] = {}
175-
self.ip6_frag_flows: dict[tuple[Ip6Address, Ip6Address, int], dict] = {}
175+
self.ip4_frag_flows: dict[IpFragFlowId, IpFragData] = {}
176+
self.ip6_frag_flows: dict[IpFragFlowId, IpFragData] = {}
176177

177178
# Thread control
178179
self._run_thread: bool = False

‎pytcp/stack/packet_handler/packet_handler__ip4__rx.py‎

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242

4343
from pytcp import stack
4444
from pytcp.lib.inet_cksum import inet_cksum
45+
from pytcp.lib.ip_frag import IpFragData, IpFragFlowId
4546
from pytcp.lib.logger import log
4647
from pytcp.lib.packet import PacketRx
4748
from pytcp.protocols.enums import IpProto
@@ -63,7 +64,7 @@ class PacketHandlerIp4Rx(ABC):
6364

6465
packet_stats_rx: PacketStatsRx
6566
ip4_multicast: list[Ip4Address]
66-
ip4_frag_flows: dict[tuple[Ip4Address, Ip4Address, int], dict]
67+
ip4_frag_flows: dict[IpFragFlowId, IpFragData]
6768

6869
# pylint: disable=unused-argument
6970

@@ -183,7 +184,7 @@ def __defragment_ip4_packet(self, packet_rx: PacketRx) -> PacketRx | None:
183184
self.ip4_frag_flows = {
184185
flow: self.ip4_frag_flows[flow]
185186
for flow in self.ip4_frag_flows
186-
if self.ip4_frag_flows[flow]["timestamp"] - time()
187+
if self.ip4_frag_flows[flow].timestamp - time()
187188
< stack.IP4__FRAG_FLOW_TIMEOUT
188189
}
189190

@@ -194,43 +195,45 @@ def __defragment_ip4_packet(self, packet_rx: PacketRx) -> PacketRx | None:
194195
f"{'' if packet_rx.ip4.flag_mf else ', last'}",
195196
)
196197

197-
flow_id = (packet_rx.ip4.src, packet_rx.ip4.dst, packet_rx.ip4.id)
198+
flow_id = IpFragFlowId(
199+
src=packet_rx.ip4.src,
200+
dst=packet_rx.ip4.dst,
201+
id=packet_rx.ip4.id,
202+
)
198203

199204
# Update flow db
200205
if flow_id in self.ip4_frag_flows:
201-
self.ip4_frag_flows[flow_id]["payload"][
206+
self.ip4_frag_flows[flow_id].payload[
202207
packet_rx.ip4.offset
203208
] = packet_rx.ip4.payload_bytes
204209
else:
205-
self.ip4_frag_flows[flow_id] = {
206-
"header": packet_rx.ip4.header_bytes,
207-
"timestamp": time(),
208-
"last": False,
209-
"payload": {packet_rx.ip4.offset: packet_rx.ip4.payload_bytes},
210-
}
210+
self.ip4_frag_flows[flow_id] = IpFragData(
211+
header=packet_rx.ip4.header_bytes,
212+
payload={packet_rx.ip4.offset: packet_rx.ip4.payload_bytes},
213+
)
211214
if not packet_rx.ip4.flag_mf:
212-
self.ip4_frag_flows[flow_id]["last"] = True
215+
self.ip4_frag_flows[flow_id].last = True
213216

214217
# Test if we received all fragments
215-
if not self.ip4_frag_flows[flow_id]["last"]:
218+
if not self.ip4_frag_flows[flow_id].last:
216219
return None
217220
payload_len = 0
218-
for offset in sorted(self.ip4_frag_flows[flow_id]["payload"]):
221+
for offset in sorted(self.ip4_frag_flows[flow_id].payload):
219222
if offset > payload_len:
220223
return None
221224
payload_len = offset + len(
222-
self.ip4_frag_flows[flow_id]["payload"][offset]
225+
self.ip4_frag_flows[flow_id].payload[offset]
223226
)
224227

225228
# Defragment packet
226-
header = bytearray(self.ip4_frag_flows[flow_id]["header"])
229+
header = bytearray(self.ip4_frag_flows[flow_id].header)
227230
payload = bytearray(payload_len)
228-
for offset in sorted(self.ip4_frag_flows[flow_id]["payload"]):
231+
for offset in sorted(self.ip4_frag_flows[flow_id].payload):
229232
struct.pack_into(
230-
f"{len(self.ip4_frag_flows[flow_id]['payload'][offset])}s",
233+
f"{len(self.ip4_frag_flows[flow_id].payload[offset])}s",
231234
payload,
232235
offset,
233-
bytes(self.ip4_frag_flows[flow_id]["payload"][offset]),
236+
bytes(self.ip4_frag_flows[flow_id].payload[offset]),
234237
)
235238
del self.ip4_frag_flows[flow_id]
236239
header[0] = 0x45

‎pytcp/stack/packet_handler/packet_handler__ip6_frag__rx.py‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
from typing import TYPE_CHECKING
4242

4343
from pytcp import stack
44+
from pytcp.lib.ip_frag import IpFragData, IpFragFlowId
4445
from pytcp.lib.logger import log
4546
from pytcp.lib.packet import PacketRx
4647
from pytcp.protocols.ip6_frag.ip6_frag__parser import Ip6FragParser
@@ -56,7 +57,7 @@ class PacketHandlerIp6FragRx(ABC):
5657
from pytcp.lib.packet_stats import PacketStatsRx
5758

5859
packet_stats_rx: PacketStatsRx
59-
ip6_frag_flows: dict[tuple[Ip6Address, Ip6Address, int], dict]
60+
ip6_frag_flows: dict[IpFragFlowId, IpFragData]
6061

6162
# pylint: disable=unused-argument
6263

@@ -96,7 +97,7 @@ def __defragment_ip6_packet(self, packet_rx: PacketRx) -> PacketRx | None:
9697
self.ip6_frag_flows = {
9798
flow: self.ip6_frag_flows[flow]
9899
for flow in self.ip6_frag_flows
99-
if self.ip6_frag_flows[flow]["timestamp"] - time()
100+
if self.ip6_frag_flows[flow].timestamp - time()
100101
< stack.IP6__FRAG_FLOW_TIMEOUT
101102
}
102103

@@ -108,49 +109,47 @@ def __defragment_ip6_packet(self, packet_rx: PacketRx) -> PacketRx | None:
108109
f"{'' if packet_rx.ip6_frag.flag_mf else ', last'}",
109110
)
110111

111-
flow_id = (
112-
packet_rx.ip6.src,
113-
packet_rx.ip6.dst,
114-
packet_rx.ip6_frag.id,
112+
flow_id = IpFragFlowId(
113+
src=packet_rx.ip6.src,
114+
dst=packet_rx.ip6.dst,
115+
id=packet_rx.ip6_frag.id,
115116
)
116117

117118
# Update flow db
118119
if flow_id in self.ip6_frag_flows:
119-
self.ip6_frag_flows[flow_id]["payload"][
120+
self.ip6_frag_flows[flow_id].payload[
120121
packet_rx.ip6_frag.offset
121122
] = packet_rx.ip6_frag.payload_bytes
122123
else:
123-
self.ip6_frag_flows[flow_id] = {
124-
"header": packet_rx.ip6.header_bytes,
125-
"timestamp": time(),
126-
"last": False,
127-
"payload": {
124+
self.ip6_frag_flows[flow_id] = IpFragData(
125+
header=packet_rx.ip6.header_bytes,
126+
payload={
128127
packet_rx.ip6_frag.offset: packet_rx.ip6_frag.payload_bytes
129128
},
130-
}
129+
)
131130
if not packet_rx.ip6_frag.flag_mf:
132-
self.ip6_frag_flows[flow_id]["last"] = True
131+
self.ip6_frag_flows[flow_id].last = True
133132

134133
# Test if we received all fragments
135-
if not self.ip6_frag_flows[flow_id]["last"]:
134+
if not self.ip6_frag_flows[flow_id].last:
136135
return None
137136
payload_len = 0
138-
for offset in sorted(self.ip6_frag_flows[flow_id]["payload"]):
137+
for offset in sorted(self.ip6_frag_flows[flow_id].payload):
139138
if offset > payload_len:
140139
return None
141140
payload_len = offset + len(
142-
self.ip6_frag_flows[flow_id]["payload"][offset]
141+
self.ip6_frag_flows[flow_id].payload[offset]
143142
)
144143

145144
# Defragment packet
146-
header = bytearray(self.ip6_frag_flows[flow_id]["header"])
145+
header = bytearray(self.ip6_frag_flows[flow_id].header)
147146
payload = bytearray(payload_len)
148-
for offset in sorted(self.ip6_frag_flows[flow_id]["payload"]):
147+
for offset in sorted(self.ip6_frag_flows[flow_id].payload):
149148
struct.pack_into(
150-
f"{len(self.ip6_frag_flows[flow_id]['payload'][offset])}s",
149+
f"{len(self.ip6_frag_flows[flow_id].payload[offset])}s",
151150
payload,
152151
offset,
153-
self.ip6_frag_flows[flow_id]["payload"][offset],
152+
self.ip6_frag_flows[flow_id].payload[offset],
154153
)
155154
del self.ip6_frag_flows[flow_id]
156155
struct.pack_into("!H", header, 4, len(payload))

0 commit comments

Comments
 (0)